"""RFC-005 decoder using Transformers' Llama attention, RoPE and generation.

The explicit model class is required when loading: ordinary Llama does not
implement this recipe's parameter-free QK normalization or two-matrix ReLU².
"""
from dataclasses import asdict
import torch
from torch import nn
from torch.nn import functional as F
from transformers import LlamaConfig, LlamaForCausalLM
from .recipe import Architecture


class EnglishBaseConfig(LlamaConfig):
    model_type = 'fabryka_english_base'

    def __init__(self, recipe=None, qk_norm=True, **kwargs):
        architecture = Architecture(**(recipe or {}))
        controlled = dict(
            vocab_size=architecture.vocab_size, hidden_size=architecture.width,
            intermediate_size=architecture.ffn_width, num_hidden_layers=architecture.layers,
            num_attention_heads=architecture.query_heads, num_key_value_heads=architecture.kv_heads,
            head_dim=architecture.head_dim, max_position_embeddings=architecture.context,
            hidden_act='silu', attention_bias=False, mlp_bias=False,
            attention_dropout=0.0, tie_word_embeddings=architecture.tied_embeddings,
        )
        for key, value in controlled.items():
            if key in kwargs and kwargs[key] != value:
                raise ValueError(f'{key} conflicts with recorded architecture')
            kwargs[key] = value
        if type(qk_norm) is not bool:
            raise ValueError('qk_norm must be boolean')
        super().__init__(**kwargs)
        self.recipe = asdict(architecture)
        self.qk_norm = qk_norm


class HeadNormalizedLinear(nn.Linear):
    """Project, then normalize each head in FP32 before RoPE; no learned gain."""
    def __init__(self, source, head_dim, eps):
        super().__init__(source.in_features, source.out_features, bias=False,
                         device=source.weight.device, dtype=source.weight.dtype)
        self.weight = source.weight
        self.head_dim, self.eps = head_dim, eps

    def forward(self, x):
        projected = F.linear(x, self.weight)
        heads = projected.unflatten(-1, (-1, self.head_dim))
        if hasattr(F, 'rms_norm'):
            normalized = F.rms_norm(heads.float(), (self.head_dim,), eps=self.eps)
        else:
            # PyTorch 2.3 compatibility for the pinned DCLM evaluation harness.
            values = heads.float()
            normalized = values * torch.rsqrt(values.square().mean(-1, keepdim=True) + self.eps)
        return normalized.to(projected.dtype).flatten(-2)


class ReluSquaredMLP(nn.Module):
    def __init__(self, original):
        super().__init__()
        self.up_proj = original.up_proj
        self.down_proj = original.down_proj

    def forward(self, x):
        return self.down_proj(F.relu(self.up_proj(x)).square())


class EnglishBaseForCausalLM(LlamaForCausalLM):
    config_class = EnglishBaseConfig

    def __init__(self, config):
        if not isinstance(config, EnglishBaseConfig):
            raise TypeError('EnglishBaseConfig is required; byte/Llama checkpoints are distinct')
        super().__init__(config)
        for block in self.model.layers:
            if config.recipe['activation'] == 'relu2':
                block.mlp = ReluSquaredMLP(block.mlp)
            if config.qk_norm:
                for name in ('q_proj', 'k_proj'):
                    projection = getattr(block.self_attn, name)
                    setattr(block.self_attn, name, HeadNormalizedLinear(
                        projection, config.head_dim, config.rms_norm_eps))

    def parameter_report(self):
        total = sum(p.numel() for p in self.parameters())
        expected = Architecture(**self.config.recipe).parameter_count
        if total != expected:
            raise ValueError(f'Implemented parameter count {total} differs from expected {expected}')
        embedding_ids = {id(self.get_input_embeddings().weight), id(self.get_output_embeddings().weight)}
        non_embedding = sum(p.numel() for p in self.parameters() if id(p) not in embedding_ids)
        return {'total_parameters': total, 'non_embedding_parameters': non_embedding,
                'tied_embeddings': self.get_input_embeddings().weight is self.get_output_embeddings().weight,
                'qk_norm': self.config.qk_norm}


    def mark_tied_weights_as_initialized(self, loading_info):
        """Preserve parameter-tying missing keys for Transformers 5.3 loading.

        This architecture uses parameter tying only, never shared modules.
        The remote-code compatibility path in 5.3 removes initialized missing
        keys before tie_weights, including the lm_head it has just marked.
        Keep that key so the loader ties it to the loaded input embedding.
        """
        for name in self.all_tied_weights_keys:
            self.get_parameter(name)._is_hf_initialized = True