"""
Specialized models for AMD phases
Different architectures optimized for each market phase
Migrated from TradingAgent for OrbiQuant IA Platform
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import pandas as pd
from typing import Dict, List, Optional, Tuple, Any
from loguru import logger
import xgboost as xgb
from dataclasses import dataclass


@dataclass
class AMDPrediction:
    """Prediction tailored to AMD phase"""
    phase: str
    predictions: Dict[str, float]
    confidence: float
    recommended_action: str
    stop_loss: float
    take_profit: float
    position_size: float
    reasoning: List[str]


class AccumulationModel(nn.Module):
    """
    Neural network optimized for accumulation phase
    Focus: Identifying breakout potential and optimal entry points
    """

    def __init__(self, input_dim: int, hidden_dim: int = 128, num_heads: int = 4):
        super().__init__()

        # Multi-head attention for pattern recognition
        self.attention = nn.MultiheadAttention(
            embed_dim=input_dim,
            num_heads=num_heads,
            batch_first=True
        )

        # Feature extraction layers
        self.feature_net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.BatchNorm1d(hidden_dim // 2),
            nn.ReLU(),
            nn.Dropout(0.1)
        )

        # Breakout prediction head
        self.breakout_head = nn.Sequential(
            nn.Linear(hidden_dim // 2, 32),
            nn.ReLU(),
            nn.Linear(32, 3)  # [no_breakout, bullish_breakout, failed_breakout]
        )

        # Entry timing head
        self.entry_head = nn.Sequential(
            nn.Linear(hidden_dim // 2, 32),
            nn.ReLU(),
            nn.Linear(32, 2)  # [entry_score, optimal_size]
        )

        # Price target head
        self.target_head = nn.Sequential(
            nn.Linear(hidden_dim // 2, 32),
            nn.ReLU(),
            nn.Linear(32, 2)  # [target_high, confidence]
        )

    def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> Dict[str, torch.Tensor]:
        """
        Forward pass for accumulation phase prediction

        Args:
            x: Input tensor [batch, seq_len, features]
            mask: Optional attention mask

        Returns:
            Dictionary of predictions
        """
        # Apply attention
        attn_out, _ = self.attention(x, x, x, key_padding_mask=mask)

        # Global pooling
        if len(attn_out.shape) == 3:
            pooled = attn_out.mean(dim=1)
        else:
            pooled = attn_out

        # Extract features
        features = self.feature_net(pooled)

        # Generate predictions
        breakout_logits = self.breakout_head(features)
        entry_scores = self.entry_head(features)
        targets = self.target_head(features)

        return {
            'breakout_probs': F.softmax(breakout_logits, dim=-1),
            'entry_score': torch.sigmoid(entry_scores[:, 0]),
            'position_size': torch.sigmoid(entry_scores[:, 1]),
            'target_high': targets[:, 0],
            'target_confidence': torch.sigmoid(targets[:, 1])
        }


class ManipulationModel(nn.Module):
    """
    Neural network optimized for manipulation phase
    Focus: Detecting false moves and avoiding traps
    """

    def __init__(self, input_dim: int, hidden_dim: int = 128):
        super().__init__()

        # LSTM for sequence modeling
        self.lstm = nn.LSTM(
            input_size=input_dim,
            hidden_size=hidden_dim,
            num_layers=2,
            batch_first=True,
            dropout=0.3,
            bidirectional=True
        )

        # Trap detection network
        self.trap_detector = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(hidden_dim, 64),
            nn.ReLU(),
            nn.Linear(64, 4)  # [no_trap, bull_trap, bear_trap, whipsaw]
        )

        # Reversal prediction
        self.reversal_predictor = nn.Sequential(
            nn.Linear(hidden_dim * 2, 64),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(64, 3)  # [reversal_probability, reversal_direction, reversal_magnitude]
        )

        # Safe zone identifier
        self.safe_zone = nn.Sequential(
            nn.Linear(hidden_dim * 2, 32),
            nn.ReLU(),
            nn.Linear(32, 2)  # [upper_safe, lower_safe]
        )

    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
        """
        Forward pass for manipulation phase prediction

        Args:
            x: Input tensor [batch, seq_len, features]

        Returns:
            Dictionary of predictions
        """
        # LSTM encoding
        lstm_out, (hidden, _) = self.lstm(x)

        # Use last hidden state
        if len(lstm_out.shape) == 3:
            final_hidden = lstm_out[:, -1, :]
        else:
            final_hidden = lstm_out

        # Detect traps
        trap_logits = self.trap_detector(final_hidden)
        trap_probs = F.softmax(trap_logits, dim=-1)

        # Predict reversals
        reversal_features = self.reversal_predictor(final_hidden)
        reversal_prob = torch.sigmoid(reversal_features[:, 0])
        reversal_dir = torch.tanh(reversal_features[:, 1])
        reversal_mag = torch.sigmoid(reversal_features[:, 2])

        # Identify safe zones
        safe_zones = self.safe_zone(final_hidden)

        return {
            'trap_probabilities': trap_probs,
            'reversal_probability': reversal_prob,
            'reversal_direction': reversal_dir,  # -1 to 1
            'reversal_magnitude': reversal_mag,
            'safe_zone_upper': safe_zones[:, 0],
            'safe_zone_lower': safe_zones[:, 1]
        }


class DistributionModel(nn.Module):
    """
    Neural network optimized for distribution phase
    Focus: Identifying exit points and downside targets
    """

    def __init__(self, input_dim: int, hidden_dim: int = 128):
        super().__init__()

        # GRU for temporal patterns
        self.gru = nn.GRU(
            input_size=input_dim,
            hidden_size=hidden_dim,
            num_layers=2,
            batch_first=True,
            dropout=0.2
        )

        # Breakdown detection
        self.breakdown_detector = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(hidden_dim, 64),
            nn.ReLU(),
            nn.Linear(64, 3)  # [breakdown_prob, breakdown_timing, breakdown_magnitude]
        )

        # Exit signal generator
        self.exit_signal = nn.Sequential(
            nn.Linear(hidden_dim, 64),
            nn.ReLU(),
            nn.Linear(64, 4)  # [exit_urgency, exit_price, stop_loss, position_reduction]
        )

        # Downside target predictor
        self.target_predictor = nn.Sequential(
            nn.Linear(hidden_dim, 64),
            nn.ReLU(),
            nn.Linear(64, 3)  # [target_1, target_2, target_3]
        )

    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
        """
        Forward pass for distribution phase prediction

        Args:
            x: Input tensor [batch, seq_len, features]

        Returns:
            Dictionary of predictions
        """
        # GRU encoding
        gru_out, hidden = self.gru(x)

        # Use last output
        if len(gru_out.shape) == 3:
            final_out = gru_out[:, -1, :]
        else:
            final_out = gru_out

        # Breakdown detection
        breakdown_features = self.breakdown_detector(final_out)
        breakdown_prob = torch.sigmoid(breakdown_features[:, 0])
        breakdown_timing = torch.sigmoid(breakdown_features[:, 1]) * 10  # 0-10 periods
        breakdown_mag = torch.sigmoid(breakdown_features[:, 2]) * 0.2  # 0-20% move

        # Exit signals
        exit_features = self.exit_signal(final_out)
        exit_urgency = torch.sigmoid(exit_features[:, 0])
        exit_price = exit_features[:, 1]
        stop_loss = exit_features[:, 2]
        position_reduction = torch.sigmoid(exit_features[:, 3])

        # Downside targets
        targets = self.target_predictor(final_out)

        return {
            'breakdown_probability': breakdown_prob,
            'breakdown_timing': breakdown_timing,
            'breakdown_magnitude': breakdown_mag,
            'exit_urgency': exit_urgency,
            'exit_price': exit_price,
            'stop_loss': stop_loss,
            'position_reduction': position_reduction,
            'downside_targets': targets
        }


class AMDEnsemble:
    """
    Ensemble model that selects and weights predictions based on AMD phase
    """

    def __init__(self, feature_dim: int = 256):
        """
        Initialize AMD ensemble

        Args:
            feature_dim: Dimension of input features
        """
        self.feature_dim = feature_dim

        # Initialize phase-specific models
        self.accumulation_model = AccumulationModel(feature_dim)
        self.manipulation_model = ManipulationModel(feature_dim)
        self.distribution_model = DistributionModel(feature_dim)

        # XGBoost models for each phase
        self.accumulation_xgb = None
        self.manipulation_xgb = None
        self.distribution_xgb = None

        # Model weights based on phase confidence
        self.phase_weights = {
            'accumulation': {'neural': 0.6, 'xgboost': 0.4},
            'manipulation': {'neural': 0.5, 'xgboost': 0.5},
            'distribution': {'neural': 0.6, 'xgboost': 0.4}
        }

        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self._move_models_to_device()

    def _move_models_to_device(self):
        """Move neural models to appropriate device"""
        self.accumulation_model = self.accumulation_model.to(self.device)
        self.manipulation_model = self.manipulation_model.to(self.device)
        self.distribution_model = self.distribution_model.to(self.device)

    def train_phase_models(
        self,
        X_train: pd.DataFrame,
        y_train: pd.DataFrame,
        phase: str,
        validation_data: Optional[Tuple[pd.DataFrame, pd.DataFrame]] = None
    ):
        """
        Train models for specific phase

        Args:
            X_train: Training features
            y_train: Training targets
            phase: AMD phase
            validation_data: Optional validation set
        """
        logger.info(f"Training {phase} models...")

        # Train XGBoost model
        xgb_params = self._get_xgb_params(phase)

        if phase == 'accumulation':
            self.accumulation_xgb = xgb.XGBRegressor(**xgb_params)
            self.accumulation_xgb.fit(X_train, y_train)
        elif phase == 'manipulation':
            self.manipulation_xgb = xgb.XGBRegressor(**xgb_params)
            self.manipulation_xgb.fit(X_train, y_train)
        elif phase == 'distribution':
            self.distribution_xgb = xgb.XGBRegressor(**xgb_params)
            self.distribution_xgb.fit(X_train, y_train)

        logger.info(f"Completed training for {phase} models")

    def _get_xgb_params(self, phase: str) -> Dict[str, Any]:
        """Get XGBoost parameters for specific phase"""
        base_params = {
            'n_estimators': 200,
            'learning_rate': 0.05,
            'max_depth': 6,
            'subsample': 0.8,
            'colsample_bytree': 0.8,
            'random_state': 42,
            'n_jobs': -1
        }

        if torch.cuda.is_available():
            base_params.update({
                'tree_method': 'hist',
                'device': 'cuda'
            })

        # Phase-specific adjustments
        if phase == 'accumulation':
            base_params['learning_rate'] = 0.03  # More conservative
            base_params['max_depth'] = 8  # Capture complex patterns
        elif phase == 'manipulation':
            base_params['learning_rate'] = 0.1  # Faster adaptation
            base_params['max_depth'] = 5  # Avoid overfitting to noise
            base_params['subsample'] = 0.6  # More regularization
        elif phase == 'distribution':
            base_params['learning_rate'] = 0.05
            base_params['max_depth'] = 7

        return base_params

    def predict(
        self,
        features: pd.DataFrame,
        phase: str,
        phase_confidence: float
    ) -> AMDPrediction:
        """
        Generate predictions based on detected phase

        Args:
            features: Input features
            phase: Detected AMD phase
            phase_confidence: Confidence in phase detection

        Returns:
            AMDPrediction with phase-specific recommendations
        """
        # Convert features to tensor
        X_tensor = torch.FloatTensor(features.values).to(self.device)
        if len(X_tensor.shape) == 2:
            X_tensor = X_tensor.unsqueeze(0)  # Add batch dimension

        predictions = {}
        confidence = phase_confidence

        with torch.no_grad():
            if phase == 'accumulation':
                nn_preds = self.accumulation_model(X_tensor)
                xgb_preds = None
                if self.accumulation_xgb is not None:
                    xgb_preds = self.accumulation_xgb.predict(features.iloc[-1:])
                predictions = self._combine_accumulation_predictions(nn_preds, xgb_preds)
                action, sl, tp, size, reasoning = self._get_accumulation_strategy(predictions)

            elif phase == 'manipulation':
                nn_preds = self.manipulation_model(X_tensor)
                xgb_preds = None
                if self.manipulation_xgb is not None:
                    xgb_preds = self.manipulation_xgb.predict(features.iloc[-1:])
                predictions = self._combine_manipulation_predictions(nn_preds, xgb_preds)
                action, sl, tp, size, reasoning = self._get_manipulation_strategy(predictions)

            elif phase == 'distribution':
                nn_preds = self.distribution_model(X_tensor)
                xgb_preds = None
                if self.distribution_xgb is not None:
                    xgb_preds = self.distribution_xgb.predict(features.iloc[-1:])
                predictions = self._combine_distribution_predictions(nn_preds, xgb_preds)
                action, sl, tp, size, reasoning = self._get_distribution_strategy(predictions)

            else:
                action = 'hold'
                sl = tp = size = 0
                reasoning = ['Unknown market phase']
                confidence = 0

        return AMDPrediction(
            phase=phase,
            predictions=predictions,
            confidence=confidence,
            recommended_action=action,
            stop_loss=sl,
            take_profit=tp,
            position_size=size,
            reasoning=reasoning
        )

    def _combine_accumulation_predictions(
        self,
        nn_preds: Dict[str, torch.Tensor],
        xgb_preds: Optional[np.ndarray]
    ) -> Dict[str, float]:
        """Combine neural network and XGBoost predictions for accumulation"""
        combined = {}

        combined['breakout_probability'] = float(nn_preds['breakout_probs'][0, 1].cpu())
        combined['entry_score'] = float(nn_preds['entry_score'][0].cpu())
        combined['position_size'] = float(nn_preds['position_size'][0].cpu())
        combined['target_high'] = float(nn_preds['target_high'][0].cpu())
        combined['target_confidence'] = float(nn_preds['target_confidence'][0].cpu())

        if xgb_preds is not None:
            weights = self.phase_weights['accumulation']
            combined['target_high'] = (
                combined['target_high'] * weights['neural'] +
                float(xgb_preds[0]) * weights['xgboost']
            )

        return combined

    def _combine_manipulation_predictions(
        self,
        nn_preds: Dict[str, torch.Tensor],
        xgb_preds: Optional[np.ndarray]
    ) -> Dict[str, float]:
        """Combine predictions for manipulation phase"""
        combined = {}

        trap_probs = nn_preds['trap_probabilities'][0].cpu().numpy()
        combined['bull_trap_prob'] = float(trap_probs[1])
        combined['bear_trap_prob'] = float(trap_probs[2])
        combined['whipsaw_prob'] = float(trap_probs[3])
        combined['reversal_probability'] = float(nn_preds['reversal_probability'][0].cpu())
        combined['reversal_direction'] = float(nn_preds['reversal_direction'][0].cpu())
        combined['safe_zone_upper'] = float(nn_preds['safe_zone_upper'][0].cpu())
        combined['safe_zone_lower'] = float(nn_preds['safe_zone_lower'][0].cpu())

        return combined

    def _combine_distribution_predictions(
        self,
        nn_preds: Dict[str, torch.Tensor],
        xgb_preds: Optional[np.ndarray]
    ) -> Dict[str, float]:
        """Combine predictions for distribution phase"""
        combined = {}

        combined['breakdown_probability'] = float(nn_preds['breakdown_probability'][0].cpu())
        combined['breakdown_timing'] = float(nn_preds['breakdown_timing'][0].cpu())
        combined['exit_urgency'] = float(nn_preds['exit_urgency'][0].cpu())
        combined['position_reduction'] = float(nn_preds['position_reduction'][0].cpu())

        targets = nn_preds['downside_targets'][0].cpu().numpy()
        combined['target_1'] = float(targets[0])
        combined['target_2'] = float(targets[1])
        combined['target_3'] = float(targets[2])

        return combined

    def _get_accumulation_strategy(
        self,
        predictions: Dict[str, float]
    ) -> Tuple[str, float, float, float, List[str]]:
        """Get trading strategy for accumulation phase"""
        reasoning = []

        if predictions['breakout_probability'] > 0.7:
            action = 'buy'
            sl = 0.98
            tp = predictions['target_high']
            size = min(1.0, predictions['position_size'] * 1.5)
            reasoning.append(f"High breakout probability: {predictions['breakout_probability']:.2%}")
            reasoning.append("Accumulation phase indicates institutional buying")
        elif predictions['entry_score'] > 0.6:
            action = 'buy'
            sl = 0.97
            tp = predictions['target_high'] * 0.98
            size = predictions['position_size']
            reasoning.append(f"Good entry opportunity: {predictions['entry_score']:.2f}")
            reasoning.append("Building position during accumulation")
        else:
            action = 'wait'
            sl = tp = size = 0
            reasoning.append("Waiting for better entry in accumulation phase")
            reasoning.append(f"Entry score too low: {predictions['entry_score']:.2f}")

        return action, sl, tp, size, reasoning

    def _get_manipulation_strategy(
        self,
        predictions: Dict[str, float]
    ) -> Tuple[str, float, float, float, List[str]]:
        """Get trading strategy for manipulation phase"""
        reasoning = []

        max_trap_prob = max(
            predictions['bull_trap_prob'],
            predictions['bear_trap_prob'],
            predictions['whipsaw_prob']
        )

        if max_trap_prob > 0.6:
            action = 'avoid'
            sl = tp = size = 0
            reasoning.append(f"High trap probability detected: {max_trap_prob:.2%}")
            reasoning.append("Manipulation phase - avoid new positions")
        elif predictions['reversal_probability'] > 0.7:
            if predictions['reversal_direction'] > 0:
                action = 'buy'
                sl = predictions['safe_zone_lower']
                tp = predictions['safe_zone_upper']
            else:
                action = 'sell'
                sl = predictions['safe_zone_upper']
                tp = predictions['safe_zone_lower']
            size = 0.3
            reasoning.append(f"Reversal signal: {predictions['reversal_probability']:.2%}")
            reasoning.append("Trading reversal with tight stops")
        else:
            action = 'hold'
            sl = tp = size = 0
            reasoning.append("Unclear signals in manipulation phase")
            reasoning.append("Waiting for clearer market structure")

        return action, sl, tp, size, reasoning

    def _get_distribution_strategy(
        self,
        predictions: Dict[str, float]
    ) -> Tuple[str, float, float, float, List[str]]:
        """Get trading strategy for distribution phase"""
        reasoning = []

        if predictions['exit_urgency'] > 0.8:
            action = 'sell'
            sl = 1.02
            tp = predictions['target_1']
            size = 1.0
            reasoning.append(f"High exit urgency: {predictions['exit_urgency']:.2%}")
            reasoning.append("Distribution phase - institutional selling")
        elif predictions['breakdown_probability'] > 0.6:
            action = 'sell'
            sl = 1.03
            tp = predictions['target_2']
            size = predictions['position_reduction']
            reasoning.append(f"Breakdown imminent: {predictions['breakdown_probability']:.2%}")
            reasoning.append(f"Expected timing: {predictions['breakdown_timing']:.1f} periods")
        elif predictions['position_reduction'] > 0.5:
            action = 'reduce'
            sl = tp = 0
            size = predictions['position_reduction']
            reasoning.append(f"Reduce position by {size:.0%}")
            reasoning.append("Distribution phase - protect capital")
        else:
            action = 'hold'
            sl = tp = size = 0
            reasoning.append("Monitor distribution development")
            reasoning.append(f"Breakdown probability: {predictions['breakdown_probability']:.2%}")

        return action, sl, tp, size, reasoning
