"""
Tests for feature pruning logic.
"""
import pytest
import pandas as pd
import numpy as np

# Try to import the module; skip tests if not found
try:
    from gru_sac_predictor.src import features
except ImportError:
    features = None

# --- Fixtures ---
@pytest.fixture
def sample_features_df():
    """Create a DataFrame with more columns than the whitelist."""
    data = {
        # Whitelisted
        "return_1m": np.random.randn(100),
        "return_15m": np.random.randn(100),
        "return_60m": np.random.randn(100),
        "ATR_14": np.random.rand(100) * 0.01,
        "volatility_14d": np.random.rand(100) * 0.02,
        "chaikin_AD_10": np.random.randn(100) * 1000,
        "svi_10": np.random.randn(100) * 500,
        "EMA_10": 100 + np.random.randn(100),
        "EMA_50": 100 + np.random.randn(100),
        "MACD": np.random.randn(100) * 0.1,
        "MACD_signal": np.random.randn(100) * 0.05,
        "hour_sin": np.sin(np.linspace(0, 2*np.pi, 100)),
        "hour_cos": np.cos(np.linspace(0, 2*np.pi, 100)),
        # Non-whitelisted
        "close": 100 + np.random.randn(100),
        "open": 100 + np.random.randn(100),
        "RSI_14": 50 + np.random.randn(100) * 10, # Assumed not in final whitelist
        "some_other_feature": np.random.rand(100)
    }
    return pd.DataFrame(data)

# --- Tests ---
@pytest.mark.skipif(features is None, reason="Module gru_sac_predictor.src.features not found")
def test_prune_features_uses_whitelist(sample_features_df):
    """
    Verify prune_features returns only columns present in features.minimal_whitelist.
    """
    df_in = sample_features_df
    whitelist = features.minimal_whitelist
    df_out = features.prune_features(df_in)

    print(f"\nWhitelist: {whitelist}")
    print(f"Input columns: {df_in.columns.tolist()}")
    print(f"Output columns: {df_out.columns.tolist()}")

    # Check that all output columns are in the whitelist
    assert all(col in whitelist for col in df_out.columns), \
        "Output DataFrame contains columns not in the whitelist."

    # Check that all whitelist columns present in the input are also in the output
    expected_cols = [col for col in whitelist if col in df_in.columns]
    assert sorted(df_out.columns.tolist()) == sorted(expected_cols), \
        "Output columns do not match the expected intersection of input and whitelist."

    # Check that non-whitelisted columns are removed
    assert "close" not in df_out.columns, "'close' column was not pruned."
    assert "some_other_feature" not in df_out.columns, "'some_other_feature' was not pruned." 