from __future__ import annotations

import torch
import torch.nn.functional as F
from torch import nn


class Attention(nn.Module):
    """Multi-head (self or cross) attention on top of SDPA; `key_mask` is bool (B, S), True = valid."""

    def __init__(self, d: int, heads: int, dropout: float = 0.0):
        super().__init__()
        assert d % heads == 0, "d_model must be divisible by n_heads"
        self.h, self.dh = heads, d // heads
        self.q = nn.Linear(d, d)
        self.kv = nn.Linear(d, 2 * d)
        self.o = nn.Linear(d, d)
        self.dropout = dropout

    def forward(self, x, ctx=None, key_mask=None):
        ctx = x if ctx is None else ctx
        B, T, D = x.shape
        S = ctx.shape[1]
        q = self.q(x).view(B, T, self.h, self.dh).transpose(1, 2)
        kv = self.kv(ctx).view(B, S, 2, self.h, self.dh).permute(2, 0, 3, 1, 4)
        mask = key_mask[:, None, None, :] if key_mask is not None else None
        out = F.scaled_dot_product_attention(q, kv[0], kv[1], attn_mask=mask,
                                             dropout_p=self.dropout if self.training else 0.0)
        return self.o(out.transpose(1, 2).reshape(B, T, D))


class MLP(nn.Module):
    def __init__(self, d: int, ratio: float = 4.0, dropout: float = 0.0):
        super().__init__()
        hidden = int(d * ratio)
        self.net = nn.Sequential(nn.Linear(d, hidden), nn.GELU(), nn.Linear(hidden, d), nn.Dropout(dropout))

    def forward(self, x):
        return self.net(x)


class EncoderBlock(nn.Module):
    """Pre-LN transformer block."""

    def __init__(self, d, heads, mlp_ratio=4.0, dropout=0.0):
        super().__init__()
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.attn = Attention(d, heads, dropout)
        self.mlp = MLP(d, mlp_ratio, dropout)

    def forward(self, x, key_mask=None):
        h = self.ln1(x)
        x = x + self.attn(h, key_mask=key_mask)
        return x + self.mlp(self.ln2(x))


class DecisionBlock(nn.Module):
    """Choices read the context (cross-attn), then compare with each other (self-attn).

    No positional information over choices -> the block is permutation-equivariant:
    reordering the choices reorders the scores, nothing else.
    """

    def __init__(self, d, heads, mlp_ratio=4.0, dropout=0.0):
        super().__init__()
        self.ln_q, self.ln_ctx = nn.LayerNorm(d), nn.LayerNorm(d)
        self.cross = Attention(d, heads, dropout)
        self.ln_s = nn.LayerNorm(d)
        self.self_attn = Attention(d, heads, dropout)
        self.ln_m = nn.LayerNorm(d)
        self.mlp = MLP(d, mlp_ratio, dropout)

    def forward(self, c, ctx, ctx_mask, choice_mask):
        c = c + self.cross(self.ln_q(c), self.ln_ctx(ctx), key_mask=ctx_mask)
        c = c + self.self_attn(self.ln_s(c), key_mask=choice_mask)
        return c + self.mlp(self.ln_m(c))