import numpy as np
import pandas as pd
import pytest
from pandas.testing import assert_series_equal

from panelbeater.features import features, _ticker_features, _meta_ticker_feature, _dt_features

@pytest.fixture
def sample_data():
    """
    Creates a sample 'wide' dataframe (tickers as columns) with a DatetimeIndex.
    """
    dates = pd.date_range(start="2024-01-01", periods=20, freq="D")
    data = {
        "AAPL": np.linspace(100, 200, 20),  # Linear trend for predictable testing
        "MSFT": np.random.normal(100, 10, 20)  # Random data
    }
    df = pd.DataFrame(data, index=dates)
    return df

@pytest.fixture
def params():
    return {
        "windows": [2, 5],
        "lags": [1]
    }

def test_ticker_features_logic(sample_data, params):
    """
    Verifies that SMA, PCT, and Z-Score columns are generated and calculated correctly.
    """
    df = sample_data.copy()
    windows = params["windows"]
    
    # Run the internal function
    df_result = _ticker_features(df, windows)
    
    # Check column existence
    original_cols = sample_data.columns
    for col in original_cols:
        for w in windows:
            assert f"{col}_sma_{w}" in df_result.columns
            assert f"{col}_pctchg_{w}" in df_result.columns
            assert f"{col}_z_{w}" in df_result.columns

    # Verify SMA Calculation logic (SMA / Price - 1)
    # For a window of 2 on a linear trend [100, 105.26...], manual check
    col = "AAPL"
    w = 2
    s = sample_data[col]
    expected_sma_feature = (s.rolling(w).mean() / s) - 1
    
    # We drop NaNs generated by rolling window for equality check
    assert_series_equal(
        df_result[f"{col}_sma_{w}"].dropna(), 
        expected_sma_feature.dropna(),
        check_names=False,
        obj="SMA Feature Calculation"
    )

def test_meta_ticker_feature_expansion(sample_data, params):
    """
    Verifies that lags and rolling stats are applied to ALL columns passed in.
    """
    df = sample_data.copy()
    lags = params["lags"]
    windows = params["windows"]
    
    df_result = _meta_ticker_feature(df, lags=lags, windows=windows)
    
    # Logic Check: _meta_ticker_feature concatenates the original df + lags + rolling
    # Expected column count: 
    # Original (2) + Lags (2*1) + RollingMean (2*2) + RollingStd (2*2) = 12 columns
    assert df_result.shape[1] == 2 + (2 * len(lags)) + (2 * len(windows)) + (2 * len(windows))
    
    # Verify Lag logic
    assert f"AAPL_lag{lags[0]}" in df_result.columns
    assert_series_equal(
        df_result["AAPL"].shift(1), 
        df_result[f"AAPL_lag{1}"], 
        check_names=False
    )

def test_dt_features_generation(sample_data):
    """
    Verifies feature_engine integration for datetime features.
    """
    df_result = _dt_features(sample_data.copy())
    
    # feature_engine 'all' usually generates month, year, day_of_week, etc.
    expected_partials = ["month", "quarter", "day_of_week"]
    
    # Check if at least some standard datetime features exist
    assert any(col.endswith("month") for col in df_result.columns)
    assert any(col.endswith("day_of_week") for col in df_result.columns)
    
    # Ensure index wasn't lost
    pd.testing.assert_index_equal(df_result.index, sample_data.index)

def test_features_integration_pipeline(sample_data, params):
    """
    Tests the full pipeline:
    1. Correct columns dropped (original tickers).
    2. Final shift(1) is applied.
    3. No infinite values.
    """
    df = sample_data.copy()
    original_cols = df.columns.tolist()
    
    df_final = features(df, windows=params["windows"], lags=params["lags"])
    
    # 1. Check Original Columns Dropped
    for col in original_cols:
        assert col not in df_final.columns
        
    # 2. Check Final Shift(1)
    # Because of the final .shift(1), the very first row of the *entire* dataframe must be NaN
    # (assuming the index aligns).
    assert df_final.iloc[0].isna().all()
    
    # 3. Check Expansion Volume
    # The pipeline is: _ticker_features -> _meta_ticker_feature -> _dt_features
    # NOTE: _ticker_features ADDS columns to df. 
    # _meta_ticker_feature calculates meta features on THAT expanded df.
    # This leads to a large feature explosion. This test ensures the code actually runs that explosion without error.
    assert df_final.shape[1] > 20  # Arbitrary check to ensure expansion happened

    # 4. Check Structural Features exist
    # Look for one of the new columns (e.g., _er_ or _skew_)
    # Note: Columns will be shifted/lagged, so names might be complex, 
    # but the base feature name should be present in the string.
    assert any("_er_" in col for col in df_final.columns), "Efficiency Ratio features missing"
    assert any("_skew_" in col for col in df_final.columns), "Skewness features missing"
    
    # 5. Check Datetime features persist
    assert any(str(c).endswith("month") for c in df_final.columns)

def test_inf_handling(params):
    """
    Test that the pipeline handles division by zero (producing infs) gracefully
    by converting them to NaNs as specified in _meta_ticker_feature.
    """
    # Create a constant dataframe where std dev will be 0, causing Z-score division by zero
    dates = pd.date_range(start="2024-01-01", periods=10, freq="D")
    df_const = pd.DataFrame({"CONST": [10.0] * 10}, index=dates)
    
    # This should trigger the Z-score calculation (x - mu) / sigma -> 0 / 0 or x / 0
    # Your code catches warnings, but we want to ensure the output is valid (NaN, not Inf)
    df_final = features(df_const, windows=params["windows"], lags=params["lags"])
    
    # Check for Inf
    assert not np.isinf(df_final.values).any(), "Result contains Infinite values"
