import unittest

from quant60.factor_risk import (
    FactorExposureSnapshot,
    RiskModelError,
    build_factor_risk_forecast,
    estimate_factor_returns,
    factor_stress_loss,
    portfolio_factor_variance,
)


class FactorRiskTests(unittest.TestCase):
    def _exposure(self):
        return FactorExposureSnapshot(
            as_of="2026-07-24T15:00:00+08:00",
            exposures={
                "600000.SH": {"market": 1.0, "style": 1.0},
                "000001.SZ": {"market": 1.0, "style": 1.0},
                "300750.SZ": {"market": 1.0, "style": -1.0},
            },
        )

    def _forecast(self):
        return build_factor_risk_forecast(
            factor_return_history=[
                {"market": 0.01, "style": 0.02},
                {"market": -0.01, "style": -0.02},
                {"market": 0.02, "style": 0.01},
                {"market": -0.02, "style": -0.01},
            ],
            residual_history={
                "600000.SH": [0.01, -0.01, 0.005],
                "000001.SZ": [0.01, -0.01, 0.005],
                "300750.SZ": [0.01, -0.01, 0.005],
            },
            shrinkage=0.2,
        )

    def test_factor_return_fit_and_residuals_are_complete(self):
        fit = estimate_factor_returns(
            returns={
                "600000.SH": 0.03,
                "000001.SZ": 0.02,
                "300750.SZ": -0.01,
            },
            exposure=self._exposure(),
        )
        self.assertEqual(set(fit.factor_returns), {"market", "style"})
        self.assertEqual(len(fit.residuals), 3)

    def test_correlated_pair_has_more_risk_than_style_diversified_pair(self):
        exposure = self._exposure()
        forecast = self._forecast()
        correlated = portfolio_factor_variance(
            portfolio_weights={
                "600000.SH": 0.5,
                "000001.SZ": 0.5,
            },
            exposure=exposure,
            forecast=forecast,
        )
        diversified = portfolio_factor_variance(
            portfolio_weights={
                "600000.SH": 0.5,
                "300750.SZ": 0.5,
            },
            exposure=exposure,
            forecast=forecast,
        )
        self.assertGreater(correlated, diversified)

    def test_missing_risk_input_fails_closed(self):
        with self.assertRaisesRegex(RiskModelError, "missing"):
            portfolio_factor_variance(
                portfolio_weights={"000333.SZ": 1.0},
                exposure=self._exposure(),
                forecast=self._forecast(),
            )

    def test_market_down_stress_is_positive_loss(self):
        loss = factor_stress_loss(
            portfolio_weights={"600000.SH": 0.5, "300750.SZ": 0.5},
            exposure=self._exposure(),
            factor_shocks={"market": -0.10},
        )
        self.assertAlmostEqual(loss, 0.10)


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