"""
MRD Trainer - Training Pipeline for Momentum Regime Detection
==============================================================

Implements the training pipeline for the MRD model including:
- HMM training for unsupervised regime detection
- LSTM training for sequence modeling
- XGBoost training for final ensemble
- Walk-forward validation for robust evaluation

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

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

try:
    import torch
    import torch.nn as nn
    from torch.utils.data import DataLoader, TensorDataset
    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False

try:
    from xgboost import XGBClassifier, XGBRegressor
    XGBOOST_AVAILABLE = True
except ImportError:
    XGBOOST_AVAILABLE = False

from sklearn.metrics import (
    accuracy_score, f1_score, precision_recall_fscore_support,
    mean_squared_error, mean_absolute_error, classification_report
)

from .feature_engineering import MRDFeatureEngineer, MRDFeatureConfig
from .hmm_regime import RegimeHMM, RegimeHMMConfig, MarketRegime
from .model import MRDModel, MRDModelConfig, MRDLSTMEncoder


@dataclass
class TrainerConfig:
    """Configuration for MRD trainer."""

    # Model configuration
    model_config: Optional[MRDModelConfig] = None

    # Training parameters
    lstm_epochs: int = 50
    lstm_batch_size: int = 64
    lstm_learning_rate: float = 0.001
    lstm_weight_decay: float = 1e-5
    lstm_early_stopping_patience: int = 10

    # Walk-forward parameters
    n_folds: int = 5
    test_ratio: float = 0.15
    gap_samples: int = 12  # Gap between train and test to avoid look-ahead
    min_train_samples: int = 5000
    expanding_window: bool = True  # True = expanding, False = sliding

    # XGBoost parameters
    xgb_early_stopping_rounds: int = 20

    # Output
    save_models: bool = True
    model_dir: str = "models/mrd"

    # Logging
    verbose: bool = True
    log_every: int = 10


@dataclass
class FoldResult:
    """Results from a single walk-forward fold."""
    fold_id: int
    train_start: int
    train_end: int
    test_start: int
    test_end: int
    train_samples: int
    test_samples: int
    hmm_metrics: Dict[str, Any]
    lstm_metrics: Dict[str, float]
    xgb_metrics: Dict[str, float]
    overall_metrics: Dict[str, float]
    predictions: Optional[np.ndarray] = None
    true_regimes: Optional[np.ndarray] = None

    def to_dict(self) -> Dict:
        return {
            'fold_id': self.fold_id,
            'train_range': (self.train_start, self.train_end),
            'test_range': (self.test_start, self.test_end),
            'train_samples': self.train_samples,
            'test_samples': self.test_samples,
            'metrics': self.overall_metrics
        }


@dataclass
class TrainingResults:
    """Complete training results."""
    symbol: str
    timeframe: str
    n_folds: int
    fold_results: List[FoldResult]
    avg_metrics: Dict[str, float]
    std_metrics: Dict[str, float]
    best_fold: int
    timestamp: str
    config: Dict

    def to_dict(self) -> Dict:
        return {
            'symbol': self.symbol,
            'timeframe': self.timeframe,
            'n_folds': self.n_folds,
            'folds': [f.to_dict() for f in self.fold_results],
            'avg_metrics': self.avg_metrics,
            'std_metrics': self.std_metrics,
            'best_fold': self.best_fold,
            'timestamp': self.timestamp
        }


class MRDTrainer:
    """
    Trainer for Momentum Regime Detection model.

    Handles the complete training pipeline:
    1. Feature engineering
    2. HMM training for regime detection
    3. LSTM training for sequence modeling
    4. XGBoost training for ensemble
    5. Walk-forward validation

    Usage:
        trainer = MRDTrainer()

        # Train on single split
        metrics = trainer.train(df_train, df_val)

        # Walk-forward training
        results = trainer.walk_forward_train(df, symbol='EURUSD', n_folds=5)
    """

    def __init__(self, config: Optional[TrainerConfig] = None):
        """
        Initialize the trainer.

        Args:
            config: Training configuration
        """
        self.config = config or TrainerConfig()
        self.model_config = self.config.model_config or MRDModelConfig()

        self.feature_engineer = MRDFeatureEngineer()
        self._current_model: Optional[MRDModel] = None

        # Training history
        self.history: Dict[str, List] = {
            'train_loss': [],
            'val_loss': [],
            'regime_accuracy': [],
            'continuation_accuracy': []
        }

    def train_hmm(
        self,
        data: pd.DataFrame,
        features: Optional[pd.DataFrame] = None
    ) -> Tuple[RegimeHMM, Dict[str, Any]]:
        """
        Train the HMM component.

        Args:
            data: DataFrame with OHLCV data
            features: Pre-computed features (optional)

        Returns:
            Tuple of (trained HMM, training metrics)
        """
        logger.info("Training HMM for regime detection...")

        # Compute features if not provided
        if features is None:
            features = self.feature_engineer.compute_all_features(data)

        # Initialize and train HMM
        hmm = RegimeHMM(self.model_config.hmm_config or RegimeHMMConfig())
        hmm.fit(features)

        # Get regime statistics
        stats = hmm.get_regime_statistics(features)

        metrics = {
            'regime_distribution': stats['regime_percentages'],
            'avg_durations': stats['avg_duration'],
            'transition_matrix': stats['transition_matrix']
        }

        logger.info(f"HMM trained - Regime distribution: {stats['regime_percentages']}")

        return hmm, metrics

    def train_lstm(
        self,
        data: pd.DataFrame,
        regimes: np.ndarray,
        features: Optional[pd.DataFrame] = None,
        val_data: Optional[pd.DataFrame] = None,
        val_regimes: Optional[np.ndarray] = None,
        val_features: Optional[pd.DataFrame] = None
    ) -> Tuple[MRDLSTMEncoder, Dict[str, float]]:
        """
        Train the LSTM component.

        Args:
            data: Training DataFrame
            regimes: Regime labels for training data
            features: Pre-computed features (optional)
            val_data: Validation DataFrame (optional)
            val_regimes: Validation regime labels
            val_features: Validation features

        Returns:
            Tuple of (trained LSTM, training metrics)
        """
        if not TORCH_AVAILABLE:
            logger.warning("PyTorch not available, skipping LSTM training")
            return None, {}

        logger.info("Training LSTM encoder...")

        # Compute features if needed
        if features is None:
            features = self.feature_engineer.compute_all_features(data)

        n_features = features.shape[1]
        seq_len = self.model_config.lstm_sequence_length
        device = torch.device(self.model_config.device)

        # Initialize LSTM
        lstm = MRDLSTMEncoder(
            input_size=n_features,
            hidden_size=self.model_config.lstm_hidden_size,
            num_layers=self.model_config.lstm_num_layers,
            dropout=self.model_config.lstm_dropout,
            bidirectional=self.model_config.lstm_bidirectional,
            n_regimes=3
        ).to(device)

        # Prepare training data
        durations = self._compute_durations(regimes)
        X_train, y_regime, y_duration = self._prepare_sequences(
            features.values, regimes, durations, seq_len
        )

        # Continuation labels
        y_continuation = self._compute_continuation_labels(regimes, seq_len)

        # Convert to tensors
        X_tensor = torch.FloatTensor(X_train).to(device)
        regime_tensor = torch.LongTensor(y_regime).to(device)
        duration_tensor = torch.FloatTensor(y_duration).to(device)
        continuation_tensor = torch.FloatTensor(y_continuation).to(device)

        # Create data loader
        train_dataset = TensorDataset(
            X_tensor, regime_tensor, duration_tensor, continuation_tensor
        )
        train_loader = DataLoader(
            train_dataset,
            batch_size=self.config.lstm_batch_size,
            shuffle=True
        )

        # Validation data
        val_loader = None
        if val_features is not None and val_regimes is not None:
            val_durations = self._compute_durations(val_regimes)
            X_val, y_val_regime, y_val_duration = self._prepare_sequences(
                val_features.values, val_regimes, val_durations, seq_len
            )
            y_val_cont = self._compute_continuation_labels(val_regimes, seq_len)

            val_dataset = TensorDataset(
                torch.FloatTensor(X_val).to(device),
                torch.LongTensor(y_val_regime).to(device),
                torch.FloatTensor(y_val_duration).to(device),
                torch.FloatTensor(y_val_cont).to(device)
            )
            val_loader = DataLoader(val_dataset, batch_size=self.config.lstm_batch_size)

        # Training setup
        optimizer = torch.optim.AdamW(
            lstm.parameters(),
            lr=self.config.lstm_learning_rate,
            weight_decay=self.config.lstm_weight_decay
        )

        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            optimizer, mode='min', factor=0.5, patience=5
        )

        regime_criterion = nn.CrossEntropyLoss()
        continuation_criterion = nn.BCELoss()

        # Training loop
        best_val_loss = float('inf')
        patience_counter = 0
        best_state = None

        for epoch in range(self.config.lstm_epochs):
            # Training phase
            lstm.train()
            train_loss = 0
            train_regime_correct = 0
            train_total = 0

            for batch in train_loader:
                batch_x, batch_regime, batch_duration, batch_cont = batch

                optimizer.zero_grad()

                regime_logits, cont_prob, _, _ = lstm(
                    batch_x, batch_regime, batch_duration
                )

                regime_loss = regime_criterion(regime_logits, batch_regime)
                cont_loss = continuation_criterion(cont_prob.squeeze(), batch_cont)
                loss = regime_loss + cont_loss

                loss.backward()
                torch.nn.utils.clip_grad_norm_(lstm.parameters(), 1.0)
                optimizer.step()

                train_loss += loss.item()
                train_regime_correct += (regime_logits.argmax(1) == batch_regime).sum().item()
                train_total += batch_regime.size(0)

            avg_train_loss = train_loss / len(train_loader)
            train_accuracy = train_regime_correct / train_total

            # Validation phase
            val_loss = 0
            val_accuracy = 0

            if val_loader is not None:
                lstm.eval()
                val_regime_correct = 0
                val_total = 0

                with torch.no_grad():
                    for batch in val_loader:
                        batch_x, batch_regime, batch_duration, batch_cont = batch

                        regime_logits, cont_prob, _, _ = lstm(
                            batch_x, batch_regime, batch_duration
                        )

                        regime_loss = regime_criterion(regime_logits, batch_regime)
                        cont_loss = continuation_criterion(cont_prob.squeeze(), batch_cont)
                        val_loss += (regime_loss + cont_loss).item()

                        val_regime_correct += (regime_logits.argmax(1) == batch_regime).sum().item()
                        val_total += batch_regime.size(0)

                val_loss = val_loss / len(val_loader)
                val_accuracy = val_regime_correct / val_total

                scheduler.step(val_loss)

                # Early stopping
                if val_loss < best_val_loss:
                    best_val_loss = val_loss
                    best_state = lstm.state_dict().copy()
                    patience_counter = 0
                else:
                    patience_counter += 1

                if patience_counter >= self.config.lstm_early_stopping_patience:
                    logger.info(f"Early stopping at epoch {epoch + 1}")
                    break

            # Logging
            if self.config.verbose and (epoch + 1) % self.config.log_every == 0:
                logger.info(
                    f"Epoch {epoch+1}/{self.config.lstm_epochs} - "
                    f"Train Loss: {avg_train_loss:.4f}, Train Acc: {train_accuracy:.4f}, "
                    f"Val Loss: {val_loss:.4f}, Val Acc: {val_accuracy:.4f}"
                )

        # Load best state
        if best_state is not None:
            lstm.load_state_dict(best_state)

        lstm.eval()

        metrics = {
            'final_train_loss': avg_train_loss,
            'final_train_accuracy': train_accuracy,
            'best_val_loss': best_val_loss,
            'final_val_accuracy': val_accuracy,
            'epochs_trained': epoch + 1
        }

        logger.info(f"LSTM trained - Best val loss: {best_val_loss:.4f}")

        return lstm, metrics

    def train_xgboost(
        self,
        lstm_outputs: np.ndarray,
        targets: np.ndarray,
        regimes: np.ndarray,
        durations: np.ndarray,
        val_lstm_outputs: Optional[np.ndarray] = None,
        val_targets: Optional[np.ndarray] = None,
        val_regimes: Optional[np.ndarray] = None,
        val_durations: Optional[np.ndarray] = None
    ) -> Tuple[Any, Any, Dict[str, float]]:
        """
        Train XGBoost models for final prediction.

        Args:
            lstm_outputs: LSTM hidden features
            targets: Target continuation labels
            regimes: Regime labels
            durations: Duration values
            val_*: Validation data (optional)

        Returns:
            Tuple of (continuation model, regime model, metrics)
        """
        if not XGBOOST_AVAILABLE:
            logger.warning("XGBoost not available")
            return None, None, {}

        logger.info("Training XGBoost ensemble...")

        # Combine features
        X_train = np.hstack([
            lstm_outputs,
            regimes.reshape(-1, 1),
            durations.reshape(-1, 1)
        ])

        # Initialize models
        xgb_params = {
            'n_estimators': self.model_config.xgb_n_estimators,
            'max_depth': self.model_config.xgb_max_depth,
            'learning_rate': self.model_config.xgb_learning_rate,
            'subsample': self.model_config.xgb_subsample,
            'colsample_bytree': self.model_config.xgb_colsample_bytree,
            'tree_method': 'hist',
            'random_state': 42
        }

        xgb_continuation = XGBRegressor(**xgb_params)
        xgb_regime = XGBClassifier(**xgb_params, num_class=3)

        # Prepare validation data
        eval_set = None
        if val_lstm_outputs is not None:
            X_val = np.hstack([
                val_lstm_outputs,
                val_regimes.reshape(-1, 1),
                val_durations.reshape(-1, 1)
            ])
            eval_set = [(X_val, val_targets)]

        # Train continuation model
        xgb_continuation.fit(
            X_train, targets,
            eval_set=eval_set,
            verbose=False
        )

        # Train regime model
        xgb_regime.fit(
            X_train, regimes,
            eval_set=[(X_val, val_regimes)] if eval_set else None,
            verbose=False
        )

        # Calculate metrics
        cont_pred = xgb_continuation.predict(X_train)
        regime_pred = xgb_regime.predict(X_train)

        train_metrics = {
            'continuation_mse': float(mean_squared_error(targets, cont_pred)),
            'regime_accuracy': float(accuracy_score(regimes, regime_pred))
        }

        if eval_set:
            val_cont_pred = xgb_continuation.predict(X_val)
            val_regime_pred = xgb_regime.predict(X_val)
            train_metrics['val_continuation_mse'] = float(mean_squared_error(val_targets, val_cont_pred))
            train_metrics['val_regime_accuracy'] = float(accuracy_score(val_regimes, val_regime_pred))

        logger.info(f"XGBoost trained - Regime Acc: {train_metrics['regime_accuracy']:.4f}")

        return xgb_continuation, xgb_regime, train_metrics

    def walk_forward_train(
        self,
        data: pd.DataFrame,
        symbol: str = 'UNKNOWN',
        timeframe: str = '5min',
        n_folds: Optional[int] = None
    ) -> TrainingResults:
        """
        Perform walk-forward training and validation.

        Args:
            data: Complete DataFrame with OHLCV data
            symbol: Symbol being trained
            timeframe: Timeframe of the data
            n_folds: Number of folds (overrides config)

        Returns:
            TrainingResults with metrics for all folds
        """
        n_folds = n_folds or self.config.n_folds
        n_samples = len(data)

        logger.info(f"Starting walk-forward training: {n_folds} folds on {n_samples} samples")

        # Calculate fold boundaries
        folds = self._create_walk_forward_splits(n_samples, n_folds)
        fold_results = []

        for fold_idx, (train_start, train_end, test_start, test_end) in enumerate(folds):
            logger.info(f"\n{'='*60}")
            logger.info(f"Fold {fold_idx + 1}/{n_folds}")
            logger.info(f"Train: [{train_start}:{train_end}] ({train_end - train_start} samples)")
            logger.info(f"Test: [{test_start}:{test_end}] ({test_end - test_start} samples)")
            logger.info(f"{'='*60}")

            # Split data
            df_train = data.iloc[train_start:train_end].copy()
            df_test = data.iloc[test_start:test_end].copy()

            # Train model
            fold_result = self._train_single_fold(
                df_train, df_test, fold_idx + 1,
                train_start, train_end, test_start, test_end
            )
            fold_results.append(fold_result)

            logger.info(f"Fold {fold_idx + 1} - Regime Accuracy: {fold_result.overall_metrics.get('regime_accuracy', 0):.4f}")

        # Aggregate results
        results = self._aggregate_results(
            fold_results, symbol, timeframe, n_folds
        )

        # Save results if configured
        if self.config.save_models:
            self._save_results(results)

        logger.info(f"\nWalk-Forward Training Complete")
        logger.info(f"Average Regime Accuracy: {results.avg_metrics['regime_accuracy']:.4f} "
                    f"(+/- {results.std_metrics['regime_accuracy']:.4f})")

        return results

    def _create_walk_forward_splits(
        self,
        n_samples: int,
        n_folds: int
    ) -> List[Tuple[int, int, int, int]]:
        """Create walk-forward validation splits."""
        folds = []
        test_size = int(n_samples * self.config.test_ratio / n_folds)

        for i in range(n_folds):
            if self.config.expanding_window:
                train_start = 0
            else:
                train_start = i * test_size

            train_end = int(n_samples * (1 - self.config.test_ratio) + i * test_size / n_folds)
            test_start = train_end + self.config.gap_samples
            test_end = min(test_start + test_size, n_samples)

            if train_end - train_start >= self.config.min_train_samples:
                folds.append((train_start, train_end, test_start, test_end))

        return folds

    def _train_single_fold(
        self,
        df_train: pd.DataFrame,
        df_test: pd.DataFrame,
        fold_id: int,
        train_start: int,
        train_end: int,
        test_start: int,
        test_end: int
    ) -> FoldResult:
        """Train model for a single fold."""
        # Compute features
        train_features = self.feature_engineer.compute_all_features(df_train)
        test_features = self.feature_engineer.compute_all_features(df_test)

        # Train HMM
        hmm, hmm_metrics = self.train_hmm(df_train, train_features)

        # Get regimes
        train_regimes = hmm.predict_regime(train_features)
        test_regimes = hmm.predict_regime(test_features)

        train_durations = self._compute_durations(train_regimes)
        test_durations = self._compute_durations(test_regimes)

        # Train LSTM
        lstm, lstm_metrics = self.train_lstm(
            df_train, train_regimes, train_features,
            df_test, test_regimes, test_features
        )

        # Get LSTM features
        xgb_metrics = {}
        if lstm is not None:
            seq_len = self.model_config.lstm_sequence_length

            train_lstm_out = self._get_lstm_features(lstm, train_features, seq_len)
            test_lstm_out = self._get_lstm_features(lstm, test_features, seq_len)

            # Continuation labels
            train_cont = self._compute_continuation_labels(train_regimes, seq_len)
            test_cont = self._compute_continuation_labels(test_regimes, seq_len)

            # Train XGBoost
            xgb_cont, xgb_regime, xgb_metrics = self.train_xgboost(
                train_lstm_out, train_cont, train_regimes[seq_len:], train_durations[seq_len:],
                test_lstm_out, test_cont, test_regimes[seq_len:], test_durations[seq_len:]
            )

        # Calculate overall metrics
        overall_metrics = self._calculate_fold_metrics(
            test_regimes, test_features, hmm, lstm
        )

        return FoldResult(
            fold_id=fold_id,
            train_start=train_start,
            train_end=train_end,
            test_start=test_start,
            test_end=test_end,
            train_samples=train_end - train_start,
            test_samples=test_end - test_start,
            hmm_metrics=hmm_metrics,
            lstm_metrics=lstm_metrics,
            xgb_metrics=xgb_metrics,
            overall_metrics=overall_metrics,
            predictions=None,
            true_regimes=test_regimes
        )

    def _get_lstm_features(
        self,
        lstm: MRDLSTMEncoder,
        features: pd.DataFrame,
        seq_len: int
    ) -> np.ndarray:
        """Extract LSTM features."""
        device = torch.device(self.model_config.device)
        regimes = np.zeros(len(features))
        durations = np.zeros(len(features))

        X, _, _ = self._prepare_sequences(features.values, regimes, durations, seq_len)
        X_tensor = torch.FloatTensor(X).to(device)

        lstm.eval()
        with torch.no_grad():
            _, _, lstm_features, _ = lstm(X_tensor)

        return lstm_features.cpu().numpy()

    def _calculate_fold_metrics(
        self,
        true_regimes: np.ndarray,
        features: pd.DataFrame,
        hmm: RegimeHMM,
        lstm: Optional[MRDLSTMEncoder]
    ) -> Dict[str, float]:
        """Calculate metrics for a fold."""
        # HMM regime prediction
        pred_regimes = hmm.predict_regime(features)

        # Basic metrics
        regime_accuracy = accuracy_score(true_regimes, pred_regimes)
        regime_f1 = f1_score(true_regimes, pred_regimes, average='weighted')

        metrics = {
            'regime_accuracy': float(regime_accuracy),
            'regime_f1': float(regime_f1)
        }

        # Per-class metrics
        precision, recall, f1, _ = precision_recall_fscore_support(
            true_regimes, pred_regimes, average=None
        )
        for i, name in enumerate(['trend_down', 'range', 'trend_up']):
            if i < len(f1):
                metrics[f'{name}_f1'] = float(f1[i])

        return metrics

    def _aggregate_results(
        self,
        fold_results: List[FoldResult],
        symbol: str,
        timeframe: str,
        n_folds: int
    ) -> TrainingResults:
        """Aggregate results from all folds."""
        # Collect metrics
        metrics_keys = fold_results[0].overall_metrics.keys()
        all_metrics = {key: [] for key in metrics_keys}

        for fold in fold_results:
            for key in metrics_keys:
                all_metrics[key].append(fold.overall_metrics.get(key, 0))

        # Calculate average and std
        avg_metrics = {key: np.mean(values) for key, values in all_metrics.items()}
        std_metrics = {key: np.std(values) for key, values in all_metrics.items()}

        # Find best fold
        best_fold = max(
            range(len(fold_results)),
            key=lambda i: fold_results[i].overall_metrics.get('regime_accuracy', 0)
        )

        return TrainingResults(
            symbol=symbol,
            timeframe=timeframe,
            n_folds=n_folds,
            fold_results=fold_results,
            avg_metrics=avg_metrics,
            std_metrics=std_metrics,
            best_fold=best_fold + 1,
            timestamp=datetime.now().isoformat(),
            config=self.config.__dict__
        )

    def _save_results(self, results: TrainingResults):
        """Save training results."""
        save_dir = Path(self.config.model_dir) / results.symbol
        save_dir.mkdir(parents=True, exist_ok=True)

        # Save results
        results_path = save_dir / f"results_{results.timestamp.replace(':', '-')}.joblib"
        joblib.dump(results.to_dict(), results_path)
        logger.info(f"Saved results to {results_path}")

    def _compute_durations(self, regimes: np.ndarray) -> np.ndarray:
        """Compute regime durations."""
        durations = np.ones_like(regimes, dtype=np.float32)
        current_duration = 1

        for i in range(1, len(regimes)):
            if regimes[i] == regimes[i - 1]:
                current_duration += 1
            else:
                current_duration = 1
            durations[i] = current_duration

        return durations

    def _compute_continuation_labels(
        self,
        regimes: np.ndarray,
        seq_len: int
    ) -> np.ndarray:
        """Compute continuation labels."""
        n = len(regimes) - seq_len
        labels = np.zeros(n, dtype=np.float32)

        for i in range(n):
            idx = seq_len + i
            if idx < len(regimes) - 1:
                labels[i] = float(regimes[idx] == regimes[idx + 1])

        return labels

    def _prepare_sequences(
        self,
        features: np.ndarray,
        regimes: np.ndarray,
        durations: np.ndarray,
        seq_len: int
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """Prepare sequences for LSTM."""
        n_samples = len(features) - seq_len

        X = np.zeros((n_samples, seq_len, features.shape[1]))
        y_regime = np.zeros(n_samples, dtype=np.int64)
        y_duration = np.zeros(n_samples, dtype=np.float32)

        for i in range(n_samples):
            X[i] = features[i:i + seq_len]
            y_regime[i] = regimes[i + seq_len - 1]
            y_duration[i] = durations[i + seq_len - 1]

        return X, y_regime, y_duration


if __name__ == "__main__":
    # Test the trainer
    np.random.seed(42)
    n_samples = 5000

    # Create sample data
    dates = pd.date_range('2024-01-01', periods=n_samples, freq='5min')
    price = np.zeros(n_samples)
    price[0] = 2000

    for i in range(1, n_samples):
        phase = (i // 500) % 3
        if phase == 0:
            drift = 0.0003
        elif phase == 1:
            drift = 0
        else:
            drift = -0.0003
        price[i] = price[i-1] * (1 + drift + np.random.randn() * 0.0005)

    df = pd.DataFrame({
        'open': price * (1 + np.random.randn(n_samples) * 0.0002),
        'high': price * (1 + np.abs(np.random.randn(n_samples)) * 0.0005),
        'low': price * (1 - np.abs(np.random.randn(n_samples)) * 0.0005),
        'close': price,
        'volume': np.random.randint(100, 10000, n_samples)
    }, index=dates)

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

    # Create trainer
    config = TrainerConfig(
        n_folds=3,
        lstm_epochs=20,
        save_models=False,
        verbose=True
    )

    trainer = MRDTrainer(config)

    # Walk-forward training
    print(f"\n{'='*60}")
    print("Testing MRD Trainer - Walk Forward")
    print(f"{'='*60}")

    results = trainer.walk_forward_train(df, symbol='TEST', n_folds=3)

    print(f"\nFinal Results:")
    print(f"  Average Regime Accuracy: {results.avg_metrics['regime_accuracy']:.4f}")
    print(f"  Std Regime Accuracy: {results.std_metrics['regime_accuracy']:.4f}")
    print(f"  Best Fold: {results.best_fold}")

    print("\nTest complete!")
