"""Tests for eapctf.ctf.predict — ctf_expanding_window_oos()."""

from __future__ import annotations

import numpy as np
import pandas as pd
import pytest

from eapctf.ctf.predict import ctf_expanding_window_oos
from eapctf.predict import make_predictor


# ---------------------------------------------------------------------------
# Shared synthetic CTF data builder
# ---------------------------------------------------------------------------


def _make_ctf_data(
    n_stocks: int = 20,
    n_months: int = 24,
    n_features: int = 5,
    n_test_months: int = 6,
    seed: int = 0,
) -> 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


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------


class TestCtfExpandingWindowOos:
    def test_returns_oos_result(self) -> None:
        """ctf_expanding_window_oos returns an OOSResult."""
        from eapctf.predict.oos import OOSResult

        chars, features = _make_ctf_data(n_stocks=20, n_months=24, n_test_months=6)
        models = [make_predictor("ridge")]
        oos = ctf_expanding_window_oos(
            chars, features,
            models=models,
            window_length=12,
        )
        assert isinstance(oos, OOSResult)

    def test_predictions_columns(self) -> None:
        """Predictions DataFrame has required columns."""
        chars, features = _make_ctf_data(n_stocks=20, n_months=24, n_test_months=6)
        oos = ctf_expanding_window_oos(
            chars, features,
            models=[make_predictor("ridge")],
            window_length=12,
        )
        preds = oos.predictions
        assert "date" in preds.columns
        assert "permno" in preds.columns
        assert "ret" in preds.columns
        # At least one pred_* column
        pred_cols = [c for c in preds.columns if c.startswith("pred_")]
        assert len(pred_cols) >= 1

    def test_test_only_filters_to_test_dates(self) -> None:
        """With test_only=True, predictions include only ctff_test dates."""
        chars, features = _make_ctf_data(n_stocks=20, n_months=24, n_test_months=6)
        test_dates = set(
            chars.loc[chars["ctff_test"].astype(bool), "eom"].unique()
        )
        oos = ctf_expanding_window_oos(
            chars, features,
            models=[make_predictor("ridge")],
            window_length=12,
            test_only=True,
        )
        pred_dates = set(oos.predictions["date"].astype(str).unique())
        # All predicted dates must be test dates
        assert pred_dates.issubset(test_dates)

    def test_test_only_false_includes_more_dates(self) -> None:
        """With test_only=False, predictions include non-test-period dates too."""
        chars, features = _make_ctf_data(n_stocks=20, n_months=24, n_test_months=6)
        oos_all = ctf_expanding_window_oos(
            chars, features,
            models=[make_predictor("ridge")],
            window_length=12,
            test_only=False,
        )
        oos_test = ctf_expanding_window_oos(
            chars, features,
            models=[make_predictor("ridge")],
            window_length=12,
            test_only=True,
        )
        # test_only=False should have >= rows than test_only=True
        assert len(oos_all.predictions) >= len(oos_test.predictions)

    def test_rank_normalize_false_runs(self) -> None:
        """rank_normalize=False should still produce valid OOSResult."""
        chars, features = _make_ctf_data(n_stocks=20, n_months=24, n_test_months=6)
        oos = ctf_expanding_window_oos(
            chars, features,
            models=[make_predictor("ridge")],
            window_length=12,
            rank_normalize=False,
        )
        assert not oos.predictions.empty

    def test_feature_subset_jkp153_empty_falls_back(self) -> None:
        """subset='jkp153' with no JKP-153 features returns empty or raises gracefully."""
        chars, features = _make_ctf_data(n_stocks=20, n_months=24, n_features=3)
        # None of f0, f1, f2 are in JKP_153 — predictions may be empty
        oos = ctf_expanding_window_oos(
            chars, features,
            models=[make_predictor("ridge")],
            window_length=12,
            feature_subset="jkp153",
        )
        # Should return OOSResult regardless (empty predictions is acceptable)
        from eapctf.predict.oos import OOSResult
        assert isinstance(oos, OOSResult)

    def test_feature_subset_all(self) -> None:
        """feature_subset='all' uses all available features."""
        chars, features = _make_ctf_data(n_stocks=20, n_months=24, n_features=5)
        oos = ctf_expanding_window_oos(
            chars, features,
            models=[make_predictor("ridge")],
            window_length=12,
            feature_subset="all",
        )
        assert not oos.predictions.empty

    def test_feature_subset_available(self) -> None:
        """feature_subset='available' selects features with sufficient coverage."""
        chars, features = _make_ctf_data(n_stocks=20, n_months=24, n_features=5)
        oos = ctf_expanding_window_oos(
            chars, features,
            models=[make_predictor("ridge")],
            window_length=12,
            feature_subset="available",
        )
        from eapctf.predict.oos import OOSResult
        assert isinstance(oos, OOSResult)

    def test_invalid_feature_subset_raises(self) -> None:
        """Invalid feature_subset raises ValueError."""
        chars, features = _make_ctf_data()
        with pytest.raises(ValueError, match="feature_subset must be one of"):
            ctf_expanding_window_oos(
                chars, features,
                models=[make_predictor("ridge")],
                window_length=12,
                feature_subset="invalid",
            )

    def test_to_ctf_weights_pipeline(self) -> None:
        """End-to-end: ctf_expanding_window_oos -> to_ctf_weights produces valid output."""
        chars, features = _make_ctf_data(n_stocks=20, n_months=24, n_test_months=6)
        oos = ctf_expanding_window_oos(
            chars, features,
            models=[make_predictor("ridge")],
            window_length=12,
        )
        # Use a model-specific column (varies cross-sectionally) rather than
        # pred_HistMean (constant per date) to avoid degenerate z-score input.
        model_pred_cols = [
            c for c in oos.predictions.columns
            if c.startswith("pred_") and c != "pred_HistMean"
        ]
        if not model_pred_cols or oos.predictions.empty:
            pytest.skip("No predictions generated with this synthetic data size")

        weights = oos.to_ctf_weights(signal_col=model_pred_cols[0], mode="zscore")
        assert set(weights.columns) == {"id", "eom", "w"}
        assert weights["id"].dtype in (int, "int64", "int32")
        assert not weights[["id", "eom"]].duplicated().any(), "Duplicate (id, eom) pairs found"
