"""Tests for eapctf.ctf.conformal — ctf_conformal_oos() and helpers."""

from __future__ import annotations

import numpy as np
import pandas as pd
import pytest

# Skip entire module if mapie is not installed
pytest.importorskip("mapie", reason="mapie not installed — skip conformal tests")

from eapctf.ctf.conformal import (  # noqa: E402
    AdaptiveConformalInference,
    _signal_to_weights,
    ctf_conformal_oos,
)
from eapctf.predict import make_predictor  # noqa: E402


# ---------------------------------------------------------------------------
# Synthetic data builder
# ---------------------------------------------------------------------------


def _make_ctf_data(
    n_stocks: int = 20,
    n_months: int = 30,
    n_features: int = 4,
    n_test_months: int = 6,
    seed: int = 42,
) -> tuple[pd.DataFrame, pd.DataFrame]:
    """Build minimal synthetic (chars, features) for testing."""
    rng = np.random.default_rng(seed)
    dates = pd.date_range("2010-01-31", periods=n_months, freq="ME")
    ids = list(range(1, n_stocks + 1))
    feat_names = [f"f{i}" for i in range(n_features)]
    test_cutoff = dates[-n_test_months]

    rows = []
    for d in dates:
        for sid in ids:
            row: dict = {
                "id": sid,
                "eom": d.date().isoformat(),
                "ret_exc_lead1m": float(rng.normal(0.005, 0.02)),
                "ctff_test": d >= test_cutoff,
            }
            for f in feat_names:
                row[f] = float(rng.normal(0, 1))
            rows.append(row)

    chars = pd.DataFrame(rows)
    features_df = pd.DataFrame({"features": feat_names})
    return chars, features_df


# ---------------------------------------------------------------------------
# AdaptiveConformalInference
# ---------------------------------------------------------------------------


class TestAdaptiveConformalInference:
    def test_initial_alpha_t_equals_alpha(self) -> None:
        aci = AdaptiveConformalInference(alpha=0.1, gamma=0.05)
        assert aci.alpha_t == 0.1

    def test_update_increases_alpha_when_over_covered(self) -> None:
        """If all stocks are covered (err < alpha), alpha_t should increase."""
        aci = AdaptiveConformalInference(alpha=0.1, gamma=0.1)
        y = np.array([0.01, 0.02, -0.01])
        lo = np.array([-1.0, -1.0, -1.0])
        hi = np.array([1.0, 1.0, 1.0])
        new_alpha = aci.update(y, lo, hi)
        # err_t = 0 (all covered), update = 0.1 + 0.1 * (0.1 - 0) = 0.11
        assert new_alpha > 0.1
        assert aci.alpha_t == new_alpha

    def test_update_decreases_alpha_when_under_covered(self) -> None:
        """If nothing is covered (err > alpha), alpha_t should decrease."""
        aci = AdaptiveConformalInference(alpha=0.1, gamma=0.1)
        y = np.array([5.0, 5.0, 5.0])
        lo = np.array([-0.01, -0.01, -0.01])
        hi = np.array([0.01, 0.01, 0.01])
        new_alpha = aci.update(y, lo, hi)
        # err_t = 1.0 (none covered), update = 0.1 + 0.1 * (0.1 - 1.0) = 0.01
        assert new_alpha < 0.1

    def test_alpha_t_clipped_to_valid_range(self) -> None:
        aci = AdaptiveConformalInference(alpha=0.1, gamma=1.0)
        # Force many under-covered updates
        for _ in range(20):
            aci.update(
                np.array([5.0]),
                np.array([-0.01]),
                np.array([0.01]),
            )
        assert aci.alpha_t >= 0.01

        aci2 = AdaptiveConformalInference(alpha=0.1, gamma=1.0)
        for _ in range(20):
            aci2.update(
                np.array([0.0]),
                np.array([-1.0]),
                np.array([1.0]),
            )
        assert aci2.alpha_t <= 0.99


# ---------------------------------------------------------------------------
# _signal_to_weights
# ---------------------------------------------------------------------------


class TestSignalToWeights:
    def test_zscore_sum_abs_one(self) -> None:
        sig = np.array([1.0, 2.0, 3.0, -1.0])
        w = _signal_to_weights(sig, "zscore")
        assert abs(np.abs(w).sum() - 1.0) < 1e-9

    def test_zscore_sum_zero(self) -> None:
        sig = np.array([1.0, 2.0, 3.0, -1.0])
        w = _signal_to_weights(sig, "zscore")
        assert abs(w.sum()) < 1e-9

    def test_rank_weighted_sum_abs_one(self) -> None:
        sig = np.array([1.0, 3.0, 2.0])
        w = _signal_to_weights(sig, "rank_weighted")
        assert abs(np.abs(w).sum() - 1.0) < 1e-9

    def test_long_short_sum_zero(self) -> None:
        sig = np.linspace(-1, 1, 10)
        w = _signal_to_weights(sig, "long_short", n_stocks=3)
        assert abs(w.sum()) < 1e-12

    def test_invalid_mode_raises(self) -> None:
        with pytest.raises(ValueError, match="weight_mode"):
            _signal_to_weights(np.array([1.0, 2.0]), "invalid_mode")

    def test_empty_signal_returns_empty(self) -> None:
        w = _signal_to_weights(np.array([]), "zscore")
        assert len(w) == 0


# ---------------------------------------------------------------------------
# ctf_conformal_oos
# ---------------------------------------------------------------------------


class TestCtfConformalOos:
    """Integration tests for the full conformal OOS pipeline."""

    @pytest.fixture(scope="class")
    def small_data(self) -> tuple[pd.DataFrame, pd.DataFrame]:
        return _make_ctf_data(n_stocks=15, n_months=20, n_features=3, n_test_months=4)

    def test_returns_dataframe(
        self, small_data: tuple[pd.DataFrame, pd.DataFrame]
    ) -> None:
        chars, features = small_data
        result = ctf_conformal_oos(
            chars, features,
            model=make_predictor("ridge"),
            window_length=10,
            n_conformal_folds=3,
        )
        assert isinstance(result, pd.DataFrame)

    def test_output_columns(
        self, small_data: tuple[pd.DataFrame, pd.DataFrame]
    ) -> None:
        chars, features = small_data
        result = ctf_conformal_oos(
            chars, features,
            model=make_predictor("ridge"),
            window_length=10,
            n_conformal_folds=3,
        )
        expected = {"id", "eom", "w", "w_baseline", "interval_width", "alpha_t"}
        assert expected.issubset(result.columns)

    def test_test_only_filters_to_test_dates(
        self, small_data: tuple[pd.DataFrame, pd.DataFrame]
    ) -> None:
        chars, features = small_data
        result = ctf_conformal_oos(
            chars, features,
            model=make_predictor("ridge"),
            window_length=10,
            n_conformal_folds=3,
            test_only=True,
        )
        test_eoms = set(
            chars.loc[chars["ctff_test"].astype(bool), "eom"].unique()
        )
        assert set(result["eom"].unique()).issubset(test_eoms)

    def test_test_only_false_includes_all_dates(
        self, small_data: tuple[pd.DataFrame, pd.DataFrame]
    ) -> None:
        chars, features = small_data
        result_all = ctf_conformal_oos(
            chars, features,
            model=make_predictor("ridge"),
            window_length=10,
            n_conformal_folds=3,
            test_only=False,
        )
        result_test = ctf_conformal_oos(
            chars, features,
            model=make_predictor("ridge"),
            window_length=10,
            n_conformal_folds=3,
            test_only=True,
        )
        assert len(result_all) >= len(result_test)

    def test_interval_width_nonnegative(
        self, small_data: tuple[pd.DataFrame, pd.DataFrame]
    ) -> None:
        chars, features = small_data
        result = ctf_conformal_oos(
            chars, features,
            model=make_predictor("ridge"),
            window_length=10,
            n_conformal_folds=3,
        )
        assert (result["interval_width"] >= 0).all()

    def test_alpha_t_in_valid_range(
        self, small_data: tuple[pd.DataFrame, pd.DataFrame]
    ) -> None:
        chars, features = small_data
        result = ctf_conformal_oos(
            chars, features,
            model=make_predictor("ridge"),
            window_length=10,
            n_conformal_folds=3,
        )
        assert (result["alpha_t"] >= 0.01).all()
        assert (result["alpha_t"] <= 0.99).all()

    def test_zscore_weights_sum_zero_per_date(
        self, small_data: tuple[pd.DataFrame, pd.DataFrame]
    ) -> None:
        """zscore mode: sum(w) ≈ 0 per date."""
        chars, features = small_data
        result = ctf_conformal_oos(
            chars, features,
            model=make_predictor("ridge"),
            window_length=10,
            n_conformal_folds=3,
            weight_mode="zscore",
            test_only=False,
        )
        for _, grp in result.groupby("eom"):
            assert abs(grp["w"].sum()) < 1e-6, f"sum(w) ≠ 0 at {grp['eom'].iloc[0]}"

    def test_lambda_zero_w_equals_w_baseline(
        self, small_data: tuple[pd.DataFrame, pd.DataFrame]
    ) -> None:
        """With lambda_t=0, mu_tilde = point_pred → w should equal w_baseline."""
        chars, features = small_data
        result = ctf_conformal_oos(
            chars, features,
            model=make_predictor("ridge"),
            window_length=10,
            n_conformal_folds=3,
            lambda_t=0.0,
        )
        if not result.empty:
            np.testing.assert_allclose(
                result["w"].to_numpy(),
                result["w_baseline"].to_numpy(),
                atol=1e-9,
            )

    def test_invalid_feature_subset_raises(
        self, small_data: tuple[pd.DataFrame, pd.DataFrame]
    ) -> None:
        chars, features = small_data
        with pytest.raises(ValueError, match="feature_subset"):
            ctf_conformal_oos(
                chars, features,
                model=make_predictor("ridge"),
                feature_subset="bad_subset",
            )

    def test_invalid_weight_mode_raises(
        self, small_data: tuple[pd.DataFrame, pd.DataFrame]
    ) -> None:
        chars, features = small_data
        with pytest.raises(ValueError, match="weight_mode"):
            ctf_conformal_oos(
                chars, features,
                model=make_predictor("ridge"),
                window_length=10,
                n_conformal_folds=3,
                weight_mode="bad_mode",
            )

    def test_invalid_refit_freq_raises(
        self, small_data: tuple[pd.DataFrame, pd.DataFrame]
    ) -> None:
        chars, features = small_data
        with pytest.raises(ValueError, match="refit_freq"):
            ctf_conformal_oos(
                chars, features,
                model=make_predictor("ridge"),
                window_length=10,
                n_conformal_folds=3,
                refit_freq="weekly",
            )

    def test_rank_normalize_false(
        self, small_data: tuple[pd.DataFrame, pd.DataFrame]
    ) -> None:
        chars, features = small_data
        result = ctf_conformal_oos(
            chars, features,
            model=make_predictor("ridge"),
            window_length=10,
            n_conformal_folds=3,
            rank_normalize=False,
        )
        assert isinstance(result, pd.DataFrame)
        assert not result.empty

    def test_monthly_refit_freq(
        self, small_data: tuple[pd.DataFrame, pd.DataFrame]
    ) -> None:
        chars, features = small_data
        result = ctf_conformal_oos(
            chars, features,
            model=make_predictor("ridge"),
            window_length=10,
            n_conformal_folds=3,
            refit_freq="monthly",
        )
        assert isinstance(result, pd.DataFrame)

    def test_no_duplicate_id_eom(
        self, small_data: tuple[pd.DataFrame, pd.DataFrame]
    ) -> None:
        chars, features = small_data
        result = ctf_conformal_oos(
            chars, features,
            model=make_predictor("ridge"),
            window_length=10,
            n_conformal_folds=3,
        )
        dupes = result.duplicated(subset=["id", "eom"])
        assert not dupes.any(), "Duplicate (id, eom) pairs found"
