"""Tests for eapctf.utils: preprocessing, validation, data utilities."""

from __future__ import annotations

import math

import numpy as np
import pandas as pd
import pytest

from eapctf.utils import (
    EAPPanel,
    TemporalLeakageError,
    align_returns_chars,
    check_no_lookahead,
    check_panel_dims,
    classify,
    lag_panel,
    rank_normalize,
    standardize,
    winsorize,
)


# ---------------------------------------------------------------------------
# rank_normalize
# ---------------------------------------------------------------------------


class TestRankNormalize:
    def test_output_shape_and_range(self, small_panel: pd.DataFrame) -> None:
        """rank_normalize output is in [-1, 1] for all char cols."""
        char_cols = ["bm", "op", "inv"]
        result = rank_normalize(small_panel, char_cols=char_cols)
        assert result.shape == small_panel.shape
        for col in char_cols:
            vals = result[col].dropna()
            assert vals.min() >= -1.0 - 1e-9, f"{col} min below -1"
            assert vals.max() <= 1.0 + 1e-9, f"{col} max above 1"

    def test_non_char_cols_unchanged(self, small_panel: pd.DataFrame) -> None:
        """Columns not in char_cols should be identical to input."""
        result = rank_normalize(small_panel, char_cols=["bm"])
        pd.testing.assert_series_equal(result["ret"], small_panel["ret"])

    def test_single_stock(self) -> None:
        """Single stock panel: rank_normalize should run without error."""
        df = pd.DataFrame(
            {"permno": [1001, 1001, 1001], "date": [1, 2, 3], "bm": [0.5, 1.0, 1.5]}
        )
        result = rank_normalize(df, char_cols=["bm"])
        assert result.shape == df.shape

    def test_missing_values_propagate(self, small_panel: pd.DataFrame) -> None:
        """NaN in input stays NaN after rank normalization."""
        panel = small_panel.copy()
        panel.loc[panel.index[0], "bm"] = np.nan
        result = rank_normalize(panel, char_cols=["bm"])
        assert pd.isna(result.loc[panel.index[0], "bm"])


# ---------------------------------------------------------------------------
# winsorize
# ---------------------------------------------------------------------------


class TestWinsorize:
    def test_bounds_respected(self, small_panel: pd.DataFrame) -> None:
        """After winsorizing, no values should exceed the quantile bounds."""
        result = winsorize(small_panel, cols=["bm"], q_low=0.05, q_high=0.95)
        assert result.shape == small_panel.shape
        # Values should stay within 1st and 99th percentile (approximately)
        assert result["bm"].isna().sum() == small_panel["bm"].isna().sum()

    def test_global_winsorize(self, small_panel: pd.DataFrame) -> None:
        """by_date=False should winsorize over full sample."""
        result = winsorize(small_panel, cols=["ret"], by_date=False)
        assert result.shape == small_panel.shape


# ---------------------------------------------------------------------------
# standardize
# ---------------------------------------------------------------------------


class TestStandardize:
    def test_cross_sectional_mean_near_zero(self, small_panel: pd.DataFrame) -> None:
        """Each date's cross-sectional mean should be approximately zero."""
        result = standardize(small_panel, cols=["bm"])
        group_means = result.groupby("date")["bm"].mean().abs()
        assert group_means.max() < 1e-10

    def test_global_standardize(self, small_panel: pd.DataFrame) -> None:
        """by_date=False should standardize over full sample."""
        result = standardize(small_panel, cols=["bm"], by_date=False)
        assert abs(result["bm"].mean()) < 1e-10
        assert abs(result["bm"].std() - 1.0) < 0.01


# ---------------------------------------------------------------------------
# classify
# ---------------------------------------------------------------------------


class TestClassify:
    def test_n_portfolios(self, small_panel: pd.DataFrame) -> None:
        """classify with n_portfolios=5 should produce labels 0..4."""
        labels = classify(small_panel, col="bm", n_portfolios=5, by_nyse=False)
        valid_labels = labels.dropna().astype(int)
        assert set(valid_labels.unique()).issubset({0, 1, 2, 3, 4})

    def test_splits(self, small_panel: pd.DataFrame) -> None:
        """classify with splits=[0.3, 0.7] should produce 3 groups."""
        labels = classify(small_panel, col="bm", splits=[0.3, 0.7], by_nyse=False)
        valid_labels = labels.dropna().astype(int)
        assert len(valid_labels.unique()) <= 3

    def test_mutual_exclusive_error(self, small_panel: pd.DataFrame) -> None:
        """Passing both n_portfolios and splits should raise ValueError."""
        with pytest.raises(ValueError):
            classify(small_panel, col="bm", n_portfolios=5, splits=[0.5])

    def test_neither_error(self, small_panel: pd.DataFrame) -> None:
        """Passing neither n_portfolios nor splits should raise ValueError."""
        with pytest.raises(ValueError):
            classify(small_panel, col="bm")


# ---------------------------------------------------------------------------
# check_no_lookahead / TemporalLeakageError
# ---------------------------------------------------------------------------


class TestCheckNoLookahead:
    def test_valid_split_passes(self, small_panel: pd.DataFrame) -> None:
        """Train ending before test beginning should not raise."""
        train = small_panel[small_panel["date"] < 60]
        test = small_panel[small_panel["date"] >= 60]
        # Should not raise
        check_no_lookahead(train, test)

    def test_overlap_raises(self, small_panel: pd.DataFrame) -> None:
        """Overlapping train and test dates should raise TemporalLeakageError."""
        train = small_panel[small_panel["date"] <= 60]
        test = small_panel[small_panel["date"] >= 60]
        with pytest.raises(TemporalLeakageError):
            check_no_lookahead(train, test)

    def test_same_dates_raises(self, small_panel: pd.DataFrame) -> None:
        """Same max train date and min test date should raise."""
        train = small_panel[small_panel["date"] <= 50]
        test = small_panel[small_panel["date"] >= 50]
        with pytest.raises(TemporalLeakageError):
            check_no_lookahead(train, test)


# ---------------------------------------------------------------------------
# check_panel_dims
# ---------------------------------------------------------------------------


class TestCheckPanelDims:
    def test_returns_correct_keys(self, small_panel: pd.DataFrame) -> None:
        """check_panel_dims should return dict with expected keys."""
        result = check_panel_dims(small_panel)
        assert set(result.keys()) == {"n_entities", "n_dates", "n_obs", "pct_missing", "unbalanced"}

    def test_balanced_panel(self, small_panel: pd.DataFrame) -> None:
        """small_panel fixture is balanced (50 stocks x 120 months)."""
        result = check_panel_dims(small_panel)
        assert result["n_entities"] == 50
        assert result["n_dates"] == 120
        assert result["unbalanced"] == 0

    def test_missing_required_cols_raises(self, small_panel: pd.DataFrame) -> None:
        """Missing required columns should raise ValueError."""
        with pytest.raises(ValueError):
            check_panel_dims(small_panel, required_cols=["nonexistent_col"])


# ---------------------------------------------------------------------------
# lag_panel
# ---------------------------------------------------------------------------


class TestLagPanel:
    def test_first_period_is_nan(self, small_panel: pd.DataFrame) -> None:
        """After lag_panel(n=1), each entity's first period should be NaN."""
        result = lag_panel(small_panel, cols=["bm"], n_periods=1)
        first_dates = result.groupby("permno").apply(lambda g: g.sort_values("date").iloc[0])
        assert first_dates["bm"].isna().all()

    def test_shape_preserved(self, small_panel: pd.DataFrame) -> None:
        """lag_panel should not change the shape of the DataFrame."""
        result = lag_panel(small_panel, cols=["bm", "op"])
        assert result.shape == small_panel.shape


# ---------------------------------------------------------------------------
# align_returns_chars
# ---------------------------------------------------------------------------


class TestAlignReturnsChars:
    def test_merge_produces_correct_shape(self, small_panel: pd.DataFrame) -> None:
        """align_returns_chars should produce a merged panel with correct columns."""
        chars = small_panel[["permno", "date", "bm", "op"]]
        returns = small_panel[["permno", "date", "ret"]]
        merged = align_returns_chars(chars, returns, lag=1)
        # Merged result should contain both ret and characteristic columns
        assert "bm" in merged.columns
        assert "op" in merged.columns
        assert "ret" in merged.columns
        # First period's chars (date=1) are NaN after lag=1 shift;
        # the inner merge drops these NaN rows, so length <= small_panel
        assert len(merged) <= len(small_panel)


# ---------------------------------------------------------------------------
# EAPPanel
# ---------------------------------------------------------------------------


class TestEAPPanel:
    def test_construction(self, small_panel: pd.DataFrame) -> None:
        """EAPPanel should construct from a flat DataFrame."""
        panel = EAPPanel(small_panel)
        info = panel.inspect()
        assert info["n_entities"] == 50
        assert info["n_dates"] == 120

    def test_shift_creates_nans(self, small_panel: pd.DataFrame) -> None:
        """shift() should create NaN in first period per entity."""
        panel = EAPPanel(small_panel)
        shifted = panel.shift(["bm"], n=1)
        long_df = shifted.to_long()
        first_obs = long_df.groupby("permno").apply(lambda g: g.sort_values("date").iloc[0])
        assert first_obs["bm"].isna().all()

    def test_populate_ffill(self, small_panel: pd.DataFrame) -> None:
        """populate with ffill should fill NaN values within entities."""
        panel_with_nan = small_panel.copy()
        # Introduce NaN in a few rows
        panel_with_nan.loc[panel_with_nan["date"] == 5, "bm"] = np.nan
        panel = EAPPanel(panel_with_nan)
        filled = panel.populate(method="ffill")
        long_df = filled.to_long()
        # After forward-fill, date-5 NaN should be filled (from date=4)
        date5_vals = long_df[long_df["date"] == 5]["bm"]
        assert date5_vals.isna().sum() == 0

    def test_to_long_roundtrip(self, small_panel: pd.DataFrame) -> None:
        """to_long() should recover a flat DataFrame."""
        panel = EAPPanel(small_panel)
        long_df = panel.to_long()
        assert set(["permno", "date"]).issubset(set(long_df.columns))
        assert len(long_df) == len(small_panel)

    def test_single_entity(self) -> None:
        """EAPPanel with a single entity should not crash."""
        df = pd.DataFrame(
            {"permno": [1001] * 10, "date": list(range(1, 11)), "ret": np.random.normal(size=10)}
        )
        panel = EAPPanel(df)
        info = panel.inspect()
        assert info["n_entities"] == 1
        assert info["n_dates"] == 10
