"""
MSA (Market Structure Analysis) - Model
========================================
XGBoost-based model for market structure predictions.

Predictions:
- Next BOS direction (bullish/bearish/neutral)
- POI reaction probability
- Structure continuation probability

Optional GNN component for swing point relationships.

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
from loguru import logger
import joblib
from enum import IntEnum

try:
    from xgboost import XGBClassifier
    HAS_XGBOOST = True
except ImportError:
    HAS_XGBOOST = False
    logger.warning("XGBoost not available, MSAModel will have limited functionality")

try:
    import torch
    import torch.nn as nn
    import torch.nn.functional as F
    HAS_TORCH = True
except ImportError:
    HAS_TORCH = False
    logger.warning("PyTorch not available, GNN component disabled")

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


class BOSDirection(IntEnum):
    """Break of Structure direction"""
    NEUTRAL = 0
    BULLISH = 1
    BEARISH = 2


@dataclass
class MSAPrediction:
    """MSA model prediction result"""
    # BOS prediction
    next_bos_direction: str  # 'bullish', 'bearish', 'neutral'
    bos_confidence: float
    bos_probabilities: Dict[str, float]

    # POI reaction
    poi_reaction_prob: float
    poi_reaction_direction: str  # 'bullish', 'bearish', 'none'

    # Structure continuation
    structure_continuation_prob: float
    expected_structure: str  # 'bullish_continuation', 'bearish_continuation', 'reversal'

    # Metadata
    timestamp: Optional[Any] = None
    features_used: List[str] = field(default_factory=list)

    def to_dict(self) -> Dict[str, Any]:
        return {
            'next_bos_direction': self.next_bos_direction,
            'bos_confidence': float(self.bos_confidence),
            'bos_probabilities': self.bos_probabilities,
            'poi_reaction_prob': float(self.poi_reaction_prob),
            'poi_reaction_direction': self.poi_reaction_direction,
            'structure_continuation_prob': float(self.structure_continuation_prob),
            'expected_structure': self.expected_structure,
            'timestamp': str(self.timestamp) if self.timestamp else None
        }

    @property
    def trading_bias(self) -> str:
        """Get overall trading bias"""
        if self.bos_confidence >= 0.7:
            if self.next_bos_direction == 'bullish':
                return 'LONG_BIAS'
            elif self.next_bos_direction == 'bearish':
                return 'SHORT_BIAS'
        return 'NEUTRAL'

    @property
    def signal_strength(self) -> float:
        """Calculate overall signal strength 0-1"""
        # Combine BOS confidence, POI reaction, and continuation
        bos_factor = self.bos_confidence if self.next_bos_direction != 'neutral' else 0
        poi_factor = self.poi_reaction_prob * 0.5
        cont_factor = self.structure_continuation_prob * 0.3

        return min(1.0, bos_factor + poi_factor + cont_factor)


@dataclass
class MSAMetrics:
    """Metrics for MSA model evaluation"""
    # BOS direction metrics
    bos_accuracy: float = 0.0
    bos_macro_f1: float = 0.0
    bos_weighted_f1: float = 0.0
    bos_per_class_f1: Dict[str, float] = field(default_factory=dict)
    bos_confusion_matrix: Optional[np.ndarray] = None

    # POI reaction metrics
    poi_accuracy: float = 0.0
    poi_precision: float = 0.0
    poi_recall: float = 0.0
    poi_f1: float = 0.0

    # Structure continuation metrics
    continuation_accuracy: float = 0.0
    continuation_f1: float = 0.0

    n_samples: int = 0


class SwingPointGNN(nn.Module):
    """
    Graph Neural Network for swing point relationships.

    Models the relationships between swing points to capture
    market structure patterns.
    """

    def __init__(
        self,
        input_dim: int = 8,
        hidden_dim: int = 32,
        output_dim: int = 16,
        n_layers: int = 2,
        dropout: float = 0.1
    ):
        """
        Initialize Swing Point GNN.

        Args:
            input_dim: Input feature dimension per node
            hidden_dim: Hidden layer dimension
            output_dim: Output embedding dimension
            n_layers: Number of GNN layers
            dropout: Dropout probability
        """
        super().__init__()

        if not HAS_TORCH:
            raise ImportError("PyTorch is required for SwingPointGNN")

        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.output_dim = output_dim
        self.n_layers = n_layers

        # Node embedding
        self.node_embed = nn.Linear(input_dim, hidden_dim)

        # GNN layers (simple message passing)
        self.gnn_layers = nn.ModuleList([
            nn.Linear(hidden_dim * 2, hidden_dim)
            for _ in range(n_layers)
        ])

        # Layer norms
        self.layer_norms = nn.ModuleList([
            nn.LayerNorm(hidden_dim)
            for _ in range(n_layers)
        ])

        # Output projection
        self.output_proj = nn.Linear(hidden_dim, output_dim)

        self.dropout = nn.Dropout(dropout)
        self.activation = nn.GELU()

    def forward(
        self,
        node_features: torch.Tensor,
        edge_index: torch.Tensor
    ) -> torch.Tensor:
        """
        Forward pass.

        Args:
            node_features: [n_nodes, input_dim] node feature matrix
            edge_index: [2, n_edges] edge index tensor

        Returns:
            [n_nodes, output_dim] node embeddings
        """
        # Initial embedding
        x = self.node_embed(node_features)
        x = self.activation(x)

        # Message passing layers
        for i, (gnn_layer, layer_norm) in enumerate(zip(self.gnn_layers, self.layer_norms)):
            # Aggregate neighbor features
            row, col = edge_index
            neighbor_features = x[col]

            # Create edge features by concatenating source and target
            source_features = x[row]
            edge_features = torch.cat([source_features, neighbor_features], dim=-1)

            # Transform
            messages = gnn_layer(edge_features)
            messages = self.activation(messages)

            # Aggregate messages per node
            aggregated = torch.zeros_like(x)
            aggregated.index_add_(0, row, messages)

            # Combine with residual
            x = x + self.dropout(aggregated)
            x = layer_norm(x)

        # Output projection
        return self.output_proj(x)


class MSAModel:
    """
    Market Structure Analysis Model.

    Uses XGBoost as the main classifier with optional GNN
    for swing point relationship modeling.

    Predictions:
    - next_bos_direction: Direction of next BOS (bullish/bearish/neutral)
    - poi_reaction_prob: Probability of price reacting at current POI
    - structure_continuation: Probability of structure continuation vs reversal
    """

    BOS_CLASSES = {0: 'neutral', 1: 'bullish', 2: 'bearish'}

    def __init__(
        self,
        config: Optional[Dict[str, Any]] = None,
        use_gpu: bool = True,
        use_gnn: bool = False,
        gnn_config: Optional[Dict[str, Any]] = None
    ):
        """
        Initialize MSA Model.

        Args:
            config: XGBoost configuration
            use_gpu: Use GPU acceleration for XGBoost
            use_gnn: Enable GNN component for swing point relationships
            gnn_config: GNN configuration (if use_gnn=True)
        """
        if not HAS_XGBOOST:
            raise ImportError("XGBoost is required for MSAModel")

        self.config = config or self._default_config()
        self.use_gpu = use_gpu
        self.use_gnn = use_gnn and HAS_TORCH

        # Initialize XGBoost classifiers
        self.bos_classifier: Optional[XGBClassifier] = None
        self.poi_classifier: Optional[XGBClassifier] = None
        self.continuation_classifier: Optional[XGBClassifier] = None

        # GNN component
        self.gnn: Optional[SwingPointGNN] = None
        if self.use_gnn:
            gnn_cfg = gnn_config or {}
            self.gnn = SwingPointGNN(**gnn_cfg)
            logger.info("GNN component enabled for swing point relationships")

        # Feature columns
        self.feature_columns: List[str] = []

        # Training state
        self._is_trained = False
        self.metrics: Optional[MSAMetrics] = None

        self._init_classifiers()
        logger.info(f"MSAModel initialized: use_gpu={use_gpu}, use_gnn={use_gnn}")

    def _default_config(self) -> Dict[str, Any]:
        """Default XGBoost configuration"""
        return {
            'n_estimators': 200,
            'max_depth': 6,
            'learning_rate': 0.05,
            'subsample': 0.8,
            'colsample_bytree': 0.8,
            'min_child_weight': 3,
            'gamma': 0.1,
            'reg_alpha': 0.1,
            'reg_lambda': 1.0,
            'tree_method': 'hist',
            'random_state': 42
        }

    def _init_classifiers(self):
        """Initialize XGBoost classifiers"""
        params = self.config.copy()

        # GPU configuration
        if self.use_gpu:
            try:
                import torch
                if torch.cuda.is_available():
                    params['device'] = 'cuda'
                    params['tree_method'] = 'gpu_hist'
                    logger.info("GPU acceleration enabled for MSAModel")
            except ImportError:
                pass

        # BOS direction classifier (3 classes)
        bos_params = params.copy()
        bos_params['objective'] = 'multi:softprob'
        bos_params['num_class'] = 3
        self.bos_classifier = XGBClassifier(**bos_params)

        # POI reaction classifier (binary)
        poi_params = params.copy()
        poi_params['objective'] = 'binary:logistic'
        poi_params.pop('num_class', None)
        self.poi_classifier = XGBClassifier(**poi_params)

        # Structure continuation classifier (binary)
        cont_params = params.copy()
        cont_params['objective'] = 'binary:logistic'
        cont_params.pop('num_class', None)
        self.continuation_classifier = XGBClassifier(**cont_params)

    def forward(
        self,
        X: Union[pd.DataFrame, np.ndarray],
        swing_graph: Optional[Tuple[torch.Tensor, torch.Tensor]] = None
    ) -> List[MSAPrediction]:
        """
        Generate predictions for input features.

        Args:
            X: Feature matrix [n_samples, n_features]
            swing_graph: Optional tuple of (node_features, edge_index) for GNN

        Returns:
            List of MSAPrediction objects
        """
        if not self._is_trained:
            raise RuntimeError("Model must be trained before prediction")

        if isinstance(X, pd.DataFrame):
            X_arr = X.values
        else:
            X_arr = X

        # Get classifier predictions
        bos_proba = self.bos_classifier.predict_proba(X_arr)
        poi_proba = self.poi_classifier.predict_proba(X_arr)[:, 1]
        cont_proba = self.continuation_classifier.predict_proba(X_arr)[:, 1]

        # Optional: Incorporate GNN features
        if self.use_gnn and swing_graph is not None and self.gnn is not None:
            node_features, edge_index = swing_graph
            with torch.no_grad():
                gnn_embeddings = self.gnn(node_features, edge_index)
                # Could use GNN embeddings to adjust predictions
                # For now, just log
                logger.debug(f"GNN embeddings shape: {gnn_embeddings.shape}")

        # Generate predictions
        predictions = []
        for i in range(len(X_arr)):
            # BOS direction
            bos_class = int(np.argmax(bos_proba[i]))
            bos_direction = self.BOS_CLASSES[bos_class]
            bos_conf = float(bos_proba[i].max())

            # POI reaction
            poi_prob = float(poi_proba[i])
            if poi_prob > 0.5:
                # Determine direction based on structure features
                poi_direction = 'bullish' if bos_direction == 'bullish' else 'bearish'
            else:
                poi_direction = 'none'

            # Structure continuation
            cont_prob = float(cont_proba[i])
            if cont_prob > 0.5:
                expected = f"{bos_direction}_continuation" if bos_direction != 'neutral' else 'neutral'
            else:
                expected = 'reversal'

            pred = MSAPrediction(
                next_bos_direction=bos_direction,
                bos_confidence=bos_conf,
                bos_probabilities={
                    'neutral': float(bos_proba[i][0]),
                    'bullish': float(bos_proba[i][1]),
                    'bearish': float(bos_proba[i][2])
                },
                poi_reaction_prob=poi_prob,
                poi_reaction_direction=poi_direction,
                structure_continuation_prob=cont_prob,
                expected_structure=expected,
                timestamp=X.index[i] if isinstance(X, pd.DataFrame) else None,
                features_used=self.feature_columns
            )
            predictions.append(pred)

        return predictions

    def fit(
        self,
        X_train: Union[pd.DataFrame, np.ndarray],
        y_bos: np.ndarray,
        y_poi: np.ndarray,
        y_continuation: np.ndarray,
        X_val: Optional[Union[pd.DataFrame, np.ndarray]] = None,
        y_bos_val: Optional[np.ndarray] = None,
        y_poi_val: Optional[np.ndarray] = None,
        y_continuation_val: Optional[np.ndarray] = None,
        sample_weight: Optional[np.ndarray] = None,
        verbose: bool = True
    ) -> MSAMetrics:
        """
        Train all MSA classifiers.

        Args:
            X_train: Training features
            y_bos: BOS direction labels (0=neutral, 1=bullish, 2=bearish)
            y_poi: POI reaction labels (0/1)
            y_continuation: Structure continuation labels (0/1)
            X_val: Validation features (optional)
            y_bos_val: Validation BOS labels
            y_poi_val: Validation POI labels
            y_continuation_val: Validation continuation labels
            sample_weight: Sample weights for training
            verbose: Print training progress

        Returns:
            Training metrics
        """
        if isinstance(X_train, pd.DataFrame):
            self.feature_columns = X_train.columns.tolist()
            X_train_arr = X_train.values
        else:
            X_train_arr = X_train
            self.feature_columns = [f"feature_{i}" for i in range(X_train.shape[1])]

        logger.info(f"Training MSAModel on {len(X_train_arr)} samples with {X_train_arr.shape[1]} features")

        # Prepare fit parameters
        fit_params = {}
        if sample_weight is not None:
            fit_params['sample_weight'] = sample_weight

        # Add validation set if provided
        eval_set = None
        if X_val is not None:
            X_val_arr = X_val.values if isinstance(X_val, pd.DataFrame) else X_val
            eval_set = [(X_val_arr, None)]  # Placeholder, updated per classifier

        # Train BOS classifier
        logger.info("Training BOS direction classifier...")
        bos_fit_params = fit_params.copy()
        if X_val is not None and y_bos_val is not None:
            bos_fit_params['eval_set'] = [(X_val.values if isinstance(X_val, pd.DataFrame) else X_val, y_bos_val)]
        self.bos_classifier.fit(X_train_arr, y_bos, **bos_fit_params)

        # Train POI classifier
        logger.info("Training POI reaction classifier...")
        poi_fit_params = fit_params.copy()
        if X_val is not None and y_poi_val is not None:
            poi_fit_params['eval_set'] = [(X_val.values if isinstance(X_val, pd.DataFrame) else X_val, y_poi_val)]
        self.poi_classifier.fit(X_train_arr, y_poi, **poi_fit_params)

        # Train continuation classifier
        logger.info("Training structure continuation classifier...")
        cont_fit_params = fit_params.copy()
        if X_val is not None and y_continuation_val is not None:
            cont_fit_params['eval_set'] = [(X_val.values if isinstance(X_val, pd.DataFrame) else X_val, y_continuation_val)]
        self.continuation_classifier.fit(X_train_arr, y_continuation, **cont_fit_params)

        self._is_trained = True

        # Calculate metrics on training data
        self.metrics = self._calculate_metrics(
            X_train_arr, y_bos, y_poi, y_continuation
        )

        if verbose:
            self._print_metrics(self.metrics, "Training")

        # Calculate validation metrics if provided
        if X_val is not None:
            val_metrics = self._calculate_metrics(
                X_val.values if isinstance(X_val, pd.DataFrame) else X_val,
                y_bos_val, y_poi_val, y_continuation_val
            )
            if verbose:
                self._print_metrics(val_metrics, "Validation")

        return self.metrics

    def predict(
        self,
        X: Union[pd.DataFrame, np.ndarray]
    ) -> List[MSAPrediction]:
        """
        Generate predictions (alias for forward).

        Args:
            X: Feature matrix

        Returns:
            List of MSAPrediction objects
        """
        return self.forward(X)

    def predict_single(
        self,
        X: Union[pd.DataFrame, np.ndarray]
    ) -> MSAPrediction:
        """
        Predict for the last row of features.

        Args:
            X: Feature matrix (uses last row)

        Returns:
            Single MSAPrediction
        """
        if isinstance(X, pd.DataFrame):
            X = X.tail(1)
        else:
            X = X[-1:, :]

        predictions = self.forward(X)
        return predictions[-1] if predictions else None

    def _calculate_metrics(
        self,
        X: np.ndarray,
        y_bos: np.ndarray,
        y_poi: np.ndarray,
        y_continuation: np.ndarray
    ) -> MSAMetrics:
        """Calculate metrics for all classifiers"""
        # BOS metrics
        bos_pred = self.bos_classifier.predict(X)
        bos_precision, bos_recall, bos_f1, _ = precision_recall_fscore_support(
            y_bos, bos_pred, average=None, zero_division=0
        )

        bos_per_class_f1 = {}
        for i, name in self.BOS_CLASSES.items():
            if i < len(bos_f1):
                bos_per_class_f1[name] = float(bos_f1[i])

        # POI metrics
        poi_pred = self.poi_classifier.predict(X)
        poi_precision, poi_recall, poi_f1, _ = precision_recall_fscore_support(
            y_poi, poi_pred, average='binary', zero_division=0
        )

        # Continuation metrics
        cont_pred = self.continuation_classifier.predict(X)
        cont_f1 = f1_score(y_continuation, cont_pred, zero_division=0)

        return MSAMetrics(
            bos_accuracy=accuracy_score(y_bos, bos_pred),
            bos_macro_f1=f1_score(y_bos, bos_pred, average='macro', zero_division=0),
            bos_weighted_f1=f1_score(y_bos, bos_pred, average='weighted', zero_division=0),
            bos_per_class_f1=bos_per_class_f1,
            bos_confusion_matrix=confusion_matrix(y_bos, bos_pred),
            poi_accuracy=accuracy_score(y_poi, poi_pred),
            poi_precision=float(poi_precision),
            poi_recall=float(poi_recall),
            poi_f1=float(poi_f1),
            continuation_accuracy=accuracy_score(y_continuation, cont_pred),
            continuation_f1=float(cont_f1),
            n_samples=len(X)
        )

    def _print_metrics(self, metrics: MSAMetrics, prefix: str = ""):
        """Print metrics summary"""
        logger.info(f"\n{prefix} Metrics (n={metrics.n_samples}):")
        logger.info(f"  BOS Direction:")
        logger.info(f"    Accuracy: {metrics.bos_accuracy:.2%}")
        logger.info(f"    Macro F1: {metrics.bos_macro_f1:.4f}")
        logger.info(f"    Per-class F1: {metrics.bos_per_class_f1}")
        logger.info(f"  POI Reaction:")
        logger.info(f"    Accuracy: {metrics.poi_accuracy:.2%}")
        logger.info(f"    F1: {metrics.poi_f1:.4f}")
        logger.info(f"  Structure Continuation:")
        logger.info(f"    Accuracy: {metrics.continuation_accuracy:.2%}")
        logger.info(f"    F1: {metrics.continuation_f1:.4f}")

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

        Args:
            top_n: Number of top features to return

        Returns:
            Dictionary with feature importance per classifier
        """
        if not self._is_trained:
            return {}

        result = {}

        # BOS classifier importance
        bos_imp = dict(zip(self.feature_columns, self.bos_classifier.feature_importances_))
        result['bos'] = dict(sorted(bos_imp.items(), key=lambda x: x[1], reverse=True)[:top_n])

        # POI classifier importance
        poi_imp = dict(zip(self.feature_columns, self.poi_classifier.feature_importances_))
        result['poi'] = dict(sorted(poi_imp.items(), key=lambda x: x[1], reverse=True)[:top_n])

        # Continuation classifier importance
        cont_imp = dict(zip(self.feature_columns, self.continuation_classifier.feature_importances_))
        result['continuation'] = dict(sorted(cont_imp.items(), key=lambda x: x[1], reverse=True)[:top_n])

        return result

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

        # Save classifiers
        joblib.dump(self.bos_classifier, path / 'bos_classifier.joblib')
        joblib.dump(self.poi_classifier, path / 'poi_classifier.joblib')
        joblib.dump(self.continuation_classifier, path / 'continuation_classifier.joblib')

        # Save GNN if present
        if self.use_gnn and self.gnn is not None:
            torch.save(self.gnn.state_dict(), path / 'gnn.pt')

        # Save metadata
        metadata = {
            'config': self.config,
            'feature_columns': self.feature_columns,
            'use_gnn': self.use_gnn,
            'metrics': {
                'bos_accuracy': self.metrics.bos_accuracy if self.metrics else None,
                'bos_macro_f1': self.metrics.bos_macro_f1 if self.metrics else None,
                'poi_f1': self.metrics.poi_f1 if self.metrics else None,
                'continuation_f1': self.metrics.continuation_f1 if self.metrics else None
            },
            'version': '1.0.0'
        }
        joblib.dump(metadata, path / 'metadata.joblib')

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

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

        # Load classifiers
        self.bos_classifier = joblib.load(path / 'bos_classifier.joblib')
        self.poi_classifier = joblib.load(path / 'poi_classifier.joblib')
        self.continuation_classifier = joblib.load(path / 'continuation_classifier.joblib')

        # Load GNN if present
        if self.use_gnn and (path / 'gnn.pt').exists():
            self.gnn.load_state_dict(torch.load(path / 'gnn.pt'))

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

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


if __name__ == "__main__":
    # Test MSAModel
    np.random.seed(42)

    # Create sample data
    n_samples = 1000
    n_features = 50

    X = np.random.randn(n_samples, n_features)
    y_bos = np.random.randint(0, 3, n_samples)  # 0=neutral, 1=bullish, 2=bearish
    y_poi = np.random.randint(0, 2, n_samples)  # 0=no reaction, 1=reaction
    y_continuation = np.random.randint(0, 2, n_samples)  # 0=reversal, 1=continuation

    # Split data
    train_size = 800
    X_train, X_val = X[:train_size], X[train_size:]
    y_bos_train, y_bos_val = y_bos[:train_size], y_bos[train_size:]
    y_poi_train, y_poi_val = y_poi[:train_size], y_poi[train_size:]
    y_cont_train, y_cont_val = y_continuation[:train_size], y_continuation[train_size:]

    # Create and train model
    model = MSAModel(use_gpu=False, use_gnn=False)
    metrics = model.fit(
        X_train, y_bos_train, y_poi_train, y_cont_train,
        X_val, y_bos_val, y_poi_val, y_cont_val
    )

    # Generate predictions
    predictions = model.predict(X_val[:5])

    print("\n=== Sample Predictions ===")
    for i, pred in enumerate(predictions):
        print(f"Sample {i}: BOS={pred.next_bos_direction} ({pred.bos_confidence:.2f}), "
              f"POI_react={pred.poi_reaction_prob:.2f}, "
              f"Continuation={pred.structure_continuation_prob:.2f}, "
              f"Bias={pred.trading_bias}")

    # Feature importance
    print("\n=== Top Features (BOS) ===")
    importance = model.get_feature_importance(10)
    for feat, imp in list(importance['bos'].items())[:5]:
        print(f"  {feat}: {imp:.4f}")
