"""Tests for the factor-weight learning loop (P4)."""

from __future__ import annotations

import datetime as dt
import unittest

from app import weight_learning as wl


class SpearmanICTest(unittest.TestCase):
    def test_perfect_positive(self):
        # Factor and forward returns perfectly rank-aligned -> IC = 1.
        fv = {"A": 1.0, "B": 2.0, "C": 3.0, "D": 4.0}
        fr = {"A": 0.1, "B": 0.2, "C": 0.3, "D": 0.4}
        ic = wl.spearman_ic(fv, fr)
        self.assertIsNotNone(ic)
        self.assertTrue(ic is not None and abs(ic - 1.0) < 1e-5)

    def test_inverse_is_negative_one(self):
        fv = {"A": 1.0, "B": 2.0, "C": 3.0, "D": 4.0}
        fr = {"A": 0.4, "B": 0.3, "C": 0.2, "D": 0.1}
        ic = wl.spearman_ic(fv, fr)
        self.assertIsNotNone(ic)
        self.assertTrue(ic is not None and abs(ic + 1.0) < 1e-5)

    def test_too_few_symbols_returns_none(self):
        self.assertIsNone(wl.spearman_ic({"A": 1.0}, {"A": 0.1}))

    def test_invalid_symbols_excluded(self):
        fv = {"A": 1.0, "B": 2.0, "C": 3.0, "D": 4.0, "E": float("nan")}
        fr = {"A": 0.1, "B": 0.2, "C": 0.3, "D": 0.4, "E": 0.5}
        ic = wl.spearman_ic(fv, fr)
        self.assertIsNotNone(ic)
        self.assertTrue(ic is not None and abs(ic - 1.0) < 1e-5)


class WeightUpdateTest(unittest.TestCase):
    def test_positive_ic_raises_weight(self):
        l = wl.FactorLearning("f", n_periods=12, ic_mean=0.3,
                              old_weight=1.0)
        wl.apply_weight_update(l, shrink=0.5)
        self.assertIsNotNone(l.new_weight)
        self.assertTrue(l.new_weight is not None and l.new_weight > 1.0)

    def test_negative_ic_lowers_weight(self):
        l = wl.FactorLearning("f", n_periods=12, ic_mean=-0.4,
                              old_weight=1.0)
        wl.apply_weight_update(l, shrink=0.5)
        self.assertIsNotNone(l.new_weight)
        self.assertTrue(l.new_weight is not None and l.new_weight < 1.0)

    def test_clamped_to_bounds(self):
        l = wl.FactorLearning("f", n_periods=12, ic_mean=10.0, old_weight=1.0)
        wl.apply_weight_update(l, shrink=1.0, max_w=3.0)
        self.assertEqual(l.new_weight, 3.0)

    def test_blocked_keeps_weight(self):
        l = wl.FactorLearning("f", n_periods=0, ic_mean=None,
                              old_weight=1.0, blocked=True)
        wl.apply_weight_update(l)
        self.assertEqual(l.new_weight, 1.0)

    def test_no_old_weight_returns(self):
        l = wl.FactorLearning("f", n_periods=12, ic_mean=0.2, old_weight=None)
        wl.apply_weight_update(l)
        self.assertIsNone(l.new_weight)


class MomentumLearningTest(unittest.TestCase):
    @staticmethod
    def _series():
        # A trends up strongly (positive momentum), B flat.
        def bars(base, drift):
            out = []
            for i in range(500):
                d = (dt.date(2024, 1, 1) + dt.timedelta(days=i)).isoformat()
                out.append({"date": d, "adjusted_close": base + drift * i})
            return out
        return {"A": {"bars": bars(10.0, 0.05)}, "B": {"bars": bars(20.0, 0.0)}}

    def test_learn_runs_without_error(self):
        res = wl.learn_momentum(self._series(), ["A", "B"],
                                "2025-06-01", "2026-06-01")
        self.assertIsInstance(res, wl.FactorLearning)
        self.assertGreaterEqual(res.n_periods, 0)


class AddMonthsTest(unittest.TestCase):
    def test_clamps_day_to_end_of_month(self):
        # Jan 31 + 1 month must clamp to Feb 28/29, not raise.
        self.assertEqual(wl._add_months(dt.date(2026, 1, 31), 1), dt.date(2026, 2, 28))
        self.assertEqual(wl._add_months(dt.date(2026, 1, 31), 2), dt.date(2026, 3, 31))
        self.assertEqual(wl._add_months(dt.date(2024, 1, 31), 1), dt.date(2024, 2, 29))  # leap


class LearnFactorSeriesTest(unittest.TestCase):
    def test_aggregates_ics(self):
        res = wl.learn_factor_series([0.1, 0.2, 0.3, 0.4])
        self.assertEqual(res.n_periods, 4)
        self.assertIsNotNone(res.ic_mean)
        self.assertTrue(res.ic_mean is not None and abs(res.ic_mean - 0.25) < 1e-6)
        self.assertIsNotNone(res.ic_tstat)

    def test_empty_is_blocked_like(self):
        res = wl.learn_factor_series([])
        self.assertEqual(res.n_periods, 0)
        self.assertIsNone(res.ic_mean)


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