#!/usr/bin/env python3
"""
MTS Feature Engineering - Multi-Timeframe Synthesis Features
=============================================================
Feature engineering for multi-timeframe analysis. Aggregates base timeframe (5m)
to higher timeframes (15m, 1h, 4h) and computes alignment/conflict scores.

Timeframe Hierarchy:
- 5m (base): Raw market data
- 15m (3x): Short-term trends
- 1h (12x): Intraday trends
- 4h (48x): Swing trends

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

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


@dataclass
class TimeframeDefinition:
    """Definition of a timeframe for aggregation."""
    name: str
    aggregation_factor: int
    feature_windows: List[int] = field(default_factory=list)
    weight: float = 1.0


class MTSFeatureEngineer:
    """
    Multi-Timeframe Synthesis Feature Engineer.

    Aggregates 5-minute base data to higher timeframes and computes
    features for each timeframe along with alignment and conflict scores.

    Features computed per timeframe:
    - Price returns (close-to-close)
    - Range features (high-low, body size)
    - Momentum indicators (RSI, MACD-like)
    - Volatility measures (ATR, std)
    - Volume features (if available)
    - Trend strength indicators

    Cross-timeframe features:
    - Alignment score: How aligned are signals across TFs
    - Conflict score: How conflicting are signals across TFs
    - Dominant TF: Which TF has the strongest signal
    """

    TIMEFRAMES = {
        '5m': TimeframeDefinition(name='5m', aggregation_factor=1, feature_windows=[12, 24, 48], weight=0.15),
        '15m': TimeframeDefinition(name='15m', aggregation_factor=3, feature_windows=[4, 8, 16], weight=0.25),
        '1h': TimeframeDefinition(name='1h', aggregation_factor=12, feature_windows=[4, 8, 24], weight=0.35),
        '4h': TimeframeDefinition(name='4h', aggregation_factor=48, feature_windows=[3, 6, 12], weight=0.25),
    }

    def __init__(
        self,
        timeframes: Optional[List[str]] = None,
        include_volume: bool = True,
        feature_lookback: int = 100
    ):
        """
        Initialize MTS Feature Engineer.

        Args:
            timeframes: List of timeframes to use (default: all)
            include_volume: Whether to include volume features
            feature_lookback: Lookback period for feature calculation
        """
        self.timeframes = timeframes or list(self.TIMEFRAMES.keys())
        self.include_volume = include_volume
        self.feature_lookback = feature_lookback

        # Validate timeframes
        for tf in self.timeframes:
            if tf not in self.TIMEFRAMES:
                raise ValueError(f"Unknown timeframe: {tf}")

        logger.info(f"Initialized MTSFeatureEngineer with timeframes: {self.timeframes}")

    def aggregate_to_timeframe(
        self,
        df_5m: pd.DataFrame,
        target_tf: str
    ) -> pd.DataFrame:
        """
        Aggregate 5-minute data to a target timeframe.

        Args:
            df_5m: DataFrame with 5-minute OHLCV data
            target_tf: Target timeframe ('5m', '15m', '1h', '4h')

        Returns:
            Aggregated DataFrame
        """
        if target_tf not in self.TIMEFRAMES:
            raise ValueError(f"Unknown timeframe: {target_tf}")

        tf_def = self.TIMEFRAMES[target_tf]
        factor = tf_def.aggregation_factor

        if factor == 1:
            return df_5m.copy()

        # Ensure we have the required columns
        required_cols = ['open', 'high', 'low', 'close']
        for col in required_cols:
            if col not in df_5m.columns:
                raise ValueError(f"Missing required column: {col}")

        n_rows = len(df_5m)
        n_aggregated = n_rows // factor

        if n_aggregated == 0:
            logger.warning(f"Not enough data to aggregate to {target_tf}")
            return pd.DataFrame()

        # Truncate to exact multiple of factor
        df_truncated = df_5m.iloc[-(n_aggregated * factor):].copy()

        # Reshape for aggregation
        idx = np.arange(len(df_truncated)) // factor

        aggregated_data = {
            'open': df_truncated.groupby(idx)['open'].first().values,
            'high': df_truncated.groupby(idx)['high'].max().values,
            'low': df_truncated.groupby(idx)['low'].min().values,
            'close': df_truncated.groupby(idx)['close'].last().values,
        }

        # Volume if available
        if 'volume' in df_5m.columns and self.include_volume:
            aggregated_data['volume'] = df_truncated.groupby(idx)['volume'].sum().values

        # Create timestamp index
        if isinstance(df_truncated.index, pd.DatetimeIndex):
            timestamps = df_truncated.groupby(idx).apply(lambda x: x.index[-1]).values
            df_agg = pd.DataFrame(aggregated_data, index=timestamps)
        else:
            df_agg = pd.DataFrame(aggregated_data)

        return df_agg

    def compute_tf_features(
        self,
        df: pd.DataFrame,
        tf: str
    ) -> pd.DataFrame:
        """
        Compute standard features for a single timeframe.

        Args:
            df: OHLCV DataFrame
            tf: Timeframe name (for naming features)

        Returns:
            DataFrame with computed features
        """
        if len(df) < 10:
            return pd.DataFrame()

        tf_def = self.TIMEFRAMES.get(tf, TimeframeDefinition(name=tf, aggregation_factor=1, feature_windows=[4, 8, 16]))
        windows = tf_def.feature_windows

        features = {}

        # Basic price features
        close = df['close'].values
        high = df['high'].values
        low = df['low'].values
        open_price = df['open'].values

        # Returns
        returns = np.zeros(len(close))
        returns[1:] = (close[1:] - close[:-1]) / close[:-1]
        features[f'{tf}_return'] = returns

        # Log returns
        log_returns = np.zeros(len(close))
        log_returns[1:] = np.log(close[1:] / close[:-1])
        features[f'{tf}_log_return'] = log_returns

        # Range features
        range_hl = (high - low) / close
        features[f'{tf}_range_hl'] = range_hl

        body_size = np.abs(close - open_price) / close
        features[f'{tf}_body_size'] = body_size

        # Direction (1 = bullish, -1 = bearish)
        direction = np.sign(close - open_price)
        features[f'{tf}_direction'] = direction

        # Upper/lower wick ratios
        upper_wick = np.where(close >= open_price, high - close, high - open_price)
        lower_wick = np.where(close >= open_price, open_price - low, close - low)
        total_range = high - low + 1e-10

        features[f'{tf}_upper_wick_ratio'] = upper_wick / total_range
        features[f'{tf}_lower_wick_ratio'] = lower_wick / total_range

        # Window-based features
        for window in windows:
            if len(close) < window:
                continue

            # SMA
            sma = self._rolling_mean(close, window)
            features[f'{tf}_sma_{window}'] = sma

            # Price distance from SMA
            sma_dist = (close - sma) / (sma + 1e-10)
            features[f'{tf}_sma_dist_{window}'] = sma_dist

            # Standard deviation
            std = self._rolling_std(close, window)
            features[f'{tf}_std_{window}'] = std / (close + 1e-10)

            # ATR
            atr = self._compute_atr(high, low, close, window)
            features[f'{tf}_atr_{window}'] = atr / (close + 1e-10)

            # RSI
            rsi = self._compute_rsi(close, window)
            features[f'{tf}_rsi_{window}'] = rsi

            # Momentum
            momentum = np.zeros(len(close))
            momentum[window:] = (close[window:] - close[:-window]) / close[:-window]
            features[f'{tf}_momentum_{window}'] = momentum

            # Rate of Change (ROC)
            roc = np.zeros(len(close))
            roc[window:] = 100 * (close[window:] - close[:-window]) / close[:-window]
            features[f'{tf}_roc_{window}'] = roc

            # Return sum (cumulative return over window)
            return_sum = self._rolling_sum(returns, window)
            features[f'{tf}_return_sum_{window}'] = return_sum

            # Highest high / lowest low
            rolling_high = self._rolling_max(high, window)
            rolling_low = self._rolling_min(low, window)

            # Price position in range
            range_position = (close - rolling_low) / (rolling_high - rolling_low + 1e-10)
            features[f'{tf}_range_position_{window}'] = range_position

        # MACD-like features (fast EMA - slow EMA)
        if len(close) >= max(windows):
            fast_window = min(windows)
            slow_window = max(windows)

            fast_ema = self._compute_ema(close, fast_window)
            slow_ema = self._compute_ema(close, slow_window)

            macd = (fast_ema - slow_ema) / close
            features[f'{tf}_macd'] = macd

            # MACD signal line
            macd_signal = self._compute_ema(macd, min(9, len(macd) // 2))
            features[f'{tf}_macd_signal'] = macd_signal

            # MACD histogram
            features[f'{tf}_macd_hist'] = macd - macd_signal

        # Trend strength (ADX-like)
        adx = self._compute_adx(high, low, close, min(14, len(close) // 2))
        features[f'{tf}_adx'] = adx

        # Volume features
        if 'volume' in df.columns and self.include_volume:
            volume = df['volume'].values

            features[f'{tf}_volume_norm'] = volume / (self._rolling_mean(volume, 20) + 1e-10)

            for window in windows:
                if len(volume) >= window:
                    vol_sma = self._rolling_mean(volume, window)
                    features[f'{tf}_vol_sma_{window}'] = vol_sma / (self._rolling_mean(volume, max(windows)) + 1e-10)

        # Convert to DataFrame
        df_features = pd.DataFrame(features, index=df.index if hasattr(df, 'index') else None)

        return df_features

    def compute_tf_alignment_score(
        self,
        features_dict: Dict[str, pd.DataFrame]
    ) -> np.ndarray:
        """
        Compute alignment score across timeframes.

        Alignment is high when all timeframes agree on direction/trend.

        Args:
            features_dict: Dictionary of {timeframe: features_df}

        Returns:
            Array of alignment scores (0 to 1)
        """
        if len(features_dict) < 2:
            return np.ones(len(list(features_dict.values())[0]))

        # Get minimum length across all timeframes
        min_len = min(len(df) for df in features_dict.values() if len(df) > 0)

        if min_len == 0:
            return np.array([])

        # Collect direction signals from each timeframe
        directions = []
        weights = []

        for tf, df in features_dict.items():
            if len(df) == 0:
                continue

            tf_def = self.TIMEFRAMES.get(tf, TimeframeDefinition(name=tf, aggregation_factor=1, weight=0.25))

            # Get direction feature
            dir_col = f'{tf}_direction'
            if dir_col in df.columns:
                # Align to minimum length (take last min_len values)
                dir_values = df[dir_col].values[-min_len:]
                directions.append(dir_values)
                weights.append(tf_def.weight)
            else:
                # Use momentum as proxy
                momentum_cols = [c for c in df.columns if 'momentum' in c or 'return' in c]
                if momentum_cols:
                    momentum = df[momentum_cols[0]].values[-min_len:]
                    directions.append(np.sign(momentum))
                    weights.append(tf_def.weight)

        if len(directions) < 2:
            return np.ones(min_len)

        directions = np.array(directions)
        weights = np.array(weights)
        weights = weights / weights.sum()

        # Compute alignment as weighted agreement
        alignment_scores = np.zeros(min_len)

        for i in range(min_len):
            dir_at_i = directions[:, i]

            # Count agreements
            bullish = np.sum(weights[dir_at_i > 0])
            bearish = np.sum(weights[dir_at_i < 0])

            # Alignment is max of bullish/bearish agreement
            alignment_scores[i] = max(bullish, bearish)

        return alignment_scores

    def compute_conflict_score(
        self,
        features_dict: Dict[str, pd.DataFrame]
    ) -> np.ndarray:
        """
        Compute conflict score across timeframes.

        Conflict is high when timeframes disagree strongly.

        Args:
            features_dict: Dictionary of {timeframe: features_df}

        Returns:
            Array of conflict scores (0 to 1)
        """
        alignment = self.compute_tf_alignment_score(features_dict)

        # Conflict is inverse of alignment, but normalized
        # High alignment = low conflict
        conflict = 1.0 - alignment

        # Scale to emphasize strong conflicts
        conflict = np.power(conflict, 0.5)

        return conflict

    def compute_dominant_tf(
        self,
        features_dict: Dict[str, pd.DataFrame]
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Compute the dominant timeframe at each point.

        The dominant TF is the one with the strongest directional signal.

        Args:
            features_dict: Dictionary of {timeframe: features_df}

        Returns:
            Tuple of (dominant_tf_indices, signal_strengths)
        """
        if len(features_dict) == 0:
            return np.array([]), np.array([])

        min_len = min(len(df) for df in features_dict.values() if len(df) > 0)

        if min_len == 0:
            return np.array([]), np.array([])

        tf_names = list(features_dict.keys())
        signal_matrix = np.zeros((len(tf_names), min_len))

        for i, (tf, df) in enumerate(features_dict.items()):
            if len(df) == 0:
                continue

            # Compute signal strength using multiple indicators
            strength = np.zeros(min_len)
            count = 0

            # RSI deviation from neutral
            rsi_cols = [c for c in df.columns if 'rsi' in c]
            for col in rsi_cols:
                rsi = df[col].values[-min_len:]
                strength += np.abs(rsi - 50) / 50
                count += 1

            # Momentum strength
            momentum_cols = [c for c in df.columns if 'momentum' in c]
            for col in momentum_cols:
                mom = df[col].values[-min_len:]
                # Normalize momentum
                mom_norm = np.abs(mom) / (np.std(mom) + 1e-10)
                strength += np.clip(mom_norm, 0, 3) / 3
                count += 1

            # MACD histogram
            macd_hist_col = f'{tf}_macd_hist'
            if macd_hist_col in df.columns:
                macd_hist = df[macd_hist_col].values[-min_len:]
                macd_norm = np.abs(macd_hist) / (np.std(macd_hist) + 1e-10)
                strength += np.clip(macd_norm, 0, 3) / 3
                count += 1

            # ADX (trend strength)
            adx_col = f'{tf}_adx'
            if adx_col in df.columns:
                adx = df[adx_col].values[-min_len:]
                strength += adx / 100
                count += 1

            if count > 0:
                signal_matrix[i] = strength / count

        # Apply timeframe weights
        for i, tf in enumerate(tf_names):
            tf_def = self.TIMEFRAMES.get(tf, TimeframeDefinition(name=tf, aggregation_factor=1, weight=0.25))
            signal_matrix[i] *= tf_def.weight

        # Find dominant TF at each point
        dominant_indices = np.argmax(signal_matrix, axis=0)
        signal_strengths = np.max(signal_matrix, axis=0)

        return dominant_indices, signal_strengths

    def prepare_hierarchical_input(
        self,
        dfs_dict: Dict[str, pd.DataFrame]
    ) -> Dict[str, np.ndarray]:
        """
        Prepare input for hierarchical attention model.

        Args:
            dfs_dict: Dictionary of {timeframe: ohlcv_df}

        Returns:
            Dictionary with features, alignment, conflict, and dominant TF
        """
        # Compute features for each timeframe
        features_dict = {}
        for tf, df in dfs_dict.items():
            if tf in self.timeframes:
                features_dict[tf] = self.compute_tf_features(df, tf)

        # Find common length (use last N samples where all TFs have data)
        min_len = min(len(df) for df in features_dict.values() if len(df) > 0)

        # Align all features to same length
        aligned_features = {}
        for tf, df in features_dict.items():
            aligned_features[tf] = df.iloc[-min_len:].values

        # Compute cross-TF features
        alignment = self.compute_tf_alignment_score(features_dict)
        conflict = self.compute_conflict_score(features_dict)
        dominant_tf, signal_strength = self.compute_dominant_tf(features_dict)

        return {
            'tf_features': aligned_features,
            'alignment': alignment[-min_len:] if len(alignment) > min_len else alignment,
            'conflict': conflict[-min_len:] if len(conflict) > min_len else conflict,
            'dominant_tf': dominant_tf[-min_len:] if len(dominant_tf) > min_len else dominant_tf,
            'signal_strength': signal_strength[-min_len:] if len(signal_strength) > min_len else signal_strength,
            'feature_names': {tf: list(features_dict[tf].columns) for tf in features_dict},
            'length': min_len
        }

    def prepare_training_data(
        self,
        df_5m: pd.DataFrame,
        target_horizon: int = 12,
        return_threshold: float = 0.001
    ) -> Tuple[Dict[str, np.ndarray], np.ndarray, np.ndarray]:
        """
        Prepare training data with features and targets.

        Args:
            df_5m: Base 5-minute OHLCV data
            target_horizon: Forward horizon in 5m bars for target
            return_threshold: Threshold for direction classification

        Returns:
            Tuple of (hierarchical_input, direction_targets, magnitude_targets)
        """
        # Aggregate to all timeframes
        dfs_dict = {}
        for tf in self.timeframes:
            dfs_dict[tf] = self.aggregate_to_timeframe(df_5m, tf)

        # Prepare hierarchical input
        hierarchical_input = self.prepare_hierarchical_input(dfs_dict)

        # Compute targets from 5m data
        close = df_5m['close'].values

        # Forward returns
        forward_returns = np.zeros(len(close))
        if target_horizon < len(close):
            forward_returns[:-target_horizon] = (
                close[target_horizon:] - close[:-target_horizon]
            ) / close[:-target_horizon]

        # Align targets to feature length
        min_len = hierarchical_input['length']
        forward_returns = forward_returns[-min_len:]

        # Direction targets (1=up, 0=down)
        direction_targets = (forward_returns > return_threshold).astype(np.float32)

        # Magnitude targets (absolute return)
        magnitude_targets = np.abs(forward_returns).astype(np.float32)

        return hierarchical_input, direction_targets, magnitude_targets

    # ==================== Helper Methods ====================

    def _rolling_mean(self, arr: np.ndarray, window: int) -> np.ndarray:
        """Compute rolling mean."""
        result = np.full(len(arr), np.nan)
        if len(arr) >= window:
            cumsum = np.cumsum(np.insert(arr, 0, 0))
            result[window-1:] = (cumsum[window:] - cumsum[:-window]) / window
            result[:window-1] = np.nan
        return result

    def _rolling_std(self, arr: np.ndarray, window: int) -> np.ndarray:
        """Compute rolling standard deviation."""
        result = np.full(len(arr), np.nan)
        if len(arr) >= window:
            for i in range(window - 1, len(arr)):
                result[i] = np.std(arr[i - window + 1:i + 1])
        return result

    def _rolling_sum(self, arr: np.ndarray, window: int) -> np.ndarray:
        """Compute rolling sum."""
        result = np.full(len(arr), np.nan)
        if len(arr) >= window:
            cumsum = np.cumsum(np.insert(arr, 0, 0))
            result[window-1:] = cumsum[window:] - cumsum[:-window]
        return result

    def _rolling_max(self, arr: np.ndarray, window: int) -> np.ndarray:
        """Compute rolling maximum."""
        result = np.full(len(arr), np.nan)
        for i in range(window - 1, len(arr)):
            result[i] = np.max(arr[i - window + 1:i + 1])
        return result

    def _rolling_min(self, arr: np.ndarray, window: int) -> np.ndarray:
        """Compute rolling minimum."""
        result = np.full(len(arr), np.nan)
        for i in range(window - 1, len(arr)):
            result[i] = np.min(arr[i - window + 1:i + 1])
        return result

    def _compute_ema(self, arr: np.ndarray, span: int) -> np.ndarray:
        """Compute exponential moving average."""
        alpha = 2 / (span + 1)
        result = np.zeros(len(arr))
        result[0] = arr[0]
        for i in range(1, len(arr)):
            result[i] = alpha * arr[i] + (1 - alpha) * result[i - 1]
        return result

    def _compute_atr(
        self,
        high: np.ndarray,
        low: np.ndarray,
        close: np.ndarray,
        window: int
    ) -> np.ndarray:
        """Compute Average True Range."""
        tr = np.zeros(len(close))
        tr[0] = high[0] - low[0]

        for i in range(1, len(close)):
            hl = high[i] - low[i]
            hc = abs(high[i] - close[i - 1])
            lc = abs(low[i] - close[i - 1])
            tr[i] = max(hl, hc, lc)

        return self._rolling_mean(tr, window)

    def _compute_rsi(self, close: np.ndarray, window: int) -> np.ndarray:
        """Compute Relative Strength Index."""
        delta = np.diff(close)
        gain = np.where(delta > 0, delta, 0)
        loss = np.where(delta < 0, -delta, 0)

        avg_gain = np.zeros(len(close))
        avg_loss = np.zeros(len(close))

        if len(close) > window:
            avg_gain[window] = np.mean(gain[:window])
            avg_loss[window] = np.mean(loss[:window])

            for i in range(window + 1, len(close)):
                avg_gain[i] = (avg_gain[i - 1] * (window - 1) + gain[i - 1]) / window
                avg_loss[i] = (avg_loss[i - 1] * (window - 1) + loss[i - 1]) / window

        rs = avg_gain / (avg_loss + 1e-10)
        rsi = 100 - (100 / (1 + rs))
        rsi[:window] = 50  # Neutral for initial period

        return rsi

    def _compute_adx(
        self,
        high: np.ndarray,
        low: np.ndarray,
        close: np.ndarray,
        window: int
    ) -> np.ndarray:
        """Compute Average Directional Index."""
        if len(close) < window * 2:
            return np.full(len(close), 25.0)  # Neutral ADX

        # Plus and Minus Directional Movement
        plus_dm = np.zeros(len(close))
        minus_dm = np.zeros(len(close))

        for i in range(1, len(close)):
            up_move = high[i] - high[i - 1]
            down_move = low[i - 1] - low[i]

            if up_move > down_move and up_move > 0:
                plus_dm[i] = up_move
            if down_move > up_move and down_move > 0:
                minus_dm[i] = down_move

        # True Range
        tr = np.zeros(len(close))
        tr[0] = high[0] - low[0]
        for i in range(1, len(close)):
            tr[i] = max(
                high[i] - low[i],
                abs(high[i] - close[i - 1]),
                abs(low[i] - close[i - 1])
            )

        # Smoothed values
        atr = self._compute_ema(tr, window)
        plus_di = 100 * self._compute_ema(plus_dm, window) / (atr + 1e-10)
        minus_di = 100 * self._compute_ema(minus_dm, window) / (atr + 1e-10)

        # DX and ADX
        dx = 100 * np.abs(plus_di - minus_di) / (plus_di + minus_di + 1e-10)
        adx = self._compute_ema(dx, window)

        return adx


if __name__ == "__main__":
    # Test MTSFeatureEngineer
    np.random.seed(42)

    # Create sample 5m OHLCV data
    n_samples = 1000

    # Simulate price data with trend and noise
    base_price = 100.0
    returns = np.random.randn(n_samples) * 0.001 + 0.0001  # Slight upward bias
    close = base_price * np.cumprod(1 + returns)

    high = close * (1 + np.abs(np.random.randn(n_samples) * 0.002))
    low = close * (1 - np.abs(np.random.randn(n_samples) * 0.002))
    open_price = close * (1 + np.random.randn(n_samples) * 0.001)
    volume = np.random.randint(1000, 10000, n_samples)

    df_5m = pd.DataFrame({
        'open': open_price,
        'high': high,
        'low': low,
        'close': close,
        'volume': volume
    })

    # Initialize feature engineer
    engineer = MTSFeatureEngineer()

    # Test aggregation
    print("Testing timeframe aggregation:")
    for tf in ['5m', '15m', '1h', '4h']:
        df_agg = engineer.aggregate_to_timeframe(df_5m, tf)
        print(f"  {tf}: {len(df_agg)} bars")

    # Test feature computation
    print("\nTesting feature computation:")
    for tf in ['5m', '15m', '1h']:
        df_agg = engineer.aggregate_to_timeframe(df_5m, tf)
        features = engineer.compute_tf_features(df_agg, tf)
        print(f"  {tf}: {features.shape[1]} features")

    # Test hierarchical input preparation
    print("\nTesting hierarchical input:")
    dfs_dict = {tf: engineer.aggregate_to_timeframe(df_5m, tf) for tf in ['5m', '15m', '1h', '4h']}
    hierarchical = engineer.prepare_hierarchical_input(dfs_dict)

    print(f"  Aligned length: {hierarchical['length']}")
    print(f"  Alignment score range: [{hierarchical['alignment'].min():.3f}, {hierarchical['alignment'].max():.3f}]")
    print(f"  Conflict score range: [{hierarchical['conflict'].min():.3f}, {hierarchical['conflict'].max():.3f}]")
    print(f"  Features per TF: {[len(v) for v in hierarchical['feature_names'].values()]}")

    print("\nMTSFeatureEngineer test complete!")
