#!/usr/bin/env python3
"""
MTS Model - Multi-Timeframe Synthesis Complete Model
=====================================================
Complete MTS model combining Hierarchical Attention encoder with
XGBoost for final prediction. Predicts unified direction, confidence
based on alignment, and optimal entry timeframe.

Architecture:
1. MTSFeatureEngineer: Compute multi-TF features
2. HierarchicalAttention: Learn TF relationships -> unified representation
3. XGBoost: Final prediction (direction, confidence, entry TF)

Outputs:
- unified_direction: -1 (bearish) to 1 (bullish)
- confidence_by_alignment: 0 to 1 based on TF alignment
- optimal_entry_tf: Best timeframe for entry (0=5m, 1=15m, 2=1h, 3=4h)

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

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
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:
    from xgboost import XGBClassifier, XGBRegressor
    HAS_XGBOOST = True
except ImportError:
    HAS_XGBOOST = False
    logger.warning("XGBoost not available - using PyTorch fallback")

from .feature_engineering import MTSFeatureEngineer
from .hierarchical_attention import (
    HierarchicalAttention,
    HierarchicalAttentionWithAlignment,
    HierarchicalAttentionConfig
)


@dataclass
class MTSConfig:
    """Configuration for MTS Model."""
    # Feature engineering
    timeframes: Tuple[str, ...] = ('5m', '15m', '1h', '4h')
    include_volume: bool = True
    feature_lookback: int = 100

    # Hierarchical attention
    d_model: int = 128
    n_heads: int = 4
    d_ff: int = 256
    n_layers: int = 2
    dropout: float = 0.1
    max_seq_len: int = 256
    use_cross_tf_attention: bool = True
    use_learnable_tf_weights: bool = True

    # XGBoost
    xgb_n_estimators: int = 200
    xgb_max_depth: int = 6
    xgb_learning_rate: float = 0.05
    xgb_subsample: float = 0.8
    xgb_colsample_bytree: float = 0.8
    xgb_min_child_weight: int = 5
    xgb_reg_alpha: float = 0.1
    xgb_reg_lambda: float = 1.0

    # Prediction thresholds
    direction_threshold: float = 0.6
    confidence_threshold: float = 0.65


@dataclass
class MTSPrediction:
    """MTS Model prediction output."""
    unified_direction: float  # -1 to 1
    direction_class: str  # 'bullish', 'bearish', 'neutral'
    confidence: float  # 0 to 1
    confidence_by_alignment: float  # Based on TF alignment
    optimal_entry_tf: str  # '5m', '15m', '1h', '4h'
    tf_contributions: Dict[str, float]  # Contribution of each TF
    signal_strength: float  # 0 to 1
    recommended_action: str  # 'buy', 'sell', 'hold'

    def to_dict(self) -> Dict[str, Any]:
        return {
            'unified_direction': float(self.unified_direction),
            'direction_class': self.direction_class,
            'confidence': float(self.confidence),
            'confidence_by_alignment': float(self.confidence_by_alignment),
            'optimal_entry_tf': self.optimal_entry_tf,
            'tf_contributions': self.tf_contributions,
            'signal_strength': float(self.signal_strength),
            'recommended_action': self.recommended_action
        }


class MTSHead(nn.Module):
    """
    Prediction head for MTS model.

    Takes unified representation from HierarchicalAttention and produces:
    - Direction logits
    - Confidence score
    - Entry TF logits
    """

    def __init__(self, d_model: int, n_timeframes: int, dropout: float = 0.1):
        super().__init__()

        self.d_model = d_model
        self.n_timeframes = n_timeframes

        # Shared backbone
        self.backbone = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_model, d_model // 2),
            nn.GELU(),
            nn.Dropout(dropout)
        )

        # Direction head (3 classes: bearish, neutral, bullish)
        self.direction_head = nn.Sequential(
            nn.Linear(d_model // 2, d_model // 4),
            nn.GELU(),
            nn.Linear(d_model // 4, 3)
        )

        # Confidence head (regression 0-1)
        self.confidence_head = nn.Sequential(
            nn.Linear(d_model // 2, d_model // 4),
            nn.GELU(),
            nn.Linear(d_model // 4, 1),
            nn.Sigmoid()
        )

        # Entry TF head (classification)
        self.entry_tf_head = nn.Sequential(
            nn.Linear(d_model // 2, d_model // 4),
            nn.GELU(),
            nn.Linear(d_model // 4, n_timeframes)
        )

    def forward(
        self,
        x: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Forward pass.

        Args:
            x: Unified representation (batch, d_model)

        Returns:
            direction_logits: (batch, 3)
            confidence: (batch, 1)
            entry_tf_logits: (batch, n_timeframes)
        """
        features = self.backbone(x)

        direction_logits = self.direction_head(features)
        confidence = self.confidence_head(features)
        entry_tf_logits = self.entry_tf_head(features)

        return direction_logits, confidence, entry_tf_logits


class MTSModel:
    """
    Multi-Timeframe Synthesis Model.

    Complete model for multi-timeframe analysis combining:
    1. Feature engineering across timeframes
    2. Hierarchical attention for unified representation
    3. Prediction heads (or XGBoost) for final output

    Usage:
        model = MTSModel(MTSConfig())
        model.train(symbol='EURUSD')
        prediction = model.predict(df_5m)
    """

    def __init__(
        self,
        config: Optional[MTSConfig] = None,
        use_xgboost: bool = True,
        device: Optional[str] = None
    ):
        """
        Initialize MTS Model.

        Args:
            config: Model configuration
            use_xgboost: Whether to use XGBoost for final prediction
            device: Device for PyTorch ('cuda', 'cpu', or None for auto)
        """
        self.config = config or MTSConfig()
        self.use_xgboost = use_xgboost and HAS_XGBOOST
        self.device = device or ('cuda' if torch.cuda.is_available() else 'cpu')

        # Initialize feature engineer
        self.feature_engineer = MTSFeatureEngineer(
            timeframes=list(self.config.timeframes),
            include_volume=self.config.include_volume,
            feature_lookback=self.config.feature_lookback
        )

        # Store input dimensions after first feature computation
        self.input_dims: Optional[Dict[str, int]] = None

        # Initialize attention model (will be built on first train/predict)
        self.attention_model: Optional[HierarchicalAttentionWithAlignment] = None

        # Initialize prediction components
        if self.use_xgboost:
            self._init_xgboost()
        else:
            self.prediction_head: Optional[MTSHead] = None

        self._is_trained = False

        logger.info(f"Initialized MTSModel with device={self.device}, xgboost={self.use_xgboost}")

    def _init_xgboost(self):
        """Initialize XGBoost models."""
        xgb_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,
            'min_child_weight': self.config.xgb_min_child_weight,
            'reg_alpha': self.config.xgb_reg_alpha,
            'reg_lambda': self.config.xgb_reg_lambda,
            'random_state': 42,
            'n_jobs': -1
        }

        # Check for GPU
        try:
            if torch.cuda.is_available():
                xgb_params['tree_method'] = 'gpu_hist'
                xgb_params['device'] = 'cuda'
        except Exception:
            pass

        # Direction classifier (3 classes)
        self.xgb_direction = XGBClassifier(
            **xgb_params,
            objective='multi:softprob',
            num_class=3
        )

        # Confidence regressor
        self.xgb_confidence = XGBRegressor(
            **xgb_params,
            objective='reg:squarederror'
        )

        # Entry TF classifier
        self.xgb_entry_tf = XGBClassifier(
            **xgb_params,
            objective='multi:softprob',
            num_class=len(self.config.timeframes)
        )

    def _build_attention_model(self, input_dims: Dict[str, int]):
        """Build hierarchical attention model."""
        attention_config = HierarchicalAttentionConfig(
            d_model=self.config.d_model,
            n_heads=self.config.n_heads,
            d_ff=self.config.d_ff,
            n_layers=self.config.n_layers,
            dropout=self.config.dropout,
            max_seq_len=self.config.max_seq_len,
            timeframes=self.config.timeframes,
            use_cross_tf_attention=self.config.use_cross_tf_attention,
            use_learnable_tf_weights=self.config.use_learnable_tf_weights
        )

        self.attention_model = HierarchicalAttentionWithAlignment(
            attention_config,
            input_dims
        ).to(self.device)

        if not self.use_xgboost:
            self.prediction_head = MTSHead(
                self.config.d_model,
                len(self.config.timeframes),
                self.config.dropout
            ).to(self.device)

        self.input_dims = input_dims

    def _prepare_features(
        self,
        df_5m: pd.DataFrame
    ) -> Tuple[Dict[str, np.ndarray], np.ndarray, np.ndarray]:
        """
        Prepare features from 5m data.

        Returns:
            tf_features: Dictionary of {timeframe: features}
            alignment: Alignment scores
            conflict: Conflict scores
        """
        # Aggregate to all timeframes
        dfs_dict = {}
        for tf in self.config.timeframes:
            dfs_dict[tf] = self.feature_engineer.aggregate_to_timeframe(df_5m, tf)

        # Prepare hierarchical input
        hierarchical = self.feature_engineer.prepare_hierarchical_input(dfs_dict)

        return (
            hierarchical['tf_features'],
            hierarchical['alignment'],
            hierarchical['conflict']
        )

    def _extract_unified_representation(
        self,
        tf_features: Dict[str, np.ndarray],
        alignment: np.ndarray,
        conflict: np.ndarray
    ) -> np.ndarray:
        """
        Extract unified representation using attention model.

        Args:
            tf_features: Per-TF features
            alignment: Alignment scores
            conflict: Conflict scores

        Returns:
            Unified representation array
        """
        # Determine input dimensions if not set
        if self.input_dims is None:
            input_dims = {tf: features.shape[-1] for tf, features in tf_features.items()}
            self._build_attention_model(input_dims)

        self.attention_model.eval()

        with torch.no_grad():
            # Convert to tensors
            tf_tensors = {}
            for tf, features in tf_features.items():
                tensor = torch.tensor(features, dtype=torch.float32, device=self.device)
                if tensor.dim() == 2:
                    tensor = tensor.unsqueeze(0)
                tf_tensors[tf] = tensor

            alignment_tensor = torch.tensor(alignment, dtype=torch.float32, device=self.device)
            conflict_tensor = torch.tensor(conflict, dtype=torch.float32, device=self.device)

            if alignment_tensor.dim() == 1:
                alignment_tensor = alignment_tensor.unsqueeze(0)
            if conflict_tensor.dim() == 1:
                conflict_tensor = conflict_tensor.unsqueeze(0)

            # Get unified representation
            unified = self.attention_model(
                tf_tensors,
                alignment_tensor,
                conflict_tensor
            )

            return unified.cpu().numpy()

    def forward(
        self,
        tf_features: Dict[str, torch.Tensor],
        alignment: torch.Tensor,
        conflict: torch.Tensor
    ) -> Dict[str, torch.Tensor]:
        """
        Full forward pass (PyTorch mode).

        Args:
            tf_features: Per-TF feature tensors
            alignment: Alignment scores
            conflict: Conflict scores

        Returns:
            Dictionary with predictions
        """
        if self.attention_model is None:
            raise RuntimeError("Model not initialized. Call train() first.")

        # Get unified representation
        unified = self.attention_model(tf_features, alignment, conflict)

        if self.use_xgboost:
            # Can't use XGBoost in PyTorch forward
            raise RuntimeError("Use predict() for XGBoost mode")

        # Use PyTorch prediction head
        direction_logits, confidence, entry_tf_logits = self.prediction_head(unified)

        return {
            'direction_logits': direction_logits,
            'confidence': confidence,
            'entry_tf_logits': entry_tf_logits,
            'unified_representation': unified
        }

    def predict(
        self,
        df_5m: pd.DataFrame
    ) -> MTSPrediction:
        """
        Make prediction from 5m data.

        Args:
            df_5m: 5-minute OHLCV DataFrame

        Returns:
            MTSPrediction object
        """
        if not self._is_trained:
            raise RuntimeError("Model must be trained before prediction")

        # Prepare features
        tf_features, alignment, conflict = self._prepare_features(df_5m)

        # Get unified representation
        unified_rep = self._extract_unified_representation(tf_features, alignment, conflict)

        if self.use_xgboost:
            # XGBoost predictions
            direction_probs = self.xgb_direction.predict_proba(unified_rep)[0]
            confidence = float(self.xgb_confidence.predict(unified_rep)[0])
            entry_tf_probs = self.xgb_entry_tf.predict_proba(unified_rep)[0]

            # Direction class (0=bearish, 1=neutral, 2=bullish)
            direction_class_idx = int(np.argmax(direction_probs))

            # Unified direction (-1 to 1)
            unified_direction = direction_probs[2] - direction_probs[0]

            # Optimal entry TF
            optimal_tf_idx = int(np.argmax(entry_tf_probs))
        else:
            # PyTorch predictions
            self.attention_model.eval()
            self.prediction_head.eval()

            with torch.no_grad():
                tf_tensors = {
                    tf: torch.tensor(features, dtype=torch.float32, device=self.device).unsqueeze(0)
                    for tf, features in tf_features.items()
                }
                alignment_tensor = torch.tensor(alignment, dtype=torch.float32, device=self.device).unsqueeze(0)
                conflict_tensor = torch.tensor(conflict, dtype=torch.float32, device=self.device).unsqueeze(0)

                unified = self.attention_model(tf_tensors, alignment_tensor, conflict_tensor)
                direction_logits, conf, entry_logits = self.prediction_head(unified)

                direction_probs = torch.softmax(direction_logits, dim=-1)[0].cpu().numpy()
                confidence = float(conf[0].cpu().numpy())
                entry_tf_probs = torch.softmax(entry_logits, dim=-1)[0].cpu().numpy()

                direction_class_idx = int(np.argmax(direction_probs))
                unified_direction = float(direction_probs[2] - direction_probs[0])
                optimal_tf_idx = int(np.argmax(entry_tf_probs))

        # Get TF contributions
        tf_contributions = self.get_tf_contributions()

        # Confidence by alignment
        confidence_by_alignment = float(np.mean(alignment)) * confidence

        # Map direction class
        direction_classes = ['bearish', 'neutral', 'bullish']
        direction_class = direction_classes[direction_class_idx]

        # Optimal entry TF
        optimal_entry_tf = self.config.timeframes[optimal_tf_idx]

        # Signal strength
        signal_strength = abs(unified_direction) * confidence

        # Recommended action
        if unified_direction > self.config.direction_threshold and confidence > self.config.confidence_threshold:
            recommended_action = 'buy'
        elif unified_direction < -self.config.direction_threshold and confidence > self.config.confidence_threshold:
            recommended_action = 'sell'
        else:
            recommended_action = 'hold'

        return MTSPrediction(
            unified_direction=unified_direction,
            direction_class=direction_class,
            confidence=confidence,
            confidence_by_alignment=confidence_by_alignment,
            optimal_entry_tf=optimal_entry_tf,
            tf_contributions=tf_contributions,
            signal_strength=signal_strength,
            recommended_action=recommended_action
        )

    def get_tf_contributions(self) -> Dict[str, float]:
        """Get the contribution of each timeframe."""
        if self.attention_model is None:
            # Return equal weights
            return {tf: 1.0 / len(self.config.timeframes) for tf in self.config.timeframes}

        return self.attention_model.get_tf_contributions()

    def train_attention(
        self,
        train_data: List[pd.DataFrame],
        val_data: Optional[List[pd.DataFrame]] = None,
        epochs: int = 50,
        batch_size: int = 32,
        lr: float = 0.001
    ):
        """
        Train attention model (PyTorch phase).

        Args:
            train_data: List of 5m DataFrames
            val_data: Optional validation data
            epochs: Number of training epochs
            batch_size: Batch size
            lr: Learning rate
        """
        logger.info("Training attention model...")

        # Prepare training features
        all_features = []
        all_alignments = []
        all_conflicts = []
        all_targets = []

        for df in train_data:
            tf_features, alignment, conflict = self._prepare_features(df)

            # Prepare targets (direction based on future return)
            _, direction_targets, _ = self.feature_engineer.prepare_training_data(df)

            all_features.append(tf_features)
            all_alignments.append(alignment)
            all_conflicts.append(conflict)
            all_targets.append(direction_targets)

        # Determine input dims from first sample
        if self.input_dims is None:
            input_dims = {tf: features.shape[-1] for tf, features in all_features[0].items()}
            self._build_attention_model(input_dims)

        # Training loop
        self.attention_model.train()
        optimizer = torch.optim.AdamW(self.attention_model.parameters(), lr=lr)
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, epochs)

        for epoch in range(epochs):
            total_loss = 0.0

            for i in range(len(all_features)):
                tf_tensors = {
                    tf: torch.tensor(features, dtype=torch.float32, device=self.device).unsqueeze(0)
                    for tf, features in all_features[i].items()
                }
                alignment = torch.tensor(all_alignments[i], dtype=torch.float32, device=self.device).unsqueeze(0)
                conflict = torch.tensor(all_conflicts[i], dtype=torch.float32, device=self.device).unsqueeze(0)

                optimizer.zero_grad()

                # Forward pass
                unified = self.attention_model(tf_tensors, alignment, conflict)

                # Simple contrastive loss based on alignment
                alignment_mean = alignment.mean()
                loss = -alignment_mean * unified.norm() + (1 - alignment_mean) * unified.norm()

                loss.backward()
                optimizer.step()
                total_loss += loss.item()

            scheduler.step()

            if (epoch + 1) % 10 == 0:
                logger.info(f"Epoch {epoch + 1}/{epochs}, Loss: {total_loss / len(all_features):.4f}")

        logger.info("Attention model training complete")

    def train_xgboost(
        self,
        X_train: np.ndarray,
        y_direction: np.ndarray,
        y_confidence: np.ndarray,
        y_entry_tf: np.ndarray,
        X_val: Optional[np.ndarray] = None,
        y_val_direction: Optional[np.ndarray] = None
    ):
        """
        Train XGBoost models on unified representations.

        Args:
            X_train: Unified representations (n_samples, d_model)
            y_direction: Direction labels (0, 1, 2)
            y_confidence: Confidence targets (0-1)
            y_entry_tf: Entry TF labels (0 to n_timeframes-1)
            X_val: Validation features
            y_val_direction: Validation direction labels
        """
        if not self.use_xgboost:
            raise RuntimeError("XGBoost not enabled")

        logger.info(f"Training XGBoost models on {len(X_train)} samples...")

        # Train direction classifier
        eval_set = [(X_val, y_val_direction)] if X_val is not None and y_val_direction is not None else None
        self.xgb_direction.fit(
            X_train, y_direction,
            eval_set=eval_set,
            verbose=False
        )

        # Train confidence regressor
        self.xgb_confidence.fit(X_train, y_confidence, verbose=False)

        # Train entry TF classifier
        self.xgb_entry_tf.fit(X_train, y_entry_tf, verbose=False)

        self._is_trained = True
        logger.info("XGBoost training complete")

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

        # Save config
        joblib.dump(self.config, path / 'config.joblib')

        # Save attention model
        if self.attention_model is not None:
            torch.save(self.attention_model.state_dict(), path / 'attention_model.pt')

        # Save prediction head
        if not self.use_xgboost and self.prediction_head is not None:
            torch.save(self.prediction_head.state_dict(), path / 'prediction_head.pt')

        # Save XGBoost models
        if self.use_xgboost:
            joblib.dump(self.xgb_direction, path / 'xgb_direction.joblib')
            joblib.dump(self.xgb_confidence, path / 'xgb_confidence.joblib')
            joblib.dump(self.xgb_entry_tf, path / 'xgb_entry_tf.joblib')

        # Save metadata
        metadata = {
            'input_dims': self.input_dims,
            'is_trained': self._is_trained,
            'use_xgboost': self.use_xgboost,
            'device': self.device
        }
        joblib.dump(metadata, path / 'metadata.joblib')

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

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

        # Load config
        self.config = joblib.load(path / 'config.joblib')

        # Load metadata
        metadata = joblib.load(path / 'metadata.joblib')
        self.input_dims = metadata['input_dims']
        self._is_trained = metadata['is_trained']
        self.use_xgboost = metadata['use_xgboost']

        # Rebuild attention model
        if self.input_dims is not None:
            self._build_attention_model(self.input_dims)

            # Load weights
            if (path / 'attention_model.pt').exists():
                self.attention_model.load_state_dict(
                    torch.load(path / 'attention_model.pt', map_location=self.device)
                )

        # Load prediction head
        if not self.use_xgboost and (path / 'prediction_head.pt').exists():
            self.prediction_head.load_state_dict(
                torch.load(path / 'prediction_head.pt', map_location=self.device)
            )

        # Load XGBoost models
        if self.use_xgboost:
            self.xgb_direction = joblib.load(path / 'xgb_direction.joblib')
            self.xgb_confidence = joblib.load(path / 'xgb_confidence.joblib')
            self.xgb_entry_tf = joblib.load(path / 'xgb_entry_tf.joblib')

        logger.info(f"Loaded MTSModel from {path}")


if __name__ == "__main__":
    # Test MTSModel
    np.random.seed(42)
    torch.manual_seed(42)

    # Create sample 5m data
    n_samples = 500

    base_price = 100.0
    returns = np.random.randn(n_samples) * 0.001 + 0.0001
    close = base_price * np.cumprod(1 + returns)

    df_5m = pd.DataFrame({
        'open': close * (1 + np.random.randn(n_samples) * 0.001),
        'high': close * (1 + np.abs(np.random.randn(n_samples) * 0.002)),
        'low': close * (1 - np.abs(np.random.randn(n_samples) * 0.002)),
        'close': close,
        'volume': np.random.randint(1000, 10000, n_samples)
    })

    # Initialize model
    config = MTSConfig(
        d_model=64,
        n_heads=2,
        n_layers=1
    )

    model = MTSModel(config, use_xgboost=True)

    # Prepare features
    tf_features, alignment, conflict = model._prepare_features(df_5m)

    print(f"Feature shapes: {[f'{tf}: {f.shape}' for tf, f in tf_features.items()]}")
    print(f"Alignment shape: {alignment.shape}")
    print(f"Conflict shape: {conflict.shape}")

    # Build attention model
    input_dims = {tf: features.shape[-1] for tf, features in tf_features.items()}
    model._build_attention_model(input_dims)

    # Get unified representation
    unified_rep = model._extract_unified_representation(tf_features, alignment, conflict)
    print(f"Unified representation shape: {unified_rep.shape}")

    # Train XGBoost with synthetic targets
    n_train = unified_rep.shape[0]
    y_direction = np.random.randint(0, 3, n_train)
    y_confidence = np.random.rand(n_train)
    y_entry_tf = np.random.randint(0, 4, n_train)

    model.train_xgboost(unified_rep, y_direction, y_confidence, y_entry_tf)

    # Make prediction
    prediction = model.predict(df_5m)

    print("\nPrediction:")
    print(f"  Direction: {prediction.direction_class} ({prediction.unified_direction:.3f})")
    print(f"  Confidence: {prediction.confidence:.3f}")
    print(f"  Confidence by Alignment: {prediction.confidence_by_alignment:.3f}")
    print(f"  Optimal Entry TF: {prediction.optimal_entry_tf}")
    print(f"  TF Contributions: {prediction.tf_contributions}")
    print(f"  Recommended Action: {prediction.recommended_action}")

    print("\nMTSModel test complete!")
