import unittest
from pathlib import Path
from unittest import mock

import pandas as pd


class USAlphaPipelineTests(unittest.TestCase):
    def test_build_equal_weight_portfolio_caps_holdings_under_ties(self):
        from research.us_alpha_pipeline import _build_equal_weight_portfolio

        dates = pd.date_range("2024-01-01", periods=2, freq="D")
        score = pd.DataFrame(
            {
                "AAA": [0.9, 0.9],
                "BBB": [0.9, 0.9],
                "CCC": [0.9, 0.9],
            },
            index=dates,
        )
        tradable_mask = pd.DataFrame(True, index=dates, columns=score.columns)
        regime = pd.Series([True, True], index=dates)

        weights = _build_equal_weight_portfolio(score, tradable_mask, regime, top_n=2)

        self.assertEqual(int((weights.iloc[-1] > 0).sum()), 2)
        self.assertAlmostEqual(float(weights.iloc[-1].sum()), 1.0)

    def test_equity_curve_uses_prior_day_weights_for_returns(self):
        from research.us_alpha_pipeline import _equity_curve

        dates = pd.date_range("2024-01-01", periods=3, freq="D")
        close = pd.DataFrame({"AAA": [1.0, 2.0, 4.0]}, index=dates)
        weights = pd.DataFrame({"AAA": [0.0, 1.0, 0.0]}, index=dates)

        equity = _equity_curve(close, weights)

        self.assertEqual(float(equity.iloc[1]), 2.0)
        self.assertEqual(float(equity.iloc[2]), 2.0)

    def test_summarize_equity_window_returns_nans_when_history_is_too_short(self):
        from research.us_alpha_report import summarize_equity_window

        dates = pd.date_range("2024-01-01", periods=10, freq="D")
        equity = pd.Series([1.0 + 0.01 * i for i in range(10)], index=dates)

        summary = summarize_equity_window(equity, "demo", window_years=1)

        self.assertTrue(pd.isna(summary["CAGR"]))
        self.assertTrue(pd.isna(summary["Sharpe"]))
        self.assertTrue(pd.isna(summary["MaxDD"]))
        self.assertTrue(pd.isna(summary["TotalRet"]))

    def test_run_alpha_pipeline_returns_expected_strategy_summary(self):
        from research.us_alpha_pipeline import run_alpha_pipeline

        dates = pd.date_range("2023-01-01", periods=400, freq="D")

        aaa_close = [50.0 + 0.20 * i for i in range(400)]
        bbb_close = [55.0 + 0.12 * i for i in range(400)]
        ccc_close = [60.0 + 0.05 * i for i in range(400)]
        close = pd.DataFrame(
            {
                "AAA": aaa_close,
                "BBB": bbb_close,
                "CCC": ccc_close,
            },
            index=dates,
        )
        high = pd.DataFrame(
            {
                "AAA": [value + 0.5 for value in aaa_close],
                "BBB": [value + 1.0 for value in bbb_close],
                "CCC": [value + 1.5 for value in ccc_close],
            },
            index=dates,
        )
        low = pd.DataFrame(
            {
                "AAA": [value - 0.5 for value in aaa_close],
                "BBB": [value - 1.0 for value in bbb_close],
                "CCC": [value - 1.5 for value in ccc_close],
            },
            index=dates,
        )
        volume = pd.DataFrame(
            {
                "AAA": [1_500_000.0] * 400,
                "BBB": [1_400_000.0] * 400,
                "CCC": [1_300_000.0] * 400,
            },
            index=dates,
        )
        volume.loc[dates[-2], "AAA"] = 4_000_000.0

        etf_close = pd.DataFrame(
            {
                "SPY": [300.0 + 0.8 * i for i in range(400)],
                "QQQ": [280.0 + 1.1 * i for i in range(400)],
                "XLF": [200.0 + 0.4 * i for i in range(400)],
            },
            index=dates,
        )

        market_data = {
            "close": close,
            "high": high,
            "low": low,
            "volume": volume,
        }

        summary = run_alpha_pipeline(
            market_data=market_data,
            etf_close=etf_close,
            pit_membership=None,
            windows=(1,),
            top_n=2,
        )

        required_columns = {"strategy", "window_years", "CAGR", "Sharpe", "MaxDD", "TotalRet"}
        self.assertTrue(required_columns.issubset(summary.columns))
        self.assertEqual(set(summary["strategy"]), {"breakout_regime", "rank_blend_regime"})
        self.assertEqual(set(summary["window_years"]), {1})
        self.assertEqual(len(summary), 2)
        self.assertTrue(summary[["CAGR", "Sharpe", "MaxDD", "TotalRet"]].notna().all().all())

    def test_run_saved_pit_alpha_pipeline_reads_saved_inputs(self):
        from research.us_alpha_pipeline import run_saved_pit_alpha_pipeline

        dates = pd.date_range("2024-01-01", periods=320, freq="D")
        close = pd.DataFrame(
            {
                "AAA": [50.0 + 0.2 * i for i in range(320)],
                "BBB": [40.0 + 0.1 * i for i in range(320)],
            },
            index=dates,
        )
        high = close + 1.0
        low = close - 1.0
        volume = pd.DataFrame({"AAA": [2_500_000.0] * 320, "BBB": [2_000_000.0] * 320}, index=dates)
        etf_close = pd.DataFrame(
            {"SPY": [300.0 + 0.8 * i for i in range(320)], "QQQ": [280.0 + 1.1 * i for i in range(320)]},
            index=dates,
        )

        with self.subTest("saved_inputs"):
            import tempfile

            with tempfile.TemporaryDirectory() as tmpdir:
                close.to_csv(Path(tmpdir) / "us_pit_close.csv")
                high.to_csv(Path(tmpdir) / "us_pit_high.csv")
                low.to_csv(Path(tmpdir) / "us_pit_low.csv")
                volume.to_csv(Path(tmpdir) / "us_pit_volume.csv")
                etf_close.to_csv(Path(tmpdir) / "us_etf.csv")

                intervals = {"AAA": [[None, None]], "BBB": [[None, None]]}
                with mock.patch("research.us_alpha_pipeline.uh.load_sp500_history", return_value=intervals):
                    summary = run_saved_pit_alpha_pipeline(data_dir=tmpdir, windows=(1,), top_n=1)

        self.assertEqual(set(summary["strategy"]), {"breakout_regime", "rank_blend_regime"})


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