"""Random-initialized decoder: pre-RMSNorm, RoPE, causal SDPA, SwiGLU, tied head."""
from __future__ import annotations
import math
import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.checkpoint import checkpoint
from .config import ModelConfig


class RMSNorm(nn.Module):
    def __init__(self, size, epsilon):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(size))
        self.epsilon = epsilon

    def forward(self, x):
        y = x.float()
        y = y * torch.rsqrt(y.square().mean(-1, keepdim=True) + self.epsilon)
        return y.to(x.dtype) * self.weight


class Attention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.heads = config.heads
        self.head_dim = config.hidden_size // config.heads
        self.qkv = nn.Linear(config.hidden_size, 3*config.hidden_size, bias=False)
        self.output = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
        self.dropout = config.dropout
        inv_freq = 1.0 / (config.rope_theta ** (torch.arange(0, self.head_dim, 2).float() / self.head_dim))
        frequencies = torch.outer(torch.arange(config.context_length).float(), inv_freq)
        self.register_buffer("cos", frequencies.cos()[None, None], persistent=False)
        self.register_buffer("sin", frequencies.sin()[None, None], persistent=False)

    def rotate(self, x):
        # Interleaved real/imaginary coordinates; RoPE computed in float32.
        even, odd = x.float()[..., 0::2], x.float()[..., 1::2]
        cos, sin = self.cos[:, :, :x.shape[-2]], self.sin[:, :, :x.shape[-2]]
        return torch.stack((even*cos-odd*sin, even*sin+odd*cos), -1).flatten(-2).to(x.dtype)

    def forward(self, x):
        b, t, h = x.shape
        q, k, v = self.qkv(x).view(b,t,3,self.heads,self.head_dim).permute(2,0,3,1,4).unbind(0)
        y = F.scaled_dot_product_attention(self.rotate(q), self.rotate(k), v, is_causal=True,
                dropout_p=self.dropout if self.training else 0.0)
        return self.output(y.transpose(1,2).contiguous().view(b,t,h))


class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        h, f = config.hidden_size, config.intermediate_size
        self.attention_norm = RMSNorm(h, config.rms_epsilon)
        self.attention = Attention(config)
        self.ffn_norm = RMSNorm(h, config.rms_epsilon)
        self.gate = nn.Linear(h,f,bias=False)
        self.up = nn.Linear(h,f,bias=False)
        self.down = nn.Linear(f,h,bias=False)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x):
        x = x + self.dropout(self.attention(self.attention_norm(x)))
        z = self.ffn_norm(x)
        return x + self.dropout(self.down(F.silu(self.gate(z))*self.up(z)))


class StasisDecoder(nn.Module):
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.config = config
        self.embedding = nn.Embedding(config.vocab_size, config.hidden_size)
        self.blocks = nn.ModuleList(Block(config) for _ in range(config.layers))
        self.norm = RMSNorm(config.hidden_size, config.rms_epsilon)
        self.apply(self._init)
        # Scale residual projections with depth, after common random initialization.
        for block in self.blocks:
            nn.init.normal_(block.attention.output.weight, std=0.02/math.sqrt(2*config.layers))
            nn.init.normal_(block.down.weight, std=0.02/math.sqrt(2*config.layers))

    @staticmethod
    def _init(module):
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, tokens, targets=None, reduction="mean"):
        if tokens.ndim != 2 or tokens.shape[1] > self.config.context_length:
            raise ValueError("Expected [batch, sequence] within configured context length")
        x = self.embedding(tokens)
        for block in self.blocks:
            x = checkpoint(block, x, use_reentrant=False) if self.training and self.config.gradient_checkpointing else block(x)
        logits = F.linear(self.norm(x), self.embedding.weight)
        loss = None
        if targets is not None:
            if targets.shape != tokens.shape:
                raise ValueError("Targets must match input shape")
            loss = F.cross_entropy(logits.float().reshape(-1, self.config.vocab_size), targets.reshape(-1), reduction=reduction)
        return logits, loss

    @torch.no_grad()
    def generate(self, tokens, max_new_tokens=64, temperature=1.0, top_k=50, eos_token_id=None):
        if max_new_tokens < 0 or temperature <= 0 or top_k < 0:
            raise ValueError("Invalid sampling controls")
        was_training = self.training
        self.eval()
        try:
            for _ in range(max_new_tokens):
                logits, _ = self(tokens[:, -self.config.context_length:])
                logits = logits[:, -1].float() / temperature
                if top_k:
                    threshold = torch.topk(logits, min(top_k, logits.size(-1))).values[:, -1,None]
                    logits = logits.masked_fill(logits < threshold, -float("inf"))
                nxt = torch.multinomial(F.softmax(logits, dim=-1), 1)
                tokens = torch.cat([tokens,nxt], dim=1)
                if eos_token_id is not None and bool((nxt == eos_token_id).all()):
                    break
            return tokens
        finally:
            self.train(was_training)