import unittest

import pandas as pd


class USFundamentalTransformsTests(unittest.TestCase):
    def test_build_quarterly_factor_pack_derives_expected_signals(self):
        from research.us_fundamentals import build_quarterly_factor_pack

        quarter_ends = pd.to_datetime(
            ["2023-03-31", "2023-06-30", "2023-09-30", "2023-12-31", "2024-03-31"]
        )
        close = pd.DataFrame(
            {"AAA": [100.0, 105.0], "BBB": [50.0, 48.0]},
            index=pd.to_datetime(["2024-06-03", "2024-06-04"]),
        )
        quarterly = {
            "net_income": pd.DataFrame(
                {"AAA": [10.0, 11.0, 12.0, 13.0, 14.0], "BBB": [4.0, 4.0, 5.0, 5.0, 5.0]},
                index=quarter_ends,
            ),
            "gross_profit": pd.DataFrame(
                {"AAA": [30.0, 31.0, 32.0, 33.0, 34.0], "BBB": [10.0, 10.0, 11.0, 11.0, 11.0]},
                index=quarter_ends,
            ),
            "equity": pd.DataFrame(
                {"AAA": [200.0, 205.0, 210.0, 215.0, 220.0], "BBB": [80.0, 81.0, 82.0, 83.0, 84.0]},
                index=quarter_ends,
            ),
            "assets": pd.DataFrame(
                {"AAA": [300.0, 305.0, 310.0, 315.0, 320.0], "BBB": [130.0, 131.0, 132.0, 133.0, 134.0]},
                index=quarter_ends,
            ),
            "shares": pd.DataFrame(
                {"AAA": [10.0, 10.0, 10.0, 10.0, 10.0], "BBB": [10.0, 10.0, 11.0, 11.0, 11.0]},
                index=quarter_ends,
            ),
        }

        factor_pack = build_quarterly_factor_pack(quarterly, close, lag_days=60)

        self.assertIn("composite", factor_pack)
        self.assertIn("book_to_market", factor_pack)
        self.assertEqual(list(factor_pack["composite"].columns), ["AAA", "BBB"])
        self.assertGreater(
            float(factor_pack["composite"].iloc[-1]["AAA"]),
            float(factor_pack["composite"].iloc[-1]["BBB"]),
        )


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