#!/usr/bin/env python3
"""
VBP Trainer - Training Module for Volatility Breakout Predictor
================================================================
Comprehensive training system for VBP model with:
- Balanced sampling for rare breakout events
- Walk-forward cross-validation
- Performance tracking and evaluation
- Hyperparameter optimization support

Key Features:
1. Oversampling of breakout events (3x default)
2. Walk-forward validation for time series
3. Comprehensive breakout detection metrics
4. Early stopping and checkpointing

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

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

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset, WeightedRandomSampler

from sklearn.metrics import (
    accuracy_score, precision_score, recall_score, f1_score,
    classification_report, confusion_matrix, roc_auc_score,
    precision_recall_curve, average_precision_score
)

from .model import VBPModel, VBPModelConfig, VBPPrediction
from .feature_engineering import VBPFeatureEngineer, VBPFeatureConfig


@dataclass
class VBPTrainerConfig:
    """Configuration for VBP trainer."""

    # Oversampling configuration
    oversample_breakouts: bool = True
    oversample_multiplier: float = 3.0  # How much to oversample breakouts

    # Walk-forward configuration
    n_folds: int = 5
    train_ratio: float = 0.8  # Ratio of training data in each fold
    gap_periods: int = 10  # Gap between train and test to avoid leakage

    # Training configuration
    epochs: int = 100
    batch_size: int = 64
    learning_rate: float = 1e-3
    weight_decay: float = 1e-4

    # Early stopping
    early_stopping_patience: int = 10
    early_stopping_min_delta: float = 1e-4

    # Checkpointing
    checkpoint_dir: str = './checkpoints/vbp'
    save_best_only: bool = True

    # Evaluation thresholds
    breakout_prob_threshold: float = 0.5
    high_confidence_threshold: float = 0.7

    # Class balancing
    use_class_weights: bool = True
    class_weight_method: str = 'balanced'  # 'balanced', 'sqrt', 'log'


@dataclass
class VBPMetrics:
    """Metrics for VBP model evaluation."""

    # Classification metrics
    accuracy: float = 0.0
    precision: float = 0.0
    recall: float = 0.0
    f1: float = 0.0

    # Per-class metrics
    precision_per_class: Dict[str, float] = field(default_factory=dict)
    recall_per_class: Dict[str, float] = field(default_factory=dict)
    f1_per_class: Dict[str, float] = field(default_factory=dict)

    # Breakout-specific metrics
    breakout_precision: float = 0.0  # Precision for breakout classes
    breakout_recall: float = 0.0  # Recall for breakout classes
    breakout_f1: float = 0.0

    # Direction accuracy (when breakout is predicted)
    direction_accuracy: float = 0.0

    # Magnitude metrics (for regression)
    magnitude_mae: float = 0.0
    magnitude_rmse: float = 0.0

    # ROC-AUC (one-vs-rest)
    roc_auc: float = 0.0

    # Confusion matrix
    confusion_matrix: Optional[np.ndarray] = None

    # Sample counts
    n_samples: int = 0
    n_breakouts: int = 0
    n_predicted_breakouts: int = 0

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            'accuracy': self.accuracy,
            'precision': self.precision,
            'recall': self.recall,
            'f1': self.f1,
            'precision_per_class': self.precision_per_class,
            'recall_per_class': self.recall_per_class,
            'f1_per_class': self.f1_per_class,
            'breakout_precision': self.breakout_precision,
            'breakout_recall': self.breakout_recall,
            'breakout_f1': self.breakout_f1,
            'direction_accuracy': self.direction_accuracy,
            'magnitude_mae': self.magnitude_mae,
            'magnitude_rmse': self.magnitude_rmse,
            'roc_auc': self.roc_auc,
            'n_samples': self.n_samples,
            'n_breakouts': self.n_breakouts,
            'n_predicted_breakouts': self.n_predicted_breakouts
        }


class VBPDataset(Dataset):
    """
    PyTorch Dataset for VBP training.

    Handles sequence creation and label alignment.
    """

    def __init__(
        self,
        features: np.ndarray,
        labels: np.ndarray,
        magnitudes: np.ndarray,
        sequence_length: int = 50
    ):
        self.features = features
        self.labels = labels
        self.magnitudes = magnitudes
        self.sequence_length = sequence_length

        # Create valid indices (need enough history for sequence)
        self.valid_indices = np.arange(sequence_length, len(features))

    def __len__(self):
        return len(self.valid_indices)

    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        actual_idx = self.valid_indices[idx]

        # Get sequence
        seq_start = actual_idx - self.sequence_length
        sequence = self.features[seq_start:actual_idx]

        # Get current features (last row of sequence for tabular)
        current_features = self.features[actual_idx - 1]

        # Get labels
        label = self.labels[actual_idx]
        magnitude = self.magnitudes[actual_idx]

        return (
            torch.FloatTensor(sequence),
            torch.FloatTensor(current_features),
            torch.LongTensor([label]),
            torch.FloatTensor([magnitude])
        )


class VBPTrainer:
    """
    Trainer for Volatility Breakout Predictor.

    Handles:
    - Data preparation with oversampling
    - Walk-forward cross-validation
    - Training with early stopping
    - Comprehensive evaluation

    Usage:
        trainer = VBPTrainer(VBPTrainerConfig())
        metrics = trainer.train(symbol, df)
        eval_metrics = trainer.evaluate_breakout_detection(y_true, y_pred)
    """

    def __init__(
        self,
        config: Optional[VBPTrainerConfig] = None,
        model_config: Optional[VBPModelConfig] = None
    ):
        """
        Initialize VBP Trainer.

        Args:
            config: Trainer configuration
            model_config: Model configuration
        """
        self.config = config or VBPTrainerConfig()
        self.model_config = model_config or VBPModelConfig()

        # Training state
        self.best_metrics: Optional[VBPMetrics] = None
        self.training_history: List[Dict[str, float]] = []
        self.fold_results: List[VBPMetrics] = []

    def compute_class_weights(
        self,
        labels: np.ndarray
    ) -> np.ndarray:
        """
        Compute class weights for imbalanced data.

        Supports multiple weighting methods:
        - 'balanced': Inverse frequency
        - 'sqrt': Square root of inverse frequency
        - 'log': Logarithm of inverse frequency

        Args:
            labels: Array of class labels

        Returns:
            Array of sample weights
        """
        class_counts = Counter(labels)
        n_samples = len(labels)
        n_classes = len(class_counts)

        if self.config.class_weight_method == 'balanced':
            # Standard inverse frequency weighting
            class_weights = {
                cls: n_samples / (n_classes * count)
                for cls, count in class_counts.items()
            }
        elif self.config.class_weight_method == 'sqrt':
            # Sqrt of inverse frequency (less aggressive)
            class_weights = {
                cls: np.sqrt(n_samples / (n_classes * count))
                for cls, count in class_counts.items()
            }
        elif self.config.class_weight_method == 'log':
            # Log of inverse frequency (even less aggressive)
            class_weights = {
                cls: np.log1p(n_samples / (n_classes * count))
                for cls, count in class_counts.items()
            }
        else:
            # Default to balanced
            class_weights = {
                cls: n_samples / (n_classes * count)
                for cls, count in class_counts.items()
            }

        # Apply weights to samples
        sample_weights = np.array([class_weights[label] for label in labels])

        # Normalize
        sample_weights = sample_weights / sample_weights.mean()

        logger.info(f"Class distribution: {dict(class_counts)}")
        logger.info(f"Class weights ({self.config.class_weight_method}): {class_weights}")

        return sample_weights

    def oversample_breakouts(
        self,
        df: pd.DataFrame,
        labels: pd.Series,
        multiplier: Optional[float] = None
    ) -> Tuple[pd.DataFrame, pd.Series]:
        """
        Oversample breakout events.

        Since breakouts are rare (~5-10%), we oversample them
        to improve model training.

        Args:
            df: DataFrame with features
            labels: Series with breakout labels
            multiplier: Oversampling multiplier (default: from config)

        Returns:
            Tuple of (oversampled_df, oversampled_labels)
        """
        multiplier = multiplier or self.config.oversample_multiplier

        # Identify breakout samples (classes 1 and 2)
        breakout_mask = labels > 0
        breakout_indices = df.index[breakout_mask]
        n_breakouts = len(breakout_indices)

        if n_breakouts == 0:
            logger.warning("No breakout samples found for oversampling")
            return df, labels

        # Calculate number of samples to add
        n_oversample = int(n_breakouts * (multiplier - 1))

        if n_oversample <= 0:
            return df, labels

        # Randomly sample breakout indices with replacement
        oversample_indices = np.random.choice(breakout_indices, size=n_oversample, replace=True)

        # Create oversampled dataframes
        df_oversampled = pd.concat([df, df.loc[oversample_indices]])
        labels_oversampled = pd.concat([labels, labels.loc[oversample_indices]])

        # Shuffle
        shuffle_idx = np.random.permutation(len(df_oversampled))
        df_oversampled = df_oversampled.iloc[shuffle_idx].reset_index(drop=True)
        labels_oversampled = labels_oversampled.iloc[shuffle_idx].reset_index(drop=True)

        original_breakout_ratio = n_breakouts / len(labels)
        new_breakout_ratio = (labels_oversampled > 0).sum() / len(labels_oversampled)

        logger.info(f"Oversampling breakouts: {n_breakouts} -> {n_breakouts + n_oversample}")
        logger.info(f"Breakout ratio: {original_breakout_ratio:.2%} -> {new_breakout_ratio:.2%}")

        return df_oversampled, labels_oversampled

    def create_walk_forward_splits(
        self,
        df: pd.DataFrame,
        n_folds: Optional[int] = None
    ) -> List[Tuple[pd.DataFrame, pd.DataFrame]]:
        """
        Create walk-forward cross-validation splits.

        Walk-forward ensures temporal ordering:
        - Train on past data, test on future data
        - Each fold moves forward in time
        - Gap between train and test to avoid leakage

        Args:
            df: Input DataFrame
            n_folds: Number of folds (default: from config)

        Returns:
            List of (train_df, test_df) tuples
        """
        n_folds = n_folds or self.config.n_folds
        n_samples = len(df)

        # Calculate fold size
        fold_size = n_samples // (n_folds + 1)  # +1 for initial training data

        splits = []

        for fold in range(n_folds):
            # Training end (expanding window)
            train_end = fold_size * (fold + 2)

            # Apply gap
            train_end_with_gap = train_end - self.config.gap_periods

            # Test start and end
            test_start = train_end
            test_end = min(train_end + fold_size, n_samples)

            if test_end <= test_start:
                continue

            train_df = df.iloc[:train_end_with_gap]
            test_df = df.iloc[test_start:test_end]

            splits.append((train_df, test_df))

            logger.info(f"Fold {fold + 1}: Train [{0}:{train_end_with_gap}], "
                       f"Test [{test_start}:{test_end}]")

        return splits

    def train(
        self,
        symbol: str,
        df: pd.DataFrame,
        verbose: bool = True
    ) -> VBPMetrics:
        """
        Train VBP model for a symbol.

        Args:
            symbol: Trading symbol (for logging/saving)
            df: OHLCV DataFrame
            verbose: Print training progress

        Returns:
            Training metrics
        """
        logger.info(f"Training VBP model for {symbol} on {len(df)} samples...")

        # Create model
        model = VBPModel(self.model_config)

        # Prepare data with oversampling if enabled
        feature_engineer = VBPFeatureEngineer(self.model_config.feature_config)
        features = feature_engineer.compute_all_features(df)
        labels = feature_engineer.label_breakouts(df)

        # Get labels
        y_class = labels['breakout_label'].values
        y_magnitude = labels['breakout_magnitude'].values

        # Handle oversampling
        if self.config.oversample_breakouts:
            # Create temporary df with features and labels aligned
            features['_label'] = y_class
            features['_magnitude'] = y_magnitude

            # Remove NaN labels
            valid_mask = ~np.isnan(features['_label'])
            features = features[valid_mask]

            # Oversample
            label_series = features['_label']
            features_clean, labels_oversampled = self.oversample_breakouts(
                features.drop(columns=['_label', '_magnitude']),
                label_series
            )

            # Update data
            df_for_training = df[valid_mask]

        # Train model
        metrics_dict = model.fit(df, verbose=verbose)

        # Create metrics object
        metrics = VBPMetrics(
            accuracy=metrics_dict.get('accuracy', 0),
            f1=metrics_dict.get('f1_macro', 0),
            n_samples=metrics_dict.get('n_samples', 0),
            n_breakouts=metrics_dict.get('n_breakouts', 0)
        )

        # Save model
        checkpoint_path = Path(self.config.checkpoint_dir) / symbol
        model.save(str(checkpoint_path))

        self.best_metrics = metrics
        return metrics

    def walk_forward_train(
        self,
        symbol: str,
        df: pd.DataFrame,
        n_folds: Optional[int] = None,
        verbose: bool = True
    ) -> VBPMetrics:
        """
        Train with walk-forward cross-validation.

        More robust evaluation for time series data.

        Args:
            symbol: Trading symbol
            df: OHLCV DataFrame
            n_folds: Number of CV folds
            verbose: Print progress

        Returns:
            Aggregated metrics across all folds
        """
        n_folds = n_folds or self.config.n_folds

        logger.info(f"Walk-forward training for {symbol} with {n_folds} folds...")

        # Create splits
        splits = self.create_walk_forward_splits(df, n_folds)

        self.fold_results = []
        all_y_true = []
        all_y_pred = []
        all_y_proba = []
        all_magnitudes_true = []
        all_magnitudes_pred = []

        for fold_idx, (train_df, test_df) in enumerate(splits):
            logger.info(f"\n{'='*50}")
            logger.info(f"Fold {fold_idx + 1}/{len(splits)}")
            logger.info(f"{'='*50}")

            # Create model for this fold
            model = VBPModel(self.model_config)

            # Train
            train_metrics = model.fit(train_df, verbose=verbose)

            # Predict on test set
            predictions = model.predict(test_df)

            # Extract predictions
            feature_engineer = VBPFeatureEngineer(self.model_config.feature_config)
            test_labels = feature_engineer.label_breakouts(test_df)

            # Get valid indices (where we have labels)
            valid_mask = ~np.isnan(test_labels['breakout_label'].values)
            valid_indices = np.where(valid_mask)[0]

            if len(valid_indices) == 0:
                logger.warning(f"No valid test samples in fold {fold_idx + 1}")
                continue

            # Collect predictions for valid samples
            for idx in valid_indices:
                if idx < len(predictions):
                    pred = predictions[idx]
                    all_y_true.append(int(test_labels['breakout_label'].iloc[idx]))
                    all_y_pred.append(pred.breakout_class)
                    all_y_proba.append(pred.breakout_probability)
                    all_magnitudes_true.append(test_labels['breakout_magnitude'].iloc[idx])
                    all_magnitudes_pred.append(pred.magnitude)

            # Evaluate fold
            if len(all_y_true) > 0:
                fold_metrics = self.evaluate_breakout_detection(
                    np.array(all_y_true[-len(valid_indices):]),
                    np.array(all_y_pred[-len(valid_indices):]),
                    np.array(all_magnitudes_true[-len(valid_indices):]),
                    np.array(all_magnitudes_pred[-len(valid_indices):])
                )
                self.fold_results.append(fold_metrics)

                if verbose:
                    logger.info(f"Fold {fold_idx + 1} Results:")
                    logger.info(f"  Accuracy: {fold_metrics.accuracy:.4f}")
                    logger.info(f"  Breakout F1: {fold_metrics.breakout_f1:.4f}")
                    logger.info(f"  Breakout Recall: {fold_metrics.breakout_recall:.4f}")

        # Aggregate results
        if len(all_y_true) > 0:
            aggregated_metrics = self.evaluate_breakout_detection(
                np.array(all_y_true),
                np.array(all_y_pred),
                np.array(all_magnitudes_true),
                np.array(all_magnitudes_pred)
            )

            logger.info(f"\n{'='*50}")
            logger.info("Walk-Forward Aggregated Results:")
            logger.info(f"{'='*50}")
            logger.info(f"  Total samples: {aggregated_metrics.n_samples}")
            logger.info(f"  Accuracy: {aggregated_metrics.accuracy:.4f}")
            logger.info(f"  Breakout Precision: {aggregated_metrics.breakout_precision:.4f}")
            logger.info(f"  Breakout Recall: {aggregated_metrics.breakout_recall:.4f}")
            logger.info(f"  Breakout F1: {aggregated_metrics.breakout_f1:.4f}")
            logger.info(f"  Magnitude MAE: {aggregated_metrics.magnitude_mae:.4f}")

            self.best_metrics = aggregated_metrics
            return aggregated_metrics

        return VBPMetrics()

    def evaluate_breakout_detection(
        self,
        y_true: np.ndarray,
        y_pred: np.ndarray,
        magnitudes_true: Optional[np.ndarray] = None,
        magnitudes_pred: Optional[np.ndarray] = None
    ) -> VBPMetrics:
        """
        Comprehensive evaluation of breakout detection.

        Computes:
        - Standard classification metrics
        - Breakout-specific metrics (combining bullish/bearish)
        - Direction accuracy
        - Magnitude prediction metrics

        Args:
            y_true: True labels (0: no breakout, 1: bullish, 2: bearish)
            y_pred: Predicted labels
            magnitudes_true: True breakout magnitudes
            magnitudes_pred: Predicted magnitudes

        Returns:
            VBPMetrics object with all metrics
        """
        metrics = VBPMetrics()

        # Basic counts
        metrics.n_samples = len(y_true)
        metrics.n_breakouts = int((y_true > 0).sum())
        metrics.n_predicted_breakouts = int((y_pred > 0).sum())

        # Standard classification metrics
        metrics.accuracy = accuracy_score(y_true, y_pred)

        # Per-class metrics
        class_names = ['no_breakout', 'bullish', 'bearish']
        precisions = precision_score(y_true, y_pred, average=None, zero_division=0)
        recalls = recall_score(y_true, y_pred, average=None, zero_division=0)
        f1s = f1_score(y_true, y_pred, average=None, zero_division=0)

        for i, name in enumerate(class_names[:len(precisions)]):
            metrics.precision_per_class[name] = float(precisions[i])
            metrics.recall_per_class[name] = float(recalls[i])
            metrics.f1_per_class[name] = float(f1s[i])

        # Macro averages
        metrics.precision = precision_score(y_true, y_pred, average='macro', zero_division=0)
        metrics.recall = recall_score(y_true, y_pred, average='macro', zero_division=0)
        metrics.f1 = f1_score(y_true, y_pred, average='macro', zero_division=0)

        # Breakout-specific metrics (treat classes 1 and 2 as "breakout")
        y_true_binary = (y_true > 0).astype(int)
        y_pred_binary = (y_pred > 0).astype(int)

        metrics.breakout_precision = precision_score(y_true_binary, y_pred_binary, zero_division=0)
        metrics.breakout_recall = recall_score(y_true_binary, y_pred_binary, zero_division=0)
        metrics.breakout_f1 = f1_score(y_true_binary, y_pred_binary, zero_division=0)

        # Direction accuracy (only when both true and predicted are breakouts)
        true_breakout_mask = y_true > 0
        pred_breakout_mask = y_pred > 0
        both_breakout_mask = true_breakout_mask & pred_breakout_mask

        if both_breakout_mask.sum() > 0:
            correct_direction = (y_true[both_breakout_mask] == y_pred[both_breakout_mask])
            metrics.direction_accuracy = correct_direction.mean()
        else:
            metrics.direction_accuracy = 0.0

        # Magnitude metrics (only for actual breakouts)
        if magnitudes_true is not None and magnitudes_pred is not None:
            breakout_mask = y_true > 0
            if breakout_mask.sum() > 0:
                mag_true = magnitudes_true[breakout_mask]
                mag_pred = magnitudes_pred[breakout_mask]

                # Remove NaN
                valid_mag = ~(np.isnan(mag_true) | np.isnan(mag_pred))
                if valid_mag.sum() > 0:
                    metrics.magnitude_mae = np.mean(np.abs(mag_true[valid_mag] - mag_pred[valid_mag]))
                    metrics.magnitude_rmse = np.sqrt(np.mean((mag_true[valid_mag] - mag_pred[valid_mag]) ** 2))

        # ROC-AUC (for binary breakout detection)
        try:
            if len(np.unique(y_true_binary)) > 1:
                metrics.roc_auc = roc_auc_score(y_true_binary, y_pred_binary)
        except Exception:
            metrics.roc_auc = 0.0

        # Confusion matrix
        metrics.confusion_matrix = confusion_matrix(y_true, y_pred)

        return metrics

    def print_evaluation_report(
        self,
        metrics: VBPMetrics,
        y_true: np.ndarray,
        y_pred: np.ndarray
    ) -> None:
        """
        Print detailed evaluation report.

        Args:
            metrics: VBPMetrics object
            y_true: True labels
            y_pred: Predicted labels
        """
        logger.info("\n" + "=" * 60)
        logger.info("VBP Model Evaluation Report")
        logger.info("=" * 60)

        logger.info(f"\nSample Statistics:")
        logger.info(f"  Total samples: {metrics.n_samples}")
        logger.info(f"  Actual breakouts: {metrics.n_breakouts} ({metrics.n_breakouts/metrics.n_samples*100:.1f}%)")
        logger.info(f"  Predicted breakouts: {metrics.n_predicted_breakouts}")

        logger.info(f"\nOverall Metrics:")
        logger.info(f"  Accuracy: {metrics.accuracy:.4f}")
        logger.info(f"  Precision (macro): {metrics.precision:.4f}")
        logger.info(f"  Recall (macro): {metrics.recall:.4f}")
        logger.info(f"  F1 (macro): {metrics.f1:.4f}")

        logger.info(f"\nBreakout Detection Metrics:")
        logger.info(f"  Precision: {metrics.breakout_precision:.4f}")
        logger.info(f"  Recall: {metrics.breakout_recall:.4f}")
        logger.info(f"  F1: {metrics.breakout_f1:.4f}")
        logger.info(f"  ROC-AUC: {metrics.roc_auc:.4f}")

        logger.info(f"\nDirection Accuracy (when breakout predicted correctly):")
        logger.info(f"  {metrics.direction_accuracy:.4f}")

        logger.info(f"\nMagnitude Prediction:")
        logger.info(f"  MAE: {metrics.magnitude_mae:.4f} ATR units")
        logger.info(f"  RMSE: {metrics.magnitude_rmse:.4f} ATR units")

        logger.info(f"\nPer-Class Metrics:")
        for cls in metrics.f1_per_class:
            logger.info(f"  {cls}:")
            logger.info(f"    Precision: {metrics.precision_per_class.get(cls, 0):.4f}")
            logger.info(f"    Recall: {metrics.recall_per_class.get(cls, 0):.4f}")
            logger.info(f"    F1: {metrics.f1_per_class.get(cls, 0):.4f}")

        logger.info(f"\nConfusion Matrix:")
        logger.info(f"  {metrics.confusion_matrix}")

        logger.info("\nClassification Report:")
        logger.info(classification_report(y_true, y_pred, target_names=['no_breakout', 'bullish', 'bearish'], zero_division=0))

    def get_training_summary(self) -> Dict[str, Any]:
        """Get training summary."""
        return {
            'best_metrics': self.best_metrics.to_dict() if self.best_metrics else None,
            'n_folds': len(self.fold_results),
            'fold_results': [m.to_dict() for m in self.fold_results],
            'config': {
                'oversample_multiplier': self.config.oversample_multiplier,
                'n_folds': self.config.n_folds,
                'class_weight_method': self.config.class_weight_method
            }
        }


if __name__ == "__main__":
    # Test the VBP trainer
    print("Testing VBP Trainer")
    print("=" * 60)

    # Create sample data
    np.random.seed(42)
    n = 2000

    dates = pd.date_range('2025-01-01', periods=n, freq='1H')

    # Simulate price with breakout patterns
    volatility = np.where(
        (np.arange(n) % 100 > 85),
        0.02,
        0.005
    )

    returns = np.random.randn(n) * volatility
    price = 2000 * (1 + returns).cumprod()

    df = pd.DataFrame({
        'open': price,
        'high': price * (1 + np.abs(np.random.randn(n)) * volatility),
        'low': price * (1 - np.abs(np.random.randn(n)) * volatility),
        'close': price * (1 + np.random.randn(n) * volatility * 0.5),
        'volume': np.random.randint(100, 1000, n)
    }, index=dates)

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

    # Initialize trainer
    trainer_config = VBPTrainerConfig(
        n_folds=3,
        oversample_multiplier=2.0
    )

    model_config = VBPModelConfig(
        sequence_length=30,
        xgb_n_estimators=50,
        use_gpu=False
    )

    trainer = VBPTrainer(trainer_config, model_config)

    print("\n1. Testing class weight computation...")
    labels = np.array([0] * 900 + [1] * 50 + [2] * 50)
    weights = trainer.compute_class_weights(labels)
    print(f"   Sample weights range: [{weights.min():.2f}, {weights.max():.2f}]")

    print("\n2. Testing oversampling...")
    feature_engineer = VBPFeatureEngineer()
    features = feature_engineer.compute_all_features(df)
    label_df = feature_engineer.label_breakouts(df)
    label_series = pd.Series(label_df['breakout_label'].values, index=features.index)

    # Remove NaN
    valid = ~label_series.isna()
    features_valid = features[valid]
    labels_valid = label_series[valid]

    features_os, labels_os = trainer.oversample_breakouts(features_valid, labels_valid)
    print(f"   Original: {len(features_valid)}, Oversampled: {len(features_os)}")

    print("\n3. Testing walk-forward splits...")
    splits = trainer.create_walk_forward_splits(df, n_folds=3)
    print(f"   Created {len(splits)} folds")

    print("\n4. Testing simple training...")
    metrics = trainer.train('TEST_SYMBOL', df, verbose=True)
    print(f"   Training completed with accuracy: {metrics.accuracy:.4f}")

    print("\n5. Testing breakout detection evaluation...")
    y_true = np.array([0, 0, 1, 2, 0, 1, 0, 2, 0, 0])
    y_pred = np.array([0, 0, 1, 1, 0, 0, 0, 2, 1, 0])
    mag_true = np.array([0, 0, 2.5, 3.0, 0, 2.0, 0, 2.8, 0, 0])
    mag_pred = np.array([0, 0, 2.3, 2.8, 0, 1.5, 0, 2.5, 2.0, 0])

    eval_metrics = trainer.evaluate_breakout_detection(y_true, y_pred, mag_true, mag_pred)
    print(f"   Breakout F1: {eval_metrics.breakout_f1:.4f}")
    print(f"   Direction Accuracy: {eval_metrics.direction_accuracy:.4f}")
    print(f"   Magnitude MAE: {eval_metrics.magnitude_mae:.4f}")

    print("\n6. Testing evaluation report...")
    trainer.print_evaluation_report(eval_metrics, y_true, y_pred)

    print("\n7. Testing training summary...")
    summary = trainer.get_training_summary()
    print(f"   Summary keys: {list(summary.keys())}")

    print("\n" + "=" * 60)
    print("All VBP trainer tests passed!")
