#!/usr/bin/env python3
"""
Neural Network Decision Fusion System
Central NN that merges all model outputs + market data for final trading decisions
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from typing import Dict, List, Optional, Any
from dataclasses import dataclass
from datetime import datetime
import logging

logger = logging.getLogger(__name__)

@dataclass
class ModelPrediction:
    """Standardized prediction from any model"""
    model_name: str
    prediction_type: str  # 'price', 'direction', 'action'
    value: float  # -1 to 1 for direction, actual price for price predictions
    confidence: float  # 0 to 1
    timestamp: datetime
    metadata: Optional[Dict[str, Any]] = None

@dataclass
class MarketContext:
    """Current market context for decision fusion"""
    symbol: str
    current_price: float
    price_change_1m: float
    price_change_5m: float
    volume_ratio: float
    volatility: float
    timestamp: datetime

@dataclass
class FusionDecision:
    """Final trading decision from fusion NN"""
    action: str  # 'BUY', 'SELL', 'HOLD'
    confidence: float  # 0 to 1
    expected_return: float  # Expected return percentage
    risk_score: float  # 0 to 1, higher = riskier
    position_size: float  # Recommended position size
    reasoning: str  # Human-readable explanation
    model_contributions: Dict[str, float]  # How much each model contributed
    timestamp: datetime

class DecisionFusionNetwork(nn.Module):
    """Small NN that fuses model predictions with market context"""
    
    def __init__(self, input_dim: int = 32, hidden_dim: int = 64):
        super().__init__()
        
        self.fusion_layers = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Linear(hidden_dim // 2, 16)
        )
        
        # Output heads
        self.action_head = nn.Linear(16, 3)  # BUY, SELL, HOLD
        self.confidence_head = nn.Linear(16, 1)
        self.return_head = nn.Linear(16, 1)
        
    def forward(self, features: torch.Tensor) -> Dict[str, torch.Tensor]:
        """Forward pass through fusion network"""
        fusion_output = self.fusion_layers(features)
        
        action_logits = self.action_head(fusion_output)
        action_probs = F.softmax(action_logits, dim=1)
        
        confidence = torch.sigmoid(self.confidence_head(fusion_output))
        expected_return = torch.tanh(self.return_head(fusion_output))
        
        return {
            'action_probs': action_probs,
            'confidence': confidence.squeeze(),
            'expected_return': expected_return.squeeze()
        }

class NeuralDecisionFusion:
    """Main NN-based decision fusion system"""
    
    def __init__(self, training_mode: bool = True):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.network = DecisionFusionNetwork().to(self.device)
        self.training_mode = training_mode
        self.registered_models = {}
        self.last_predictions = {}
        
        logger.info(f"🧠 Neural Decision Fusion initialized on {self.device}")
    
    def register_model(self, model_name: str, model_type: str, prediction_format: str):
        """Register a model that will provide predictions"""
        self.registered_models[model_name] = {
            'type': model_type,
            'format': prediction_format,
            'prediction_count': 0
        }
        logger.info(f"Registered NN model: {model_name} ({model_type})")
    
    def add_prediction(self, prediction: ModelPrediction):
        """Add a prediction from a registered model"""
        self.last_predictions[prediction.model_name] = prediction
        if prediction.model_name in self.registered_models:
            self.registered_models[prediction.model_name]['prediction_count'] += 1
        
        logger.debug(f"🔮 {prediction.model_name}: {prediction.value:.3f} "
                    f"(confidence: {prediction.confidence:.3f})")
    
    def make_decision(self, symbol: str, market_context: MarketContext, 
                     min_confidence: float = 0.25) -> Optional[FusionDecision]:
        """Make NN-driven trading decision"""
        try:
            if len(self.last_predictions) < 1:
                logger.debug("No NN predictions available")
                return None
            
            # Prepare features
            features = self._prepare_features(market_context)
            if features is None:
                return None
            
            # Run NN inference
            with torch.no_grad():
                self.network.eval()
                features_tensor = torch.tensor(features, dtype=torch.float32).unsqueeze(0).to(self.device)
                outputs = self.network(features_tensor)
                
                action_probs = outputs['action_probs'][0].cpu().numpy()
                confidence = outputs['confidence'].cpu().item()
                expected_return = outputs['expected_return'].cpu().item()
            
            # Determine action
            action_idx = np.argmax(action_probs)
            actions = ['BUY', 'SELL', 'HOLD']
            action = actions[action_idx]
            
            # Check confidence threshold
            if confidence < min_confidence:
                action = 'HOLD'
                logger.debug(f"Low NN confidence ({confidence:.3f}), defaulting to HOLD")
            
            # Calculate position size
            position_size = self._calculate_position_size(confidence, expected_return)
            
            # Generate reasoning
            reasoning = self._generate_reasoning(action, confidence, expected_return, action_probs)
            
            # Calculate risk score and model contributions
            risk_score = min(1.0, abs(expected_return) * 5 + (1 - confidence) * 0.5)
            model_contributions = self._calculate_model_contributions()
            
            decision = FusionDecision(
                action=action,
                confidence=confidence,
                expected_return=expected_return,
                risk_score=risk_score,
                position_size=position_size,
                reasoning=reasoning,
                model_contributions=model_contributions,
                timestamp=datetime.now()
            )
            
            logger.info(f"🧠 NN DECISION: {action} (conf: {confidence:.3f}, "
                       f"return: {expected_return:.3f}, size: {position_size:.4f})")
            
            return decision
            
        except Exception as e:
            logger.error(f"Error in NN decision making: {e}")
            return None
    
    def _prepare_features(self, context: MarketContext) -> Optional[np.ndarray]:
        """Prepare feature vector for NN"""
        try:
            features = np.zeros(32)
            
            # Model predictions (slots 0-15)
            idx = 0
            for model_name, prediction in self.last_predictions.items():
                if idx < 14:  # Leave room for other features
                    features[idx] = prediction.value
                    features[idx + 1] = prediction.confidence
                    idx += 2
            
            # Market context (slots 16-31)
            features[16] = np.tanh(context.price_change_1m * 100)  # 1m change
            features[17] = np.tanh(context.price_change_5m * 100)  # 5m change
            features[18] = np.tanh(context.volume_ratio - 1)      # Volume ratio
            features[19] = np.tanh(context.volatility * 100)      # Volatility
            features[20] = context.current_price / 10000.0        # Normalized price
            
            # Time features
            now = context.timestamp
            features[21] = now.hour / 24.0
            features[22] = now.weekday() / 7.0
            
            # Model agreement features
            if len(self.last_predictions) >= 2:
                values = [p.value for p in self.last_predictions.values()]
                features[23] = np.mean(values)  # Average prediction
                features[24] = np.std(values)   # Prediction variance
                features[25] = len(self.last_predictions)  # Model count
            
            return features
            
        except Exception as e:
            logger.error(f"Error preparing NN features: {e}")
            return None
    
    def _calculate_position_size(self, confidence: float, expected_return: float) -> float:
        """Calculate position size based on NN outputs"""
        base_size = 0.01  # 0.01 ETH base
        
        # Scale by confidence
        confidence_multiplier = max(0.1, min(2.0, confidence * 1.5))
        
        # Scale by expected return
        return_multiplier = 1.0 + abs(expected_return) * 0.5
        
        final_size = base_size * confidence_multiplier * return_multiplier
        return max(0.001, min(0.05, final_size))
    
    def _generate_reasoning(self, action: str, confidence: float, 
                          expected_return: float, action_probs: np.ndarray) -> str:
        """Generate human-readable reasoning"""
        reasons = []
        
        if action == 'BUY':
            reasons.append(f"NN suggests BUY ({action_probs[0]:.1%})")
        elif action == 'SELL':
            reasons.append(f"NN suggests SELL ({action_probs[1]:.1%})")
        else:
            reasons.append(f"NN suggests HOLD")
        
        if confidence > 0.7:
            reasons.append("High confidence")
        elif confidence > 0.5:
            reasons.append("Moderate confidence")
        else:
            reasons.append("Low confidence")
        
        if abs(expected_return) > 0.01:
            direction = "positive" if expected_return > 0 else "negative"
            reasons.append(f"Expected {direction} return: {expected_return:.2%}")
        
        reasons.append(f"Based on {len(self.last_predictions)} NN models")
        
        return " | ".join(reasons)
    
    def _calculate_model_contributions(self) -> Dict[str, float]:
        """Calculate how much each model contributed to the decision"""
        contributions = {}
        total_confidence = sum(p.confidence for p in self.last_predictions.values()) if self.last_predictions else 1.0
        
        if total_confidence > 0:
            for model_name, prediction in self.last_predictions.items():
                contributions[model_name] = prediction.confidence / total_confidence
        
        return contributions
    
    def get_status(self) -> Dict[str, Any]:
        """Get NN fusion system status"""
        return {
            'device': str(self.device),
            'training_mode': self.training_mode,
            'registered_models': len(self.registered_models),
            'recent_predictions': len(self.last_predictions),
            'model_parameters': sum(p.numel() for p in self.network.parameters())
        }