#!/usr/bin/env python3
"""
PVA (Price Variation Attention) Feature Engineering
====================================================
Computes return-based features for the PVA strategy.

This module provides feature engineering specifically designed for
price variation analysis without temporal features. All features
are derived from price returns and their derivatives.

Key Design Principles:
1. NO temporal features (hour, day, session)
2. Pure price-derived metrics
3. Multi-period return analysis
4. Statistical moments of returns

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

from typing import List, Tuple, Optional, Dict, Any
from dataclasses import dataclass, field

import numpy as np
import pandas as pd
from loguru import logger

try:
    import torch
    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False


@dataclass
class PVAFeatureConfig:
    """Configuration for PVA feature engineering."""

    # Return periods for multi-horizon analysis
    return_periods: List[int] = field(default_factory=lambda: [1, 5, 10, 20])

    # Volatility calculation window
    volatility_window: int = 20

    # Statistical moments window
    stats_window: int = 50

    # Sequence length for transformer input
    sequence_length: int = 100

    # Minimum valid data points after NaN removal
    min_valid_points: int = 200

    # Epsilon for numerical stability
    epsilon: float = 1e-8

    # Z-score clipping threshold
    zscore_clip: float = 5.0


class PVAFeatureEngineer:
    """
    Feature engineering for PVA (Price Variation Attention) strategy.

    Computes return-based features suitable for transformer models.
    NO temporal features are used - only price-derived metrics.

    Features computed:
    - Multi-period returns (1, 5, 10, 20 periods)
    - Return acceleration (derivative of returns)
    - Volatility of returns (rolling std)
    - Skewness and kurtosis of return distributions
    - Normalized price range and body
    - Momentum indicators derived from returns

    Usage:
        engineer = PVAFeatureEngineer()

        # From DataFrame
        df_features = engineer.compute_all_features(df_ohlcv)

        # Get sequences for transformer
        sequences = engineer.prepare_sequences(df_features, seq_length=100)

        # Or use convenience method
        X, y = engineer.prepare_training_data(df, target_horizon=12)
    """

    def __init__(self, config: Optional[PVAFeatureConfig] = None):
        """
        Initialize the feature engineer.

        Args:
            config: Configuration object (uses defaults if None)
        """
        self.config = config or PVAFeatureConfig()
        self._feature_names: List[str] = []

        logger.info("PVAFeatureEngineer initialized")
        logger.info(f"  Return periods: {self.config.return_periods}")
        logger.info(f"  Volatility window: {self.config.volatility_window}")
        logger.info(f"  Stats window: {self.config.stats_window}")

    def compute_returns(
        self,
        df: pd.DataFrame,
        periods: Optional[List[int]] = None
    ) -> pd.DataFrame:
        """
        Compute returns at multiple time horizons.

        Returns are computed as percentage change: (P_t - P_{t-n}) / P_{t-n}

        Args:
            df: DataFrame with 'close' column
            periods: List of periods for return calculation
                    Defaults to config.return_periods

        Returns:
            DataFrame with return columns added
        """
        if periods is None:
            periods = self.config.return_periods

        df = df.copy()
        eps = self.config.epsilon

        if 'close' not in df.columns:
            raise ValueError("DataFrame must have 'close' column")

        close = df['close']

        for period in periods:
            # Simple returns
            returns = (close - close.shift(period)) / (close.shift(period) + eps)
            df[f'return_{period}'] = returns

            # Log returns (more stable for ML)
            log_returns = np.log((close + eps) / (close.shift(period) + eps))
            df[f'log_return_{period}'] = log_returns

        # Primary return (1-period) as main feature
        df['return'] = df['return_1']
        df['log_return'] = df['log_return_1']

        logger.debug(f"Computed returns for periods: {periods}")
        return df

    def compute_acceleration(
        self,
        df: pd.DataFrame,
        periods: Optional[List[int]] = None
    ) -> pd.DataFrame:
        """
        Compute return acceleration (second derivative of price).

        Acceleration measures how fast returns are changing.
        Positive acceleration = returns are increasing
        Negative acceleration = returns are decreasing

        Args:
            df: DataFrame with return columns
            periods: Periods to compute acceleration for

        Returns:
            DataFrame with acceleration columns added
        """
        if periods is None:
            periods = self.config.return_periods

        df = df.copy()

        for period in periods:
            return_col = f'return_{period}'
            if return_col not in df.columns:
                continue

            # Acceleration = change in returns
            returns = df[return_col]
            acceleration = returns - returns.shift(1)
            df[f'acceleration_{period}'] = acceleration

            # Jerk (third derivative) for higher-order dynamics
            jerk = acceleration - acceleration.shift(1)
            df[f'jerk_{period}'] = jerk

        logger.debug(f"Computed acceleration for periods: {periods}")
        return df

    def compute_volatility_returns(
        self,
        df: pd.DataFrame,
        window: Optional[int] = None
    ) -> pd.DataFrame:
        """
        Compute volatility of returns (rolling standard deviation).

        Volatility is a key feature for attention mechanisms as it
        indicates periods of high uncertainty vs stability.

        Args:
            df: DataFrame with return columns
            window: Rolling window size (defaults to config.volatility_window)

        Returns:
            DataFrame with volatility columns added
        """
        if window is None:
            window = self.config.volatility_window

        df = df.copy()

        # Volatility of main return
        if 'log_return' in df.columns:
            df['volatility'] = df['log_return'].rolling(window=window).std()
            df['volatility_pct'] = df['volatility'] * np.sqrt(252 * 12)  # Annualized (5m)

            # Volatility ratio (current vs historical)
            long_window = window * 2
            long_vol = df['log_return'].rolling(window=long_window).std()
            df['volatility_ratio'] = df['volatility'] / (long_vol + self.config.epsilon)

            # Volatility change (momentum in volatility)
            df['volatility_change'] = df['volatility'] - df['volatility'].shift(window // 2)

        # Range-based volatility (Parkinson)
        if all(col in df.columns for col in ['high', 'low']):
            log_hl = np.log(df['high'] / df['low'] + self.config.epsilon)
            parkinson = log_hl ** 2 / (4 * np.log(2))
            df['parkinson_vol'] = np.sqrt(parkinson.rolling(window=window).mean())

        logger.debug(f"Computed volatility with window={window}")
        return df

    def compute_skewness_kurtosis(
        self,
        df: pd.DataFrame,
        window: Optional[int] = None
    ) -> pd.DataFrame:
        """
        Compute rolling skewness and kurtosis of returns.

        These higher-order moments capture:
        - Skewness: asymmetry of return distribution
        - Kurtosis: tail heaviness (fat tails = more extreme events)

        Args:
            df: DataFrame with return columns
            window: Rolling window size (defaults to config.stats_window)

        Returns:
            DataFrame with skewness and kurtosis columns
        """
        if window is None:
            window = self.config.stats_window

        df = df.copy()

        if 'log_return' not in df.columns:
            logger.warning("log_return column not found, skipping stats computation")
            return df

        returns = df['log_return']

        # Rolling mean and std for standardization
        rolling_mean = returns.rolling(window=window).mean()
        rolling_std = returns.rolling(window=window).std() + self.config.epsilon

        # Standardized returns
        standardized = (returns - rolling_mean) / rolling_std

        # Rolling skewness (3rd moment)
        skewness = standardized.pow(3).rolling(window=window).mean()
        df['return_skewness'] = skewness

        # Rolling kurtosis (4th moment) - excess kurtosis
        kurtosis = standardized.pow(4).rolling(window=window).mean() - 3.0
        df['return_kurtosis'] = kurtosis

        # Normalized skewness (bounded)
        max_skew = self.config.zscore_clip
        df['return_skewness_norm'] = np.clip(skewness, -max_skew, max_skew) / max_skew

        # Normalized kurtosis (bounded)
        max_kurt = self.config.zscore_clip * 2
        df['return_kurtosis_norm'] = np.clip(kurtosis, -max_kurt, max_kurt) / max_kurt

        logger.debug(f"Computed skewness/kurtosis with window={window}")
        return df

    def compute_price_features(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        Compute price-derived features (non-temporal).

        These features capture candle characteristics:
        - Range percentage (volatility proxy)
        - Body percentage (direction strength)
        - Wick ratios (rejection/absorption)

        Args:
            df: DataFrame with OHLC columns

        Returns:
            DataFrame with price features added
        """
        df = df.copy()
        eps = self.config.epsilon

        required_cols = ['open', 'high', 'low', 'close']
        if not all(col in df.columns for col in required_cols):
            logger.warning(f"Missing OHLC columns, skipping price features")
            return df

        open_p = df['open']
        high = df['high']
        low = df['low']
        close = df['close']

        # Range as percentage of close (intrabar volatility)
        df['range_pct'] = (high - low) / (close + eps)

        # Body as percentage of range (direction strength)
        body = np.abs(close - open_p)
        range_val = high - low + eps
        df['body_pct'] = body / range_val

        # Direction: 1 for bullish, -1 for bearish, scaled by body
        direction = np.sign(close - open_p)
        df['direction_strength'] = direction * df['body_pct']

        # Upper wick ratio
        body_high = np.maximum(close, open_p)
        df['upper_wick_pct'] = (high - body_high) / range_val

        # Lower wick ratio
        body_low = np.minimum(close, open_p)
        df['lower_wick_pct'] = (body_low - low) / range_val

        # Wick imbalance (positive = more upper wick)
        df['wick_imbalance'] = df['upper_wick_pct'] - df['lower_wick_pct']

        # Close position in range (0 = low, 1 = high)
        df['close_position'] = (close - low) / range_val

        # Gap from previous close (if available)
        prev_close = close.shift(1)
        df['gap_pct'] = (open_p - prev_close) / (prev_close + eps)

        logger.debug("Computed price-derived features")
        return df

    def compute_momentum_features(
        self,
        df: pd.DataFrame,
        periods: Optional[List[int]] = None
    ) -> pd.DataFrame:
        """
        Compute momentum-based features from returns.

        Momentum features capture trend strength and persistence.

        Args:
            df: DataFrame with return columns
            periods: Lookback periods for momentum

        Returns:
            DataFrame with momentum features
        """
        if periods is None:
            periods = self.config.return_periods

        df = df.copy()

        if 'close' not in df.columns:
            return df

        close = df['close']

        for period in periods:
            # Rate of Change (ROC)
            roc = (close - close.shift(period)) / (close.shift(period) + self.config.epsilon)
            df[f'roc_{period}'] = roc

            # Momentum (simple difference, normalized)
            momentum = close - close.shift(period)
            momentum_std = momentum.rolling(window=period * 2).std() + self.config.epsilon
            df[f'momentum_{period}'] = momentum / momentum_std

        # Trend consistency (ratio of positive returns)
        if 'return_1' in df.columns:
            window = self.config.volatility_window
            positive_ratio = (df['return_1'] > 0).rolling(window=window).mean()
            df['trend_consistency'] = 2 * positive_ratio - 1  # Scale to [-1, 1]

            # Consecutive direction count (normalized)
            direction = np.sign(df['return_1'])
            direction_change = (direction != direction.shift(1)).astype(int)

            # Reset counter on direction change
            groups = direction_change.cumsum()
            consecutive = df.groupby(groups).cumcount() + 1
            max_consecutive = window
            df['consecutive_direction'] = np.clip(consecutive, 0, max_consecutive) / max_consecutive
            df['consecutive_direction'] *= direction  # Keep sign

        logger.debug(f"Computed momentum features for periods: {periods}")
        return df

    def compute_all_features(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        Compute all PVA features from OHLCV data.

        This is the main entry point for feature computation.
        Applies all feature transformations in the correct order.

        Args:
            df: DataFrame with OHLCV columns (indexed by timestamp)

        Returns:
            DataFrame with all computed features
        """
        logger.info(f"Computing all PVA features for {len(df)} rows")

        # Step 1: Returns at multiple horizons
        df = self.compute_returns(df)

        # Step 2: Acceleration (second derivative)
        df = self.compute_acceleration(df)

        # Step 3: Volatility of returns
        df = self.compute_volatility_returns(df)

        # Step 4: Statistical moments (skewness, kurtosis)
        df = self.compute_skewness_kurtosis(df)

        # Step 5: Price-derived features
        df = self.compute_price_features(df)

        # Step 6: Momentum features
        df = self.compute_momentum_features(df)

        # Store feature names (exclude OHLCV and timestamp)
        exclude_cols = {'open', 'high', 'low', 'close', 'volume', 'vwap', 'timestamp'}
        self._feature_names = [col for col in df.columns if col not in exclude_cols]

        logger.info(f"Computed {len(self._feature_names)} features")
        logger.debug(f"Feature names: {self._feature_names}")

        return df

    def prepare_sequences(
        self,
        df: pd.DataFrame,
        seq_length: Optional[int] = None,
        feature_cols: Optional[List[str]] = None,
        stride: int = 1
    ) -> np.ndarray:
        """
        Prepare sequences of features for transformer input.

        Converts the DataFrame to overlapping sequences suitable
        for sequence-to-sequence models like transformers.

        Args:
            df: DataFrame with computed features
            seq_length: Sequence length (defaults to config.sequence_length)
            feature_cols: Feature columns to include (defaults to all computed)
            stride: Step size between sequences

        Returns:
            numpy array of shape (n_sequences, seq_length, n_features)
        """
        if seq_length is None:
            seq_length = self.config.sequence_length

        if feature_cols is None:
            feature_cols = self._feature_names
            if not feature_cols:
                # Auto-detect feature columns
                exclude_cols = {'open', 'high', 'low', 'close', 'volume', 'vwap', 'timestamp'}
                feature_cols = [col for col in df.columns if col not in exclude_cols]

        # Extract feature matrix
        features = df[feature_cols].values.astype(np.float32)

        # Handle NaN values
        features = np.nan_to_num(features, nan=0.0, posinf=0.0, neginf=0.0)

        n_samples = len(features)
        n_features = len(feature_cols)

        # Calculate number of valid sequences
        n_sequences = (n_samples - seq_length) // stride + 1

        if n_sequences <= 0:
            logger.warning(f"Not enough data for sequences. Need {seq_length}, have {n_samples}")
            return np.array([]).reshape(0, seq_length, n_features)

        # Create sequences using strides
        sequences = np.zeros((n_sequences, seq_length, n_features), dtype=np.float32)

        for i in range(n_sequences):
            start_idx = i * stride
            end_idx = start_idx + seq_length
            sequences[i] = features[start_idx:end_idx]

        logger.info(f"Prepared {n_sequences} sequences of shape ({seq_length}, {n_features})")
        return sequences

    def prepare_training_data(
        self,
        df: pd.DataFrame,
        target_horizon: int = 12,
        target_type: str = 'return',
        seq_length: Optional[int] = None,
        val_ratio: float = 0.0
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Prepare complete training data with features and targets.

        Convenience method that computes features, creates sequences,
        and generates targets in one call.

        Args:
            df: DataFrame with OHLCV data
            target_horizon: Number of periods ahead for target
            target_type: 'return' for regression, 'direction' for classification
            seq_length: Sequence length for transformer
            val_ratio: If > 0, returns (X_train, y_train, X_val, y_val)

        Returns:
            Tuple of (X, y) numpy arrays
            X shape: (n_samples, seq_length, n_features)
            y shape: (n_samples,) for single target
        """
        if seq_length is None:
            seq_length = self.config.sequence_length

        # Compute all features
        df = self.compute_all_features(df)

        # Compute target
        if target_type == 'return':
            target = df['close'].pct_change(target_horizon).shift(-target_horizon)
        elif target_type == 'direction':
            future_return = df['close'].pct_change(target_horizon).shift(-target_horizon)
            target = (future_return > 0).astype(np.float32)
        elif target_type == 'magnitude':
            target = np.abs(df['close'].pct_change(target_horizon).shift(-target_horizon))
        else:
            raise ValueError(f"Unknown target_type: {target_type}")

        df['_target'] = target

        # Drop rows with NaN target (end of series)
        valid_mask = ~df['_target'].isna()
        df_valid = df[valid_mask].copy()

        # Get feature columns
        feature_cols = self._feature_names

        # Prepare sequences
        X = self.prepare_sequences(df_valid, seq_length=seq_length, feature_cols=feature_cols)

        # Extract targets aligned with sequences
        targets = df_valid['_target'].values.astype(np.float32)
        y = targets[seq_length - 1:]  # Align with last position of each sequence

        # Ensure X and y have same length
        min_len = min(len(X), len(y))
        X = X[:min_len]
        y = y[:min_len]

        logger.info(f"Prepared training data: X={X.shape}, y={y.shape}")

        if val_ratio > 0:
            split_idx = int(len(X) * (1 - val_ratio))
            X_train, X_val = X[:split_idx], X[split_idx:]
            y_train, y_val = y[:split_idx], y[split_idx:]
            return X_train, y_train, X_val, y_val

        return X, y

    def get_feature_names(self) -> List[str]:
        """Get list of computed feature names."""
        return self._feature_names.copy()

    def get_n_features(self) -> int:
        """Get number of features."""
        return len(self._feature_names)

    def to_torch(
        self,
        X: np.ndarray,
        y: Optional[np.ndarray] = None,
        device: str = 'cpu'
    ) -> Tuple:
        """
        Convert numpy arrays to PyTorch tensors.

        Args:
            X: Feature sequences
            y: Optional targets
            device: Device to place tensors ('cpu' or 'cuda')

        Returns:
            Tuple of (X_tensor, y_tensor) or just X_tensor if y is None
        """
        if not TORCH_AVAILABLE:
            raise ImportError("PyTorch not available. Install with: pip install torch")

        X_tensor = torch.from_numpy(X).to(device)

        if y is not None:
            y_tensor = torch.from_numpy(y).to(device)
            return X_tensor, y_tensor

        return X_tensor


def compute_pva_features(
    df: pd.DataFrame,
    return_periods: Optional[List[int]] = None,
    volatility_window: int = 20,
    stats_window: int = 50
) -> pd.DataFrame:
    """
    Convenience function to compute PVA features.

    Args:
        df: DataFrame with OHLCV data
        return_periods: Periods for return calculation
        volatility_window: Window for volatility computation
        stats_window: Window for statistical moments

    Returns:
        DataFrame with all PVA features
    """
    config = PVAFeatureConfig(
        return_periods=return_periods or [1, 5, 10, 20],
        volatility_window=volatility_window,
        stats_window=stats_window
    )

    engineer = PVAFeatureEngineer(config)
    return engineer.compute_all_features(df)


if __name__ == "__main__":
    # Test the feature engineer
    print("Testing PVAFeatureEngineer...")

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

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

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

    # Initialize engineer
    engineer = PVAFeatureEngineer()

    # Compute all features
    df_features = engineer.compute_all_features(df)
    print(f"\nFeatures computed: {len(engineer.get_feature_names())}")
    print(f"Feature names: {engineer.get_feature_names()[:10]}...")

    # Prepare sequences
    sequences = engineer.prepare_sequences(df_features, seq_length=100)
    print(f"\nSequences shape: {sequences.shape}")

    # Prepare training data
    X, y = engineer.prepare_training_data(df, target_horizon=12, target_type='return')
    print(f"\nTraining data:")
    print(f"  X shape: {X.shape}")
    print(f"  y shape: {y.shape}")
    print(f"  y mean: {y.mean():.6f}")
    print(f"  y std: {y.std():.6f}")

    # Test torch conversion
    if TORCH_AVAILABLE:
        X_tensor, y_tensor = engineer.to_torch(X, y)
        print(f"\nTorch tensors:")
        print(f"  X: {X_tensor.shape}, dtype={X_tensor.dtype}")
        print(f"  y: {y_tensor.shape}, dtype={y_tensor.dtype}")

    print("\nTest complete!")
