"""Tests for eapctf.predict: OOS prediction, models, leakage detection."""

from __future__ import annotations

import numpy as np
import pandas as pd
import pytest

from eapctf.predict import (
    LassoPredictor,
    OOSResult,
    RidgePredictor,
    eval_oos_result,
    expanding_window_oos,
    portfolio_eval,
    prediction_eval,
)
from eapctf.utils import TemporalLeakageError


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_small_panel(
    n_stocks: int = 20,
    n_months: int = 40,
    seed: int = 42,
) -> pd.DataFrame:
    """Create a small panel for fast OOS tests.

    Uses integer dates and has 2 characteristics.
    """
    np.random.seed(seed)
    rows = []
    for date in range(1, n_months + 1):
        for permno in range(1001, 1001 + n_stocks):
            rows.append(
                {
                    "permno": permno,
                    "date": date,
                    "ret": np.random.normal(0.005, 0.05),
                    "bm": np.random.normal(0.0, 1.0),
                    "mom": np.random.normal(0.0, 1.0),
                }
            )
    return pd.DataFrame(rows)


# ---------------------------------------------------------------------------
# LassoPredictor / RidgePredictor (unit)
# ---------------------------------------------------------------------------


class TestLassoPredictor:
    def test_fit_predict(self) -> None:
        """LassoPredictor should fit and return predictions with correct index."""
        np.random.seed(0)
        X = pd.DataFrame(np.random.normal(size=(100, 3)), columns=["a", "b", "c"])
        y = pd.Series(np.random.normal(size=100))
        model = LassoPredictor(alpha=0.01)
        model.fit(X, y)
        preds = model.predict(X)
        assert isinstance(preds, pd.Series)
        assert len(preds) == len(X)

    def test_predict_before_fit_raises(self) -> None:
        """Calling predict before fit should raise RuntimeError."""
        model = LassoPredictor()
        X = pd.DataFrame({"a": [1.0, 2.0]})
        with pytest.raises(RuntimeError):
            model.predict(X)


class TestRidgePredictor:
    def test_fit_predict(self) -> None:
        """RidgePredictor should fit and return predictions with correct index."""
        np.random.seed(1)
        X = pd.DataFrame(np.random.normal(size=(100, 3)), columns=["a", "b", "c"])
        y = pd.Series(np.random.normal(size=100))
        model = RidgePredictor(alpha=0.01)
        model.fit(X, y)
        preds = model.predict(X)
        assert isinstance(preds, pd.Series)
        assert len(preds) == len(X)


# ---------------------------------------------------------------------------
# expanding_window_oos
# ---------------------------------------------------------------------------


class TestExpandingWindowOOS:
    def test_happy_path_returns_oos_result(self) -> None:
        """expanding_window_oos should return an OOSResult."""
        panel = _make_small_panel(n_stocks=10, n_months=30)
        result = expanding_window_oos(
            panel,
            ret_col="ret",
            char_cols=["bm", "mom"],
            models=[LassoPredictor()],
            train_min_periods=15,
        )
        assert isinstance(result, OOSResult)

    def test_predictions_contain_ret_column(self) -> None:
        """predictions DataFrame must contain the actual return column 'ret'."""
        panel = _make_small_panel(n_stocks=10, n_months=30)
        result = expanding_window_oos(
            panel,
            ret_col="ret",
            char_cols=["bm", "mom"],
            models=[LassoPredictor(alpha=0.001)],
            train_min_periods=15,
        )
        assert "ret" in result.predictions.columns
        assert len(result.predictions) > 0

    def test_models_history_has_entry_per_test_date(self) -> None:
        """models_history should have an entry for every test date."""
        panel = _make_small_panel(n_stocks=10, n_months=30)
        result = expanding_window_oos(
            panel,
            ret_col="ret",
            char_cols=["bm"],
            models=[LassoPredictor(), RidgePredictor()],
            train_min_periods=15,
        )
        # Number of test periods = n_months - train_min_periods
        expected_test_dates = sorted(panel["date"].unique())[15:]
        assert set(result.models_history.keys()) == set(expected_test_dates)
        # Each entry should have both model labels
        for date, model_dict in result.models_history.items():
            assert "LassoPredictor" in model_dict
            assert "RidgePredictor" in model_dict

    def test_predictions_dataframe_structure(self) -> None:
        """predictions DataFrame should have date, permno, ret, and pred columns."""
        panel = _make_small_panel(n_stocks=10, n_months=25)
        result = expanding_window_oos(
            panel,
            ret_col="ret",
            char_cols=["bm", "mom"],
            models=[LassoPredictor()],
            train_min_periods=12,
        )
        if len(result.predictions) > 0:
            assert "date" in result.predictions.columns
            assert "permno" in result.predictions.columns
            assert "ret" in result.predictions.columns
            pred_cols = [c for c in result.predictions.columns if c.startswith("pred_")]
            assert len(pred_cols) >= 1

    def test_oos_no_lookahead(self) -> None:
        """expanding_window_oos should call check_no_lookahead internally.

        We verify this by patching and confirming the leakage guard triggers
        when train and test data overlap.
        """
        from eapctf.utils.validation import check_no_lookahead, TemporalLeakageError

        # Build a panel where we manually verify check_no_lookahead logic:
        # train_data max_date < test_date min_date must hold for each step.
        panel = _make_small_panel(n_stocks=5, n_months=20)
        # This should complete without raising TemporalLeakageError
        result = expanding_window_oos(
            panel,
            ret_col="ret",
            char_cols=["bm"],
            models=[RidgePredictor()],
            train_min_periods=10,
        )
        assert isinstance(result, OOSResult)

        # Now test that check_no_lookahead itself raises on overlapping data
        train = panel[panel["date"] <= 10]
        test = panel[panel["date"] >= 10]
        with pytest.raises(TemporalLeakageError):
            check_no_lookahead(train, test)

    def test_insufficient_data_returns_empty(self) -> None:
        """If not enough periods, OOSResult should have empty predictions."""
        panel = _make_small_panel(n_stocks=5, n_months=5)
        result = expanding_window_oos(
            panel,
            ret_col="ret",
            char_cols=["bm"],
            models=[LassoPredictor()],
            train_min_periods=10,  # more than available months
        )
        assert isinstance(result, OOSResult)
        assert len(result.predictions) == 0

    def test_oos_r2_computable_from_predictions(self) -> None:
        """OOS R² should be computable from predictions['ret'] and pred_* columns."""
        panel = _make_small_panel(n_stocks=10, n_months=25)
        result = expanding_window_oos(
            panel,
            ret_col="ret",
            char_cols=["bm"],
            models=[LassoPredictor()],
            train_min_periods=12,
        )
        df = result.predictions
        assert len(df) > 0
        y_actual = df["ret"]
        y_pred = df["pred_LassoPredictor"]
        sse = ((y_actual - y_pred) ** 2).sum()
        sst = ((y_actual - y_actual.mean()) ** 2).sum()
        r2 = 1.0 - sse / sst
        assert isinstance(float(r2), float)
        assert r2 <= 1.0

    def test_hist_mean_benchmark_stored(self) -> None:
        """pred_HistMean column should be present in predictions."""
        panel = _make_small_panel(n_stocks=10, n_months=25)
        result = expanding_window_oos(
            panel,
            ret_col="ret",
            char_cols=["bm"],
            models=[LassoPredictor()],
            train_min_periods=12,
        )
        assert "pred_HistMean" in result.predictions.columns

    def test_invalid_retrain_freq_raises(self) -> None:
        """Invalid retrain_freq should raise ValueError."""
        panel = _make_small_panel()
        with pytest.raises(ValueError):
            expanding_window_oos(
                panel,
                char_cols=["bm"],
                models=[LassoPredictor()],
                retrain_freq="weekly",
            )


# ---------------------------------------------------------------------------
# prediction_eval
# ---------------------------------------------------------------------------


class TestPredictionEval:
    def _run_oos(self) -> OOSResult:
        panel = _make_small_panel(n_stocks=15, n_months=30)
        return expanding_window_oos(
            panel,
            ret_col="ret",
            char_cols=["bm", "mom"],
            models=[LassoPredictor(), RidgePredictor()],
            train_min_periods=15,
        )

    def test_returns_expected_keys(self) -> None:
        result = self._run_oos()
        metrics = prediction_eval(result.predictions, pred_col="pred_LassoPredictor")
        assert set(metrics.keys()) == {"mse", "mae", "mape", "oos_r2", "hit_ratio"}

    def test_mse_mae_non_negative(self) -> None:
        result = self._run_oos()
        m = prediction_eval(result.predictions, pred_col="pred_LassoPredictor")
        assert m["mse"] >= 0.0
        assert m["mae"] >= 0.0

    def test_hit_ratio_in_unit_interval(self) -> None:
        result = self._run_oos()
        m = prediction_eval(result.predictions, pred_col="pred_LassoPredictor")
        assert 0.0 <= m["hit_ratio"] <= 1.0

    def test_oos_r2_uses_hist_mean_by_default(self) -> None:
        """OOS R² should use pred_HistMean as benchmark when benchmark_col is None."""
        result = self._run_oos()
        m_default = prediction_eval(result.predictions, pred_col="pred_LassoPredictor")
        m_explicit = prediction_eval(
            result.predictions,
            pred_col="pred_LassoPredictor",
            benchmark_col="pred_HistMean",
        )
        assert abs(m_default["oos_r2"] - m_explicit["oos_r2"]) < 1e-12

    def test_benchmark_col_another_model(self) -> None:
        """benchmark_col can point to another model for head-to-head OOS R²."""
        result = self._run_oos()
        m = prediction_eval(
            result.predictions,
            pred_col="pred_LassoPredictor",
            benchmark_col="pred_RidgePredictor",
        )
        # R² can be any real number; just verify it is finite
        assert np.isfinite(m["oos_r2"])

    def test_invalid_benchmark_col_raises(self) -> None:
        result = self._run_oos()
        with pytest.raises(ValueError, match="benchmark_col"):
            prediction_eval(
                result.predictions,
                pred_col="pred_LassoPredictor",
                benchmark_col="nonexistent_col",
            )

    def test_eval_oos_result_shape(self) -> None:
        """eval_oos_result should return one row per model."""
        result = self._run_oos()
        table = eval_oos_result(result)
        assert set(table.index) == {"LassoPredictor", "RidgePredictor"}
        assert set(table.columns) == {"mse", "mae", "mape", "oos_r2", "hit_ratio"}

    def test_eval_oos_result_empty(self) -> None:
        """eval_oos_result on empty OOSResult returns empty DataFrame."""
        empty = OOSResult(predictions=pd.DataFrame())
        table = eval_oos_result(empty)
        assert len(table) == 0


# ---------------------------------------------------------------------------
# portfolio_eval
# ---------------------------------------------------------------------------


class TestPortfolioEval:
    def _make_returns(self, seed: int = 0) -> pd.Series:
        np.random.seed(seed)
        idx = pd.date_range("2000-01-01", periods=120, freq="ME")
        return pd.Series(np.random.normal(0.01, 0.05, 120), index=idx)

    def test_returns_expected_keys(self) -> None:
        r = self._make_returns()
        m = portfolio_eval(r)
        assert set(m.keys()) == {
            "mean_ret", "vol", "sharpe", "information_ratio", "hit_ratio", "n_periods"
        }

    def test_sharpe_sign(self) -> None:
        """Sharpe should have same sign as mean return."""
        r = self._make_returns()
        m = portfolio_eval(r)
        assert np.sign(m["sharpe"]) == np.sign(m["mean_ret"])

    def test_hit_ratio_in_unit_interval(self) -> None:
        r = self._make_returns()
        m = portfolio_eval(r)
        assert 0.0 <= m["hit_ratio"] <= 1.0

    def test_information_ratio_nan_without_benchmark(self) -> None:
        r = self._make_returns()
        m = portfolio_eval(r)
        assert np.isnan(m["information_ratio"])

    def test_information_ratio_computed_with_benchmark(self) -> None:
        r = self._make_returns(seed=0)
        b = self._make_returns(seed=1)
        m = portfolio_eval(r, benchmark=b)
        assert np.isfinite(m["information_ratio"])

    def test_annualization_factor(self) -> None:
        """Daily data (freq=252) should produce higher annualized vol than monthly (freq=12)."""
        np.random.seed(5)
        daily = pd.Series(np.random.normal(0.001, 0.01, 252))
        monthly = pd.Series(np.random.normal(0.01, 0.05, 12))
        m_daily = portfolio_eval(daily, freq=252)
        m_monthly = portfolio_eval(monthly, freq=12)
        # Daily vol annualized from 1% daily std should be ~16%; monthly from 5% std ~17%
        # Just verify vol is positive and finite
        assert m_daily["vol"] > 0
        assert m_monthly["vol"] > 0

    def test_empty_series(self) -> None:
        r = pd.Series([], dtype=float)
        m = portfolio_eval(r)
        assert np.isnan(m["sharpe"])
        assert m["n_periods"] == 0
