import unittest

import pandas as pd


class USComboSweepTests(unittest.TestCase):
    def test_apply_filter_threshold_masks_names_below_rank_cutoff(self):
        from research.us_combo_sweep import apply_filter_threshold

        index = pd.DatetimeIndex([pd.Timestamp("2024-01-31")])
        score = pd.DataFrame({"AAA": [0.9], "BBB": [0.8], "CCC": [0.7]}, index=index)
        filter_rank = pd.DataFrame({"AAA": [0.2], "BBB": [0.6], "CCC": [0.9]}, index=index)

        filtered = apply_filter_threshold(score, filter_rank, min_rank=0.5)

        self.assertTrue(pd.isna(filtered.iloc[0]["AAA"]))
        self.assertEqual(float(filtered.iloc[0]["BBB"]), 0.8)
        self.assertEqual(float(filtered.iloc[0]["CCC"]), 0.7)

    def test_run_combo_backtests_returns_candidates_and_yearly_summary(self):
        from research.us_combo_sweep import run_combo_backtests

        dates = pd.date_range("2022-01-01", periods=800, freq="D")
        close = pd.DataFrame(
            {
                "AAA": [50.0 + 0.12 * i for i in range(800)],
                "BBB": [40.0 + 0.08 * i for i in range(800)],
                "CCC": [35.0 + 0.06 * i for i in range(800)],
                "DDD": [30.0 + 0.04 * i for i in range(800)],
                "EEE": [25.0 + 0.03 * i for i in range(800)],
                "FFF": [20.0 + 0.02 * i for i in range(800)],
                "GGG": [18.0 + 0.015 * i for i in range(800)],
                "HHH": [16.0 + 0.010 * i for i in range(800)],
                "III": [14.0 + 0.008 * i for i in range(800)],
                "JJJ": [12.0 + 0.005 * i for i in range(800)],
                "SPY": [300.0 + 0.20 * i for i in range(800)],
            },
            index=dates,
        )
        fundamental_score = pd.DataFrame(
            {
                "AAA": [0.95] * 800,
                "BBB": [0.90] * 800,
                "CCC": [0.85] * 800,
                "DDD": [0.80] * 800,
                "EEE": [0.75] * 800,
                "FFF": [0.70] * 800,
                "GGG": [0.65] * 800,
                "HHH": [0.60] * 800,
                "III": [0.55] * 800,
                "JJJ": [0.50] * 800,
            },
            index=dates,
        )

        yearly, summary = run_combo_backtests(close, fundamental_score, top_n=3)

        self.assertIn("Recovery+Mom Top10", yearly.columns)
        self.assertIn("rm_fund_tilt_20", yearly.columns)
        self.assertIn("rm_fund_filter_50", yearly.columns)
        self.assertIn("mega_quality_fund", set(summary["strategy"]))


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