"""LSTM + Multi-Head Attention 多时间尺度预警模型"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from src.utils.config import HIDDEN_DIM, LSTM_LAYERS, ATTENTION_HEADS, DROPOUT


class MultiHeadSelfAttention(nn.Module):
    """多头自注意力层"""
    def __init__(self, embed_dim: int, num_heads: int = ATTENTION_HEADS,
                 dropout: float = DROPOUT):
        super().__init__()
        assert embed_dim % num_heads == 0, f"embed_dim {embed_dim} must be divisible by num_heads {num_heads}"
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads

        self.qkv = nn.Linear(embed_dim, 3 * embed_dim)
        self.out_proj = nn.Linear(embed_dim, embed_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, D = x.shape
        qkv = self.qkv(x).reshape(B, T, 3, self.num_heads, self.head_dim)
        qkv = qkv.permute(2, 0, 3, 1, 4)  # (3, B, heads, T, head_dim)
        q, k, v = qkv[0], qkv[1], qkv[2]

        scale = self.head_dim ** -0.5
        attn = (q @ k.transpose(-2, -1)) * scale
        attn = F.softmax(attn, dim=-1)
        attn = self.dropout(attn)

        out = attn @ v  # (B, heads, T, head_dim)
        out = out.permute(0, 2, 1, 3).reshape(B, T, D)
        return self.out_proj(out)


class HeatRiskPredictor(nn.Module):
    """LSTM-Attention 多时间尺度高温风险预测模型

    Input: (B, T, input_dim) sequence of meteorological features
    Output: dict with 'short', 'medium', 'long' keys, each (B, 4) logits
    """
    def __init__(self, input_dim: int, hidden_dim: int = HIDDEN_DIM,
                 num_layers: int = LSTM_LAYERS, num_classes: int = 4):
        super().__init__()
        self.input_proj = nn.Linear(input_dim, hidden_dim)
        self.lstm = nn.LSTM(
            hidden_dim, hidden_dim, num_layers,
            batch_first=True, bidirectional=True, dropout=DROPOUT if num_layers > 1 else 0,
        )
        lstm_out_dim = hidden_dim * 2  # bidirectional
        self.attention = MultiHeadSelfAttention(lstm_out_dim)
        self.lstm_proj = nn.Linear(lstm_out_dim, hidden_dim)

        self.head_short = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(), nn.Dropout(DROPOUT),
            nn.Linear(hidden_dim // 2, num_classes),
        )
        self.head_medium = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(), nn.Dropout(DROPOUT),
            nn.Linear(hidden_dim // 2, num_classes),
        )
        self.head_long = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(), nn.Dropout(DROPOUT),
            nn.Linear(hidden_dim // 2, num_classes),
        )

    def forward(self, x: torch.Tensor) -> dict:
        x = self.input_proj(x)
        lstm_out, _ = self.lstm(x)
        attn_out = self.attention(lstm_out)
        last_hidden = self.lstm_proj(attn_out[:, -1, :])

        return {
            "short": self.head_short(last_hidden),
            "medium": self.head_medium(last_hidden),
            "long": self.head_long(last_hidden),
        }
