"""Multi-task neural network for entry signal and SL/TP prediction.

Architecture: LSTM/Transformer encoder with multiple output heads:
- Entry signal: binary classification (should enter?)
- SL distance: regression (stop-loss in ATR units)
- TP distance: regression (take-profit in ATR units)
- Confidence: regression (model confidence score)
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from typing import Dict, Optional

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from loguru import logger

from domain import DEFAULT_DROPOUT


class MultiTaskOutput:
    """Container for multi-task model outputs.
    
    Has TWO separate entry heads: one for LONG, one for SHORT.
    This fixes the fundamental issue where a single shared entry_proba
    couldn't distinguish between LONG and SHORT entries.
    """
    
    def __init__(
        self,
        entry_long_logits: torch.Tensor,
        entry_short_logits: torch.Tensor,
        sl_distance: torch.Tensor,
        tp_distance: torch.Tensor,
        confidence: torch.Tensor,
    ):
        self.entry_long_logits = entry_long_logits
        self.entry_short_logits = entry_short_logits
        self.sl_distance = sl_distance
        self.tp_distance = tp_distance
        self.confidence = confidence
    
    @property
    def entry_long_proba(self) -> torch.Tensor:
        return torch.sigmoid(self.entry_long_logits)
    
    @property
    def entry_short_proba(self) -> torch.Tensor:
        return torch.sigmoid(self.entry_short_logits)
    
    @property
    def entry_proba(self) -> torch.Tensor:
        """Legacy: returns LONG prob for backward compatibility."""
        return self.entry_long_proba
    
    def to_dict(self) -> Dict[str, torch.Tensor]:
        return {
            "entry_long_logits": self.entry_long_logits,
            "entry_short_logits": self.entry_short_logits,
            "entry_long_proba": self.entry_long_proba,
            "entry_short_proba": self.entry_short_proba,
            "sl_distance": self.sl_distance,
            "tp_distance": self.tp_distance,
            "confidence": self.confidence,
        }
    
    def to_numpy(self) -> Dict[str, np.ndarray]:
        return {
            "entry_long_proba": self.entry_long_proba.detach().cpu().numpy(),
            "entry_short_proba": self.entry_short_proba.detach().cpu().numpy(),
            "sl_distance": self.sl_distance.detach().cpu().numpy(),
            "tp_distance": self.tp_distance.detach().cpu().numpy(),
            "confidence": self.confidence.detach().cpu().numpy(),
        }


class BaseMultiTaskModel(ABC, nn.Module):
    """Abstract base class for multi-task trading models."""
    
    SL_MIN = 0.5
    SL_MAX = 3.0
    TP_MIN = 0.5
    TP_MAX = 5.0
    
    def __init__(
        self,
        input_size: int,
        context_window: int = 24,
        dropout: float = DEFAULT_DROPOUT,
    ):
        super().__init__()
        self.input_size = input_size
        self.context_window = context_window
        self.dropout = dropout
    
    @abstractmethod
    def encode(self, x: torch.Tensor) -> torch.Tensor:
        """Encode input sequence into latent representation."""
        pass
    
    def forward(self, x: torch.Tensor) -> MultiTaskOutput:
        """Forward pass through the network.
        
        Returns TWO entry probabilities: LONG and SHORT.
        SL/TP distances are shared (direction-independent ATR values).
        
        Args:
            x: Input tensor of shape (batch, features) or (batch, seq_len, features)
            
        Returns:
            MultiTaskOutput with all predictions
        """
        encoded = self.encode(x)
        
        entry_long_logits = self.entry_long_head(encoded).squeeze(-1)
        entry_short_logits = self.entry_short_head(encoded).squeeze(-1)
        
        # Gradient-preserving bounded output: sigmoid scales to [SL_MIN, SL_MAX]
        sl_raw = self.sl_head(encoded).squeeze(-1)
        sl_normalized = torch.sigmoid(sl_raw)
        sl_distance = self.SL_MIN + (self.SL_MAX - self.SL_MIN) * sl_normalized
        
        # Gradient-preserving bounded output: sigmoid scales to [TP_MIN, TP_MAX]
        tp_raw = self.tp_head(encoded).squeeze(-1)
        tp_normalized = torch.sigmoid(tp_raw)
        tp_distance = self.TP_MIN + (self.TP_MAX - self.TP_MIN) * tp_normalized
        
        confidence = torch.sigmoid(self.confidence_head(encoded)).squeeze(-1)
        
        return MultiTaskOutput(
            entry_long_logits=entry_long_logits,
            entry_short_logits=entry_short_logits,
            sl_distance=sl_distance,
            tp_distance=tp_distance,
            confidence=confidence,
        )
    
    def predict(self, x: np.ndarray, entry_threshold: float = 0.5) -> Dict[str, np.ndarray]:
        """Make predictions on numpy array.
        
        Returns LONG and SHORT probabilities from a single forward pass.
        Also computes entry_signal for both sides.
        
        Args:
            x: Input features array
            entry_threshold: Threshold for entry signal classification
            
        Returns:
            Dictionary with predictions (includes entry_long_proba, entry_short_proba)
        """
        self.eval()
        with torch.no_grad():
            x_tensor = torch.FloatTensor(x).to(next(self.parameters()).device)
            output = self.forward(x_tensor)
            result = output.to_numpy()
            result["entry_long_signal"] = (result["entry_long_proba"] >= entry_threshold).astype(int)
            result["entry_short_signal"] = (result["entry_short_proba"] >= entry_threshold).astype(int)
            # Backward compatibility
            result["entry_signal"] = result["entry_long_signal"]
            result["entry_proba"] = result["entry_long_proba"]
        return result
    
    def _init_weights(self, module: nn.Module) -> None:
        """Initialize weights using Xavier initialization."""
        if isinstance(module, (nn.Linear, nn.Conv1d)):
            if module.weight is not None:
                nn.init.xavier_uniform_(module.weight)
            if module.bias is not None:
                nn.init.zeros_(module.bias)


class LSTMMultiTaskModel(BaseMultiTaskModel):
    """LSTM-based multi-task model for sequence processing.
    
    The input feature vector has a dual-path layout (see ai.features.extract_feature_vector):
    - First 168 values (LSTM_N_SEQ): per-bar sequence (24 bars × 7 features)
    - Last 58 values (LSTM_N_GLOBAL): global features (summary stats + scalars)
    
    LSTM processes the per-bar sequence, then concatenates global features for the heads.
    This fixes the previous "broken pseudo-sequence reshaping" issue where features
    were chunked arbitrarily rather than as meaningful per-bar time steps.
    """
    
    def __init__(
        self,
        input_size: int,
        context_window: int = 24,
        hidden_size: int = 128,
        num_layers: int = 2,
        dropout: float = DEFAULT_DROPOUT,
    ):
        super().__init__(input_size, context_window, dropout)
        
        # Dual-path layout: per-bar seq + global features
        # Per-bar: Open, High, Low, Close, Volume, returns, log_returns (7 per bar)
        # Global: summary stats (7) + FEATURE_COLS (45) + scalars (6) = 58
        self.n_per_bar = 7  # always fixed
        self.n_seq = context_window * self.n_per_bar
        self.n_global = input_size - self.n_seq
        
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        
        # Scale head sizes with hidden_size for balanced capacity:
        # hidden=32→head=16, hidden=64→head=32, hidden=128→head=64
        head_size = max(16, hidden_size // 2)
        conf_size = max(8, hidden_size // 4)
        
        self.lstm = nn.LSTM(
            input_size=self.n_per_bar,  # 7 per-bar features
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,
            dropout=dropout if num_layers > 1 else 0,
        )
        
        self.attention = nn.Sequential(
            nn.Linear(hidden_size, hidden_size // 2),
            nn.Tanh(),
            nn.Linear(hidden_size // 2, 1),
        )
        
        encoder_output_size = hidden_size + self.n_global
        
        self.entry_long_head = nn.Sequential(
            nn.Linear(encoder_output_size, head_size),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(head_size, 1),
        )
        
        self.entry_short_head = nn.Sequential(
            nn.Linear(encoder_output_size, head_size),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(head_size, 1),
        )
        
        self.sl_head = nn.Sequential(
            nn.Linear(encoder_output_size, head_size),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(head_size, 1),
        )
        
        self.tp_head = nn.Sequential(
            nn.Linear(encoder_output_size, head_size),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(head_size, 1),
        )
        
        self.confidence_head = nn.Sequential(
            nn.Linear(encoder_output_size, conf_size),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(conf_size, 1),
        )
        
        self._init_lstm_weights()
    
    def _init_lstm_weights(self) -> None:
        """Initialize LSTM weights with orthogonal initialization."""
        for name, param in self.lstm.named_parameters():
            if "weight_ih" in name:
                nn.init.xavier_uniform_(param)
            elif "weight_hh" in name:
                nn.init.orthogonal_(param)
            elif "bias" in name:
                nn.init.zeros_(param)
                n = param.size(0)
                param.data[n // 4 : n // 2].fill_(1.0)
    
    def encode(self, x: torch.Tensor) -> torch.Tensor:
        """Encode input using LSTM with attention + global features.
        
        Splits the flat input into:
          - Per-bar sequence: (batch, context_window, 7)
          - Global features:  (batch, n_global)
        
        Args:
            x: Input tensor of shape (batch, input_size)
            
        Returns:
            Encoded representation (batch, hidden_size + n_global)
        """
        batch_size = x.size(0)
        
        x_seq = x[:, :self.n_seq].view(batch_size, self.context_window, self.n_per_bar)
        x_global = x[:, self.n_seq:]
        
        lstm_out, _ = self.lstm(x_seq)
        
        attn_weights = self.attention(lstm_out)
        attn_weights = F.softmax(attn_weights, dim=1)
        context = torch.sum(lstm_out * attn_weights, dim=1)
        
        encoded = torch.cat([context, x_global], dim=1)
        
        return encoded


class TransformerMultiTaskModel(BaseMultiTaskModel):
    """Transformer-based multi-task model.
    
    Same dual-path layout as LSTM: per-bar sequence (24×7) + global features (58).
    Uses transformer encoder with positional encoding for the sequence.
    """
    
    def __init__(
        self,
        input_size: int,
        context_window: int = 24,
        d_model: int = 128,
        n_heads: int = 4,
        n_layers: int = 2,
        dropout: float = DEFAULT_DROPOUT,
    ):
        super().__init__(input_size, context_window, dropout)
        
        self.d_model = d_model
        self.n_heads = n_heads
        self.n_layers = n_layers
        
        # Dual-path layout: per-bar seq + global features
        self.n_per_bar = 7
        self.n_seq = context_window * self.n_per_bar
        self.n_global = input_size - self.n_seq
        
        self.input_projection = nn.Linear(self.n_per_bar, d_model)
        
        self.pos_encoding = nn.Parameter(
            torch.randn(1, context_window, d_model) * 0.02
        )
        
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=n_heads,
            dim_feedforward=d_model * 4,
            dropout=dropout,
            batch_first=True,
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=n_layers)
        
        encoder_output_size = d_model + self.n_global
        
        self.entry_long_head = nn.Sequential(
            nn.Linear(encoder_output_size, 32),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(32, 1),
        )
        
        self.entry_short_head = nn.Sequential(
            nn.Linear(encoder_output_size, 32),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(32, 1),
        )
        
        self.sl_head = nn.Sequential(
            nn.Linear(encoder_output_size, 32),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(32, 1),
        )
        
        self.tp_head = nn.Sequential(
            nn.Linear(encoder_output_size, 32),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(32, 1),
        )
        
        self.confidence_head = nn.Sequential(
            nn.Linear(encoder_output_size, 16),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(16, 1),
        )
        
        self._init_weights(self.input_projection)
        for module in self.entry_long_head.modules():
            self._init_weights(module)
        for module in self.entry_short_head.modules():
            self._init_weights(module)
        for module in self.sl_head.modules():
            self._init_weights(module)
        for module in self.tp_head.modules():
            self._init_weights(module)
        for module in self.confidence_head.modules():
            self._init_weights(module)
    
    def encode(self, x: torch.Tensor) -> torch.Tensor:
        """Encode input using Transformer.
        
        Splits input into per-bar sequence + global features.
        Applies transformer + mean pooling, then concatenates global features.
        
        Args:
            x: Input tensor of shape (batch, input_size)
            
        Returns:
            Encoded representation
        """
        batch_size = x.size(0)
        
        x_seq = x[:, :self.n_seq].view(batch_size, self.context_window, self.n_per_bar)
        x_global = x[:, self.n_seq:]
        
        x_proj = self.input_projection(x_seq)
        x_proj = x_proj + self.pos_encoding
        
        encoded_seq = self.transformer(x_proj)
        encoded = encoded_seq.mean(dim=1)
        
        encoded = torch.cat([encoded, x_global], dim=1)
        
        return encoded


class MLPMultiTaskModel(BaseMultiTaskModel):
    """Simple MLP-based multi-task model (baseline)."""
    
    def __init__(
        self,
        input_size: int,
        context_window: int = 24,
        hidden_sizes: Optional[list] = None,
        dropout: float = DEFAULT_DROPOUT,
    ):
        super().__init__(input_size, context_window, dropout)
        
        hidden_sizes = hidden_sizes or [256, 128, 64]
        
        layers = []
        prev_size = input_size
        for hidden_size in hidden_sizes:
            layers.extend([
                nn.Linear(prev_size, hidden_size),
                nn.BatchNorm1d(hidden_size),
                nn.ReLU(),
                nn.Dropout(dropout),
            ])
            prev_size = hidden_size
        
        self.encoder = nn.Sequential(*layers)
        
        self.entry_long_head = nn.Sequential(
            nn.Linear(prev_size, 32),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(32, 1),
        )
        
        self.entry_short_head = nn.Sequential(
            nn.Linear(prev_size, 32),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(32, 1),
        )
        
        self.sl_head = nn.Sequential(
            nn.Linear(prev_size, 32),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(32, 1),
        )
        
        self.tp_head = nn.Sequential(
            nn.Linear(prev_size, 32),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(32, 1),
        )
        
        self.confidence_head = nn.Sequential(
            nn.Linear(prev_size, 16),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(16, 1),
        )
        
        for module in self.modules():
            self._init_weights(module)
    
    def encode(self, x: torch.Tensor) -> torch.Tensor:
        """Encode input using MLP.
        
        Args:
            x: Input tensor of shape (batch, features)
            
        Returns:
            Encoded representation
        """
        return self.encoder(x)


def get_multi_task_model(
    model_type: str,
    input_size: int,
    context_window: int = 24,
    **kwargs,
) -> BaseMultiTaskModel:
    """Factory function to create multi-task model instances.
    
    Args:
        model_type: Type of model ('lstm', 'transformer', 'mlp')
        input_size: Input feature dimension
        context_window: Number of bars in context window
        **kwargs: Additional model-specific arguments
        
    Returns:
        Initialized model instance
    """
    model_type = model_type.lower()
    
    if model_type == "lstm":
        return LSTMMultiTaskModel(
            input_size=input_size,
            context_window=context_window,
            **kwargs,
        )
    elif model_type == "transformer":
        return TransformerMultiTaskModel(
            input_size=input_size,
            context_window=context_window,
            **kwargs,
        )
    elif model_type == "mlp":
        return MLPMultiTaskModel(
            input_size=input_size,
            context_window=context_window,
            **kwargs,
        )
    else:
        logger.warning(f"Неизвестный тип модели {model_type}, используется LSTM")
        return LSTMMultiTaskModel(
            input_size=input_size,
            context_window=context_window,
            **kwargs,
        )