import config as cfg
import torch
import torch.nn as nn
import torch.nn.functional as F


class SelfAttention(nn.Module):
    def __init__(self, hidden_dim: int, num_heads: int = cfg.DEFAULT_NUM_HEADS):
        super().__init__()
        self.attention = nn.MultiheadAttention(hidden_dim, num_heads, batch_first=True, dropout=0.1)
        self.norm = nn.LayerNorm(hidden_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        attn_out, _ = self.attention(x, x, x)
        return self.norm(x + attn_out)


class ConvBlock(nn.Module):
    def __init__(self, hidden_dim: int, kernel_size: int = cfg.DEFAULT_KERNEL_SIZE):
        super().__init__()
        self.conv = nn.Conv1d(hidden_dim, hidden_dim, kernel_size, padding=kernel_size // 2)
        self.bn = nn.BatchNorm1d(hidden_dim)
        self.relu = nn.ReLU()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x.transpose(1, 2)
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        return x.transpose(1, 2)


class FocalLoss(nn.Module):
    """Focal Loss с динамической балансировкой классов.

    Параметры
    ---------
    pos_weight : float или None
        Вес положительного класса. Если None → вычисляется автоматически
        из распределения меток в батче (или фиксированная alpha).
    gamma : float
        Фокусирующий параметр: чем больше, тем сильнее фокус на трудных примерах.
    alpha : float
        Базовый вес (используется если pos_weight=None).
    """
    def __init__(self, pos_weight: float = None, gamma: float = 2.0,
                 alpha: float = 0.5):
        super().__init__()
        self.gamma = gamma
        self.alpha = alpha
        # pos_weight может быть задан позже через set_pos_weight()
        self.pos_weight = pos_weight

    def set_pos_weight(self, pw: float):
        """Установить вес положительного класса (балансировка)."""
        self.pos_weight = pw

    def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        if targets.numel() == 0:
            return torch.tensor(0.0, device=inputs.device, requires_grad=False)

        if self.pos_weight is None:
            pos_count = targets.sum().item()
            neg_count = len(targets) - pos_count
            pw = neg_count / pos_count if pos_count > 0 and neg_count > 0 else 1.0
        else:
            pw = self.pos_weight

        bce = F.binary_cross_entropy_with_logits(
            inputs, targets,
            pos_weight=torch.tensor(pw, device=inputs.device),
            reduction='none'
        )
        pt = torch.exp(-bce)
        focal = (1.0 - pt) ** self.gamma * bce
        return focal.mean()


class DualHeadLSTMModel(nn.Module):
    """Улучшенная архитектура: Conv1D + Self-Attention + LSTM."""
    def __init__(self, input_dim: int, hidden_dim: int = 48, num_layers: int = cfg.DEFAULT_NUM_LAYERS, 
                 dropout: float = cfg.DEFAULT_DROPOUT, use_conv: bool = True, use_attention: bool = True):
        super().__init__()
        self.use_conv = use_conv
        self.use_attention = use_attention
        
        self.input_proj = nn.Linear(input_dim, hidden_dim)
        
        if use_conv:
            self.conv_block = ConvBlock(hidden_dim, kernel_size=cfg.DEFAULT_KERNEL_SIZE)
        
        if use_attention:
            self.attention = SelfAttention(hidden_dim, num_heads=cfg.DEFAULT_NUM_HEADS)
        
        self.lstm = nn.LSTM(hidden_dim, hidden_dim, num_layers, batch_first=True, 
                           dropout=dropout if num_layers > 1 else 0.0)
        self.bn = nn.BatchNorm1d(hidden_dim)
        
        self.fc_long = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim // 2, 1)
        )
        self.fc_short = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim // 2, 1)
        )

    def forward(self, x: torch.Tensor):
        x = self.input_proj(x)
        
        if self.use_conv:
            x = self.conv_block(x)
        
        if self.use_attention:
            x = self.attention(x)
        
        out, _ = self.lstm(x)
        last_step = out[:, -1, :]
        features = self.bn(last_step)
        
        logits_long = self.fc_long(features).squeeze(-1)
        logits_short = self.fc_short(features).squeeze(-1)
        return logits_long, logits_short