from dataclasses import dataclass, asdict
from typing import NamedTuple
import torch
from torch import Tensor, nn
import torch.nn.functional as F
MOODS=()
ROLES=()
INSTRUMENTS=()
FRAME_FEATURES=()
CONTROL_FEATURES=()

@dataclass(frozen=True)
class A02bConfig:
    width: int = 96
    layers: int = 4
    heads: int = 4
    dropout: float = 0.1
    frames_per_beat: int = 16
    instrument_embedding_dim: int = 8
    use_convolution: bool = True
    use_relative_bias: bool = True
    use_stop_head: bool = False
    relative_transitions_beats: tuple[float, ...] = (2 / 16, 3 / 16, 4 / 16, 5 / 16, 6 / 16, 7 / 16, 8 / 16, 3 / 4, 1, 2, 4, 8, 16, 32)

    def __post_init__(self) -> None:
        if self.instrument_embedding_dim <= 0:
            raise ValueError('instrument_embedding_dim must be positive')
        if self.width <= 0 or self.heads <= 0 or self.width % self.heads:
            raise ValueError('width must be positive and divisible by heads')
        if self.layers <= 0 or not 0 <= self.dropout < 1:
            raise ValueError('layers must be positive; dropout must be in [0,1)')
        if self.frames_per_beat not in (16, 48):
            raise ValueError('A02b supports separately configured rates 16 and 48')
        edges = self.relative_transitions_beats
        if len(edges) != 14 or not all((x > 0 for x in edges)):
            raise ValueError('32 signed buckets require 14 positive transitions')
        if not all((a < b for a, b in zip(edges, edges[1:]))):
            raise ValueError('relative distance transitions must increase')
        if not all((abs(x * self.frames_per_beat - round(x * self.frames_per_beat)) < 1e-08 for x in edges)):
            raise ValueError('bucket boundaries must align with the frame grid')

class A02bOutput(NamedTuple):
    normalized_cc1: Tensor
    stop_logits: Tensor | None

def zero_padding(h: Tensor, mask: Tensor) -> Tensor:
    return h.masked_fill(~mask[..., None], 0.0)

class SignedRelativeBias(nn.Module):

    def __init__(self, config: A02bConfig) -> None:
        super().__init__()
        self.register_buffer('transitions_frames', torch.tensor([round(x * config.frames_per_beat) for x in config.relative_transitions_beats], dtype=torch.long))
        self.table = nn.Embedding(32, config.heads)
        nn.init.zeros_(self.table.weight)

    def bucket_ids(self, length: int) -> Tensor:
        positions = torch.arange(length, device=self.transitions_frames.device)
        offset = positions[None, :] - positions[:, None]
        magnitude_bucket = torch.ones_like(offset)
        for transition in self.transitions_frames.unbind():
            magnitude_bucket = magnitude_bucket + (offset.abs() >= transition)
        signed_bucket = magnitude_bucket + (offset > 0).long() * 16
        return torch.where(offset == 0, 0, signed_bucket)

    def forward(self, length: int) -> Tensor:
        return self.table(self.bucket_ids(length)).permute(2, 0, 1)[None]

class FiLM(nn.Module):

    def __init__(self, width: int) -> None:
        super().__init__()
        self.modulation = nn.Linear(64, 2 * width)
        nn.init.zeros_(self.modulation.weight)
        nn.init.zeros_(self.modulation.bias)

    def forward(self, u: Tensor, c: Tensor) -> Tensor:
        delta_gamma, beta = self.modulation(c).chunk(2, dim=-1)
        return (1 + delta_gamma[:, None]) * u + beta[:, None]

class BidirectionalAttention(nn.Module):

    def __init__(self, config: A02bConfig) -> None:
        super().__init__()
        self.heads = config.heads
        self.head_dim = config.width // config.heads
        self.dropout = config.dropout
        self.qkv = nn.Linear(config.width, 3 * config.width)
        self.output = nn.Linear(config.width, config.width)

    def forward(self, u: Tensor, attention_bias: Tensor) -> Tensor:
        batch, length, width = u.shape
        qkv = self.qkv(u).reshape(batch, length, 3, self.heads, self.head_dim)
        q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0)
        result = F.scaled_dot_product_attention(q, k, v, attn_mask=attention_bias, dropout_p=self.dropout if self.training else 0.0, is_causal=False)
        result = result.transpose(1, 2).reshape(batch, length, width)
        return self.output(result)

class EncoderBlock(nn.Module):

    def __init__(self, config: A02bConfig) -> None:
        super().__init__()
        d = config.width
        self.norm_attention = nn.LayerNorm(d)
        self.norm_ff = nn.LayerNorm(d)
        self.film_attention = FiLM(d)
        self.film_ff = FiLM(d)
        self.attention = BidirectionalAttention(config)
        self.ff = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, h: Tensor, c: Tensor, mask: Tensor, bias: Tensor) -> Tensor:
        u = self.film_attention(self.norm_attention(h), c)
        h = zero_padding(h + self.dropout(self.attention(u, bias)), mask)
        u = self.film_ff(self.norm_ff(h), c)
        return zero_padding(h + self.dropout(self.ff(u)), mask)

class AireA02b(nn.Module):

    def __init__(self, config: A02bConfig | None=None) -> None:
        super().__init__()
        self.config = config or A02bConfig()
        cfg = self.config
        d = cfg.width
        self.mood_embedding = nn.Embedding(len(MOODS), 16)
        self.role_embedding = nn.Embedding(len(ROLES), 8)
        self.instrument_embedding = nn.Embedding(len(INSTRUMENTS), cfg.instrument_embedding_dim)
        self.control_mlp = nn.Sequential(nn.Linear(30 + cfg.instrument_embedding_dim, 64), nn.GELU(), nn.Linear(64, 64))
        self.lead_embedding = nn.Embedding(129, d)
        self.pitch_projection = nn.Linear(128, d, bias=False)
        self.onset_projection = nn.Linear(128, d, bias=False)
        self.offset_projection = nn.Linear(128, d, bias=False)
        self.numeric_projection = nn.Linear(25, d)
        self.control_projection = nn.Linear(64, d)
        self.convolutions = nn.ModuleList([nn.Conv1d(d, d, kernel_size=5, padding=2) for _ in range(2)] if cfg.use_convolution else [])
        self.relative_bias = SignedRelativeBias(cfg) if cfg.use_relative_bias else None
        self.blocks = nn.ModuleList([EncoderBlock(cfg) for _ in range(cfg.layers)])
        self.final_norm = nn.LayerNorm(d)
        self.cc_head = nn.Linear(d, 1)
        self.stop_head = nn.Linear(d, 1) if cfg.use_stop_head else None

    def manifest(self) -> dict:
        return {'model': 'aire-A02b', 'input_schema': 'A02b-25-instrument-v1', 'config': asdict(self.config), 'moods': list(MOODS), 'roles': list(ROLES), 'frame_features': list(FRAME_FEATURES), 'control_features': list(CONTROL_FEATURES), 'instruments': list(INSTRUMENTS), 'velocity_output': 'preserve-input-only', 'rest_pitch_index': 128, 'relative_bucket_rule': 'signed-piecewise-beats-v1; zero=0; past=1:15; future=17:31'}

    def forward(self, *, lead_pitch: Tensor, pitch_roll: Tensor, onset_roll: Tensor, offset_roll: Tensor, frame_numeric: Tensor, mood_id: Tensor, role_id: Tensor, instrument_id: Tensor, ctrl_numeric: Tensor, mask: Tensor) -> A02bOutput:
        if mask.ndim != 2 or mask.dtype != torch.bool:
            raise ValueError('mask must be bool [B,T]')
        batch, length = mask.shape
        inputs = {'lead_pitch': (lead_pitch, (batch, length), torch.long), 'pitch_roll': (pitch_roll, (batch, length, 128), self.cc_head.weight.dtype), 'onset_roll': (onset_roll, (batch, length, 128), self.cc_head.weight.dtype), 'offset_roll': (offset_roll, (batch, length, 128), self.cc_head.weight.dtype), 'frame_numeric': (frame_numeric, (batch, length, 25), self.cc_head.weight.dtype), 'mood_id': (mood_id, (batch,), torch.long), 'role_id': (role_id, (batch,), torch.long), 'instrument_id': (instrument_id, (batch,), torch.long), 'ctrl_numeric': (ctrl_numeric, (batch, 6), self.cc_head.weight.dtype)}
        device = self.cc_head.weight.device
        if mask.device != device:
            raise ValueError('mask and model must be on the same device')
        for name, (tensor, shape, dtype) in inputs.items():
            if tuple(tensor.shape) != shape or tensor.dtype != dtype or tensor.device != device:
                raise ValueError(f'{name} must have shape {shape}, dtype {dtype}, device {device}')
        if batch == 0 or length == 0:
            empty = self.cc_head.weight.new_zeros((batch, length))
            return A02bOutput(empty, empty if self.stop_head is not None else None)
        for name, tensor, count in (('mood_id', mood_id, len(MOODS)), ('role_id', role_id, len(ROLES)), ('instrument_id', instrument_id, len(INSTRUMENTS))):
            if bool(((tensor < 0) | (tensor >= count)).any()):
                raise ValueError(f'{name} is outside the frozen category range')
        if not bool(mask.any(dim=1).all()):
            raise ValueError('all-padding samples must bypass the model')
        if bool((~mask[:, :-1] & mask[:, 1:]).any()):
            raise ValueError('mask must consist of a valid prefix and right padding')
        lead_pitch = lead_pitch.masked_fill(~mask, 128)
        pitch_roll = zero_padding(pitch_roll, mask)
        onset_roll = zero_padding(onset_roll, mask)
        offset_roll = zero_padding(offset_roll, mask)
        numeric = zero_padding(frame_numeric, mask)
        velocity_enabled = ctrl_numeric[:, 5:6] != 0
        gate = velocity_enabled[:, None, :]
        numeric = torch.cat([numeric[..., :3], torch.where(gate, numeric[..., 3:4], 0.0), numeric[..., 4:24], torch.where(gate, numeric[..., 24:25], 0.0)], dim=-1)
        c = self.control_mlp(torch.cat([self.mood_embedding(mood_id), self.role_embedding(role_id), self.instrument_embedding(instrument_id), ctrl_numeric], dim=-1))
        h = zero_padding(self.lead_embedding(lead_pitch) + self.pitch_projection(pitch_roll) + self.onset_projection(onset_roll) + self.offset_projection(offset_roll) + self.numeric_projection(numeric) + self.control_projection(c)[:, None], mask)
        if self.convolutions:
            residual = h
            for convolution in self.convolutions:
                h = F.gelu(convolution(h.transpose(1, 2))).transpose(1, 2)
                h = zero_padding(h, mask)
            h = zero_padding(residual + h, mask)
        bias = h.new_zeros((batch, 1, 1, length)).masked_fill(~mask[:, None, None, :], float('-inf'))
        if self.relative_bias is not None:
            bias = bias + self.relative_bias(length).to(h.dtype)
        for block in self.blocks:
            h = block(h, c, mask, bias)
        h = self.final_norm(h)
        y = torch.sigmoid(self.cc_head(h).squeeze(-1)).masked_fill(~mask, 0.0)
        stop = None
        if self.stop_head is not None:
            stop = self.stop_head(h).squeeze(-1).masked_fill(~mask, 0.0)
        return A02bOutput(y, stop)