#!/usr/bin/env python3
"""
Reduced Features Configuration
==============================
Defines the reduced feature set for improved model training.

The 14 core features selected are:
- OHLCV: open, high, low, close, volume (5 features)
- Indicators: ATR, SAR, RSI, MFI, OBV, AD, CMF (7 features)
- Volume derived: volume_z, volume_anomaly (2 features)

Total: 14 features (reduced from 50+)

This configuration reduces overfitting and improves generalization
while maintaining key market information.

Author: ML-Specialist (NEXUS v4.0)
Version: 1.0.0
Created: 2026-01-05
"""

from typing import List, Dict, Optional
from dataclasses import dataclass, field
import numpy as np
import pandas as pd
from loguru import logger

try:
    import pandas_ta as ta
    HAS_PANDAS_TA = True
except ImportError:
    HAS_PANDAS_TA = False
    logger.warning("pandas_ta not available, using manual calculations")


# ==============================================================================
# Core Feature Configuration
# ==============================================================================

# The 14 core columns for training (reduced from 50+)
COLUMNS_TO_TRAIN: List[str] = [
    # OHLCV - 5 base features
    "open",
    "high",
    "low",
    "close",
    "volume",
    # Technical Indicators - 7 features
    "ATR",      # Average True Range (volatility)
    "SAR",      # Parabolic SAR (trend)
    "RSI",      # Relative Strength Index (momentum)
    "MFI",      # Money Flow Index (volume-weighted momentum)
    "OBV",      # On-Balance Volume (volume trend)
    "AD",       # Accumulation/Distribution (volume-price)
    "CMF",      # Chaikin Money Flow (volume-price)
    # Volume derived - 2 features
    "volume_z",       # Volume z-score
    "volume_anomaly"  # Volume anomaly flag
]

# Lowercase version for compatibility
COLUMNS_TO_TRAIN_LOWER: List[str] = [col.lower() for col in COLUMNS_TO_TRAIN]

# Feature groups for reference
FEATURE_GROUPS: Dict[str, List[str]] = {
    "ohlcv": ["open", "high", "low", "close", "volume"],
    "volatility": ["ATR"],
    "trend": ["SAR"],
    "momentum": ["RSI", "MFI"],
    "volume_flow": ["OBV", "AD", "CMF"],
    "volume_derived": ["volume_z", "volume_anomaly"]
}

# Indicator parameters
INDICATOR_PARAMS: Dict[str, Dict] = {
    "ATR": {"length": 14},
    "SAR": {"af0": 0.02, "af": 0.02, "max_af": 0.2},
    "RSI": {"length": 14},
    "MFI": {"length": 14},
    "OBV": {},  # No parameters
    "AD": {},   # No parameters
    "CMF": {"length": 20},
}


@dataclass
class ReducedFeatureConfig:
    """Configuration for reduced feature generation"""

    # ATR parameters
    atr_length: int = 14

    # SAR parameters
    sar_af0: float = 0.02
    sar_af: float = 0.02
    sar_max_af: float = 0.2

    # RSI parameters
    rsi_length: int = 14

    # MFI parameters
    mfi_length: int = 14

    # CMF parameters
    cmf_length: int = 20

    # Volume z-score parameters
    volume_z_window: int = 20

    # Volume anomaly threshold (z-score)
    volume_anomaly_threshold: float = 2.0

    # Normalization
    normalize_ohlcv: bool = False  # Keep raw for now
    normalize_volume: bool = False


# ==============================================================================
# Feature Generation Functions
# ==============================================================================

def calculate_atr(
    df: pd.DataFrame,
    length: int = 14
) -> pd.Series:
    """
    Calculate Average True Range.

    True Range = max(high-low, abs(high-prev_close), abs(low-prev_close))
    ATR = EMA(True Range, length)
    """
    high = df['high'] if 'high' in df.columns else df['High']
    low = df['low'] if 'low' in df.columns else df['Low']
    close = df['close'] if 'close' in df.columns else df['Close']

    tr1 = high - low
    tr2 = abs(high - close.shift(1))
    tr3 = abs(low - close.shift(1))
    tr = pd.concat([tr1, tr2, tr3], axis=1).max(axis=1)

    atr = tr.ewm(span=length, adjust=False).mean()
    return atr


def calculate_sar(
    df: pd.DataFrame,
    af0: float = 0.02,
    af: float = 0.02,
    max_af: float = 0.2
) -> pd.Series:
    """
    Calculate Parabolic SAR.

    Uses pandas_ta if available, otherwise manual calculation.
    """
    high = df['high'] if 'high' in df.columns else df['High']
    low = df['low'] if 'low' in df.columns else df['Low']
    close = df['close'] if 'close' in df.columns else df['Close']

    if HAS_PANDAS_TA:
        sar = ta.psar(high, low, close, af0=af0, af=af, max_af=max_af)
        if sar is not None:
            # Get the SAR values (first column)
            return sar.iloc[:, 0].fillna(close)

    # Manual calculation (simplified)
    n = len(df)
    sar = pd.Series(index=df.index, dtype=float)

    # Initialize
    trend = 1  # 1 for uptrend, -1 for downtrend
    ep = high.iloc[0]  # Extreme point
    sar_val = low.iloc[0]
    current_af = af0

    for i in range(1, n):
        if trend == 1:  # Uptrend
            sar_val = sar_val + current_af * (ep - sar_val)
            sar_val = min(sar_val, low.iloc[i-1], low.iloc[i-2] if i > 1 else low.iloc[i-1])

            if high.iloc[i] > ep:
                ep = high.iloc[i]
                current_af = min(current_af + af, max_af)

            if low.iloc[i] < sar_val:
                trend = -1
                sar_val = ep
                ep = low.iloc[i]
                current_af = af0
        else:  # Downtrend
            sar_val = sar_val + current_af * (ep - sar_val)
            sar_val = max(sar_val, high.iloc[i-1], high.iloc[i-2] if i > 1 else high.iloc[i-1])

            if low.iloc[i] < ep:
                ep = low.iloc[i]
                current_af = min(current_af + af, max_af)

            if high.iloc[i] > sar_val:
                trend = 1
                sar_val = ep
                ep = high.iloc[i]
                current_af = af0

        sar.iloc[i] = sar_val

    return sar


def calculate_rsi(
    df: pd.DataFrame,
    length: int = 14
) -> pd.Series:
    """
    Calculate Relative Strength Index.

    RSI = 100 - 100 / (1 + RS)
    RS = avg_gain / avg_loss
    """
    close = df['close'] if 'close' in df.columns else df['Close']

    if HAS_PANDAS_TA:
        return ta.rsi(close, length=length)

    delta = close.diff()
    gain = delta.where(delta > 0, 0)
    loss = -delta.where(delta < 0, 0)

    avg_gain = gain.ewm(span=length, adjust=False).mean()
    avg_loss = loss.ewm(span=length, adjust=False).mean()

    rs = avg_gain / (avg_loss + 1e-10)
    rsi = 100 - 100 / (1 + rs)

    return rsi


def calculate_mfi(
    df: pd.DataFrame,
    length: int = 14
) -> pd.Series:
    """
    Calculate Money Flow Index.

    Typical Price = (High + Low + Close) / 3
    Raw Money Flow = Typical Price * Volume
    MFI = 100 - 100 / (1 + Money Ratio)
    """
    high = df['high'] if 'high' in df.columns else df['High']
    low = df['low'] if 'low' in df.columns else df['Low']
    close = df['close'] if 'close' in df.columns else df['Close']
    volume = df['volume'] if 'volume' in df.columns else df['Volume']

    if HAS_PANDAS_TA:
        return ta.mfi(high, low, close, volume, length=length)

    typical_price = (high + low + close) / 3
    raw_money_flow = typical_price * volume

    tp_diff = typical_price.diff()
    positive_flow = raw_money_flow.where(tp_diff > 0, 0)
    negative_flow = raw_money_flow.where(tp_diff < 0, 0)

    positive_mf = positive_flow.rolling(length).sum()
    negative_mf = negative_flow.rolling(length).sum()

    money_ratio = positive_mf / (negative_mf + 1e-10)
    mfi = 100 - 100 / (1 + money_ratio)

    return mfi


def calculate_obv(df: pd.DataFrame) -> pd.Series:
    """
    Calculate On-Balance Volume.

    OBV = cumsum(sign(close_change) * volume)
    """
    close = df['close'] if 'close' in df.columns else df['Close']
    volume = df['volume'] if 'volume' in df.columns else df['Volume']

    if HAS_PANDAS_TA:
        return ta.obv(close, volume)

    sign = np.sign(close.diff())
    obv = (sign * volume).cumsum()

    return obv


def calculate_ad(df: pd.DataFrame) -> pd.Series:
    """
    Calculate Accumulation/Distribution.

    CLV = ((close - low) - (high - close)) / (high - low)
    AD = cumsum(CLV * volume)
    """
    high = df['high'] if 'high' in df.columns else df['High']
    low = df['low'] if 'low' in df.columns else df['Low']
    close = df['close'] if 'close' in df.columns else df['Close']
    volume = df['volume'] if 'volume' in df.columns else df['Volume']

    if HAS_PANDAS_TA:
        return ta.ad(high, low, close, volume)

    clv = ((close - low) - (high - close)) / (high - low + 1e-10)
    ad = (clv * volume).cumsum()

    return ad


def calculate_cmf(
    df: pd.DataFrame,
    length: int = 20
) -> pd.Series:
    """
    Calculate Chaikin Money Flow.

    CMF = sum(AD_volume, length) / sum(volume, length)
    """
    high = df['high'] if 'high' in df.columns else df['High']
    low = df['low'] if 'low' in df.columns else df['Low']
    close = df['close'] if 'close' in df.columns else df['Close']
    volume = df['volume'] if 'volume' in df.columns else df['Volume']

    if HAS_PANDAS_TA:
        return ta.cmf(high, low, close, volume, length=length)

    clv = ((close - low) - (high - close)) / (high - low + 1e-10)
    ad_volume = clv * volume

    cmf = ad_volume.rolling(length).sum() / (volume.rolling(length).sum() + 1e-10)

    return cmf


def calculate_volume_z(
    df: pd.DataFrame,
    window: int = 20
) -> pd.Series:
    """
    Calculate Volume Z-Score.

    z = (volume - mean) / std
    """
    volume = df['volume'] if 'volume' in df.columns else df['Volume']

    vol_mean = volume.rolling(window).mean()
    vol_std = volume.rolling(window).std()
    vol_std = vol_std.replace(0, 1)

    z = (volume - vol_mean) / vol_std

    return z


def calculate_volume_anomaly(
    df: pd.DataFrame,
    window: int = 20,
    threshold: float = 2.0
) -> pd.Series:
    """
    Calculate Volume Anomaly Flag.

    Returns 1 if volume z-score exceeds threshold, 0 otherwise.
    """
    z = calculate_volume_z(df, window)
    anomaly = (z.abs() > threshold).astype(float)

    return anomaly


# ==============================================================================
# Main Feature Generation
# ==============================================================================

def generate_reduced_features(
    df: pd.DataFrame,
    config: ReducedFeatureConfig = None
) -> pd.DataFrame:
    """
    Generate the reduced feature set (14 features).

    Args:
        df: DataFrame with OHLCV data
        config: Feature configuration

    Returns:
        DataFrame with the 14 core features
    """
    config = config or ReducedFeatureConfig()

    logger.info("Generating reduced feature set (14 features)...")

    # Start with OHLCV
    result = df[['open', 'high', 'low', 'close', 'volume']].copy()

    # Normalize column names to lowercase
    result.columns = [c.lower() for c in result.columns]

    # Add ATR
    result['ATR'] = calculate_atr(result, length=config.atr_length)

    # Add SAR
    result['SAR'] = calculate_sar(
        result,
        af0=config.sar_af0,
        af=config.sar_af,
        max_af=config.sar_max_af
    )

    # Add RSI
    result['RSI'] = calculate_rsi(result, length=config.rsi_length)

    # Add MFI
    result['MFI'] = calculate_mfi(result, length=config.mfi_length)

    # Add OBV
    result['OBV'] = calculate_obv(result)

    # Add AD
    result['AD'] = calculate_ad(result)

    # Add CMF
    result['CMF'] = calculate_cmf(result, length=config.cmf_length)

    # Add volume z-score
    result['volume_z'] = calculate_volume_z(result, window=config.volume_z_window)

    # Add volume anomaly
    result['volume_anomaly'] = calculate_volume_anomaly(
        result,
        window=config.volume_z_window,
        threshold=config.volume_anomaly_threshold
    )

    # Fill NaN
    result = result.fillna(method='ffill').fillna(0)

    # Verify we have all columns
    missing = set(COLUMNS_TO_TRAIN) - set(result.columns)
    if missing:
        logger.warning(f"Missing columns: {missing}")

    logger.info(f"Generated {len(result.columns)} features: {list(result.columns)}")

    return result


def get_feature_columns() -> List[str]:
    """
    Get the list of feature columns for training.

    Returns:
        List of feature column names
    """
    return COLUMNS_TO_TRAIN.copy()


def get_feature_columns_without_ohlcv() -> List[str]:
    """
    Get feature columns excluding OHLCV.

    Returns:
        List of indicator feature names
    """
    ohlcv = ["open", "high", "low", "close", "volume"]
    return [col for col in COLUMNS_TO_TRAIN if col.lower() not in ohlcv]


# ==============================================================================
# Test
# ==============================================================================

if __name__ == "__main__":
    print("Testing Reduced Features Configuration")
    print("=" * 60)

    # Create sample data
    np.random.seed(42)
    n = 1000

    dates = pd.date_range('2025-01-01', periods=n, freq='5min')
    price = 2650 + np.cumsum(np.random.randn(n) * 0.5)

    df = pd.DataFrame({
        'open': price + np.random.randn(n) * 0.5,
        'high': price + np.abs(np.random.randn(n)) * 3,
        'low': price - np.abs(np.random.randn(n)) * 3,
        'close': price + np.random.randn(n) * 0.5,
        'volume': np.random.randint(100, 1000, n)
    }, index=dates)

    # Ensure high > low
    df['high'] = df[['open', 'high', 'close']].max(axis=1)
    df['low'] = df[['open', 'low', 'close']].min(axis=1)

    print(f"\nInput data shape: {df.shape}")
    print(f"Input columns: {list(df.columns)}")

    # Generate features
    features = generate_reduced_features(df)

    print(f"\nOutput data shape: {features.shape}")
    print(f"Output columns ({len(features.columns)}): {list(features.columns)}")

    print(f"\nExpected columns ({len(COLUMNS_TO_TRAIN)}): {COLUMNS_TO_TRAIN}")

    # Check sample values
    print("\nSample values (last 5 rows):")
    print(features.tail())

    # Statistics
    print("\nFeature statistics:")
    print(features.describe())

    print("\n" + "=" * 60)
    print("Test completed!")
