"""Tests for eapctf.ctf.prep — CTF feature preprocessing utilities."""

from __future__ import annotations

import numpy as np
import pandas as pd
import pytest

from eapctf.ctf.prep import ctf_rank_normalize, ctf_select_features
from eapctf.utils.constants import JKP_153

# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------

N_STOCKS = 50
N_MONTHS = 12
N_FEATURES = 10
FEAT_NAMES = [f"feat_{i:02d}" for i in range(N_FEATURES)]
# A few extra non-feature columns that must survive normalization unchanged
NON_FEAT_COLS = ["id", "eom", "ret_exc_lead1m", "ctff_test"]


@pytest.fixture
def features_df() -> pd.DataFrame:
    """Small features DataFrame listing FEAT_NAMES."""
    return pd.DataFrame({"features": FEAT_NAMES})


@pytest.fixture
def chars_df() -> pd.DataFrame:
    """Synthetic characteristics panel: 50 stocks × 12 months × 10 features.

    Some feature values are intentionally set to NaN to test NaN handling.
    Non-feature columns are set to deterministic sentinel values so that
    tests can verify they are not modified.
    """
    rng = np.random.default_rng(42)
    dates = pd.date_range("2020-01-31", periods=N_MONTHS, freq="ME")

    records = []
    for t, eom in enumerate(dates):
        for i in range(N_STOCKS):
            row: dict = {
                "id": i,
                "eom": eom,
                "ret_exc_lead1m": rng.normal(0.01, 0.05),
                "ctff_test": t * N_STOCKS + i,  # unique sentinel per obs
            }
            for feat in FEAT_NAMES:
                val = rng.normal()
                # Introduce ~10% missingness per feature
                if rng.random() < 0.10:
                    val = np.nan
                row[feat] = val
            records.append(row)

    return pd.DataFrame(records)


@pytest.fixture
def chars_high_missing(chars_df: pd.DataFrame, features_df: pd.DataFrame) -> pd.DataFrame:
    """chars_df with feat_07, feat_08, feat_09 set to >50% NaN."""
    df = chars_df.copy()
    rng = np.random.default_rng(99)
    high_miss_feats = ["feat_07", "feat_08", "feat_09"]
    for feat in high_miss_feats:
        mask = rng.random(len(df)) < 0.70  # 70% missing
        df.loc[mask, feat] = np.nan
    return df


# ---------------------------------------------------------------------------
# Tests for ctf_rank_normalize
# ---------------------------------------------------------------------------


class TestCtfRankNormalize:
    def test_normalized_range(self, chars_df: pd.DataFrame, features_df: pd.DataFrame) -> None:
        """After normalization, all feature values must lie in [-scale, +scale]."""
        scale = 0.5
        result = ctf_rank_normalize(chars_df, features_df, scale=scale)

        for feat in FEAT_NAMES:
            vals = result[feat].dropna()
            assert (vals >= -scale).all(), f"{feat}: value below -scale"
            assert (vals <= scale).all(), f"{feat}: value above +scale"

    def test_custom_scale(self, chars_df: pd.DataFrame, features_df: pd.DataFrame) -> None:
        """Custom scale=1.0 should bound normalized values to [-1, 1]."""
        scale = 1.0
        result = ctf_rank_normalize(chars_df, features_df, scale=scale)
        for feat in FEAT_NAMES:
            vals = result[feat]
            assert (vals >= -scale).all(), f"{feat}: value below -scale=1.0"
            assert (vals <= scale).all(), f"{feat}: value above +scale=1.0"

    def test_nan_filled_with_fill_value(
        self, chars_df: pd.DataFrame, features_df: pd.DataFrame
    ) -> None:
        """NaN positions in the original data must be filled with fill_value."""
        fill_value = 0.0
        result = ctf_rank_normalize(chars_df, features_df, fill_value=fill_value)

        # There should be no remaining NaN in feature columns after fill
        for feat in FEAT_NAMES:
            assert result[feat].isna().sum() == 0, f"{feat}: NaN remains after fill"

    def test_custom_fill_value(
        self, chars_df: pd.DataFrame, features_df: pd.DataFrame
    ) -> None:
        """Custom fill_value should replace NaN positions from original data."""
        # Identify positions that were NaN in input
        fill_value = -99.0
        original_nan_mask = chars_df[FEAT_NAMES].isna()

        result = ctf_rank_normalize(chars_df, features_df, fill_value=fill_value)

        for feat in FEAT_NAMES:
            nan_positions = original_nan_mask[feat]
            if nan_positions.any():
                filled_vals = result.loc[nan_positions, feat]
                assert (filled_vals == fill_value).all(), (
                    f"{feat}: NaN positions not filled with fill_value={fill_value}"
                )

    def test_non_feature_columns_unchanged(
        self, chars_df: pd.DataFrame, features_df: pd.DataFrame
    ) -> None:
        """id, eom, ret_exc_lead1m, ctff_test must not be modified."""
        result = ctf_rank_normalize(chars_df, features_df)
        for col in NON_FEAT_COLS:
            pd.testing.assert_series_equal(
                result[col].reset_index(drop=True),
                chars_df[col].reset_index(drop=True),
                check_names=False,
                obj=f"column '{col}'",
            )

    def test_does_not_modify_input(
        self, chars_df: pd.DataFrame, features_df: pd.DataFrame
    ) -> None:
        """The input DataFrames must not be mutated."""
        chars_copy = chars_df.copy()
        feat_copy = features_df.copy()
        ctf_rank_normalize(chars_df, features_df)
        pd.testing.assert_frame_equal(chars_df, chars_copy)
        pd.testing.assert_frame_equal(features_df, feat_copy)

    def test_cross_sectional_application(
        self, chars_df: pd.DataFrame, features_df: pd.DataFrame
    ) -> None:
        """Normalization should be independent across eom dates.

        Within a single eom cross-section the normalized values should have
        mean approximately 0 (up to NaN fill) and span the expected range.
        """
        result = ctf_rank_normalize(chars_df, features_df, fill_value=np.nan)
        # For a single date, non-filled values should span a symmetric range
        for eom, grp in result.groupby("eom"):
            for feat in FEAT_NAMES[:3]:  # check a few features
                valid = grp[feat].dropna()
                if len(valid) > 1:
                    assert valid.max() <= 0.5 + 1e-9
                    assert valid.min() >= -0.5 - 1e-9

    def test_first_column_fallback(
        self, chars_df: pd.DataFrame
    ) -> None:
        """Features DataFrame whose first column is not named 'features' still works."""
        features_alt = pd.DataFrame({"signal": FEAT_NAMES})
        result = ctf_rank_normalize(chars_df, features_alt)
        # Verify normalization occurred (values bounded)
        for feat in FEAT_NAMES:
            assert result[feat].max() <= 0.5 + 1e-9


# ---------------------------------------------------------------------------
# Tests for ctf_select_features
# ---------------------------------------------------------------------------


class TestCtfSelectFeatures:
    def test_jkp153_subset(
        self, chars_df: pd.DataFrame, features_df: pd.DataFrame
    ) -> None:
        """jkp153 mode: only JKP_153 features that exist in chars are kept."""
        # Our synthetic features (feat_00 … feat_09) are NOT in JKP_153,
        # so the expected selection is empty.
        filtered, selected = ctf_select_features(chars_df, features_df, subset="jkp153")
        jkp_set = set(JKP_153)
        assert all(f in jkp_set for f in selected), "non-JKP feature selected in jkp153 mode"
        # Non-feature cols must survive
        for col in NON_FEAT_COLS:
            assert col in filtered.columns

    def test_jkp153_returns_jkp_when_present(self) -> None:
        """jkp153 mode returns JKP_153 features that actually exist in chars."""
        jkp_sample = JKP_153[:5]
        non_feat = ["id", "eom"]
        df = pd.DataFrame(
            np.random.default_rng(0).standard_normal((20, len(jkp_sample))),
            columns=jkp_sample,
        )
        df["id"] = range(20)
        df["eom"] = pd.Timestamp("2020-01-31")
        features = pd.DataFrame({"features": JKP_153})

        filtered, selected = ctf_select_features(df, features, subset="jkp153")
        assert set(selected) == set(jkp_sample)
        for col in non_feat:
            assert col in filtered.columns

    def test_all_subset(
        self, chars_df: pd.DataFrame, features_df: pd.DataFrame
    ) -> None:
        """all mode: every feature listed in features that exists in chars is selected."""
        filtered, selected = ctf_select_features(chars_df, features_df, subset="all")
        assert set(selected) == set(FEAT_NAMES)
        assert set(FEAT_NAMES).issubset(filtered.columns)
        for col in NON_FEAT_COLS:
            assert col in filtered.columns

    def test_all_subset_excludes_missing_columns(self) -> None:
        """all mode: features not present in chars are silently excluded."""
        df = pd.DataFrame({"id": [1], "eom": [pd.Timestamp("2020-01-31")], "feat_00": [1.0]})
        features = pd.DataFrame({"features": ["feat_00", "feat_99"]})  # feat_99 absent
        _, selected = ctf_select_features(df, features, subset="all")
        assert selected == ["feat_00"]

    def test_available_subset_drops_high_missing(
        self,
        chars_high_missing: pd.DataFrame,
        features_df: pd.DataFrame,
    ) -> None:
        """available mode: features with >= 70% missingness should be dropped at min_coverage=0.5."""
        _, selected = ctf_select_features(
            chars_high_missing, features_df, subset="available", min_coverage=0.5
        )
        # feat_07, feat_08, feat_09 have ~70% NaN; they must be excluded
        for feat in ["feat_07", "feat_08", "feat_09"]:
            assert feat not in selected, f"{feat} should be dropped (high missingness)"

    def test_available_subset_keeps_low_missing(
        self,
        chars_high_missing: pd.DataFrame,
        features_df: pd.DataFrame,
    ) -> None:
        """available mode: features with ~10% missingness should be retained at min_coverage=0.5."""
        _, selected = ctf_select_features(
            chars_high_missing, features_df, subset="available", min_coverage=0.5
        )
        # feat_00 … feat_06 have ~10% NaN and should pass coverage threshold
        for feat in [f"feat_{i:02d}" for i in range(7)]:
            assert feat in selected, f"{feat} should be retained (low missingness)"

    def test_available_non_feat_cols_preserved(
        self,
        chars_df: pd.DataFrame,
        features_df: pd.DataFrame,
    ) -> None:
        """Non-feature columns always appear in filtered DataFrame for available mode."""
        filtered, _ = ctf_select_features(
            chars_df, features_df, subset="available", min_coverage=0.5
        )
        for col in NON_FEAT_COLS:
            assert col in filtered.columns

    def test_invalid_subset_raises(
        self, chars_df: pd.DataFrame, features_df: pd.DataFrame
    ) -> None:
        """Unsupported subset string must raise ValueError."""
        with pytest.raises(ValueError, match="subset must be one of"):
            ctf_select_features(chars_df, features_df, subset="unknown")

    def test_does_not_modify_input(
        self, chars_df: pd.DataFrame, features_df: pd.DataFrame
    ) -> None:
        """ctf_select_features must not mutate its inputs."""
        chars_copy = chars_df.copy()
        feat_copy = features_df.copy()
        ctf_select_features(chars_df, features_df, subset="all")
        pd.testing.assert_frame_equal(chars_df, chars_copy)
        pd.testing.assert_frame_equal(features_df, feat_copy)

    def test_returned_filtered_shape(
        self, chars_df: pd.DataFrame, features_df: pd.DataFrame
    ) -> None:
        """Filtered DataFrame has same number of rows as input chars."""
        filtered, selected = ctf_select_features(chars_df, features_df, subset="all")
        assert len(filtered) == len(chars_df)
        assert list(filtered.columns) == NON_FEAT_COLS + selected
