"""Shared building blocks. Attention is written out by hand so it exports cleanly to ONNX and runs on MPS."""

import math

import torch
from torch import nn


class ConvBNAct(nn.Sequential):
    def __init__(self, cin, cout, k=3, s=1, p=None, groups=1, act=True):
        if p is None:
            p = (k // 2) if isinstance(k, int) else tuple(x // 2 for x in k)
        layers = [nn.Conv2d(cin, cout, k, s, p, groups=groups, bias=False), nn.BatchNorm2d(cout)]
        if act:
            layers.append(nn.ReLU(inplace=True))
        super().__init__(*layers)


class SelfAttention(nn.Module):
    def __init__(self, d, heads, dropout=0.0):
        super().__init__()
        assert d % heads == 0
        self.h, self.dk = heads, d // heads
        self.qkv = nn.Linear(d, 3 * d)
        self.proj = nn.Linear(d, d)
        self.drop = nn.Dropout(dropout)

    def forward(self, x, bias=None):
        """x: (B, T, D). bias: additive attention bias broadcastable to (B, H, T, T)."""
        B, T, D = x.shape
        qkv = self.qkv(x).reshape(B, T, 3, self.h, self.dk).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(self.dk))
        if bias is not None:
            att = att + bias
        att = self.drop(att.softmax(-1))
        out = (att @ v).transpose(1, 2).reshape(B, T, D)
        return self.proj(out)


class TransformerBlock(nn.Module):
    def __init__(self, d, heads, ff, dropout=0.1):
        super().__init__()
        self.ln1 = nn.LayerNorm(d)
        self.attn = SelfAttention(d, heads, dropout)
        self.ln2 = nn.LayerNorm(d)
        self.ff = nn.Sequential(nn.Linear(d, ff), nn.GELU(), nn.Dropout(dropout), nn.Linear(ff, d))
        self.drop = nn.Dropout(dropout)

    def forward(self, x, bias=None):
        x = x + self.drop(self.attn(self.ln1(x), bias))
        return x + self.drop(self.ff(self.ln2(x)))


def sinusoid_table(n, d):
    pos = torch.arange(n, dtype=torch.float32)[:, None]
    div = torch.exp(torch.arange(0, d, 2, dtype=torch.float32) * (-math.log(10000.0) / d))
    pe = torch.zeros(n, d)
    pe[:, 0::2] = torch.sin(pos * div)
    pe[:, 1::2] = torch.cos(pos * div)
    return pe


def padding_bias(pad_mask, dtype):
    """pad_mask: (B, T) bool, True = padding -> additive bias (B, 1, 1, T)."""
    return pad_mask[:, None, None, :].to(dtype) * -1e4


def count_params(module: nn.Module) -> int:
    return sum(p.numel() for p in module.parameters())