"""Fixed V5 capacity allocation: the V4 recipe and 8K tokenizer scaled to d512/L8.

Same operations as V2/V4 (RMSNorm, SwiGLU, bias-free, tied head, learned
positions, head width 64); only width, depth and FFN width change. V4 stays frozen.
"""
import math
from dataclasses import dataclass, fields

import torch

from vimeml.training.model_v2 import GPTV2Config, TinyGPTV2
from vimeml.training.model_v4 import KVCacheV4

PARAMETERS = 28_385_792
LAYER_PARAMETERS = 3_015_680
HEAD_WIDTH = 64
CACHE_SHAPE = (8, 1, 8, 128, HEAD_WIDTH)


@dataclass(frozen=True)
class GPTV5Config(GPTV2Config):
    vocab_size: int = 8192
    context_length: int = 128
    d_model: int = 512
    n_heads: int = 8
    n_layers: int = 8
    d_ff: int = 1280
    dropout: float = 0.0

    def __post_init__(self):
        super().__post_init__()
        for field in fields(self):
            if type(getattr(self, field.name)) is not type(field.default) or getattr(self, field.name) != field.default:
                raise ValueError(f"V5 {field.name} is fixed to {field.default}")


class TinyGPTV5(TinyGPTV2):
    """V2 operations at the fixed V5 size, with the FP32 research KV path of V4."""
    def __init__(self, config=GPTV5Config()):
        if type(config) is not GPTV5Config:
            raise ValueError("V5 requires its fixed GPTV5Config")
        super().__init__(config)
        if self.parameter_count() != PARAMETERS or any(
                sum(p.numel() for p in block.parameters()) != LAYER_PARAMETERS for block in self.blocks):
            raise ValueError("V5 parameter budget does not match its contract")

    def cache_shape(self, batch_size=1):
        layers, _, heads, length, width = CACHE_SHAPE
        return (layers, batch_size, heads, length, width)

    def new_cache(self, batch_size=1):
        if type(batch_size) is not int or batch_size < 1:
            raise ValueError("Positive integer cache batch required")
        shape = self.cache_shape(batch_size)
        device = self.token_embedding.weight.device
        return KVCacheV4(torch.zeros(shape, device=device, dtype=torch.float32),
                         torch.zeros(shape, device=device, dtype=torch.float32))

    @torch.inference_mode()
    def forward_cached(self, input_ids, cache):
        if self.training or self.token_embedding.weight.dtype != torch.float32:
            raise ValueError("Research KV path requires eval mode and FP32 weights")
        if input_ids.ndim != 2 or input_ids.dtype not in (torch.int32, torch.int64):
            raise ValueError("Expected integer [batch,time] tokens")
        batch, size = input_ids.shape
        shape = self.cache_shape(batch)
        heads, width, context = shape[2], shape[4], shape[3]
        if (cache.key.shape != shape or cache.value.shape != shape
                or cache.key.dtype != torch.float32 or cache.value.dtype != torch.float32
                or cache.key.device != input_ids.device or cache.value.device != input_ids.device
                or input_ids.device != self.token_embedding.weight.device
                or type(cache.length) is not int or cache.length < 0
                or size < 1 or cache.length + size > context):
            raise ValueError("Invalid KV cache or context overflow")
        if int(input_ids.min()) < 0 or int(input_ids.max()) >= self.config.vocab_size:
            raise ValueError("Token outside fixed vocabulary")
        start, end = cache.length, cache.length + size
        positions = torch.arange(start, end, device=input_ids.device)
        hidden = self.token_embedding(input_ids) + self.position_embedding(positions)
        causal = positions[:, None] >= torch.arange(end, device=input_ids.device)[None, :]
        for layer, block in enumerate(self.blocks):
            normalized = block.attention_norm(hidden)
            qkv = block.attention.qkv(normalized).reshape(batch, size, 3, heads, width)
            query, key, value = qkv.permute(2, 0, 3, 1, 4).unbind(0)
            cache.key[layer, :, :, start:end].copy_(key)
            cache.value[layer, :, :, start:end].copy_(value)
            scores = torch.matmul(query, cache.key[layer, :, :, :end].transpose(-2, -1)) / math.sqrt(width)
            probabilities = torch.softmax(scores.masked_fill(~causal, -torch.inf), dim=-1)
            attended = torch.matmul(probabilities, cache.value[layer, :, :, :end])
            attended = attended.transpose(1, 2).reshape(batch, size, self.config.d_model)
            hidden = hidden + block.attention.projection(attended)
            hidden = hidden + block.ffn(block.ffn_norm(hidden))
        cache.length = end
        return self.lm_head(self.final_norm(hidden))