from datetime import datetime, timedelta

import numpy as np

from ml import optimizer


def _row(date, returns=10.0, score_200w=10, score_drawdown=0):
    row = {
        "date": date,
        "price": 100.0,
        "fwd_365d": returns,
        "score_puell_multiple": 0,
        "score_mvrv_zscore": 0,
        "score_reserve_risk": 0,
        "score_rhodl_ratio": 0,
        "score_nupl": 0,
        "score_fear_greed": 0,
        "score_drawdown": score_drawdown,
        "score_pct_above_200w_sma": score_200w,
        "score_pct_above_lth_rp": 0,
    }
    return row


def test_purged_time_series_splits_remove_overlapping_forward_label_windows():
    start = datetime(2020, 1, 1)
    rows = [_row((start + timedelta(days=i)).strftime("%Y-%m-%d")) for i in range(900)]

    splits = list(
        optimizer.purged_time_series_splits(
            rows,
            n_splits=3,
            label_horizon_days=365,
            embargo_days=0,
        )
    )

    assert splits, "expected at least one viable split"
    for train_idx, val_idx in splits:
        val_start = datetime.strptime(rows[val_idx[0]]["date"], "%Y-%m-%d")
        latest_allowed_train_date = val_start - timedelta(days=365)
        assert len(train_idx) > 0, "purging should keep non-overlapping expanding-window training rows"
        for idx in train_idx:
            train_date = datetime.strptime(rows[idx]["date"], "%Y-%m-%d")
            assert train_date <= latest_allowed_train_date


def test_run_out_of_sample_comparison_scores_only_validation_rows_with_fold_weights():
    rows = [
        _row("2020-01-01", returns=-10, score_200w=0, score_drawdown=10),
        _row("2020-01-02", returns=-5, score_200w=0, score_drawdown=10),
        _row("2020-01-03", returns=100, score_200w=10, score_drawdown=10),
        _row("2020-01-04", returns=120, score_200w=10, score_drawdown=10),
    ]
    fold_results = [
        {
            "fold": 1,
            "val_idx": [2, 3],
            "weights": {"pct_above_200w_sma": 1.0, "drawdown": 0.0},
        }
    ]

    comparison = optimizer.run_out_of_sample_comparison(rows, fold_results)

    assert comparison["validation_days"] == 2
    assert comparison["folds"] == 1
    assert sum(bucket["days"] for bucket in comparison["ml_weighted"]) == 2
    assert sum(bucket["days"] for bucket in comparison["equal_weight"]) == 2

    extreme_ml = next(bucket for bucket in comparison["ml_weighted"] if bucket["label"] == "EXTREME ACCUMULATION ZONE")
    assert extreme_ml["days"] == 2
    assert extreme_ml["avg_365d"] == 110.0

    caution_equal = next(bucket for bucket in comparison["equal_weight"] if bucket["label"] == "CAUTION — OVERHEATED")
    assert caution_equal["days"] == 2


def test_classification_splits_skip_training_windows_with_one_class():
    y = np.array([1, 1, 1, 0, 1, 0])
    splits = [
        (np.array([0, 1]), np.array([2, 3])),
        (np.array([0, 1, 3, 4]), np.array([5])),
    ]

    viable = list(optimizer.viable_classification_splits(y, splits))

    assert len(viable) == 1
    assert viable[0][0].tolist() == [0, 1, 3, 4]