"""Feature engineering for ML models.

Generates 90+ technical features from OHLCV data.
v2: Added microstructure, Hurst exponent, wavelet features, Fourier features.
Adapted from kriptokamu/src/features/feature_engineering.py
"""
from __future__ import annotations
import logging
import numpy as np
import pandas as pd

logger = logging.getLogger(__name__)


class FeatureEngineer:
    """Creates ML-ready features from OHLCV DataFrame."""

    def __init__(self, timeframe_minutes: int = 60):
        self.tf_minutes = timeframe_minutes

    # Maximum warmup rows to drop (longest rolling window in all features)
    WARMUP_ROWS = 60  # Covers: SMA50, Hurst50, ADX26, wavelet, Fourier, etc.

    def generate(self, df: pd.DataFrame) -> pd.DataFrame:
        """Generate all features from OHLCV DataFrame.

        Args:
            df: DataFrame with columns: open, high, low, close, volume

        Returns:
            DataFrame with all features added, warmup rows removed,
            remaining NaN forward-filled then back-filled.
        """
        df = df.copy()
        df = self._price_features(df)
        df = self._momentum_features(df)
        df = self._trend_features(df)
        df = self._volatility_features(df)
        df = self._volume_features(df)
        df = self._lag_features(df)
        df = self._rolling_stats(df)
        df = self._microstructure_features(df)
        df = self._wavelet_features(df)
        df = self._fourier_features(df)

        # Replace inf with NaN
        df = df.replace([np.inf, -np.inf], np.nan)

        # Drop warmup rows (rolling windows need history to initialize)
        warmup = min(self.WARMUP_ROWS, len(df) - 50)  # Keep at least 50 rows
        if warmup > 0:
            df = df.iloc[warmup:].reset_index(drop=True)

        # Forward-fill then back-fill remaining NaN
        # (microstructure features like flow_imbalance can be sparse —
        #  ffill is valid because they represent rolling market state)
        df = df.ffill().bfill()

        # Fill any remaining all-NaN columns with 0 (neutral value)
        # This handles edge cases where a feature is entirely NaN
        # (e.g., amihud_illiq_20 when volume=0 is prevalent)
        df = df.fillna(0)

        # Final safety: drop any row that still has NaN (shouldn't happen)
        df = df.dropna()

        return df

    def _price_features(self, df: pd.DataFrame) -> pd.DataFrame:
        c, o, h, l = df['close'], df['open'], df['high'], df['low']
        df['price_change'] = c.pct_change()
        df['price_range'] = (h - l) / c.replace(0, np.nan)
        df['body_size'] = abs(c - o) / (h - l).replace(0, np.nan)
        df['upper_wick'] = (h - df[['close', 'open']].max(axis=1)) / (h - l).replace(0, np.nan)
        df['lower_wick'] = (df[['close', 'open']].min(axis=1) - l) / (h - l).replace(0, np.nan)
        df['is_bullish'] = (c > o).astype(int)
        return df

    def _momentum_features(self, df: pd.DataFrame) -> pd.DataFrame:
        c = df['close']

        # RSI (multiple periods)
        for period in [7, 14, 21]:
            delta = c.diff()
            gain = delta.clip(lower=0)
            loss = (-delta.clip(upper=0))
            avg_gain = gain.rolling(period, min_periods=period).mean()
            avg_loss = loss.rolling(period, min_periods=period).mean()
            rs = avg_gain / avg_loss.replace(0, np.nan)
            df[f'rsi_{period}'] = 100 - (100 / (1 + rs))

        # Rate of Change
        df['roc_10'] = c.pct_change(10) * 100

        # Williams %R
        h14 = df['high'].rolling(14).max()
        l14 = df['low'].rolling(14).min()
        df['williams_r'] = -100 * (h14 - c) / (h14 - l14).replace(0, np.nan)

        # Stochastic Oscillator
        df['stoch_k'] = 100 * (c - l14) / (h14 - l14).replace(0, np.nan)
        df['stoch_d'] = df['stoch_k'].rolling(3).mean()

        return df

    def _trend_features(self, df: pd.DataFrame) -> pd.DataFrame:
        c = df['close']

        # SMA & EMA
        for period in [7, 14, 21, 50]:
            df[f'sma_{period}'] = c.rolling(period).mean()
            df[f'ema_{period}'] = c.ewm(span=period, adjust=False).mean()

        # MA crossover signals
        df['sma_cross_7_21'] = (df['sma_7'] > df['sma_21']).astype(int)
        df['ema_cross_7_21'] = (df['ema_7'] > df['ema_21']).astype(int)
        df['price_above_sma21'] = (c > df['sma_21']).astype(int)
        df['price_above_ema21'] = (c > df['ema_21']).astype(int)

        # MACD
        ema12 = c.ewm(span=12, adjust=False).mean()
        ema26 = c.ewm(span=26, adjust=False).mean()
        df['macd'] = ema12 - ema26
        df['macd_signal'] = df['macd'].ewm(span=9, adjust=False).mean()
        df['macd_diff'] = df['macd'] - df['macd_signal']

        # ADX (simplified)
        plus_dm = df['high'].diff().clip(lower=0)
        minus_dm = (-df['low'].diff()).clip(lower=0)
        tr = pd.concat([
            df['high'] - df['low'],
            (df['high'] - c.shift(1)).abs(),
            (df['low'] - c.shift(1)).abs()
        ], axis=1).max(axis=1)
        atr14 = tr.rolling(14).mean()
        plus_di = 100 * plus_dm.rolling(14).mean() / atr14.replace(0, np.nan)
        minus_di = 100 * minus_dm.rolling(14).mean() / atr14.replace(0, np.nan)
        dx = 100 * (plus_di - minus_di).abs() / (plus_di + minus_di).replace(0, np.nan)
        df['adx'] = dx.rolling(14).mean()
        df['adx_plus'] = plus_di
        df['adx_minus'] = minus_di

        return df

    def _volatility_features(self, df: pd.DataFrame) -> pd.DataFrame:
        c = df['close']

        # Bollinger Bands
        sma20 = c.rolling(20).mean()
        std20 = c.rolling(20).std()
        df['bb_upper'] = sma20 + 2 * std20
        df['bb_lower'] = sma20 - 2 * std20
        df['bb_width'] = (df['bb_upper'] - df['bb_lower']) / sma20.replace(0, np.nan)
        df['bb_pct_b'] = (c - df['bb_lower']) / (df['bb_upper'] - df['bb_lower']).replace(0, np.nan)

        # ATR
        tr = pd.concat([
            df['high'] - df['low'],
            (df['high'] - c.shift(1)).abs(),
            (df['low'] - c.shift(1)).abs()
        ], axis=1).max(axis=1)
        df['atr_14'] = tr.rolling(14).mean()
        df['atr_pct'] = df['atr_14'] / c.replace(0, np.nan) * 100

        return df

    def _volume_features(self, df: pd.DataFrame) -> pd.DataFrame:
        v = df['volume']

        # OBV
        obv = (np.sign(df['close'].diff()) * v).fillna(0).cumsum()
        df['obv'] = obv

        # Volume ratios
        vol_ma20 = v.rolling(20).mean()
        df['vol_ratio'] = v / vol_ma20.replace(0, np.nan)
        df['vol_momentum'] = v.pct_change(5)

        # VWAP
        cum_vol = v.cumsum()
        cum_vp = (df['close'] * v).cumsum()
        vwap = cum_vp / cum_vol.replace(0, np.nan)
        df['vwap'] = vwap
        df['vwap_ratio'] = df['close'] / vwap.replace(0, np.nan)

        return df

    def _lag_features(self, df: pd.DataFrame) -> pd.DataFrame:
        c = df['close']
        # Lag periods adapted to timeframe
        lags = [1, 2, 4, 8, 16, 24]
        for lag in lags:
            df[f'close_lag_{lag}'] = c.shift(lag)
            df[f'return_lag_{lag}'] = c.pct_change(lag)
        return df

    def _rolling_stats(self, df: pd.DataFrame) -> pd.DataFrame:
        c = df['close']
        windows = [4, 8, 16, 24]
        for w in windows:
            roll = c.rolling(w)
            mean = roll.mean()
            std = roll.std()
            df[f'roll_{w}_mean'] = mean
            df[f'roll_{w}_std'] = std
            df[f'roll_{w}_min'] = roll.min()
            df[f'roll_{w}_max'] = roll.max()
            df[f'roll_{w}_zscore'] = (c - mean) / std.replace(0, np.nan)
        return df

    # --- v2 Advanced Features ---

    def _microstructure_features(self, df: pd.DataFrame) -> pd.DataFrame:
        """Market microstructure features."""
        c = df['close']
        v = df['volume']
        h = df['high']
        l = df['low']

        # Amihud illiquidity ratio: |return| / volume
        abs_ret = c.pct_change().abs()
        df['amihud_illiq'] = abs_ret / v.replace(0, np.nan)
        df['amihud_illiq_20'] = df['amihud_illiq'].rolling(20, min_periods=5).mean()

        # Bid-ask spread estimate (Corwin-Schultz)
        # Using high-low as proxy
        log_hl = np.log(h / l.replace(0, np.nan))
        log_hl_sq = log_hl ** 2
        beta = log_hl_sq.rolling(2).sum()
        gamma = np.log(
            df[['high']].rolling(2).max().values.flatten() /
            df[['low']].rolling(2).min().values.flatten().clip(min=1e-10)
        ) ** 2
        gamma_s = pd.Series(gamma, index=df.index)
        alpha = (np.sqrt(2 * beta) - np.sqrt(beta)) / (3 - 2 * np.sqrt(2))
        alpha = alpha.clip(lower=0)
        df['spread_estimate'] = 2 * (np.exp(alpha) - 1) / (1 + np.exp(alpha))

        # Volume-weighted price momentum
        vwpm = (c.pct_change() * v).rolling(10, min_periods=3).sum() / v.rolling(10, min_periods=3).sum().replace(0, np.nan)
        df['vw_price_momentum'] = vwpm

        # Trade flow imbalance (buy/sell pressure proxy)
        # Approximation: close > (high+low)/2 = buy pressure
        mid = (h + l) / 2
        buy_vol = v.where(c > mid, 0)
        sell_vol = v.where(c <= mid, 0)
        total_vol = (buy_vol + sell_vol).replace(0, np.nan)
        df['flow_imbalance'] = (buy_vol - sell_vol) / total_vol
        df['flow_imbalance_10'] = df['flow_imbalance'].rolling(10, min_periods=3).mean()

        # Hurst exponent (rolling)
        df['hurst'] = c.rolling(50, min_periods=30).apply(self._hurst_exponent, raw=True)

        return df

    @staticmethod
    def _hurst_exponent(series: np.ndarray, max_lag: int = 15) -> float:
        """Compute Hurst exponent for a price series."""
        if len(series) < max_lag * 2:
            return 0.5
        lags = range(2, max_lag + 1)
        tau = []
        for lag in lags:
            diffs = series[lag:] - series[:-lag]
            std = np.std(diffs)
            if std > 0:
                tau.append(std)
            else:
                tau.append(1e-10)
        if len(tau) < 2:
            return 0.5
        try:
            coeffs = np.polyfit(np.log(list(lags[:len(tau)])), np.log(tau), 1)
            return float(coeffs[0])
        except Exception:
            return 0.5

    def _wavelet_features(self, df: pd.DataFrame) -> pd.DataFrame:
        """Wavelet decomposition features for multi-scale analysis."""
        try:
            import pywt
        except ImportError:
            logger.debug('PyWavelets not installed, skipping wavelet features')
            return df

        c = df['close'].values
        if len(c) < 32:
            return df

        try:
            # Haar wavelet decomposition (3 levels)
            coeffs = pywt.wavedec(c, 'haar', level=3)

            # Reconstruct at each level
            for level in range(1, min(4, len(coeffs))):
                # Zero out all coefficients except this level
                detail = [np.zeros_like(co) for co in coeffs]
                detail[level] = coeffs[level]
                reconstructed = pywt.waverec(detail, 'haar')
                # Align length
                reconstructed = reconstructed[:len(c)]
                if len(reconstructed) == len(df):
                    df[f'wavelet_detail_{level}'] = reconstructed
                    # Energy at this level
                    df[f'wavelet_energy_{level}'] = pd.Series(
                        reconstructed, index=df.index
                    ).rolling(10).apply(lambda x: np.sum(x**2), raw=True)

            # Trend component (approximation)
            approx = [np.zeros_like(co) for co in coeffs]
            approx[0] = coeffs[0]
            trend = pywt.waverec(approx, 'haar')[:len(c)]
            if len(trend) == len(df):
                df['wavelet_trend'] = trend
                df['wavelet_detrended'] = c - trend

        except Exception as e:
            logger.debug(f'Wavelet feature generation failed: {e}')

        return df

    def _fourier_features(self, df: pd.DataFrame) -> pd.DataFrame:
        """Fourier transform features for cycle detection."""
        c = df['close'].values
        if len(c) < 32:
            return df

        try:
            returns = np.diff(c) / c[:-1]
            # Use last 64 data points for FFT
            window = min(64, len(returns))
            fft_input = returns[-window:]

            fft_vals = np.fft.fft(fft_input)
            fft_mag = np.abs(fft_vals[:window // 2])
            fft_freq = np.fft.fftfreq(window)[:window // 2]

            # Dominant frequency (excluding DC component)
            if len(fft_mag) > 1:
                dominant_idx = np.argmax(fft_mag[1:]) + 1
                dominant_freq = float(fft_freq[dominant_idx])
                dominant_power = float(fft_mag[dominant_idx])
                spectral_entropy = float(
                    -np.sum(
                        (fft_mag[1:] / fft_mag[1:].sum().clip(min=1e-10)) *
                        np.log(fft_mag[1:].clip(min=1e-10) / fft_mag[1:].sum().clip(min=1e-10))
                    )
                )

                df['fft_dominant_freq'] = dominant_freq
                df['fft_dominant_power'] = dominant_power
                df['fft_spectral_entropy'] = spectral_entropy

                # Dominant period (in candles)
                if dominant_freq > 0:
                    df['fft_dominant_period'] = 1.0 / dominant_freq
                else:
                    df['fft_dominant_period'] = 0.0

        except Exception as e:
            logger.debug(f'Fourier feature generation failed: {e}')

        return df