import math
import unittest

from quant60.research import (
    DeterministicRidge,
    NumpyRidge,
    pearson_correlation,
    purged_walk_forward,
    spearman_correlation,
)
from quant60.risk import check_target_weights, diagonal_variance


class ResearchRiskTests(unittest.TestCase):
    def test_purged_walk_forward_is_time_ordered(self):
        folds = list(
            purged_walk_forward(
                40,
                train_size=20,
                test_size=5,
                purge=2,
                embargo=1,
            )
        )
        self.assertTrue(folds)
        for fold in folds:
            self.assertLess(max(fold.train), min(fold.purged))
            self.assertLess(max(fold.purged), min(fold.test))
            self.assertTrue(set(fold.train).isdisjoint(fold.test))

    def test_risk_check_rejects_concentration(self):
        result = check_target_weights(
            {"600000.SH": 0.6, "000001.SZ": 0.2},
            max_single_weight=0.5,
            max_gross_weight=0.9,
        )
        self.assertFalse(result.approved)
        self.assertIn("SINGLE_NAME_CAP", {item.code for item in result.violations})

    def test_diagonal_variance_floor(self):
        result = diagonal_variance(
            {"600000.SH": [0.01], "000001.SZ": [0.0, 0.0, 0.0]}
        )
        self.assertGreater(result["600000.XSHG"], 0)
        self.assertGreater(result["000001.XSHE"], 0)

    def test_non_finite_weights_fail_closed(self):
        for value in (float("nan"), float("inf"), float("-inf")):
            with self.subTest(value=value):
                result = check_target_weights(
                    {"600000.SH": value},
                    max_single_weight=0.5,
                    max_gross_weight=0.9,
                )
                self.assertFalse(result.approved)
                self.assertIn(
                    "NON_FINITE_WEIGHT",
                    {item.code for item in result.violations},
                )
                self.assertTrue(math.isinf(result.gross_weight))
                self.assertTrue(math.isinf(result.one_way_turnover))

    def test_non_finite_returns_are_rejected(self):
        for value in (float("nan"), float("inf"), float("-inf")):
            with self.subTest(value=value):
                with self.assertRaisesRegex(ValueError, "returns must be finite"):
                    diagonal_variance({"600000.SH": [0.01, value]})

    def test_optional_numpy_ridge_when_numpy_is_present(self):
        try:
            model = NumpyRidge(alpha=0.1).fit([[0], [1], [2]], [1, 3, 5])
        except Exception as exc:
            if type(exc).__name__ == "ResearchDependencyUnavailable":
                self.skipTest(str(exc))
            raise
        prediction = model.predict([[3]])
        self.assertGreater(float(prediction[0]), 6.0)

    def test_dependency_free_ridge_fits_linear_relationship(self):
        model = DeterministicRidge(alpha=0.001).fit(
            [[0], [1], [2], [3]],
            [1, 3, 5, 7],
        )
        prediction = model.predict([[4]])[0]
        self.assertAlmostEqual(prediction, 9, places=2)

    def test_rank_and_linear_correlation(self):
        self.assertAlmostEqual(
            pearson_correlation([1, 2, 3], [2, 4, 6]),
            1.0,
        )
        self.assertAlmostEqual(
            spearman_correlation([10, 30, 20], [1, 3, 2]),
            1.0,
        )


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