"""
MRD Model - Complete Momentum Regime Detection Model
====================================================

Full MRD model that combines:
1. HMM for regime detection
2. LSTM for sequence modeling
3. XGBoost for final prediction

The model predicts:
- Current market regime (trend up, range, trend down)
- Regime duration
- Probability of regime continuation

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

try:
    import torch
    import torch.nn as nn
    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False
    logger.warning("PyTorch not available for LSTM component")

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

from .feature_engineering import MRDFeatureEngineer, MRDFeatureConfig
from .hmm_regime import RegimeHMM, RegimeHMMConfig, MarketRegime


@dataclass
class MRDModelConfig:
    """Configuration for the complete MRD model."""

    # Feature engineering config
    feature_config: Optional[MRDFeatureConfig] = None

    # HMM config
    hmm_config: Optional[RegimeHMMConfig] = None

    # LSTM parameters
    lstm_hidden_size: int = 128
    lstm_num_layers: int = 2
    lstm_dropout: float = 0.2
    lstm_sequence_length: int = 60
    lstm_bidirectional: bool = False

    # XGBoost parameters
    xgb_n_estimators: int = 200
    xgb_max_depth: int = 5
    xgb_learning_rate: float = 0.05
    xgb_subsample: float = 0.8
    xgb_colsample_bytree: float = 0.8

    # Training parameters
    batch_size: int = 64
    learning_rate: float = 0.001
    weight_decay: float = 1e-5

    # Device
    device: str = 'cuda' if TORCH_AVAILABLE and torch.cuda.is_available() else 'cpu'


@dataclass
class MRDPrediction:
    """Complete MRD prediction output."""
    regime: int
    regime_name: str
    regime_probabilities: Dict[str, float]
    regime_confidence: float
    duration: int
    continuation_prob: float
    reversal_prob: float
    lstm_hidden_state: Optional[np.ndarray] = None
    feature_importance: Optional[Dict[str, float]] = None

    def to_dict(self) -> Dict:
        return {
            'regime': self.regime,
            'regime_name': self.regime_name,
            'regime_probabilities': self.regime_probabilities,
            'regime_confidence': self.regime_confidence,
            'duration': self.duration,
            'continuation_prob': self.continuation_prob,
            'reversal_prob': self.reversal_prob
        }

    @property
    def is_trending(self) -> bool:
        """Check if currently in a trending regime."""
        return self.regime != MarketRegime.RANGE

    @property
    def is_bullish(self) -> bool:
        """Check if bullish trend."""
        return self.regime == MarketRegime.TREND_UP and self.regime_confidence > 0.6

    @property
    def is_bearish(self) -> bool:
        """Check if bearish trend."""
        return self.regime == MarketRegime.TREND_DOWN and self.regime_confidence > 0.6

    @property
    def should_trade(self) -> bool:
        """Check if conditions favor trading."""
        return self.is_trending and self.continuation_prob > 0.5


class MRDLSTMEncoder(nn.Module):
    """
    LSTM encoder for sequence modeling in MRD.

    Takes momentum features as input and produces:
    - Hidden state representation
    - Regime continuation prediction
    """

    def __init__(
        self,
        input_size: int,
        hidden_size: int = 128,
        num_layers: int = 2,
        dropout: float = 0.2,
        bidirectional: bool = False,
        n_regimes: int = 3
    ):
        """
        Initialize LSTM encoder.

        Args:
            input_size: Number of input features
            hidden_size: LSTM hidden size
            num_layers: Number of LSTM layers
            dropout: Dropout probability
            bidirectional: Whether to use bidirectional LSTM
            n_regimes: Number of regime classes
        """
        super().__init__()

        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.bidirectional = bidirectional
        self.n_directions = 2 if bidirectional else 1

        # Input projection
        self.input_proj = nn.Linear(input_size, hidden_size)
        self.input_norm = nn.LayerNorm(hidden_size)

        # LSTM layers
        self.lstm = nn.LSTM(
            input_size=hidden_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,
            dropout=dropout if num_layers > 1 else 0,
            bidirectional=bidirectional
        )

        # Output dimension
        output_dim = hidden_size * self.n_directions

        # Regime head (predicts regime from hidden state)
        self.regime_head = nn.Sequential(
            nn.Linear(output_dim, hidden_size),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_size, n_regimes)
        )

        # Continuation head (predicts continuation probability)
        self.continuation_head = nn.Sequential(
            nn.Linear(output_dim + n_regimes + 1, hidden_size // 2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_size // 2, 1),
            nn.Sigmoid()
        )

        # Feature extraction for XGBoost
        self.feature_extractor = nn.Linear(output_dim, hidden_size // 2)

    def forward(
        self,
        x: torch.Tensor,
        regime_labels: Optional[torch.Tensor] = None,
        durations: Optional[torch.Tensor] = None
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Forward pass.

        Args:
            x: Input features (batch_size, seq_len, input_size)
            regime_labels: Optional regime labels (batch_size,)
            durations: Optional regime durations (batch_size,)

        Returns:
            Tuple of:
            - regime_logits: (batch_size, n_regimes)
            - continuation_prob: (batch_size, 1)
            - lstm_features: (batch_size, hidden_size//2)
            - hidden_state: Final hidden state
        """
        batch_size = x.size(0)

        # Input projection
        x = self.input_proj(x)
        x = self.input_norm(x)

        # LSTM forward
        lstm_out, (hidden, cell) = self.lstm(x)

        # Take last output
        last_output = lstm_out[:, -1, :]  # (batch_size, hidden_size * n_directions)

        # Regime prediction
        regime_logits = self.regime_head(last_output)

        # Prepare continuation input
        if regime_labels is not None and durations is not None:
            # One-hot encode regime
            regime_onehot = torch.zeros(batch_size, 3, device=x.device)
            regime_onehot.scatter_(1, regime_labels.unsqueeze(1), 1)
            duration_norm = durations.unsqueeze(1) / 100.0  # Normalize duration
        else:
            # Use predicted regime probabilities
            regime_onehot = torch.softmax(regime_logits, dim=1)
            duration_norm = torch.zeros(batch_size, 1, device=x.device)

        continuation_input = torch.cat([last_output, regime_onehot, duration_norm], dim=1)
        continuation_prob = self.continuation_head(continuation_input)

        # Extract features for XGBoost
        lstm_features = self.feature_extractor(last_output)

        # Reshape hidden for return
        if self.bidirectional:
            hidden = hidden.view(self.num_layers, 2, batch_size, self.hidden_size)
            hidden = torch.cat([hidden[:, 0, :, :], hidden[:, 1, :, :]], dim=2)
        hidden_state = hidden[-1]  # Last layer

        return regime_logits, continuation_prob, lstm_features, hidden_state


class MRDModel:
    """
    Complete Momentum Regime Detection Model.

    Combines:
    1. MRDFeatureEngineer: Computes momentum indicators
    2. RegimeHMM: Detects market regimes
    3. MRDLSTMEncoder: Models sequences and predicts continuation
    4. XGBoost: Final ensemble prediction

    The model outputs:
    - Current regime (trend up, range, trend down)
    - Regime duration
    - Probability of regime continuation

    Usage:
        model = MRDModel()
        model.fit(df_train)

        prediction = model.predict(df_test)
        print(f"Regime: {prediction.regime_name}")
        print(f"Continuation prob: {prediction.continuation_prob:.2f}")
    """

    def __init__(self, config: Optional[MRDModelConfig] = None):
        """
        Initialize the MRD model.

        Args:
            config: Model configuration
        """
        self.config = config or MRDModelConfig()

        # Initialize components
        self.feature_engineer = MRDFeatureEngineer(
            self.config.feature_config or MRDFeatureConfig()
        )

        self.hmm = RegimeHMM(
            self.config.hmm_config or RegimeHMMConfig()
        )

        self.lstm: Optional[MRDLSTMEncoder] = None
        self.xgb_continuation: Optional[XGBRegressor] = None
        self.xgb_regime: Optional[XGBClassifier] = None

        self._is_fitted = False
        self._feature_columns: List[str] = []
        self._n_features: int = 0

        self.device = torch.device(self.config.device) if TORCH_AVAILABLE else None

    def _init_lstm(self, n_features: int):
        """Initialize LSTM model."""
        if not TORCH_AVAILABLE:
            logger.warning("PyTorch not available, LSTM will be skipped")
            return

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

    def _init_xgboost(self):
        """Initialize XGBoost models."""
        if not XGBOOST_AVAILABLE:
            logger.warning("XGBoost not available")
            return

        base_params = {
            'n_estimators': self.config.xgb_n_estimators,
            'max_depth': self.config.xgb_max_depth,
            'learning_rate': self.config.xgb_learning_rate,
            'subsample': self.config.xgb_subsample,
            'colsample_bytree': self.config.xgb_colsample_bytree,
            'tree_method': 'hist',
            'random_state': 42
        }

        # Regression for continuation probability
        self.xgb_continuation = XGBRegressor(**base_params)

        # Classification for regime
        self.xgb_regime = XGBClassifier(
            **base_params,
            objective='multi:softprob',
            num_class=3
        )

    def fit(
        self,
        df: pd.DataFrame,
        df_val: Optional[pd.DataFrame] = None,
        epochs: int = 50,
        verbose: bool = True
    ) -> Dict[str, Any]:
        """
        Fit the complete MRD model.

        Args:
            df: Training DataFrame with OHLCV data
            df_val: Optional validation DataFrame
            epochs: Number of LSTM training epochs
            verbose: Print training progress

        Returns:
            Dictionary with training metrics
        """
        logger.info(f"Fitting MRD Model on {len(df)} samples")

        # Step 1: Compute features
        logger.info("Step 1: Computing features...")
        features = self.feature_engineer.compute_all_features(df)
        self._feature_columns = features.columns.tolist()
        self._n_features = len(self._feature_columns)

        # Step 2: Fit HMM for regime detection
        logger.info("Step 2: Fitting HMM...")
        self.hmm.fit(features)

        # Get regime labels
        regimes = self.hmm.predict_regime(features)
        durations = self.hmm.get_regime_duration(regimes)

        # Step 3: Initialize and train LSTM
        logger.info("Step 3: Training LSTM...")
        self._init_lstm(self._n_features)

        lstm_metrics = {}
        if self.lstm is not None:
            lstm_metrics = self._train_lstm(
                features, regimes, durations, epochs, verbose
            )

        # Step 4: Train XGBoost ensemble
        logger.info("Step 4: Training XGBoost ensemble...")
        self._init_xgboost()

        xgb_metrics = {}
        if self.xgb_continuation is not None:
            xgb_metrics = self._train_xgboost(features, regimes, durations)

        self._is_fitted = True

        # Validation metrics
        val_metrics = {}
        if df_val is not None:
            val_metrics = self._validate(df_val)

        metrics = {
            'lstm': lstm_metrics,
            'xgb': xgb_metrics,
            'validation': val_metrics,
            'hmm': self.hmm.get_regime_statistics(features)
        }

        logger.info("MRD Model training complete")
        return metrics

    def _train_lstm(
        self,
        features: pd.DataFrame,
        regimes: np.ndarray,
        durations: np.ndarray,
        epochs: int,
        verbose: bool
    ) -> Dict[str, float]:
        """Train the LSTM component."""
        if self.lstm is None:
            return {}

        # Prepare sequences
        X, y_regime, y_duration = self._prepare_sequences(
            features.values, regimes, durations
        )

        # Create continuation labels (1 if same regime next period, 0 otherwise)
        y_continuation = (regimes[self.config.lstm_sequence_length:] ==
                          regimes[self.config.lstm_sequence_length - 1:-1]).astype(np.float32)
        y_continuation = y_continuation[:len(X)]

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

        # Create dataset and loader
        dataset = torch.utils.data.TensorDataset(
            X_tensor, regime_tensor, duration_tensor, continuation_tensor
        )
        loader = torch.utils.data.DataLoader(
            dataset, batch_size=self.config.batch_size, shuffle=True
        )

        # Optimizer
        optimizer = torch.optim.AdamW(
            self.lstm.parameters(),
            lr=self.config.learning_rate,
            weight_decay=self.config.weight_decay
        )

        # Loss functions
        regime_criterion = nn.CrossEntropyLoss()
        continuation_criterion = nn.BCELoss()

        # Training loop
        self.lstm.train()
        best_loss = float('inf')
        history = {'loss': [], 'regime_loss': [], 'continuation_loss': []}

        for epoch in range(epochs):
            epoch_loss = 0
            epoch_regime_loss = 0
            epoch_cont_loss = 0

            for batch_x, batch_regime, batch_duration, batch_cont in loader:
                optimizer.zero_grad()

                # Forward pass
                regime_logits, cont_prob, _, _ = self.lstm(
                    batch_x, batch_regime, batch_duration
                )

                # Losses
                regime_loss = regime_criterion(regime_logits, batch_regime)
                cont_loss = continuation_criterion(cont_prob.squeeze(), batch_cont)

                # Combined loss
                loss = regime_loss + cont_loss

                # Backward pass
                loss.backward()
                torch.nn.utils.clip_grad_norm_(self.lstm.parameters(), 1.0)
                optimizer.step()

                epoch_loss += loss.item()
                epoch_regime_loss += regime_loss.item()
                epoch_cont_loss += cont_loss.item()

            # Average losses
            n_batches = len(loader)
            avg_loss = epoch_loss / n_batches
            avg_regime = epoch_regime_loss / n_batches
            avg_cont = epoch_cont_loss / n_batches

            history['loss'].append(avg_loss)
            history['regime_loss'].append(avg_regime)
            history['continuation_loss'].append(avg_cont)

            if avg_loss < best_loss:
                best_loss = avg_loss

            if verbose and (epoch + 1) % 10 == 0:
                logger.info(
                    f"Epoch {epoch+1}/{epochs} - Loss: {avg_loss:.4f} "
                    f"(Regime: {avg_regime:.4f}, Cont: {avg_cont:.4f})"
                )

        self.lstm.eval()

        return {
            'final_loss': history['loss'][-1],
            'best_loss': best_loss,
            'regime_loss': history['regime_loss'][-1],
            'continuation_loss': history['continuation_loss'][-1]
        }

    def _train_xgboost(
        self,
        features: pd.DataFrame,
        regimes: np.ndarray,
        durations: np.ndarray
    ) -> Dict[str, float]:
        """Train XGBoost ensemble."""
        if self.xgb_continuation is None:
            return {}

        # Prepare features with LSTM outputs if available
        X = features.values
        seq_start = self.config.lstm_sequence_length

        # Get LSTM features if available
        if self.lstm is not None:
            lstm_features = self._get_lstm_features(features)
            # Combine with raw features
            X_combined = np.hstack([X[seq_start:], lstm_features])
        else:
            X_combined = X[seq_start:]

        # Add regime and duration as features
        regime_features = regimes[seq_start:].reshape(-1, 1)
        duration_features = durations[seq_start:].reshape(-1, 1)
        X_combined = np.hstack([X_combined, regime_features, duration_features])

        # Create targets
        # Continuation: 1 if same regime continues, 0 otherwise
        y_cont = (regimes[seq_start:] == regimes[seq_start - 1:-1]).astype(int)
        y_cont = np.append(y_cont, 0)  # Pad last element

        # Regime targets
        y_regime = regimes[seq_start:]

        # Train continuation model
        valid_mask = ~np.isnan(X_combined).any(axis=1)
        X_valid = X_combined[valid_mask]
        y_cont_valid = y_cont[valid_mask]
        y_regime_valid = y_regime[valid_mask]

        self.xgb_continuation.fit(X_valid, y_cont_valid)
        self.xgb_regime.fit(X_valid, y_regime_valid)

        # Calculate training metrics
        cont_pred = self.xgb_continuation.predict(X_valid)
        regime_pred = self.xgb_regime.predict(X_valid)

        cont_accuracy = np.mean((cont_pred > 0.5) == y_cont_valid)
        regime_accuracy = np.mean(regime_pred == y_regime_valid)

        return {
            'continuation_accuracy': float(cont_accuracy),
            'regime_accuracy': float(regime_accuracy)
        }

    def _prepare_sequences(
        self,
        features: np.ndarray,
        regimes: np.ndarray,
        durations: np.ndarray
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """Prepare sequences for LSTM training."""
        seq_len = self.config.lstm_sequence_length
        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

    def _get_lstm_features(self, features: pd.DataFrame) -> np.ndarray:
        """Extract LSTM hidden features for XGBoost."""
        if self.lstm is None:
            return np.zeros((len(features) - self.config.lstm_sequence_length, 64))

        self.lstm.eval()

        X, _, _ = self._prepare_sequences(
            features.values,
            np.zeros(len(features)),  # Dummy regimes
            np.zeros(len(features))   # Dummy durations
        )

        X_tensor = torch.FloatTensor(X).to(self.device)

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

        return lstm_features.cpu().numpy()

    def _validate(self, df_val: pd.DataFrame) -> Dict[str, float]:
        """Validate on held-out data."""
        predictions = self.predict_batch(df_val)

        # Calculate metrics
        features = self.feature_engineer.compute_all_features(df_val)
        true_regimes = self.hmm.predict_regime(features)
        true_durations = self.hmm.get_regime_duration(true_regimes)

        pred_regimes = np.array([p.regime for p in predictions])
        regime_accuracy = np.mean(
            pred_regimes == true_regimes[self.config.lstm_sequence_length:]
        )

        return {
            'regime_accuracy': float(regime_accuracy),
            'n_predictions': len(predictions)
        }

    def forward(
        self,
        features: pd.DataFrame
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        Forward pass through all components.

        Args:
            features: DataFrame with computed features

        Returns:
            Tuple of (regime, duration, continuation_prob)
        """
        if not self._is_fitted:
            raise RuntimeError("Model must be fitted before forward pass")

        # HMM regime detection
        regimes = self.hmm.predict_regime(features)
        durations = self.hmm.get_regime_duration(regimes)
        regime_probs = self.hmm.get_regime_probabilities(features)

        # For samples with enough history, use full pipeline
        seq_len = self.config.lstm_sequence_length
        n_samples = len(features)
        continuation_probs = np.zeros(n_samples)

        if n_samples >= seq_len and self.lstm is not None:
            X, _, _ = self._prepare_sequences(
                features.values, regimes, durations
            )

            X_tensor = torch.FloatTensor(X).to(self.device)

            with torch.no_grad():
                self.lstm.eval()
                regime_tensor = torch.LongTensor(regimes[seq_len - 1:-1]).to(self.device)
                duration_tensor = torch.FloatTensor(durations[seq_len - 1:-1]).to(self.device)

                _, cont_prob, _, _ = self.lstm(
                    X_tensor, regime_tensor, duration_tensor
                )

                continuation_probs[seq_len:] = cont_prob.squeeze().cpu().numpy()

        # For early samples, use HMM transition probabilities
        for i in range(min(seq_len, n_samples)):
            trans_matrix = self.hmm._get_ordered_transition_matrix()
            continuation_probs[i] = trans_matrix[regimes[i], regimes[i]]

        return regimes, durations, continuation_probs

    def predict(self, df: pd.DataFrame) -> MRDPrediction:
        """
        Make a single prediction for the most recent data point.

        Args:
            df: DataFrame with OHLCV data (needs enough history)

        Returns:
            MRDPrediction object
        """
        predictions = self.predict_batch(df.tail(self.config.lstm_sequence_length + 10))
        return predictions[-1] if predictions else None

    def predict_batch(self, df: pd.DataFrame) -> List[MRDPrediction]:
        """
        Make predictions for all data points.

        Args:
            df: DataFrame with OHLCV data

        Returns:
            List of MRDPrediction objects
        """
        if not self._is_fitted:
            raise RuntimeError("Model must be fitted before prediction")

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

        # Get predictions
        regimes, durations, continuation_probs = self.forward(features)
        regime_probs = self.hmm.get_regime_probabilities(features)

        # Build prediction objects
        predictions = []
        regime_names = {0: 'trend_down', 1: 'range', 2: 'trend_up'}

        for i in range(len(regimes)):
            regime = int(regimes[i])
            duration = int(durations[i])
            cont_prob = float(continuation_probs[i])

            probs_dict = {
                regime_names[j]: float(regime_probs[i, j])
                for j in range(3)
            }

            pred = MRDPrediction(
                regime=regime,
                regime_name=regime_names[regime],
                regime_probabilities=probs_dict,
                regime_confidence=float(regime_probs[i, regime]),
                duration=duration,
                continuation_prob=cont_prob,
                reversal_prob=1.0 - cont_prob
            )
            predictions.append(pred)

        return predictions

    def save(self, path: str):
        """Save the complete model."""
        path = Path(path)
        path.mkdir(parents=True, exist_ok=True)

        # Save HMM
        self.hmm.save(path / 'hmm')

        # Save LSTM
        if self.lstm is not None:
            torch.save(self.lstm.state_dict(), path / 'lstm.pt')

        # Save XGBoost models
        if self.xgb_continuation is not None:
            joblib.dump(self.xgb_continuation, path / 'xgb_continuation.joblib')
        if self.xgb_regime is not None:
            joblib.dump(self.xgb_regime, path / 'xgb_regime.joblib')

        # Save metadata
        metadata = {
            'config': self.config.__dict__,
            'feature_columns': self._feature_columns,
            'n_features': self._n_features,
            'version': '1.0.0'
        }
        joblib.dump(metadata, path / 'metadata.joblib')

        logger.info(f"Saved MRD Model to {path}")

    def load(self, path: str):
        """Load the complete model."""
        path = Path(path)

        # Load metadata
        metadata = joblib.load(path / 'metadata.joblib')
        self._feature_columns = metadata['feature_columns']
        self._n_features = metadata['n_features']

        # Load HMM
        self.hmm.load(path / 'hmm')

        # Load LSTM
        if (path / 'lstm.pt').exists() and TORCH_AVAILABLE:
            self._init_lstm(self._n_features)
            self.lstm.load_state_dict(torch.load(path / 'lstm.pt'))
            self.lstm.eval()

        # Load XGBoost models
        if (path / 'xgb_continuation.joblib').exists():
            self.xgb_continuation = joblib.load(path / 'xgb_continuation.joblib')
        if (path / 'xgb_regime.joblib').exists():
            self.xgb_regime = joblib.load(path / 'xgb_regime.joblib')

        self._is_fitted = True
        logger.info(f"Loaded MRD Model from {path}")


if __name__ == "__main__":
    # Test complete MRD model
    np.random.seed(42)
    n_samples = 3000

    # Create sample data with regime patterns
    dates = pd.date_range('2024-01-01', periods=n_samples, freq='5min')

    price = np.zeros(n_samples)
    price[0] = 2000

    # Create regime changes
    for i in range(1, n_samples):
        phase = (i // 500) % 3
        if phase == 0:  # Uptrend
            drift = 0.0003
        elif phase == 1:  # Range
            drift = 0
        else:  # Downtrend
            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)

    # Split data
    train_size = 2400
    df_train = df.iloc[:train_size]
    df_val = df.iloc[train_size:]

    # Train model
    print(f"\n{'='*60}")
    print("Testing MRD Model")
    print(f"{'='*60}")

    model = MRDModel()
    metrics = model.fit(df_train, df_val, epochs=30, verbose=True)

    print(f"\nTraining Metrics:")
    print(f"  LSTM Loss: {metrics['lstm'].get('final_loss', 'N/A')}")
    print(f"  XGB Regime Accuracy: {metrics['xgb'].get('regime_accuracy', 'N/A')}")
    print(f"  Validation Accuracy: {metrics['validation'].get('regime_accuracy', 'N/A')}")

    # Make predictions
    predictions = model.predict_batch(df_val.tail(100))
    print(f"\nSample Predictions (last 5):")
    for i, pred in enumerate(predictions[-5:]):
        print(f"  {i}: {pred.regime_name} (conf={pred.regime_confidence:.2f}, "
              f"dur={pred.duration}, cont={pred.continuation_prob:.2f})")

    print("\nTest complete!")
