"""
Tests for data leakage (Sec 6 of revisions.txt).
"""
import pytest
import pandas as pd
import numpy as np

# Assume test data is loaded via fixtures later
@pytest.fixture(scope="module")
def sample_data_for_leakage():
    """
    Provides sample features and target for leakage tests.
    Includes correctly shifted features, a feature with direct leakage,
    and a rolling feature calculated correctly vs incorrectly.
    """
    np.random.seed(43)
    dates = pd.date_range(start='2023-01-01', periods=500, freq='T')
    n = len(dates)
    df = pd.DataFrame(index=dates)
    df['noise'] = np.random.randn(n)
    df['close'] = 100 + np.cumsum(df['noise'] * 0.1)
    df['y_ret'] = np.log(df['close'].shift(-1) / df['close'])

    # --- Features ---
    # OK: Based on past noise
    df['feature_ok_past_noise'] = df['noise'].shift(1)
    # OK: Rolling mean on correctly shifted past data
    df['feature_ok_rolling_shifted'] = df['noise'].shift(1).rolling(10).mean()
    # LEAKY: Uses future return directly
    df['feature_leaky_direct'] = df['y_ret']
    # LEAKY: Rolling mean calculated *before* shifting target relationship
    df['feature_leaky_rolling_unaligned'] = df['close'].rolling(5).mean()

    # Drop rows with NaNs from shifts/rolls AND the last row where y_ret is NaN
    df.dropna(inplace=True)

    # Define features and target for the test
    y_target = df['y_ret']
    features_df = df.drop(columns=['close', 'y_ret', 'noise']) # Exclude raw data used for generation

    return features_df, y_target

@pytest.mark.parametrize("leakage_threshold", [0.02])
def test_feature_leakage_correlation(sample_data_for_leakage, leakage_threshold):
    """
    Verify that no feature has correlation > threshold with the correctly shifted target.
    """
    features_df, y_target = sample_data_for_leakage

    max_abs_corr = 0.0
    leaky_col = "None"
    all_corrs = {}

    print(f"\nTesting {features_df.shape[1]} features for leakage (threshold={leakage_threshold})...")
    for col in features_df.columns:
        if pd.api.types.is_numeric_dtype(features_df[col]):
            # Handle potential NaNs introduced by feature engineering (though fixture avoids it)
            temp_df = pd.concat([features_df[col], y_target], axis=1).dropna()
            if len(temp_df) < 0.5 * len(features_df):
                print(f"  Skipping {col} due to excessive NaNs after merging with target.")
                continue

            correlation = temp_df[col].corr(temp_df['y_ret'])
            all_corrs[col] = correlation
            # print(f"  Corr({col}, y_ret): {correlation:.4f}")
            if abs(correlation) > max_abs_corr:
                max_abs_corr = abs(correlation)
                leaky_col = col
        else:
            print(f"  Skipping non-numeric column: {col}")

    print(f"Correlations found: { {k: round(v, 4) for k, v in all_corrs.items()} }")
    print(f"Maximum absolute correlation found: {max_abs_corr:.4f} (feature: {leaky_col})")

    assert max_abs_corr < leakage_threshold, \
        f"Feature '{leaky_col}' has correlation {max_abs_corr:.4f} > threshold {leakage_threshold}, suggesting leakage."

@pytest.mark.skipif(features is None, reason="Module gru_sac_predictor.src.features not found")
def test_ta_feature_leakage(sample_data_for_leakage, leakage_threshold=0.02):
    """
    Specifically test TA features (EMA, MACD etc.) for leakage.
    Ensures they were calculated on shifted data.
    """
    features_df, y_target = sample_data_for_leakage
    # Add TA features using the helper (simulating pipeline)
    # We need OHLC in the input df for add_ta_features
    # Recreate a df with shifted OHLC + other features for TA calc
    np.random.seed(43) # Ensure consistent data with primary fixture
    dates = pd.date_range(start='2023-01-01', periods=500, freq='T')
    n = len(dates)
    df_ohlc = pd.DataFrame(index=dates)
    df_ohlc['close'] = 100 + np.cumsum(np.random.randn(n) * 0.1)
    df_ohlc['open'] = df_ohlc['close'].shift(1) * (1 + np.random.randn(n) * 0.001)
    df_ohlc['high'] = df_ohlc[['open','close']].max(axis=1) * (1 + np.random.rand(n) * 0.001)
    df_ohlc['low'] = df_ohlc[['open','close']].min(axis=1) * (1 - np.random.rand(n) * 0.001)
    df_ohlc['volume'] = np.random.rand(n) * 1000

    # IMPORTANT: Shift before calculating TA features
    df_shifted_ohlc = df_ohlc.shift(1)
    df_ta = features.add_ta_features(df_shifted_ohlc)

    # Align with the target (requires original non-shifted index)
    df_ta = df_ta.loc[y_target.index]

    ta_features_to_test = [col for col in features.minimal_whitelist if col in df_ta.columns and col not in ["return_1m", "return_15m", "return_60m", "hour_sin", "hour_cos"]]
    max_abs_corr = 0.0
    leaky_col = "None"
    all_corrs = {}

    print(f"\nTesting {len(ta_features_to_test)} TA features for leakage (threshold={leakage_threshold})...")
    print(f" Features: {ta_features_to_test}")

    for col in ta_features_to_test:
        if pd.api.types.is_numeric_dtype(df_ta[col]):
            temp_df = pd.concat([df_ta[col], y_target], axis=1).dropna()
            if len(temp_df) < 0.5 * len(y_target):
                print(f"  Skipping {col} due to excessive NaNs after merging.")
                continue
            correlation = temp_df[col].corr(temp_df['y_ret'])
            all_corrs[col] = correlation
            if abs(correlation) > max_abs_corr:
                max_abs_corr = abs(correlation)
                leaky_col = col
        else:
            print(f"  Skipping non-numeric TA column: {col}")

    print(f"TA Feature Correlations: { {k: round(v, 4) for k, v in all_corrs.items()} }")
    print(f"Maximum absolute TA correlation found: {max_abs_corr:.4f} (feature: {leaky_col})")

    assert max_abs_corr < leakage_threshold, \
        f"TA Feature '{leaky_col}' has correlation {max_abs_corr:.4f} > threshold {leakage_threshold}, suggesting leakage from TA calculation."

# test_label_timing is usually covered by the correlation test, so removed for brevity. 