"""Generate features over a dataframe."""

# pylint: disable=too-many-locals
import warnings

import numpy as np
import pandas as pd
from feature_engine.datetime import DatetimeFeatures


def _ticker_features(df: pd.DataFrame, windows: list[int]) -> pd.DataFrame:
    cols = df.columns.values.tolist()

    # Establish the maximum window to act as our long-term baseline for volatility
    max_window = max(windows) if windows else 1

    # --- OPTIMIZATION: Calculate all returns and the market index once ---
    all_daily_returns = df.pct_change(fill_method=None)
    market_index = all_daily_returns.mean(axis=1)

    for col in cols:
        s = df[col]

        # Slice the pre-calculated returns for this ticker
        daily_returns = all_daily_returns[col]

        # Pre-calculate baseline long-term volatility (standard deviation of returns)
        baseline_vol = daily_returns.rolling(max_window).std()

        for w in windows:
            with warnings.catch_warnings():
                warnings.simplefilter("ignore", category=pd.errors.PerformanceWarning)
                roll = s.rolling(w)

                # SMA
                sma = roll.mean()
                df[f"{col}_sma_{w}"] = sma / s - 1

                # PCT
                df[f"{col}_pctchg_{w}"] = s.pct_change(w, fill_method=None)

                # Z-Score
                mu = roll.mean()
                sigma = roll.std()
                df[f"{col}_z_{w}"] = (s - mu) / sigma

                # --- 1. Kaufman Efficiency Ratio (Fractal Efficiency) ---
                change = s.diff(w).abs()
                volatility = s.diff(1).abs().rolling(w).sum()
                er = change.div(volatility).replace([np.inf, -np.inf], 0).fillna(0)
                df[f"{col}_er_{w}"] = er

                # --- 2. Rolling Skewness & Kurtosis ---
                df[f"{col}_skew_{w}"] = roll.skew().fillna(0)
                df[f"{col}_kurt_{w}"] = roll.kurt().fillna(0)

                # --- 3. Volatility Regime Ratios ---
                current_vol = daily_returns.rolling(w).std()
                vol_ratio = (
                    current_vol.div(baseline_vol)
                    .replace([np.inf, -np.inf], 0)
                    .fillna(0)
                )
                df[f"{col}_vol_regime_{w}"] = vol_ratio

                # --- 4. Bollinger Band Width (Volatility Squeeze) ---
                # 4 sigma covers ~95% of data (2 up, 2 down)
                bb_width = (4 * sigma) / mu
                df[f"{col}_bbwidth_{w}"] = bb_width.replace([np.inf, -np.inf], 0)

                # --- 5. Rolling Autocorrelation (Serial Correlation) ---
                autocorr = daily_returns.rolling(w).corr(daily_returns.shift(1))
                df[f"{col}_autocorr_{w}"] = autocorr.fillna(0)

                # --- 6. RSI (Relative Strength Index) - Vectorized ---
                diff = s.diff(1)
                gain = diff.clip(lower=0).rolling(w).mean()
                loss = -1 * diff.clip(upper=0).rolling(w).mean()

                # Safe division for RS
                rs = gain / loss.replace(0, np.nan)  # type: ignore
                rsi = 1 - (1 / (1 + rs))
                # If loss is 0 (no down days), RSI is 1.0 (100 in standard RSI terms)
                df[f"{col}_rsi_{w}"] = rsi.fillna(1.0)

                # --- 7. Market Physics (Rolling Beta) ---
                # "Is this asset decoupled from the market?"
                beta_proxy = daily_returns.rolling(w).corr(market_index)  # type: ignore
                df[f"{col}_beta_{w}"] = beta_proxy.fillna(0)

    return df


def _meta_ticker_feature(
    df: pd.DataFrame, lags: list[int], windows: list[int]
) -> pd.DataFrame:
    dfs = [df]
    for lag in lags:
        dfs.append(df.shift(lag).add_suffix(f"_lag{lag}"))
    for window in windows:
        dfs.append(df.rolling(window).mean().add_suffix(f"_rmean{window}"))  # type: ignore
        dfs.append(df.rolling(window).std().add_suffix(f"_rstd{window}"))  # type: ignore
    return pd.concat(dfs, axis=1).replace([np.inf, -np.inf], np.nan)


def _dt_features(df: pd.DataFrame) -> pd.DataFrame:
    dtf = DatetimeFeatures(features_to_extract="all", variables="index")
    return dtf.fit_transform(df)


def _cross_sectional_features(df: pd.DataFrame) -> pd.DataFrame:
    """
    Generates relative strength features across the panel at each timestamp.
    """
    daily_returns = df.pct_change(fill_method=None)

    # 1. Cross-Sectional Rank (0.0 to 1.0)
    cs_rank = daily_returns.rank(axis=1, pct=True).add_suffix("_xs_rank")

    # 2. Cross-Sectional Z-Score (Distance from Market Mean)
    mean = daily_returns.mean(axis=1)
    std = daily_returns.std(axis=1)

    cs_z = daily_returns.sub(mean, axis=0).div(std, axis=0)
    cs_z = cs_z.add_suffix("_xs_z")

    return pd.concat([cs_rank, cs_z], axis=1)


def features(df: pd.DataFrame, windows: list[int], lags: list[int]) -> pd.DataFrame:
    """Generate features on a dataframe."""
    # 1. Deduplicate inputs to prevent column collisions
    windows = sorted(list(set(windows)))
    lags = sorted(list(set(lags)))

    original_cols = df.columns.values.tolist()

    # 2. Generate Cross-Sectional Features (Use ONLY original cols)
    df_xs = _cross_sectional_features(df[original_cols])  # pyright: ignore

    # 3. Generate Time-Series & Market Physics Features (modifies df in-place)
    df = _ticker_features(df=df, windows=windows)

    # 4. Generate Meta Features (Lags/Rolling of the expanded df)
    df = _meta_ticker_feature(df, lags=lags, windows=windows)

    # 5. Merge Cross-Sectional Features
    df = pd.concat([df, df_xs], axis=1)

    # 6. Datetime Features
    df = _dt_features(df=df)

    # 7. Final Clean up
    return df.drop(columns=original_cols).shift(1)
