#!/usr/bin/env python3
"""
Symbol-Timeframe Specific Trainer
=================================
Trains separate models for each symbol and timeframe combination.

Key Features:
1. Separate models per symbol (XAUUSD, BTCUSD, EURUSD)
2. Separate models per timeframe (5m, 15m)
3. Last year excluded from training for backtesting
4. Dynamic factor-based sample weighting
5. Walk-forward validation support

Architecture:
- Each (symbol, timeframe) pair gets its own model
- Models are organized in a hierarchy: symbol/timeframe/model_type
- Supports retraining short-term models independently

Author: Trading Strategist + ML Specialist
Version: 1.0.0
Created: 2026-01-05
"""

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, timedelta
from pathlib import Path
import joblib
from loguru import logger

try:
    from xgboost import XGBRegressor, XGBClassifier
    HAS_XGBOOST = True
except ImportError:
    HAS_XGBOOST = False
    logger.warning("XGBoost not available")

from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score


@dataclass
class SymbolConfig:
    """Configuration for a specific trading symbol"""
    symbol: str
    base_factor: float  # Base volatility factor in USD/pips
    pip_value: float = 1.0  # Value of 1 pip (1.0 for commodities, 0.0001 for forex)
    typical_spread: float = 0.0  # Typical spread for the symbol


# Default configurations for common symbols
SYMBOL_CONFIGS = {
    'XAUUSD': SymbolConfig(
        symbol='XAUUSD',
        base_factor=5.0,  # ~5 USD typical 5m range
        pip_value=0.01,
        typical_spread=0.30
    ),
    'BTCUSD': SymbolConfig(
        symbol='BTCUSD',
        base_factor=100.0,  # ~100 USD typical 5m range
        pip_value=0.01,
        typical_spread=10.0
    ),
    'EURUSD': SymbolConfig(
        symbol='EURUSD',
        base_factor=0.0005,  # ~5 pips typical 5m range
        pip_value=0.0001,
        typical_spread=0.0001
    ),
    'GBPUSD': SymbolConfig(
        symbol='GBPUSD',
        base_factor=0.0006,  # ~6 pips typical 5m range
        pip_value=0.0001,
        typical_spread=0.00012
    ),
    'USDJPY': SymbolConfig(
        symbol='USDJPY',
        base_factor=0.05,  # ~5 pips typical 5m range
        pip_value=0.01,
        typical_spread=0.012
    )
}


@dataclass
class TrainerConfig:
    """Configuration for the symbol-timeframe trainer"""

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

    # Symbols to train
    symbols: List[str] = field(default_factory=lambda: ['XAUUSD', 'BTCUSD', 'EURUSD'])

    # Horizon configuration per timeframe (bars to look ahead)
    horizons: Dict[str, int] = field(default_factory=lambda: {
        '5m': 3,   # 15 minutes
        '15m': 3,  # 45 minutes
        '1H': 3,   # 3 hours
    })

    # Training data configuration
    train_years: float = 5.0  # Years of data for training
    holdout_years: float = 1.0  # Last year excluded for backtesting

    # XGBoost parameters (more regularization for better generalization)
    xgb_params: Dict = field(default_factory=lambda: {
        'n_estimators': 150,     # Reduced from 300
        'max_depth': 4,          # Reduced from 6 - shallower trees
        'learning_rate': 0.02,   # Reduced from 0.03
        'subsample': 0.7,        # Reduced from 0.8
        'colsample_bytree': 0.7, # Reduced from 0.8
        'min_child_weight': 20,  # Increased from 10 - more regularization
        'gamma': 0.3,            # Increased from 0.1
        'reg_alpha': 0.5,        # Increased from 0.1 - L1 regularization
        'reg_lambda': 5.0,       # Increased from 1.0 - L2 regularization
        'tree_method': 'hist',
        'random_state': 42
    })

    # Sample weighting (reduced aggressiveness for better generalization)
    use_dynamic_factor_weighting: bool = True
    factor_window: int = 200
    softplus_beta: float = 2.0  # Reduced from 4.0 - less aggressive
    softplus_w_max: float = 2.0  # Reduced from 3.0 - cap weights lower

    # ATR weighting
    use_atr_weighting: bool = True
    use_session_weighting: bool = False

    # Validation
    val_split: float = 0.15

    # Minimum samples required
    min_train_samples: int = 5000

    # Attention model integration (Nivel 0 -> Nivel 1)
    use_attention_features: bool = False
    attention_model_path: str = 'models/attention'  # Path to trained attention models


@dataclass
class ModelKey:
    """Unique identifier for a model"""
    symbol: str
    timeframe: str
    target_type: str  # 'high', 'low', 'direction'
    horizon_bars: int

    def __str__(self):
        return f"{self.symbol}_{self.timeframe}_{self.target_type}_h{self.horizon_bars}"

    def to_path(self) -> str:
        return f"{self.symbol}/{self.timeframe}/{self.target_type}_h{self.horizon_bars}"


@dataclass
class TrainingResult:
    """Result of training a single model"""
    model_key: ModelKey
    mae: float = 0.0
    rmse: float = 0.0
    r2: float = 0.0
    directional_accuracy: float = 0.0
    n_train: int = 0
    n_val: int = 0
    train_date_range: str = ""
    holdout_date_range: str = ""
    feature_importance: Dict[str, float] = field(default_factory=dict)


class SymbolTimeframeTrainer:
    """
    Trains separate models for each symbol and timeframe combination.

    Key Features:
    1. Symbol-specific: Each symbol gets its own model
    2. Timeframe-specific: Each timeframe gets separate models
    3. Holdout support: Automatically excludes last year for backtesting
    4. Dynamic weighting: Uses rolling median factor with shift(1)

    Usage:
        trainer = SymbolTimeframeTrainer(TrainerConfig())

        # Train all models
        results = trainer.train_all(data_dict)

        # Or train specific symbol/timeframe
        result = trainer.train_single(df, 'XAUUSD', '15m')

        # Predict
        predictions = trainer.predict(df, 'XAUUSD', '15m')

        # Save/load
        trainer.save('models/')
        trainer.load('models/')
    """

    def __init__(self, config: TrainerConfig = None, attention_trainer=None):
        self.config = config or TrainerConfig()
        self.models: Dict[str, Any] = {}
        self.results: Dict[str, TrainingResult] = {}
        self.symbol_configs: Dict[str, SymbolConfig] = {}
        self._is_trained = False
        self.attention_trainer = attention_trainer
        self._attention_loaded = False

        # Initialize symbol configs
        for symbol in self.config.symbols:
            if symbol in SYMBOL_CONFIGS:
                self.symbol_configs[symbol] = SYMBOL_CONFIGS[symbol]
            else:
                # Create default config
                self.symbol_configs[symbol] = SymbolConfig(
                    symbol=symbol,
                    base_factor=1.0
                )

        # Load attention models if configured and not provided
        if self.config.use_attention_features and attention_trainer is None:
            self._load_attention_trainer()

        logger.info(f"Initialized SymbolTimeframeTrainer")
        logger.info(f"  Symbols: {self.config.symbols}")
        logger.info(f"  Timeframes: {self.config.timeframes}")
        logger.info(f"  Holdout: last {self.config.holdout_years} year(s)")
        logger.info(f"  Use attention features: {self.config.use_attention_features}")

    def _load_attention_trainer(self):
        """Load attention trainer from disk if use_attention_features is enabled."""
        try:
            # Import modules dynamically to avoid circular imports
            import importlib.util

            # First load the attention_score_model with consistent module name
            models_dir = Path(__file__).parent.parent / 'models'
            attention_model_path = models_dir / 'attention_score_model.py'

            if not attention_model_path.exists():
                logger.warning(f"Attention score model not found: {attention_model_path}")
                self.config.use_attention_features = False
                return

            # Use consistent module name for joblib pickle compatibility
            module_name = "models.attention_score_model"
            if module_name not in sys.modules:
                spec_model = importlib.util.spec_from_file_location(module_name, attention_model_path)
                attention_model_module = importlib.util.module_from_spec(spec_model)
                sys.modules[module_name] = attention_model_module
                spec_model.loader.exec_module(attention_model_module)
            else:
                attention_model_module = sys.modules[module_name]

            AttentionScoreModel = attention_model_module.AttentionScoreModel

            # Create a simple AttentionModelTrainer class for loading
            class AttentionTrainerConfig:
                def __init__(self, symbols=None, timeframes=None, train_years=5.0, holdout_years=1.0):
                    self.symbols = symbols or ['XAUUSD', 'EURUSD', 'BTCUSD', 'GBPUSD', 'USDJPY']
                    self.timeframes = timeframes or ['5m', '15m']
                    self.train_years = train_years
                    self.holdout_years = holdout_years

            class AttentionModelTrainer:
                def __init__(self, config=None):
                    self.config = config or AttentionTrainerConfig()
                    self.models = {}

                def predict(self, df, symbol, timeframe):
                    key = f"{symbol}_{timeframe}_attention"
                    if key not in self.models:
                        raise ValueError(f"No model for {key}")
                    return self.models[key].predict(df)

                @classmethod
                def load(cls, path):
                    import joblib as jl
                    path = Path(path)
                    metadata = jl.load(path / 'trainer_metadata.joblib')
                    config = AttentionTrainerConfig(
                        symbols=metadata['config']['symbols'],
                        timeframes=metadata['config']['timeframes'],
                        train_years=metadata['config']['train_years'],
                        holdout_years=metadata['config']['holdout_years']
                    )
                    trainer = cls(config)
                    for key in metadata['model_keys']:
                        model_path = path / key
                        if model_path.exists():
                            trainer.models[key] = AttentionScoreModel.load(str(model_path))
                    logger.info(f"Loaded {len(trainer.models)} attention models from {path}")
                    return trainer

            # Load the trainer
            attention_path = Path(self.config.attention_model_path)
            if attention_path.exists():
                self.attention_trainer = AttentionModelTrainer.load(str(attention_path))
                self._attention_loaded = True
                logger.info(f"Loaded attention trainer from {attention_path}")
            else:
                logger.warning(f"Attention model path not found: {attention_path}")
                logger.warning("Training will proceed WITHOUT attention features")
                self.config.use_attention_features = False
        except Exception as e:
            logger.warning(f"Failed to load attention trainer: {e}")
            import traceback
            traceback.print_exc()
            logger.warning("Training will proceed WITHOUT attention features")
            self.config.use_attention_features = False

    def _generate_attention_features(
        self,
        df: pd.DataFrame,
        symbol: str,
        timeframe: str
    ) -> pd.DataFrame:
        """
        Generate attention features for the DataFrame.

        Args:
            df: DataFrame with OHLCV data
            symbol: Trading symbol
            timeframe: Timeframe

        Returns:
            DataFrame with attention_score and attention_class columns added
        """
        if not self.config.use_attention_features or self.attention_trainer is None:
            return df

        try:
            # Get attention prediction
            prediction = self.attention_trainer.predict(df, symbol, timeframe)

            # Add attention features
            df_with_attention = df.copy()
            df_with_attention['attention_score'] = prediction.attention_score
            df_with_attention['attention_class'] = prediction.flow_class

            logger.info(f"Added attention features for {symbol} {timeframe}")
            logger.info(f"  attention_score mean: {prediction.attention_score.mean():.2f}")
            logger.info(f"  high_flow (class=2) ratio: {(prediction.flow_class == 2).mean():.1%}")

            return df_with_attention

        except Exception as e:
            logger.warning(f"Failed to generate attention features: {e}")
            return df

    def _get_date_split(
        self,
        df: pd.DataFrame
    ) -> Tuple[pd.Timestamp, pd.Timestamp, pd.Timestamp]:
        """
        Get date splits for training and holdout.

        Returns:
            Tuple of (train_start, train_end, holdout_end)
        """
        if isinstance(df.index, pd.DatetimeIndex):
            max_date = df.index.max()
            min_date = df.index.min()
        else:
            # Try to find datetime column
            for col in ['timestamp', 'datetime', 'date', 'time']:
                if col in df.columns:
                    max_date = pd.to_datetime(df[col]).max()
                    min_date = pd.to_datetime(df[col]).min()
                    break
            else:
                raise ValueError("No datetime index or column found")

        # Holdout = last N years
        holdout_start = max_date - timedelta(days=self.config.holdout_years * 365)

        # Training = holdout_years before holdout
        train_start = holdout_start - timedelta(days=self.config.train_years * 365)
        train_start = max(train_start, min_date)  # Can't go before data start

        return train_start, holdout_start, max_date

    def _split_train_holdout(
        self,
        df: pd.DataFrame
    ) -> Tuple[pd.DataFrame, pd.DataFrame]:
        """
        Split data into training and holdout sets.

        Holdout = last holdout_years of data (for backtesting)
        Training = everything before holdout

        Args:
            df: DataFrame with datetime index

        Returns:
            Tuple of (train_df, holdout_df)
        """
        train_start, holdout_start, max_date = self._get_date_split(df)

        if isinstance(df.index, pd.DatetimeIndex):
            train_mask = (df.index >= train_start) & (df.index < holdout_start)
            holdout_mask = df.index >= holdout_start
        else:
            # Use timestamp column
            for col in ['timestamp', 'datetime', 'date', 'time']:
                if col in df.columns:
                    timestamps = pd.to_datetime(df[col])
                    train_mask = (timestamps >= train_start) & (timestamps < holdout_start)
                    holdout_mask = timestamps >= holdout_start
                    break

        train_df = df[train_mask].copy()
        holdout_df = df[holdout_mask].copy()

        logger.info(f"Data split:")
        logger.info(f"  Training: {train_start.strftime('%Y-%m-%d')} to {holdout_start.strftime('%Y-%m-%d')} "
                   f"({len(train_df)} samples)")
        logger.info(f"  Holdout: {holdout_start.strftime('%Y-%m-%d')} to {max_date.strftime('%Y-%m-%d')} "
                   f"({len(holdout_df)} samples)")

        return train_df, holdout_df

    def _compute_atr(
        self,
        df: pd.DataFrame,
        period: int = 14
    ) -> np.ndarray:
        """
        Compute ATR with shift(1) to avoid data leakage.

        Args:
            df: DataFrame with OHLC columns
            period: ATR period

        Returns:
            Array of shifted ATR values
        """
        high_col = 'High' if 'High' in df.columns else 'high'
        low_col = 'Low' if 'Low' in df.columns else 'low'
        close_col = 'Close' if 'Close' in df.columns else 'close'

        high = df[high_col]
        low = df[low_col]
        close = df[close_col]

        # True Range components
        tr1 = high - low
        tr2 = abs(high - close.shift(1))
        tr3 = abs(low - close.shift(1))

        # True Range
        true_range = pd.concat([tr1, tr2, tr3], axis=1).max(axis=1)

        # ATR with shift(1) to avoid leakage
        atr = true_range.rolling(period, min_periods=period//2).mean().shift(1)

        return atr.values

    def _compute_targets(
        self,
        df: pd.DataFrame,
        horizon_bars: int,
        normalize: bool = True
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        Compute corrected targets for range prediction.

        Formula:
        - target_high = MAX(high[t+1:t+horizon+1]) - close[t]
        - target_low = close[t] - MIN(low[t+1:t+horizon+1])

        If normalize=True, targets are normalized by ATR(14) with shift(1).
        This puts targets in a scale of [-3, 3] ATR multiples instead of
        raw price units which can be very small for forex pairs.

        Args:
            df: DataFrame with OHLC columns
            horizon_bars: Number of bars to look ahead
            normalize: Whether to normalize by ATR (default True)

        Returns:
            Tuple of (target_high, target_low, atr) arrays
            - If normalize=True: targets are in ATR multiples
            - If normalize=False: targets are in raw price units
            - atr: ATR values for denormalization during prediction
        """
        # Handle column name variations
        close_col = 'Close' if 'Close' in df.columns else 'close'
        high_col = 'High' if 'High' in df.columns else 'high'
        low_col = 'Low' if 'Low' in df.columns else 'low'

        close = df[close_col].values
        high = df[high_col].values
        low = df[low_col].values
        n = len(df)

        target_high = np.full(n, np.nan)
        target_low = np.full(n, np.nan)

        for i in range(n - horizon_bars):
            # Future window [t+1, t+horizon]
            future_high = high[i+1:i+1+horizon_bars]
            future_low = low[i+1:i+1+horizon_bars]

            target_high[i] = np.max(future_high) - close[i]
            target_low[i] = close[i] - np.min(future_low)

        # Compute ATR for normalization
        atr = self._compute_atr(df, period=14)

        if normalize:
            # Normalize targets by ATR
            # Use epsilon to avoid division by zero
            epsilon = 1e-10
            atr_safe = np.where(atr > epsilon, atr, epsilon)

            target_high_norm = target_high / atr_safe
            target_low_norm = target_low / atr_safe

            # Clip extreme values (beyond 5 ATR is outlier)
            target_high_norm = np.clip(target_high_norm, -5, 5)
            target_low_norm = np.clip(target_low_norm, -5, 5)

            logger.info(f"Target normalization applied:")
            logger.info(f"  Raw target_high mean: {np.nanmean(target_high):.6f}")
            logger.info(f"  Normalized target_high mean: {np.nanmean(target_high_norm):.4f}")
            logger.info(f"  ATR mean: {np.nanmean(atr):.6f}")

            return target_high_norm, target_low_norm, atr

        return target_high, target_low, atr

    def _compute_sample_weights(
        self,
        df: pd.DataFrame,
        target_high: np.ndarray,
        target_low: np.ndarray
    ) -> np.ndarray:
        """
        Compute sample weights using dynamic factor approach.

        Args:
            df: DataFrame with OHLC data
            target_high: Target high values
            target_low: Target low values

        Returns:
            Array of sample weights
        """
        if not self.config.use_dynamic_factor_weighting:
            return np.ones(len(df))

        # Handle column name variations
        high_col = 'High' if 'High' in df.columns else 'high'
        low_col = 'Low' if 'Low' in df.columns else 'low'

        # Calculate candle range
        candle_range = (df[high_col] - df[low_col]).abs()

        # Rolling median with shift(1) to avoid leakage
        factor = candle_range.rolling(
            window=self.config.factor_window,
            min_periods=self.config.factor_window // 2
        ).median().shift(1)

        # Compute move multiplier from target movement
        total_target = np.abs(target_high) + np.abs(target_low)
        m = total_target / (factor.values + 1e-12)

        # Apply softplus weighting
        beta = self.config.softplus_beta
        x = beta * (m - 1.0)
        w = np.where(
            x > 20,
            x / beta,
            np.log1p(np.exp(x)) / beta
        )
        weights = np.clip(w, 0.0, self.config.softplus_w_max)

        # Handle NaN
        nan_mask = np.isnan(weights) | np.isnan(factor.values)
        weights[nan_mask] = 1.0

        # Normalize
        valid_mask = ~nan_mask
        if valid_mask.sum() > 0 and weights[valid_mask].mean() > 0:
            weights[valid_mask] = weights[valid_mask] / weights[valid_mask].mean()

        logger.info(f"Sample weights computed:")
        logger.info(f"  Mean multiplier: {np.nanmean(m):.2f}")
        logger.info(f"  High attention (w>1.5): {(weights > 1.5).sum()} samples")

        return weights

    def _prepare_features(
        self,
        df: pd.DataFrame,
        symbol: str
    ) -> pd.DataFrame:
        """
        Prepare feature DataFrame, excluding target columns.

        Args:
            df: Raw DataFrame
            symbol: Symbol for context

        Returns:
            DataFrame with only feature columns
        """
        # Exclude known non-feature columns
        exclude_patterns = [
            'target_', 'high', 'low', 'open', 'close', 'volume',
            'High', 'Low', 'Open', 'Close', 'Volume',
            'timestamp', 'datetime', 'date', 'time',
            'rr_', 'direction', 'is_valid'
        ]

        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)

        logger.info(f"Selected {len(feature_cols)} features for {symbol}")
        return df[feature_cols]

    def train_single(
        self,
        df: pd.DataFrame,
        symbol: str,
        timeframe: str,
        features_df: pd.DataFrame = None
    ) -> Dict[str, TrainingResult]:
        """
        Train models for a single symbol/timeframe combination.

        Args:
            df: OHLCV DataFrame with datetime index
            symbol: Trading symbol (e.g., 'XAUUSD')
            timeframe: Timeframe (e.g., '15m')
            features_df: Optional pre-computed features

        Returns:
            Dictionary of TrainingResult for each target type
        """
        logger.info(f"\n{'='*60}")
        logger.info(f"Training {symbol} {timeframe}")
        logger.info(f"{'='*60}")

        # Generate attention features if configured (Nivel 0 -> Nivel 1 integration)
        if self.config.use_attention_features:
            df = self._generate_attention_features(df, symbol, timeframe)

        # Split train/holdout
        train_df, holdout_df = self._split_train_holdout(df)

        if len(train_df) < self.config.min_train_samples:
            logger.warning(f"Insufficient training data: {len(train_df)} < {self.config.min_train_samples}")
            return {}

        # Get horizon
        horizon = self.config.horizons.get(timeframe, 3)

        # Compute targets (normalized by ATR for better scale)
        target_high, target_low, atr = self._compute_targets(train_df, horizon, normalize=True)

        # Compute sample weights (using un-normalized targets for proper weighting)
        # Re-compute raw targets just for weighting
        target_high_raw, target_low_raw, _ = self._compute_targets(train_df, horizon, normalize=False)
        sample_weights = self._compute_sample_weights(train_df, target_high_raw, target_low_raw)

        # Prepare features
        if features_df is not None:
            # Use provided features (aligned with train_df)
            X = features_df.loc[train_df.index].values if hasattr(features_df, 'loc') else features_df
        else:
            # Auto-detect features
            X = self._prepare_features(train_df, symbol).values

        # Remove invalid samples (NaN targets)
        valid_mask = ~(np.isnan(target_high) | np.isnan(target_low))
        X_valid = X[valid_mask]
        y_high_valid = target_high[valid_mask]
        y_low_valid = target_low[valid_mask]
        weights_valid = sample_weights[valid_mask]

        # Train/val split (time-based)
        split_idx = int(len(X_valid) * (1 - self.config.val_split))
        X_train, X_val = X_valid[:split_idx], X_valid[split_idx:]
        y_high_train, y_high_val = y_high_valid[:split_idx], y_high_valid[split_idx:]
        y_low_train, y_low_val = y_low_valid[:split_idx], y_low_valid[split_idx:]
        weights_train = weights_valid[:split_idx]

        results = {}

        # Train HIGH model
        model_key_high = ModelKey(symbol, timeframe, 'high', horizon)
        result_high = self._train_model(
            X_train, y_high_train, X_val, y_high_val,
            weights_train, model_key_high
        )
        results[str(model_key_high)] = result_high

        # Train LOW model
        model_key_low = ModelKey(symbol, timeframe, 'low', horizon)
        result_low = self._train_model(
            X_train, y_low_train, X_val, y_low_val,
            weights_train, model_key_low
        )
        results[str(model_key_low)] = result_low

        logger.info(f"\nResults for {symbol} {timeframe}:")
        logger.info(f"  HIGH: MAE={result_high.mae:.4f}, R2={result_high.r2:.4f}")
        logger.info(f"  LOW:  MAE={result_low.mae:.4f}, R2={result_low.r2:.4f}")

        return results

    def _train_model(
        self,
        X_train: np.ndarray,
        y_train: np.ndarray,
        X_val: np.ndarray,
        y_val: np.ndarray,
        sample_weights: np.ndarray,
        model_key: ModelKey
    ) -> TrainingResult:
        """Train a single XGBoost model."""
        xgb_params = self.config.xgb_params.copy()

        # Force CPU-only training (GPU requires special XGBoost build)
        # Ensure tree_method is valid for CPU
        if xgb_params.get('tree_method') == 'gpu_hist':
            xgb_params['tree_method'] = 'hist'
        if 'device' in xgb_params:
            del xgb_params['device']

        model = XGBRegressor(**xgb_params)

        # Fit with sample weights
        model.fit(
            X_train, y_train,
            sample_weight=sample_weights,
            eval_set=[(X_val, y_val)],
            verbose=False
        )

        # Store model
        self.models[str(model_key)] = model

        # Evaluate
        y_pred_train = model.predict(X_train)
        y_pred_val = model.predict(X_val)

        # Metrics
        mae = mean_absolute_error(y_val, y_pred_val)
        rmse = np.sqrt(mean_squared_error(y_val, y_pred_val))
        r2 = r2_score(y_val, y_pred_val)

        # Directional accuracy
        dir_acc = np.mean(np.sign(y_val) == np.sign(y_pred_val))

        result = TrainingResult(
            model_key=model_key,
            mae=mae,
            rmse=rmse,
            r2=r2,
            directional_accuracy=dir_acc,
            n_train=len(X_train),
            n_val=len(X_val)
        )

        self.results[str(model_key)] = result
        return result

    def train_all(
        self,
        data_dict: Dict[str, Dict[str, pd.DataFrame]]
    ) -> Dict[str, TrainingResult]:
        """
        Train all symbol/timeframe combinations.

        Args:
            data_dict: Dictionary structured as {symbol: {timeframe: df}}
                       e.g., {'XAUUSD': {'5m': df_5m, '15m': df_15m}}

        Returns:
            Dictionary of all TrainingResult objects
        """
        all_results = {}

        for symbol in self.config.symbols:
            if symbol not in data_dict:
                logger.warning(f"No data for symbol: {symbol}")
                continue

            for timeframe in self.config.timeframes:
                if timeframe not in data_dict[symbol]:
                    logger.warning(f"No {timeframe} data for {symbol}")
                    continue

                df = data_dict[symbol][timeframe]
                results = self.train_single(df, symbol, timeframe)
                all_results.update(results)

        self._is_trained = True
        return all_results

    def predict(
        self,
        X: Union[pd.DataFrame, np.ndarray],
        symbol: str,
        timeframe: str
    ) -> Dict[str, np.ndarray]:
        """
        Generate predictions for a symbol/timeframe.

        Args:
            X: Features (DataFrame or array)
            symbol: Symbol to predict
            timeframe: Timeframe

        Returns:
            Dictionary with 'high' and 'low' predictions
        """
        horizon = self.config.horizons.get(timeframe, 3)

        X_np = X.values if isinstance(X, pd.DataFrame) else X
        if X_np.ndim == 1:
            X_np = X_np.reshape(1, -1)

        key_high = f"{symbol}_{timeframe}_high_h{horizon}"
        key_low = f"{symbol}_{timeframe}_low_h{horizon}"

        if key_high not in self.models or key_low not in self.models:
            raise ValueError(f"No trained model for {symbol} {timeframe}")

        pred_high = self.models[key_high].predict(X_np)
        pred_low = self.models[key_low].predict(X_np)

        return {
            'high': pred_high,
            'low': pred_low,
            'symbol': symbol,
            'timeframe': timeframe,
            'horizon_bars': horizon
        }

    def get_holdout_data(
        self,
        df: pd.DataFrame
    ) -> pd.DataFrame:
        """Get the holdout (backtesting) portion of data."""
        _, holdout_df = self._split_train_holdout(df)
        return holdout_df

    def get_training_summary(self) -> pd.DataFrame:
        """Get summary of all trained models."""
        rows = []
        for key, result in self.results.items():
            rows.append({
                'model_key': key,
                'symbol': result.model_key.symbol,
                'timeframe': result.model_key.timeframe,
                'target': result.model_key.target_type,
                'horizon': result.model_key.horizon_bars,
                'mae': result.mae,
                'rmse': result.rmse,
                'r2': result.r2,
                'dir_accuracy': result.directional_accuracy,
                'n_train': result.n_train,
                'n_val': result.n_val
            })
        return pd.DataFrame(rows)

    def save(self, path: str):
        """Save all models to disk."""
        path = Path(path)
        path.mkdir(parents=True, exist_ok=True)

        # Save each model
        for key, model in self.models.items():
            model_path = path / f"{key}.joblib"
            joblib.dump(model, model_path)

        # Save metadata
        metadata = {
            'config': self.config,
            'symbol_configs': self.symbol_configs,
            'results': {k: {
                'mae': v.mae,
                'rmse': v.rmse,
                'r2': v.r2,
                'directional_accuracy': v.directional_accuracy,
                'n_train': v.n_train,
                'n_val': v.n_val
            } for k, v in self.results.items()},
            'model_keys': list(self.models.keys())
        }
        joblib.dump(metadata, path / 'trainer_metadata.joblib')

        logger.info(f"Saved {len(self.models)} models to {path}")

    def load(self, path: str):
        """Load models from disk."""
        path = Path(path)

        # Load metadata
        metadata = joblib.load(path / 'trainer_metadata.joblib')
        self.config = metadata['config']
        self.symbol_configs = metadata['symbol_configs']

        # Load models
        self.models = {}
        for key in metadata['model_keys']:
            model_path = path / f"{key}.joblib"
            if model_path.exists():
                self.models[key] = joblib.load(model_path)

        self._is_trained = True
        logger.info(f"Loaded {len(self.models)} models from {path}")


if __name__ == "__main__":
    # Test the module
    print("Testing SymbolTimeframeTrainer...")

    # Create sample data
    np.random.seed(42)

    # Generate 2 years of 15m data for XAUUSD
    n = 70000  # ~2 years of 15m data
    dates = pd.date_range('2023-01-01', periods=n, freq='15min')
    price = 2000 + np.cumsum(np.random.randn(n) * 0.5)

    # Simulate varying volatility
    volatility = np.where(
        (dates.hour >= 13) & (dates.hour < 16),
        5.0, 2.0
    )

    df = pd.DataFrame({
        'Open': price,
        'High': price + np.abs(np.random.randn(n)) * volatility,
        'Low': price - np.abs(np.random.randn(n)) * volatility,
        'Close': price + np.random.randn(n) * 0.5,
        'Volume': np.random.randint(100, 1000, n),
        # Add some features
        'rsi': 50 + np.random.randn(n) * 10,
        'macd': np.random.randn(n),
        'bb_width': 10 + np.random.randn(n)
    }, index=dates)

    # Test trainer
    config = TrainerConfig(
        symbols=['XAUUSD'],
        timeframes=['15m'],
        train_years=1.0,
        holdout_years=0.5,
        min_train_samples=1000
    )

    trainer = SymbolTimeframeTrainer(config)

    # Train
    results = trainer.train_single(df, 'XAUUSD', '15m')

    print("\n" + "="*60)
    print("TRAINING RESULTS")
    print("="*60)
    for key, result in results.items():
        print(f"\n{key}:")
        print(f"  MAE: {result.mae:.4f}")
        print(f"  RMSE: {result.rmse:.4f}")
        print(f"  R2: {result.r2:.4f}")
        print(f"  Dir Accuracy: {result.directional_accuracy:.2%}")

    # Test prediction
    print("\n" + "="*60)
    print("SAMPLE PREDICTIONS")
    print("="*60)
    test_features = df[['rsi', 'macd', 'bb_width']].iloc[-10:].values
    predictions = trainer.predict(test_features, 'XAUUSD', '15m')
    print(f"High predictions: {predictions['high'][:5]}")
    print(f"Low predictions: {predictions['low'][:5]}")
