"""Integrity checks for the study's forecasting and dependence helpers.

Run: python -m unittest discover -s scripts/vrp-article -p test_study.py
Loads only the helper definitions so tests need no licensed source files and do not
rewrite results. The full study separately asserts input and horizon integrity.
"""
import ast
import math
from pathlib import Path
import unittest

import numpy as np
import pandas as pd


def load_helpers():
    source = ast.parse(Path(__file__).with_name("study.py").read_text(encoding="utf-8"))
    namespace = {"np": np, "pd": pd, "math": math, "H": 21, "MINOBS": 189}
    for node in source.body:
        if isinstance(node, ast.FunctionDef) and node.name in ("hac_mean", "forecast_comparison", "past_rank", "signal_uncertainty", "qlike"):
            exec(compile(ast.Module(body=[node], type_ignores=[]), "study.py", "exec"), namespace)
    return namespace


HELPERS = load_helpers()
hac_mean, forecast_comparison = HELPERS["hac_mean"], HELPERS["forecast_comparison"]


class StudyIntegrityTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        rng = np.random.default_rng(101)
        n = 500
        iv, rv = rng.uniform(10, 40, n), rng.uniform(8, 32, n)
        cls.frame = pd.DataFrame({
            "iv": iv, "cc20": rv,
            "rvf": 3 + .7 * iv + .2 * rv + rng.normal(0, .1, n),
            "end21": np.arange(n) + 21,
        }, index=pd.bdate_range("2000-01-03", periods=n).strftime("%Y-%m-%d"))
        cls.result = forecast_comparison(cls.frame)

    def test_training_embargo_and_nonoverlap_boundaries(self):
        sample = self.result["sample"]
        self.assertEqual(sample["forecast_days"], 500 - 252 - 21)
        self.assertEqual(sample["first_forecast"], self.frame.index[273])
        self.assertEqual(sample["minimum_origin_gap_sessions"], 22)
        self.assertEqual(sample["nonoverlap_days"], 11)

    def test_unobserved_future_outcomes_cannot_change_fitted_coefficients(self):
        changed = self.frame.copy()
        changed.loc[changed.index[-22:], "rvf"] += 10000
        altered = forecast_comparison(changed)
        for model, fit in self.result["last_fit_coefficients"].items():
            np.testing.assert_allclose(fit, altered["last_fit_coefficients"][model])

    def test_missing_origins_do_not_compress_the_session_embargo(self):
        missing = self.frame.copy()
        missing.loc[missing.index[300], "iv"] = np.nan
        sample = forecast_comparison(missing)["sample"]
        self.assertEqual(sample["forecast_days"], 226)
        self.assertEqual(sample["last_training_days"], 477)
        self.assertEqual(sample["minimum_origin_gap_sessions"], 22)

    def test_longer_outcomes_must_finish_before_training_use(self):
        longer = self.frame.copy()
        longer["end_longer"] = np.arange(len(longer)) + 23
        sample = forecast_comparison(longer, end_col="end_longer")["sample"]
        self.assertEqual(sample["first_forecast"], longer.index[275])
        self.assertEqual(sample["minimum_origin_gap_sessions"], 24)

    def test_unconstrained_joint_model_recovers_known_two_input_relation(self):
        models = self.result["forward_realized_volatility"]["models"]
        self.assertLess(models["iv_plus_rv"]["mse"], .02)
        self.assertLess(models["iv_plus_rv"]["mse"], models["iv_only"]["mse"])
        fit = np.asarray(self.result["last_fit_coefficients"]["iv_plus_rv"])
        np.testing.assert_allclose(fit[:, 0], [3, .7, .2], atol=.06)

    def test_hac_matches_hand_calculation_and_keeps_missing_session_gaps(self):
        for values in ([1., 2., 4., 8.], [1., np.nan, 4., 8.]):
            x = np.asarray(values)
            finite = np.isfinite(x)
            z = np.where(finite, x - np.nanmean(x), 0)
            # With lag one the Bartlett weight is 1/2; the two covariance legs cancel it.
            expected = math.sqrt(z @ z + z[1:] @ z[:-1]) / finite.sum()
            result = hac_mean(x, 1)
            self.assertAlmostEqual(result["se"], expected)
            self.assertEqual(result["observations"], int(finite.sum()))
        compressed = hac_mean([1., 4., 8.], 1)["se"]
        with_gap = hac_mean([1., np.nan, 4., 8.], 1)["se"]
        self.assertNotAlmostEqual(compressed, with_gap)

    def test_rank_uses_only_prior_values_and_published_minimum_coverage(self):
        rank = HELPERS["past_rank"]
        window = np.r_[np.arange(189, dtype=float), np.full(63, np.nan), 10.0]
        self.assertAlmostEqual(rank(window), 10 / 189)
        window[0] = np.nan
        self.assertTrue(np.isnan(rank(window)))
        window[0], window[-1] = 0.0, np.nan
        self.assertTrue(np.isnan(rank(window)))

    def test_signal_uncertainty_is_invariant_to_common_outcome_level_shifts(self):
        calendar = pd.bdate_range("2001-01-01", periods=120)
        rng = np.random.default_rng(18)
        group = np.tile([.1, .5, .9], 40)
        frame = pd.DataFrame({"pr": group, "var_post": 10 * group + rng.normal(size=120)}, index=calendar)
        frame = frame.drop(calendar[::7])
        original = HELPERS["signal_uncertainty"](frame, calendar)
        frame.var_post += 10000
        shifted = HELPERS["signal_uncertainty"](frame, calendar)
        self.assertAlmostEqual(original["mean"], shifted["mean"], places=9)
        self.assertAlmostEqual(original["se"], shifted["se"], places=9)

    def test_qlike_is_zero_for_perfect_forecasts_and_matches_ratio_loss(self):
        qlike = HELPERS["qlike"]
        truth = np.array([1.0, 2.0, 3.0])
        self.assertEqual(qlike(truth, truth), 0.0)
        self.assertAlmostEqual(qlike(2 * truth, truth), 1 - math.log(2))


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