#!/usr/bin/env python3
"""
Hierarchical ML Pipeline
========================
Orchestrates the 3-level hierarchical architecture for predictions.

Architecture:
    Level 0: Attention Model - Determines WHEN to pay attention
    Level 1: Base Models - Symbol/timeframe specific predictions
    Level 2: Metamodel - Synthesizes 5m and 15m predictions

Usage:
    from pipelines.hierarchical_pipeline import HierarchicalPipeline

    pipeline = HierarchicalPipeline()
    pipeline.load_models('XAUUSD')

    # Single prediction
    result = pipeline.predict(df_5m, df_15m, 'XAUUSD')

    # Batch predictions
    results = pipeline.predict_batch(data_dict)

Author: ML Pipeline
Version: 1.0.0
Created: 2026-01-07
"""

import sys
import numpy as np
import pandas as pd
from typing import Dict, List, Tuple, Optional, Any, Union
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
import joblib
from loguru import logger


@dataclass
class PipelineConfig:
    """Configuration for the hierarchical pipeline."""

    # Model paths
    attention_model_path: str = 'models/attention'
    base_model_path: str = 'models/symbol_timeframe_models'
    metamodel_path: str = 'models/metamodels'
    neural_gating_path: str = 'models/metamodels_neural'

    # Supported assets
    symbols: List[str] = field(default_factory=lambda: [
        'XAUUSD', 'EURUSD', 'BTCUSD', 'GBPUSD', 'USDJPY'
    ])

    # Timeframes
    timeframes: List[str] = field(default_factory=lambda: ['5m', '15m'])

    # Trading thresholds
    attention_threshold_low: float = 0.8    # Don't trade if attention < this
    attention_threshold_high: float = 2.0   # High confidence if attention > this
    confidence_threshold: float = 0.5       # Minimum confidence to trade

    # Feature configuration
    atr_window: int = 50
    volume_window: int = 20

    # Metamodel selection per symbol (from config/models.yaml)
    # default: 'xgboost', options: 'xgboost', 'neural_gating'
    metamodel_type_default: str = 'xgboost'
    metamodel_type_per_symbol: Dict[str, str] = field(default_factory=lambda: {
        'XAUUSD': 'neural_gating',  # Neural wins
        'EURUSD': 'xgboost',        # XGBoost wins
        'BTCUSD': 'xgboost',
        'GBPUSD': 'xgboost',
        'USDJPY': 'xgboost'
    })


@dataclass
class PredictionResult:
    """Result of hierarchical prediction."""

    symbol: str
    timestamp: datetime

    # Level 0 outputs
    attention_score_5m: float
    attention_score_15m: float
    attention_class_5m: int
    attention_class_15m: int

    # Level 1 outputs
    pred_high_5m: float
    pred_low_5m: float
    pred_high_15m: float
    pred_low_15m: float

    # Level 2 outputs (final)
    delta_high_final: float
    delta_low_final: float
    confidence: bool
    confidence_proba: float

    # Trading signals
    should_trade: bool
    trade_quality: str  # 'high', 'medium', 'low', 'skip'

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            'symbol': self.symbol,
            'timestamp': self.timestamp.isoformat() if self.timestamp else None,
            'attention_score_5m': self.attention_score_5m,
            'attention_score_15m': self.attention_score_15m,
            'attention_class_5m': self.attention_class_5m,
            'attention_class_15m': self.attention_class_15m,
            'pred_high_5m': self.pred_high_5m,
            'pred_low_5m': self.pred_low_5m,
            'pred_high_15m': self.pred_high_15m,
            'pred_low_15m': self.pred_low_15m,
            'delta_high_final': self.delta_high_final,
            'delta_low_final': self.delta_low_final,
            'confidence': self.confidence,
            'confidence_proba': self.confidence_proba,
            'should_trade': self.should_trade,
            'trade_quality': self.trade_quality
        }


class HierarchicalPipeline:
    """
    Orchestrates the 3-level hierarchical ML architecture.

    This pipeline:
    1. Loads all models for a symbol
    2. Processes features through Level 0 (Attention)
    3. Augments features and runs Level 1 (Base Models)
    4. Combines predictions in Level 2 (Metamodel)
    5. Returns final trading signals

    Supports dynamic metamodel selection per symbol (Neural Gating vs XGBoost).
    """

    def __init__(self, config: PipelineConfig = None):
        self.config = config or PipelineConfig()

        # Model storage
        self.attention_models: Dict[str, Any] = {}
        self.base_models: Dict[str, Any] = {}
        self.metamodels: Dict[str, Any] = {}  # XGBoost metamodels
        self.neural_gating_models: Dict[str, Any] = {}  # Neural Gating metamodels

        # Track which metamodel type is loaded for each symbol
        self._metamodel_types: Dict[str, str] = {}

        # State
        self._models_loaded: Dict[str, bool] = {}

        # Import model classes dynamically
        self._import_model_classes()

        logger.info("Initialized HierarchicalPipeline")
        logger.info(f"  Attention path: {self.config.attention_model_path}")
        logger.info(f"  Base path: {self.config.base_model_path}")
        logger.info(f"  Metamodel XGBoost path: {self.config.metamodel_path}")
        logger.info(f"  Metamodel Neural path: {self.config.neural_gating_path}")

    def _import_model_classes(self):
        """Import model classes dynamically."""
        import importlib.util

        models_dir = Path(__file__).parent.parent / 'models'

        # Import AttentionScoreModel
        attention_path = models_dir / 'attention_score_model.py'
        if attention_path.exists():
            spec = importlib.util.spec_from_file_location(
                "attention_score_model", attention_path
            )
            module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(module)
            self._AttentionScoreModel = module.AttentionScoreModel
        else:
            logger.warning(f"AttentionScoreModel not found at {attention_path}")
            self._AttentionScoreModel = None

        # Import AssetMetamodel (XGBoost)
        metamodel_path = models_dir / 'asset_metamodel.py'
        if metamodel_path.exists():
            spec = importlib.util.spec_from_file_location(
                "asset_metamodel", metamodel_path
            )
            module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(module)
            self._AssetMetamodel = module.AssetMetamodel
        else:
            logger.warning(f"AssetMetamodel not found at {metamodel_path}")
            self._AssetMetamodel = None

        # Import NeuralGatingMetamodelWrapper (PyTorch)
        neural_gating_path = models_dir / 'neural_gating_metamodel.py'
        if neural_gating_path.exists():
            try:
                spec = importlib.util.spec_from_file_location(
                    "neural_gating_metamodel", neural_gating_path
                )
                module = importlib.util.module_from_spec(spec)
                spec.loader.exec_module(module)
                self._NeuralGatingMetamodel = module.NeuralGatingMetamodelWrapper
                logger.info("NeuralGatingMetamodel class imported successfully")
            except Exception as e:
                logger.warning(f"NeuralGatingMetamodel import failed: {e}")
                self._NeuralGatingMetamodel = None
        else:
            logger.warning(f"NeuralGatingMetamodel not found at {neural_gating_path}")
            self._NeuralGatingMetamodel = None

    def load_models(self, symbol: str) -> bool:
        """
        Load all models for a symbol.

        Args:
            symbol: Trading symbol (e.g., 'XAUUSD')

        Returns:
            True if all models loaded successfully
        """
        if symbol in self._models_loaded and self._models_loaded[symbol]:
            logger.debug(f"Models already loaded for {symbol}")
            return True

        logger.info(f"Loading models for {symbol}...")

        success = True

        # Load attention models
        for timeframe in self.config.timeframes:
            key = f"{symbol}_{timeframe}_attention"
            model_path = Path(self.config.attention_model_path) / key

            if model_path.exists() and self._AttentionScoreModel:
                try:
                    self.attention_models[key] = self._AttentionScoreModel.load(str(model_path))
                    logger.info(f"  Loaded: {key}")
                except Exception as e:
                    logger.error(f"  Failed to load {key}: {e}")
                    success = False
            else:
                logger.warning(f"  Not found: {model_path}")
                success = False

        # Load base models
        for timeframe in self.config.timeframes:
            for target in ['high', 'low']:
                key = f"{symbol}_{timeframe}_{target}_h3"
                model_file = Path(self.config.base_model_path) / f"{key}.joblib"

                if model_file.exists():
                    try:
                        self.base_models[key] = joblib.load(model_file)
                        logger.info(f"  Loaded: {key}")
                    except Exception as e:
                        logger.error(f"  Failed to load {key}: {e}")
                        success = False
                else:
                    logger.warning(f"  Not found: {model_file}")
                    success = False

        # Determine which metamodel type to use for this symbol
        metamodel_type = self.config.metamodel_type_per_symbol.get(
            symbol, self.config.metamodel_type_default
        )
        logger.info(f"  Metamodel type for {symbol}: {metamodel_type}")

        metamodel_loaded = False

        # Try to load Neural Gating if configured
        if metamodel_type == 'neural_gating' and self._NeuralGatingMetamodel:
            neural_path = Path(self.config.neural_gating_path) / symbol
            if neural_path.exists():
                try:
                    self.neural_gating_models[symbol] = self._NeuralGatingMetamodel.load(str(neural_path))
                    self._metamodel_types[symbol] = 'neural_gating'
                    logger.info(f"  Loaded: {symbol} Neural Gating metamodel")
                    metamodel_loaded = True
                except Exception as e:
                    logger.warning(f"  Failed to load Neural Gating for {symbol}: {e}")
                    logger.info(f"  Falling back to XGBoost metamodel")
            else:
                logger.warning(f"  Neural Gating not found at {neural_path}, falling back to XGBoost")

        # Fallback to XGBoost metamodel
        if not metamodel_loaded:
            metamodel_path = Path(self.config.metamodel_path) / symbol
            if metamodel_path.exists() and self._AssetMetamodel:
                try:
                    self.metamodels[symbol] = self._AssetMetamodel.load(str(metamodel_path))
                    self._metamodel_types[symbol] = 'xgboost'
                    logger.info(f"  Loaded: {symbol} XGBoost metamodel")
                    metamodel_loaded = True
                except Exception as e:
                    logger.error(f"  Failed to load XGBoost metamodel for {symbol}: {e}")
            else:
                logger.warning(f"  XGBoost metamodel not found: {metamodel_path}")

        if not metamodel_loaded:
            logger.warning(f"  No metamodel loaded for {symbol} - will use simple average")
            self._metamodel_types[symbol] = 'none'
            success = False

        self._models_loaded[symbol] = success
        return success

    def _generate_features(self, df: pd.DataFrame) -> pd.DataFrame:
        """Generate features matching base model training."""
        if len(df) == 0:
            return df

        df = df.copy()

        # Normalize column names
        col_map = {'Open': 'open', 'High': 'high', 'Low': 'low',
                   'Close': 'close', 'Volume': 'volume'}
        df.rename(columns={k: v for k, v in col_map.items() if k in df.columns}, inplace=True)

        features = pd.DataFrame(index=df.index)

        close = df['close']
        high = df['high']
        low = df['low']
        open_price = df['open']
        volume = df.get('volume', pd.Series(1, index=df.index))

        # Returns
        features['returns_1'] = close.pct_change(1)
        features['returns_3'] = close.pct_change(3)
        features['returns_5'] = close.pct_change(5)
        features['returns_10'] = close.pct_change(10)
        features['returns_20'] = close.pct_change(20)

        # Volatility
        features['volatility_5'] = close.pct_change().rolling(5).std()
        features['volatility_10'] = close.pct_change().rolling(10).std()
        features['volatility_20'] = close.pct_change().rolling(20).std()

        # Range features
        candle_range = high - low
        features['range'] = candle_range
        features['range_pct'] = candle_range / close
        features['range_ma_5'] = candle_range.rolling(5).mean()
        features['range_ma_10'] = candle_range.rolling(10).mean()
        features['range_ma_20'] = candle_range.rolling(20).mean()
        features['range_ratio_5'] = candle_range / features['range_ma_5']
        features['range_ratio_20'] = candle_range / features['range_ma_20']

        # ATR
        tr1 = high - low
        tr2 = abs(high - close.shift(1))
        tr3 = abs(low - close.shift(1))
        true_range = pd.concat([tr1, tr2, tr3], axis=1).max(axis=1)
        features['atr_5'] = true_range.rolling(5).mean()
        features['atr_14'] = true_range.rolling(14).mean()
        features['atr_20'] = true_range.rolling(20).mean()
        features['atr_ratio'] = true_range / features['atr_14']

        # Moving averages
        sma_5 = close.rolling(5).mean()
        sma_10 = close.rolling(10).mean()
        sma_20 = close.rolling(20).mean()
        sma_50 = close.rolling(50).mean()
        ema_5 = close.ewm(span=5, adjust=False).mean()
        ema_20 = close.ewm(span=20, adjust=False).mean()

        features['price_vs_sma5'] = (close - sma_5) / features['atr_14']
        features['price_vs_sma10'] = (close - sma_10) / features['atr_14']
        features['price_vs_sma20'] = (close - sma_20) / features['atr_14']
        features['price_vs_sma50'] = (close - sma_50) / features['atr_14']
        features['sma5_vs_sma20'] = (sma_5 - sma_20) / features['atr_14']
        features['ema5_vs_ema20'] = (ema_5 - ema_20) / features['atr_14']

        # RSI
        delta = close.diff()
        gain = delta.where(delta > 0, 0).rolling(14).mean()
        loss = (-delta.where(delta < 0, 0)).rolling(14).mean()
        rs = gain / (loss + 1e-10)
        features['rsi_14'] = 100 - (100 / (1 + rs))
        features['rsi_oversold'] = (features['rsi_14'] < 30).astype(float)
        features['rsi_overbought'] = (features['rsi_14'] > 70).astype(float)

        # Bollinger Bands
        bb_middle = close.rolling(20).mean()
        bb_std = close.rolling(20).std()
        bb_upper = bb_middle + 2 * bb_std
        bb_lower = bb_middle - 2 * bb_std
        features['bb_width'] = (bb_upper - bb_lower) / bb_middle
        features['bb_position'] = (close - bb_lower) / (bb_upper - bb_lower + 1e-10)

        # MACD
        ema_12 = close.ewm(span=12, adjust=False).mean()
        ema_26 = close.ewm(span=26, adjust=False).mean()
        macd = ema_12 - ema_26
        macd_signal = macd.ewm(span=9, adjust=False).mean()
        features['macd'] = macd / features['atr_14']
        features['macd_signal'] = macd_signal / features['atr_14']
        features['macd_hist'] = (macd - macd_signal) / features['atr_14']

        # Momentum
        features['momentum_5'] = (close - close.shift(5)) / features['atr_14']
        features['momentum_10'] = (close - close.shift(10)) / features['atr_14']
        features['momentum_20'] = (close - close.shift(20)) / features['atr_14']

        # Stochastic
        low_14 = low.rolling(14).min()
        high_14 = high.rolling(14).max()
        features['stoch_k'] = 100 * (close - low_14) / (high_14 - low_14 + 1e-10)
        features['stoch_d'] = features['stoch_k'].rolling(3).mean()

        # Williams %R
        features['williams_r'] = -100 * (high_14 - close) / (high_14 - low_14 + 1e-10)

        # Volume
        if volume.sum() > 0:
            vol_ma_5 = volume.rolling(5).mean()
            vol_ma_20 = volume.rolling(20).mean()
            features['volume_ratio'] = volume / (vol_ma_20 + 1)
            features['volume_trend'] = (vol_ma_5 - vol_ma_20) / (vol_ma_20 + 1)
        else:
            features['volume_ratio'] = 1.0
            features['volume_trend'] = 0.0

        # Candle patterns
        body = close - open_price
        features['body_pct'] = body / (candle_range + 1e-10)
        features['upper_shadow'] = (high - np.maximum(close, open_price)) / (candle_range + 1e-10)
        features['lower_shadow'] = (np.minimum(close, open_price) - low) / (candle_range + 1e-10)

        # Price position
        features['close_position'] = (close - low) / (candle_range + 1e-10)
        high_5 = high.rolling(5).max()
        low_5 = low.rolling(5).min()
        features['price_position_5'] = (close - low_5) / (high_5 - low_5 + 1e-10)
        high_20 = high.rolling(20).max()
        low_20 = low.rolling(20).min()
        features['price_position_20'] = (close - low_20) / (high_20 - low_20 + 1e-10)

        # Time features
        if hasattr(df.index, 'hour'):
            hour = df.index.hour
            day_of_week = df.index.dayofweek
            features['hour_sin'] = np.sin(2 * np.pi * hour / 24)
            features['hour_cos'] = np.cos(2 * np.pi * hour / 24)
            features['dow_sin'] = np.sin(2 * np.pi * day_of_week / 7)
            features['dow_cos'] = np.cos(2 * np.pi * day_of_week / 7)
            features['is_london'] = ((hour >= 8) & (hour < 16)).astype(float)
            features['is_newyork'] = ((hour >= 13) & (hour < 21)).astype(float)
            features['is_overlap'] = ((hour >= 13) & (hour < 16)).astype(float)

        # Clean
        features = features.replace([np.inf, -np.inf], np.nan)

        # Combine with OHLCV
        result = pd.concat([df[['open', 'high', 'low', 'close', 'volume']], features], axis=1)

        return result

    def _prepare_features_for_base_model(self, df: pd.DataFrame) -> np.ndarray:
        """Prepare features array matching base model training."""
        exclude_patterns = [
            'target_', 'high', 'low', 'open', 'close', 'volume',
            'High', 'Low', 'Open', 'Close', 'Volume',
            'timestamp', 'datetime', 'date', 'time',
            'rr_', 'direction', 'is_valid',
            'attention_score', 'attention_class'  # Exclude if base models weren't trained with attention
        ]

        feature_cols = []
        for col in df.columns:
            if not any(pat.lower() in col.lower() for pat in exclude_patterns):
                if df[col].dtype in [np.float64, np.float32, np.int64, np.int32, float, int]:
                    feature_cols.append(col)

        return df[feature_cols].fillna(0).values

    def _compute_context_features(self, df: pd.DataFrame) -> Tuple[float, float]:
        """Compute ATR_ratio and volume_z for metamodel."""
        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.get('volume', df.get('Volume', pd.Series(1, index=df.index)))

        # ATR ratio
        tr1 = high - low
        tr2 = abs(high - close.shift(1))
        tr3 = abs(low - close.shift(1))
        true_range = pd.concat([tr1, tr2, tr3], axis=1).max(axis=1)
        atr = true_range.rolling(14).mean()
        atr_median = atr.rolling(self.config.atr_window).median()
        atr_ratio = (atr / (atr_median + 1e-10)).iloc[-1]

        # Volume z-score
        vol_mean = volume.rolling(self.config.volume_window).mean()
        vol_std = volume.rolling(self.config.volume_window).std()
        volume_z = ((volume - vol_mean) / (vol_std + 1e-10)).iloc[-1]

        return float(atr_ratio), float(volume_z)

    def predict(
        self,
        df_5m: pd.DataFrame,
        df_15m: pd.DataFrame,
        symbol: str
    ) -> PredictionResult:
        """
        Generate prediction using full hierarchical pipeline.

        Args:
            df_5m: 5-minute OHLCV data (recent history needed for features)
            df_15m: 15-minute OHLCV data
            symbol: Trading symbol

        Returns:
            PredictionResult with all outputs
        """
        # Ensure models are loaded
        if not self._models_loaded.get(symbol, False):
            if not self.load_models(symbol):
                raise ValueError(f"Failed to load models for {symbol}")

        timestamp = df_15m.index[-1] if hasattr(df_15m.index, '__getitem__') else datetime.now()

        # ===== LEVEL 0: Attention =====
        # Generate attention for 5m
        key_5m = f"{symbol}_5m_attention"
        if key_5m in self.attention_models:
            attention_pred_5m = self.attention_models[key_5m].predict(df_5m)
            attention_score_5m = float(attention_pred_5m.attention_score[-1])
            attention_class_5m = int(attention_pred_5m.flow_class[-1])
        else:
            attention_score_5m = 1.5
            attention_class_5m = 1

        # Generate attention for 15m
        key_15m = f"{symbol}_15m_attention"
        if key_15m in self.attention_models:
            attention_pred_15m = self.attention_models[key_15m].predict(df_15m)
            attention_score_15m = float(attention_pred_15m.attention_score[-1])
            attention_class_15m = int(attention_pred_15m.flow_class[-1])
        else:
            attention_score_15m = 1.5
            attention_class_15m = 1

        # ===== LEVEL 1: Base Models =====
        # Generate features for 5m
        df_5m_features = self._generate_features(df_5m)
        df_5m_features['attention_score'] = attention_score_5m
        df_5m_features['attention_class'] = attention_class_5m
        features_5m = self._prepare_features_for_base_model(df_5m_features)

        # Predict with 5m base models
        key_high_5m = f"{symbol}_5m_high_h3"
        key_low_5m = f"{symbol}_5m_low_h3"

        if key_high_5m in self.base_models:
            pred_high_5m = float(self.base_models[key_high_5m].predict(features_5m[-1:].reshape(1, -1))[0])
        else:
            pred_high_5m = 0.0

        if key_low_5m in self.base_models:
            pred_low_5m = float(self.base_models[key_low_5m].predict(features_5m[-1:].reshape(1, -1))[0])
        else:
            pred_low_5m = 0.0

        # Generate features for 15m
        df_15m_features = self._generate_features(df_15m)
        df_15m_features['attention_score'] = attention_score_15m
        df_15m_features['attention_class'] = attention_class_15m
        features_15m = self._prepare_features_for_base_model(df_15m_features)

        # Predict with 15m base models
        key_high_15m = f"{symbol}_15m_high_h3"
        key_low_15m = f"{symbol}_15m_low_h3"

        if key_high_15m in self.base_models:
            pred_high_15m = float(self.base_models[key_high_15m].predict(features_15m[-1:].reshape(1, -1))[0])
        else:
            pred_high_15m = 0.0

        if key_low_15m in self.base_models:
            pred_low_15m = float(self.base_models[key_low_15m].predict(features_15m[-1:].reshape(1, -1))[0])
        else:
            pred_low_15m = 0.0

        # ===== LEVEL 2: Metamodel =====
        # Get context features
        atr_ratio, volume_z = self._compute_context_features(df_15m)

        # Build meta features
        meta_features = pd.DataFrame({
            'pred_high_5m': [pred_high_5m],
            'pred_low_5m': [pred_low_5m],
            'pred_high_15m': [pred_high_15m],
            'pred_low_15m': [pred_low_15m],
            'attention_5m': [attention_score_5m],
            'attention_15m': [attention_score_15m],
            'attention_class_5m': [attention_class_5m],
            'attention_class_15m': [attention_class_15m],
            'ATR_ratio': [atr_ratio],
            'volume_z': [volume_z]
        })

        # Predict with metamodel (Neural Gating or XGBoost based on symbol config)
        metamodel_type = self._metamodel_types.get(symbol, 'none')

        if metamodel_type == 'neural_gating' and symbol in self.neural_gating_models:
            # Use Neural Gating metamodel
            meta_pred = self.neural_gating_models[symbol].predict(meta_features)
            delta_high_final = float(meta_pred.delta_high_final[0])
            delta_low_final = float(meta_pred.delta_low_final[0])
            confidence = bool(meta_pred.confidence[0])
            confidence_proba = float(meta_pred.confidence_proba[0])
            logger.debug(f"  {symbol}: Using Neural Gating (alpha_h={meta_pred.alpha_high[0]:.2f}, alpha_l={meta_pred.alpha_low[0]:.2f})")

        elif metamodel_type == 'xgboost' and symbol in self.metamodels:
            # Use XGBoost metamodel
            meta_pred = self.metamodels[symbol].predict(meta_features)
            delta_high_final = float(meta_pred.delta_high_final[0])
            delta_low_final = float(meta_pred.delta_low_final[0])
            confidence = bool(meta_pred.confidence[0])
            confidence_proba = float(meta_pred.confidence_proba[0])
            logger.debug(f"  {symbol}: Using XGBoost metamodel")

        else:
            # Fallback to simple average
            delta_high_final = (pred_high_5m + pred_high_15m) / 2
            delta_low_final = (pred_low_5m + pred_low_15m) / 2
            confidence = True
            confidence_proba = 0.5
            logger.debug(f"  {symbol}: Using simple average (no metamodel)")

        # ===== Trading Signals =====
        avg_attention = (attention_score_5m + attention_score_15m) / 2

        if avg_attention < self.config.attention_threshold_low:
            should_trade = False
            trade_quality = 'skip'
        elif not confidence or confidence_proba < self.config.confidence_threshold:
            should_trade = False
            trade_quality = 'low'
        elif avg_attention >= self.config.attention_threshold_high:
            should_trade = True
            trade_quality = 'high'
        else:
            should_trade = True
            trade_quality = 'medium'

        return PredictionResult(
            symbol=symbol,
            timestamp=timestamp,
            attention_score_5m=attention_score_5m,
            attention_score_15m=attention_score_15m,
            attention_class_5m=attention_class_5m,
            attention_class_15m=attention_class_15m,
            pred_high_5m=pred_high_5m,
            pred_low_5m=pred_low_5m,
            pred_high_15m=pred_high_15m,
            pred_low_15m=pred_low_15m,
            delta_high_final=delta_high_final,
            delta_low_final=delta_low_final,
            confidence=confidence,
            confidence_proba=confidence_proba,
            should_trade=should_trade,
            trade_quality=trade_quality
        )

    def predict_batch(
        self,
        df_5m: pd.DataFrame,
        df_15m: pd.DataFrame,
        symbol: str,
        step: int = 1
    ) -> List[PredictionResult]:
        """
        Generate predictions for historical data (backtesting).

        Args:
            df_5m: Full 5-minute history
            df_15m: Full 15-minute history
            symbol: Trading symbol
            step: How many bars to step forward between predictions

        Returns:
            List of PredictionResult
        """
        results = []
        min_lookback = 100  # Minimum bars needed for features

        # Align indices
        common_start = max(df_5m.index[min_lookback], df_15m.index[min_lookback // 3])

        df_15m_filtered = df_15m[df_15m.index >= common_start]

        for i in range(0, len(df_15m_filtered), step):
            current_time = df_15m_filtered.index[i]

            # Get data up to current time
            df_5m_slice = df_5m[df_5m.index <= current_time].tail(min_lookback * 3)
            df_15m_slice = df_15m[df_15m.index <= current_time].tail(min_lookback)

            if len(df_5m_slice) < min_lookback or len(df_15m_slice) < min_lookback // 3:
                continue

            try:
                result = self.predict(df_5m_slice, df_15m_slice, symbol)
                results.append(result)
            except Exception as e:
                logger.warning(f"Prediction failed at {current_time}: {e}")
                continue

        return results

    def get_model_info(self, symbol: str) -> Dict[str, Any]:
        """Get information about loaded models for a symbol."""
        metamodel_type = self._metamodel_types.get(symbol, 'none')

        info = {
            'symbol': symbol,
            'models_loaded': self._models_loaded.get(symbol, False),
            'attention_models': [],
            'base_models': [],
            'metamodel_type': metamodel_type,
            'has_xgboost_metamodel': symbol in self.metamodels,
            'has_neural_gating_metamodel': symbol in self.neural_gating_models
        }

        for timeframe in self.config.timeframes:
            key = f"{symbol}_{timeframe}_attention"
            if key in self.attention_models:
                info['attention_models'].append(key)

        for timeframe in self.config.timeframes:
            for target in ['high', 'low']:
                key = f"{symbol}_{timeframe}_{target}_h3"
                if key in self.base_models:
                    info['base_models'].append(key)

        return info


if __name__ == "__main__":
    # Test the pipeline
    print("Testing HierarchicalPipeline...")

    np.random.seed(42)

    # Create sample data
    n_5m = 500
    n_15m = n_5m // 3

    dates_5m = pd.date_range('2024-01-01', periods=n_5m, freq='5min')
    dates_15m = pd.date_range('2024-01-01', periods=n_15m, freq='15min')

    price = 2650 + np.cumsum(np.random.randn(n_5m) * 0.5)

    df_5m = pd.DataFrame({
        'open': price,
        'high': price + np.abs(np.random.randn(n_5m)) * 3,
        'low': price - np.abs(np.random.randn(n_5m)) * 3,
        'close': price + np.random.randn(n_5m) * 0.3,
        'volume': np.random.randint(100, 1000, n_5m)
    }, index=dates_5m)

    price_15m = 2650 + np.cumsum(np.random.randn(n_15m) * 0.8)

    df_15m = pd.DataFrame({
        'open': price_15m,
        'high': price_15m + np.abs(np.random.randn(n_15m)) * 5,
        'low': price_15m - np.abs(np.random.randn(n_15m)) * 5,
        'close': price_15m + np.random.randn(n_15m) * 0.5,
        'volume': np.random.randint(300, 3000, n_15m)
    }, index=dates_15m)

    # Initialize pipeline
    pipeline = HierarchicalPipeline()

    # Test with XAUUSD
    print("\nLoading models for XAUUSD...")
    loaded = pipeline.load_models('XAUUSD')
    print(f"Models loaded: {loaded}")

    if loaded:
        print("\nGenerating prediction...")
        result = pipeline.predict(df_5m, df_15m, 'XAUUSD')

        print(f"\nPrediction Result:")
        print(f"  Timestamp: {result.timestamp}")
        print(f"  Attention 5m: {result.attention_score_5m:.2f} (class={result.attention_class_5m})")
        print(f"  Attention 15m: {result.attention_score_15m:.2f} (class={result.attention_class_15m})")
        print(f"  Pred High 5m: {result.pred_high_5m:.2f}")
        print(f"  Pred Low 5m: {result.pred_low_5m:.2f}")
        print(f"  Pred High 15m: {result.pred_high_15m:.2f}")
        print(f"  Pred Low 15m: {result.pred_low_15m:.2f}")
        print(f"  Delta High Final: {result.delta_high_final:.2f}")
        print(f"  Delta Low Final: {result.delta_low_final:.2f}")
        print(f"  Confidence: {result.confidence} ({result.confidence_proba:.2%})")
        print(f"  Should Trade: {result.should_trade}")
        print(f"  Trade Quality: {result.trade_quality}")

    print("\nTest complete!")
