import unittest

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]), 1.0)
        self.assertEqual(float(equity.iloc[2]), 2.0)

    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())


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