import unittest

import numpy as np
import pandas as pd

from research import trend_rider_robustness as robustness


class TrendRiderRobustnessTests(unittest.TestCase):
    def test_evaluate_weights_reports_core_risk_metrics(self):
        dates = pd.bdate_range("2024-01-01", periods=6)
        prices = pd.DataFrame(
            {
                "AAA": [100, 110, 105, 120, 118, 130],
                "BBB": [50, 49, 51, 50, 52, 53],
            },
            index=dates,
        )
        weights = pd.DataFrame(
            {
                "AAA": [0, 1, 1, 0, 0, 1],
                "BBB": [0, 0, 0, 1, 1, 0],
            },
            index=dates,
        )

        result = robustness.evaluate_weights("synthetic", weights, prices, transaction_cost=0.001)

        self.assertEqual(result.name, "synthetic")
        self.assertGreater(result.final_multiple, 1.0)
        self.assertLessEqual(result.max_drawdown, 0.0)
        self.assertGreater(result.switches, 0)
        self.assertGreater(result.avg_daily_turnover, 0.0)

    def test_parameter_sweep_returns_rankable_rows(self):
        dates = pd.bdate_range("2023-01-02", periods=320)
        trend = np.linspace(100, 180, len(dates))
        prices = pd.DataFrame(
            {
                "SPY": trend,
                "TQQQ": trend * 1.5,
                "UPRO": trend * 1.4,
                "GLD": np.linspace(100, 105, len(dates)),
                "DBC": np.linspace(90, 95, len(dates)),
            },
            index=dates,
        )

        sweep = robustness.parameter_sweep(
            prices,
            variants=[
                {"vol_enter": 0.14, "dd_stop": 0.05, "peak_enter": 0.02, "mom_lookback": 63},
                {"vol_enter": 0.16, "dd_stop": 0.07, "peak_enter": 0.03, "mom_lookback": 84},
            ],
            start="2023-01-02",
        )

        self.assertEqual(len(sweep), 2)
        self.assertIn("cagr", sweep.columns)
        self.assertIn("max_drawdown", sweep.columns)
        self.assertTrue(sweep["cagr"].notna().all())

    def test_candidate_weights_include_v4_and_market_benchmarks(self):
        dates = pd.bdate_range("2023-01-02", periods=320)
        trend = np.linspace(100, 180, len(dates))
        prices = pd.DataFrame(
            {
                "SPY": trend,
                "QQQ": trend * 1.1,
                "SSO": trend * 1.5,
                "QLD": trend * 1.6,
                "UPRO": trend * 2.0,
                "TQQQ": trend * 2.2,
                "SHY": np.linspace(100, 103, len(dates)),
                "IEF": np.linspace(100, 104, len(dates)),
                "TLT": np.linspace(100, 105, len(dates)),
                "GLD": np.linspace(100, 115, len(dates)),
                "DBC": np.linspace(90, 105, len(dates)),
            },
            index=dates,
        )

        candidates = robustness.candidate_weights(prices)

        self.assertIn("TrendRiderV4", candidates)
        self.assertIn("SPY Buy&Hold", candidates)
        self.assertIn("QQQ Buy&Hold", candidates)


if __name__ == "__main__":
    unittest.main()