#!/usr/bin/env python3
"""
VBP Model - Volatility Breakout Predictor Complete Model
=========================================================
Complete model for predicting volatility breakouts combining:
1. CNN encoder for temporal pattern extraction
2. XGBoost classifier for breakout detection (handles class imbalance)
3. Separate heads for breakout probability, direction, and magnitude

The model uses a hybrid approach:
- Neural network (CNN + Attention) for feature encoding
- XGBoost for final classification (better for tabular/imbalanced data)

Key Features:
- Handles class imbalance (breakouts are rare ~5-10%)
- Predicts: breakout_probability, direction, magnitude
- Supports both training and inference modes

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

import torch
import torch.nn as nn
import torch.nn.functional as F

try:
    from xgboost import XGBClassifier, XGBRegressor
    HAS_XGBOOST = True
except ImportError:
    HAS_XGBOOST = False
    logger.warning("XGBoost not available - some features will be limited")

from .feature_engineering import VBPFeatureEngineer, VBPFeatureConfig
from .cnn_encoder import VolatilityCNN, CNNEncoderConfig


@dataclass
class VBPModelConfig:
    """Configuration for the VBP model."""

    # Feature engineering config
    feature_config: VBPFeatureConfig = field(default_factory=VBPFeatureConfig)

    # CNN encoder config
    cnn_config: CNNEncoderConfig = field(default_factory=lambda: CNNEncoderConfig(
        conv_configs=[(32, 3), (64, 5), (128, 7)],
        attention_heads=4,
        dropout=0.3,
        use_attention=True,
        output_dim=128
    ))

    # Sequence length for CNN input
    sequence_length: int = 50

    # XGBoost configuration
    xgb_n_estimators: int = 200
    xgb_max_depth: int = 6
    xgb_learning_rate: float = 0.05
    xgb_min_child_weight: int = 5
    xgb_subsample: float = 0.8
    xgb_colsample_bytree: float = 0.8
    xgb_scale_pos_weight: float = 3.0  # For class imbalance

    # Use GPU if available
    use_gpu: bool = True

    # Output heads configuration
    num_classes: int = 3  # 0: no breakout, 1: bullish, 2: bearish

    # Confidence thresholds
    breakout_threshold: float = 0.5
    high_confidence_threshold: float = 0.7


@dataclass
class VBPPrediction:
    """VBP model prediction result."""

    breakout_probability: float
    breakout_class: int  # 0: no breakout, 1: bullish, 2: bearish
    direction: int  # 1: up, -1: down, 0: no breakout
    magnitude: float  # Predicted magnitude in ATR units
    confidence: float
    class_probabilities: Dict[str, float]
    is_breakout: bool
    is_high_confidence: bool

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            'breakout_probability': float(self.breakout_probability),
            'breakout_class': int(self.breakout_class),
            'direction': int(self.direction),
            'magnitude': float(self.magnitude),
            'confidence': float(self.confidence),
            'class_probabilities': self.class_probabilities,
            'is_breakout': bool(self.is_breakout),
            'is_high_confidence': bool(self.is_high_confidence)
        }

    @property
    def signal(self) -> str:
        """Get trading signal string."""
        if not self.is_breakout:
            return 'WAIT'
        if not self.is_high_confidence:
            return 'WEAK_' + ('LONG' if self.direction > 0 else 'SHORT')
        return 'STRONG_' + ('LONG' if self.direction > 0 else 'SHORT')


class NeuralHead(nn.Module):
    """
    Neural network head for regression/classification.

    Used for direction and magnitude prediction when CNN encoding
    is used directly (alternative to XGBoost).
    """

    def __init__(
        self,
        input_dim: int,
        output_dim: int,
        hidden_dims: List[int] = None,
        dropout: float = 0.3,
        activation: str = 'relu'
    ):
        super().__init__()

        hidden_dims = hidden_dims or [128, 64]

        layers = []
        in_dim = input_dim

        for hidden_dim in hidden_dims:
            layers.append(nn.Linear(in_dim, hidden_dim))
            layers.append(nn.BatchNorm1d(hidden_dim))
            if activation == 'relu':
                layers.append(nn.ReLU())
            elif activation == 'gelu':
                layers.append(nn.GELU())
            layers.append(nn.Dropout(dropout))
            in_dim = hidden_dim

        layers.append(nn.Linear(in_dim, output_dim))
        self.network = nn.Sequential(*layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.network(x)


class VBPModel:
    """
    Volatility Breakout Predictor Model.

    Combines CNN-based feature encoding with XGBoost classification
    for robust breakout prediction.

    Architecture:
    1. Feature Engineering: Extract volatility features
    2. CNN Encoder: Encode temporal patterns
    3. XGBoost Classifier: Predict breakout class
    4. XGBoost Regressor: Predict magnitude

    Usage:
        model = VBPModel(VBPModelConfig())
        model.fit(df_train)
        predictions = model.predict(df_test)
    """

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

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

        # Initialize feature engineer
        self.feature_engineer = VBPFeatureEngineer(self.config.feature_config)

        # Device selection
        self.device = torch.device('cuda' if torch.cuda.is_available() and self.config.use_gpu else 'cpu')

        # Initialize CNN encoder (will be sized after first data pass)
        self.cnn_encoder: Optional[VolatilityCNN] = None

        # Initialize XGBoost models
        self.classifier: Optional[XGBClassifier] = None
        self.regressor: Optional[XGBRegressor] = None

        # Feature info
        self.feature_columns: List[str] = []
        self.n_features: int = 0

        # Training state
        self._is_fitted = False
        self._label_encoder = None

    def _initialize_cnn(self, n_features: int):
        """Initialize CNN encoder with correct input dimension."""
        self.n_features = n_features
        self.config.cnn_config.input_channels = n_features

        self.cnn_encoder = VolatilityCNN(
            input_dim=n_features,
            config=self.config.cnn_config
        ).to(self.device)

    def _initialize_xgboost(self, n_classes: int = 3):
        """Initialize XGBoost models."""
        if not HAS_XGBOOST:
            raise ImportError("XGBoost is required for VBPModel")

        # Common parameters
        common_params = {
            'n_estimators': self.config.xgb_n_estimators,
            'max_depth': self.config.xgb_max_depth,
            'learning_rate': self.config.xgb_learning_rate,
            'min_child_weight': self.config.xgb_min_child_weight,
            'subsample': self.config.xgb_subsample,
            'colsample_bytree': self.config.xgb_colsample_bytree,
            'random_state': 42,
            'n_jobs': -1
        }

        # GPU settings
        if self.config.use_gpu and torch.cuda.is_available():
            common_params['tree_method'] = 'hist'
            common_params['device'] = 'cuda'

        # Classifier for breakout detection
        classifier_params = common_params.copy()
        classifier_params['objective'] = 'multi:softprob'
        classifier_params['num_class'] = n_classes
        classifier_params['scale_pos_weight'] = self.config.xgb_scale_pos_weight

        self.classifier = XGBClassifier(**classifier_params)

        # Regressor for magnitude prediction
        regressor_params = common_params.copy()
        regressor_params['objective'] = 'reg:squarederror'

        self.regressor = XGBRegressor(**regressor_params)

    def _prepare_features(
        self,
        df: pd.DataFrame
    ) -> Tuple[np.ndarray, np.ndarray, pd.DataFrame]:
        """
        Prepare features for training/prediction.

        Returns:
        - X_tabular: Tabular features for XGBoost (n_samples, n_features)
        - X_sequence: Sequence features for CNN (n_samples, seq_len, n_features)
        - labels_df: Label DataFrame

        Args:
            df: Input OHLCV DataFrame

        Returns:
            Tuple of (X_tabular, X_sequence, labels_df)
        """
        # Extract features
        features = self.feature_engineer.compute_all_features(df)
        labels = self.feature_engineer.label_breakouts(df)

        self.feature_columns = features.columns.tolist()

        # Tabular features (most recent values)
        X_tabular = features.values

        # Sequence features for CNN
        seq_len = self.config.sequence_length
        n_samples = len(features)
        n_features = len(self.feature_columns)

        # Create sequences (rolling window)
        X_sequence = np.zeros((n_samples, seq_len, n_features))

        for i in range(seq_len, n_samples):
            X_sequence[i] = features.iloc[i-seq_len:i].values

        # Pad beginning samples
        for i in range(min(seq_len, n_samples)):
            pad_len = seq_len - i
            if i > 0:
                X_sequence[i, pad_len:] = features.iloc[:i].values
            # First sample gets zeros (already initialized)

        return X_tabular, X_sequence, labels

    def _encode_sequences(self, X_sequence: np.ndarray) -> np.ndarray:
        """
        Encode sequences using CNN encoder.

        Args:
            X_sequence: Sequence array (n_samples, seq_len, n_features)

        Returns:
            Encoded features (n_samples, cnn_output_dim)
        """
        if self.cnn_encoder is None:
            n_features = X_sequence.shape[2]
            self._initialize_cnn(n_features)

        self.cnn_encoder.eval()

        # Process in batches to avoid memory issues
        batch_size = 512
        n_samples = len(X_sequence)
        encoded_list = []

        with torch.no_grad():
            for i in range(0, n_samples, batch_size):
                batch = X_sequence[i:i+batch_size]
                batch_tensor = torch.FloatTensor(batch).to(self.device)

                encoded, _ = self.cnn_encoder(batch_tensor)
                encoded_list.append(encoded.cpu().numpy())

        return np.vstack(encoded_list)

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

        Gives higher weight to rare breakout classes.

        Args:
            y: Label array

        Returns:
            Sample weights array
        """
        from collections import Counter

        class_counts = Counter(y)
        total = len(y)
        n_classes = len(class_counts)

        # Compute inverse frequency weights
        class_weights = {
            cls: total / (n_classes * count)
            for cls, count in class_counts.items()
        }

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

        # Normalize to mean=1
        sample_weights = sample_weights / sample_weights.mean()

        logger.info(f"Class weights: {class_weights}")
        logger.info(f"Sample weight range: [{sample_weights.min():.2f}, {sample_weights.max():.2f}]")

        return sample_weights

    def fit(
        self,
        df: pd.DataFrame,
        verbose: bool = True,
        val_df: Optional[pd.DataFrame] = None
    ) -> Dict[str, float]:
        """
        Fit the VBP model.

        Args:
            df: Training DataFrame with OHLCV data
            verbose: Print training progress
            val_df: Optional validation DataFrame

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

        # Prepare features
        X_tabular, X_sequence, labels = self._prepare_features(df)

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

        # Remove samples with NaN labels (forward-looking window at end)
        valid_mask = ~np.isnan(y_class) & ~np.isnan(y_magnitude)
        X_tabular = X_tabular[valid_mask]
        X_sequence = X_sequence[valid_mask]
        y_class = y_class[valid_mask].astype(int)
        y_magnitude = y_magnitude[valid_mask]

        logger.info(f"Valid samples after label filtering: {len(y_class)}")

        # Initialize CNN if needed
        if self.cnn_encoder is None:
            self._initialize_cnn(X_sequence.shape[2])

        # Encode sequences with CNN
        logger.info("Encoding sequences with CNN...")
        X_encoded = self._encode_sequences(X_sequence)

        # Combine tabular and encoded features
        X_combined = np.hstack([X_tabular, X_encoded])
        logger.info(f"Combined feature dimension: {X_combined.shape[1]}")

        # Compute class weights for imbalanced data
        sample_weights = self._compute_class_weights(y_class)

        # Initialize XGBoost models
        n_classes = len(np.unique(y_class))
        self._initialize_xgboost(n_classes)

        # Train classifier
        logger.info("Training breakout classifier...")
        self.classifier.fit(
            X_combined, y_class,
            sample_weight=sample_weights,
            verbose=verbose
        )

        # Train magnitude regressor (only on breakout samples)
        breakout_mask = y_class > 0
        if breakout_mask.sum() > 0:
            logger.info(f"Training magnitude regressor on {breakout_mask.sum()} breakout samples...")
            self.regressor.fit(
                X_combined[breakout_mask],
                y_magnitude[breakout_mask],
                verbose=verbose
            )
        else:
            logger.warning("No breakout samples found for magnitude regressor training")

        self._is_fitted = True

        # Calculate training metrics
        y_pred_class = self.classifier.predict(X_combined)
        y_pred_proba = self.classifier.predict_proba(X_combined)

        from sklearn.metrics import accuracy_score, f1_score, classification_report

        metrics = {
            'accuracy': accuracy_score(y_class, y_pred_class),
            'f1_macro': f1_score(y_class, y_pred_class, average='macro', zero_division=0),
            'f1_weighted': f1_score(y_class, y_pred_class, average='weighted', zero_division=0),
            'n_samples': len(y_class),
            'n_breakouts': int(breakout_mask.sum()),
            'breakout_ratio': float(breakout_mask.sum() / len(y_class))
        }

        if verbose:
            logger.info(f"\nTraining Metrics:")
            logger.info(f"  Accuracy: {metrics['accuracy']:.4f}")
            logger.info(f"  F1 (macro): {metrics['f1_macro']:.4f}")
            logger.info(f"  F1 (weighted): {metrics['f1_weighted']:.4f}")
            logger.info(f"  Breakout ratio: {metrics['breakout_ratio']:.2%}")
            logger.info(f"\n{classification_report(y_class, y_pred_class, zero_division=0)}")

        return metrics

    def predict(self, df: pd.DataFrame) -> List[VBPPrediction]:
        """
        Generate predictions for input data.

        Args:
            df: Input DataFrame with OHLCV data

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

        # Prepare features
        X_tabular, X_sequence, _ = self._prepare_features(df)

        # Encode sequences
        X_encoded = self._encode_sequences(X_sequence)

        # Combine features
        X_combined = np.hstack([X_tabular, X_encoded])

        # Get predictions
        y_pred_class = self.classifier.predict(X_combined)
        y_pred_proba = self.classifier.predict_proba(X_combined)

        # Predict magnitude for all samples
        y_pred_magnitude = self.regressor.predict(X_combined)

        # Create prediction objects
        predictions = []
        for i in range(len(X_combined)):
            pred_class = int(y_pred_class[i])
            proba = y_pred_proba[i]

            # Breakout probability (sum of class 1 and 2)
            breakout_prob = proba[1] + proba[2] if len(proba) > 2 else proba[1]

            # Direction
            if pred_class == 1:
                direction = 1  # Bullish
            elif pred_class == 2:
                direction = -1  # Bearish
            else:
                direction = 0  # No breakout

            # Confidence is the probability of predicted class
            confidence = float(proba[pred_class])

            # Is breakout?
            is_breakout = pred_class > 0 and breakout_prob > self.config.breakout_threshold

            # Is high confidence?
            is_high_confidence = confidence > self.config.high_confidence_threshold

            pred = VBPPrediction(
                breakout_probability=float(breakout_prob),
                breakout_class=pred_class,
                direction=direction,
                magnitude=float(y_pred_magnitude[i]),
                confidence=confidence,
                class_probabilities={
                    'no_breakout': float(proba[0]),
                    'bullish': float(proba[1]) if len(proba) > 1 else 0.0,
                    'bearish': float(proba[2]) if len(proba) > 2 else 0.0
                },
                is_breakout=is_breakout,
                is_high_confidence=is_high_confidence
            )
            predictions.append(pred)

        return predictions

    def predict_single(self, df: pd.DataFrame) -> VBPPrediction:
        """
        Get prediction for the most recent bar.

        Args:
            df: DataFrame with at least sequence_length bars

        Returns:
            Single VBPPrediction
        """
        # Ensure enough data for sequence
        min_length = self.config.sequence_length + 50  # Extra for feature calculation
        if len(df) < min_length:
            raise ValueError(f"Need at least {min_length} bars for prediction")

        predictions = self.predict(df.tail(min_length))
        return predictions[-1]

    def get_feature_importance(self, top_n: int = 20) -> Dict[str, float]:
        """
        Get feature importance from the classifier.

        Args:
            top_n: Number of top features to return

        Returns:
            Dictionary of feature names to importance scores
        """
        if not self._is_fitted:
            raise RuntimeError("Model must be fitted first")

        # Get importance from classifier
        importance = self.classifier.feature_importances_

        # Create feature names (tabular + encoded)
        feature_names = self.feature_columns.copy()
        for i in range(self.config.cnn_config.output_dim):
            feature_names.append(f'cnn_encoded_{i}')

        # Create importance dict
        importance_dict = dict(zip(feature_names, importance))

        # Sort by importance
        sorted_importance = dict(
            sorted(importance_dict.items(), key=lambda x: x[1], reverse=True)[:top_n]
        )

        return sorted_importance

    def save(self, path: str) -> None:
        """
        Save model to disk.

        Args:
            path: Directory path to save model
        """
        path = Path(path)
        path.mkdir(parents=True, exist_ok=True)

        # Save CNN encoder
        if self.cnn_encoder is not None:
            torch.save(
                self.cnn_encoder.state_dict(),
                path / 'cnn_encoder.pt'
            )

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

        if self.regressor is not None:
            joblib.dump(self.regressor, path / 'regressor.joblib')

        # Save metadata
        metadata = {
            'config': self.config,
            'feature_columns': self.feature_columns,
            'n_features': self.n_features,
            'is_fitted': self._is_fitted
        }
        joblib.dump(metadata, path / 'metadata.joblib')

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

    def load(self, path: str) -> None:
        """
        Load model from disk.

        Args:
            path: Directory path containing saved model
        """
        path = Path(path)

        # Load metadata
        metadata = joblib.load(path / 'metadata.joblib')
        self.config = metadata['config']
        self.feature_columns = metadata['feature_columns']
        self.n_features = metadata['n_features']
        self._is_fitted = metadata['is_fitted']

        # Load CNN encoder
        if (path / 'cnn_encoder.pt').exists():
            self._initialize_cnn(self.n_features)
            self.cnn_encoder.load_state_dict(
                torch.load(path / 'cnn_encoder.pt', map_location=self.device)
            )
            self.cnn_encoder.eval()

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

        if (path / 'regressor.joblib').exists():
            self.regressor = joblib.load(path / 'regressor.joblib')

        logger.info(f"Model loaded from {path}")


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

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

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

    # Simulate price with breakout patterns
    volatility = np.where(
        (np.arange(n) % 100 > 85),  # High vol after compression
        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)

    # Split data
    train_size = 800
    df_train = df.iloc[:train_size]
    df_test = df.iloc[train_size:]

    print(f"\n1. Initializing model...")
    config = VBPModelConfig(
        sequence_length=30,
        xgb_n_estimators=50,  # Reduced for testing
        use_gpu=False  # CPU for testing
    )
    model = VBPModel(config)

    print(f"\n2. Training model...")
    metrics = model.fit(df_train, verbose=True)

    print(f"\n3. Making predictions...")
    predictions = model.predict(df_test)
    print(f"   Generated {len(predictions)} predictions")

    # Sample predictions
    print(f"\n4. Sample predictions:")
    for i in [0, 50, 100, 150]:
        if i < len(predictions):
            pred = predictions[i]
            print(f"   Sample {i}: {pred.signal} "
                  f"(prob={pred.breakout_probability:.3f}, "
                  f"dir={pred.direction}, mag={pred.magnitude:.2f})")

    print(f"\n5. Feature importance (top 10):")
    importance = model.get_feature_importance(top_n=10)
    for name, score in importance.items():
        print(f"   {name}: {score:.4f}")

    print(f"\n6. Testing single prediction...")
    single_pred = model.predict_single(df_test)
    print(f"   Latest signal: {single_pred.signal}")
    print(f"   Prediction: {single_pred.to_dict()}")

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