"""
PyTorch port of the flaxchat GPT with FULL YatNMN: YatNMN attention + YatNMN MLP.

YatNMN Attention (no Q/K projections):
    x_heads = RoPE(x.reshape(B,T,H,D))       # reuse input as both Q and K
    dots = x_heads @ x_heads^T                 # pairwise dot products
    dist² = ||x_i||² + ||x_j||² - 2·dots      # pairwise squared distances
    scores = (dots + softplus(b))² / (dist² + softplus(ε))
    causal: strict j < i (tokens cannot attend to themselves)
    normalize: L1 norm (scores / sum + 1e-8), NOT softmax

YatNMN MLP: same as yatnmn_gpt.py (softplus bias, learnable epsilon, optional alpha)

No value embeddings (value_embeds dict is empty).
"""
from __future__ import annotations

import math
from dataclasses import dataclass, field
from typing import Optional, Tuple, List

import torch
import torch.nn as nn
import torch.nn.functional as F

try:
    from .torch_gpt import (
        rms_norm, precompute_rotary_embeddings, apply_rotary_emb,
        has_ve, compute_window_sizes,
    )
    from .yatnmn_gpt import YatNMN
except ImportError:
    from torch_port.torch_gpt import (
        rms_norm, precompute_rotary_embeddings, apply_rotary_emb,
        has_ve, compute_window_sizes,
    )
    from torch_port.yatnmn.yatnmn_gpt import YatNMN


@dataclass
class YatFullGPTConfig:
    sequence_len: int = 1024
    vocab_size: int = 32768
    n_layer: int = 12
    n_head: int = 12
    n_kv_head: int = 12
    n_embd: int = 768
    window_pattern: str = "SSSL"
    tie_embeddings: bool = True
    rope_base: float = 100000.0
    pad_vocab_size_to: int = 64

    mlp_type: str = "yatnmn-softplus"
    scalar_bias: bool = True
    softplus_bias: bool = True
    learnable_epsilon: bool = True
    epsilon_init: float = 1e-3
    constant_alpha: bool = False

    @property
    def head_dim(self) -> int:
        return self.n_embd // self.n_head

    @property
    def padded_vocab_size(self) -> int:
        v = self.vocab_size
        p = self.pad_vocab_size_to
        return ((v + p - 1) // p) * p


class YatNMNAttention(nn.Module):
    """YatNMN attention: no Q/K projections, pairwise scoring, L1 norm."""

    def __init__(self, config: YatFullGPTConfig, layer_idx: int):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        n_embd, n_head, n_kv_head = config.n_embd, config.n_head, config.n_kv_head
        head_dim = n_embd // n_head
        self.c_v = nn.Linear(n_embd, n_kv_head * head_dim, bias=False)
        self.c_proj = nn.Linear(n_embd, n_embd, bias=False)
        self.attn_bias_raw = nn.Parameter(torch.zeros(n_head))
        raw_eps = math.log(math.expm1(1e-3))
        self.attn_eps_raw = nn.Parameter(torch.full((n_head,), raw_eps))
        self._has_ve = has_ve(layer_idx, config.n_layer)
        self.ve_gate = nn.Linear(12, n_kv_head, bias=False) if self._has_ve else None

    def forward(
        self,
        x: torch.Tensor,
        ve: Optional[torch.Tensor],
        cos: torch.Tensor,
        sin: torch.Tensor,
        window_size: Tuple[int, int],
    ) -> torch.Tensor:
        B, T, C = x.shape
        n_head = self.config.n_head
        n_kv_head = self.config.n_kv_head
        head_dim = C // n_head

        x_heads = apply_rotary_emb(x.reshape(B, T, n_head, head_dim), cos, sin)
        xh = x_heads.float()
        xh_t = xh.transpose(1, 2)  # (B, H, T, D)

        dots = torch.matmul(xh_t, xh_t.transpose(-2, -1))  # (B, H, T, T)
        x_sq = (xh_t ** 2).sum(dim=-1)  # (B, H, T)
        dist_sq = torch.clamp(x_sq.unsqueeze(-1) + x_sq.unsqueeze(-2) - 2.0 * dots, min=0.0)

        b = F.softplus(self.attn_bias_raw)  # (H,)
        eps = F.softplus(self.attn_eps_raw)  # (H,)
        scores = (dots + b[None, :, None, None]) ** 2 / (dist_sq + eps[None, :, None, None])

        # Strict causal: j < i (no self-attention)
        device = x.device
        row_idx = torch.arange(T, device=device)
        col_idx = torch.arange(T, device=device)
        causal = col_idx.unsqueeze(0) < row_idx.unsqueeze(1)  # (T, T)

        window_left = window_size[0]
        if 0 < window_left < T:
            causal = causal & ((row_idx.unsqueeze(1) - col_idx.unsqueeze(0)) <= window_left)

        scores = torch.where(causal.unsqueeze(0).unsqueeze(0), scores, torch.zeros_like(scores))
        scores = scores / (scores.sum(dim=-1, keepdim=True) + 1e-8)

        # Value
        v = self.c_v(x).reshape(B, T, n_kv_head, head_dim)
        if self._has_ve and ve is not None:
            ve2 = ve.reshape(B, T, n_kv_head, head_dim)
            gate = 3.0 * torch.sigmoid(self.ve_gate(x[..., :12]))
            v = v + gate.unsqueeze(-1) * ve2

        if n_kv_head < n_head:
            v = v.repeat_interleave(n_head // n_kv_head, dim=2)

        v_t = v.transpose(1, 2)  # (B, H, T, D)
        scores = scores.to(v_t.dtype)
        y = torch.matmul(scores, v_t)
        y = y.transpose(1, 2).reshape(B, T, C)
        return self.c_proj(y)


class YatMLP(nn.Module):
    def __init__(self, config: YatFullGPTConfig):
        super().__init__()
        n, ff = config.n_embd, 4 * config.n_embd
        self.c_fc = YatNMN(
            n, ff,
            use_bias=True,
            softplus_bias=config.softplus_bias,
            scalar_bias=config.scalar_bias,
            learnable_epsilon=config.learnable_epsilon,
            epsilon_init=config.epsilon_init,
            use_alpha=True,
            constant_alpha=config.constant_alpha,
        )
        self.c_proj = nn.Linear(ff, n, bias=False)
        nn.init.zeros_(self.c_proj.weight)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.c_proj(self.c_fc(x))


class Block(nn.Module):
    def __init__(self, config: YatFullGPTConfig, layer_idx: int):
        super().__init__()
        self.attn = YatNMNAttention(config, layer_idx)
        self.mlp = YatMLP(config)

    def forward(self, x, ve, cos, sin, window_size):
        x = x + self.attn(rms_norm(x), ve, cos, sin, window_size)
        x = x + self.mlp(rms_norm(x))
        return x


class YatFull_GPT(nn.Module):
    """Full YatNMN GPT: YatNMN attention + YatNMN MLP, no value embeddings."""

    def __init__(self, config: YatFullGPTConfig):
        super().__init__()
        self.config = config
        self.window_sizes = compute_window_sizes(_ConfigShim(config))
        padded_vocab = config.padded_vocab_size
        self.padded_vocab_size = padded_vocab

        self.wte = nn.Embedding(padded_vocab, config.n_embd)
        self.blocks = nn.ModuleList([Block(config, i) for i in range(config.n_layer)])

        self.tie_embeddings = config.tie_embeddings
        self.lm_head = None if config.tie_embeddings else nn.Linear(padded_vocab, config.n_embd, bias=False)

        self.resid_lambdas = nn.Parameter(torch.ones(config.n_layer))
        self.x0_lambdas = nn.Parameter(torch.zeros(config.n_layer))

        self.smear_gate = nn.Linear(24, 1, bias=False)
        self.smear_lambda = nn.Parameter(torch.zeros(1))
        self.backout_lambda = nn.Parameter(0.2 * torch.ones(1))

        # Empty — no value embeddings in this variant
        self.value_embeds = nn.ModuleDict()

        # RoPE buffers are computed lazily in _get_rope below to avoid HF meta-init
        # zeroing them out when a checkpoint doesn't include them.
        self._rope_max_len = config.sequence_len * 10
        self._rope_head_dim = config.head_dim
        self._rope_base = config.rope_base
        # Placeholder buffers so state_dict save/load still sees them (but persistent=False
        # keeps them out of safetensors so HF won't overwrite with zeros).
        self.register_buffer("rope_cos", torch.empty(0), persistent=False)
        self.register_buffer("rope_sin", torch.empty(0), persistent=False)
        self._rope_initialized = False

    def _get_rope(self, T: int, dtype: torch.dtype, device: torch.device):
        if (not self._rope_initialized
                or self.rope_cos.numel() == 0
                or self.rope_cos.shape[1] < T):
            cos, sin = precompute_rotary_embeddings(
                max(T, self._rope_max_len), self._rope_head_dim, base=self._rope_base
            )
            self.rope_cos = cos.to(device)
            self.rope_sin = sin.to(device)
            self._rope_initialized = True
        return self.rope_cos[:, :T].to(dtype), self.rope_sin[:, :T].to(dtype)

    def forward(self, idx: torch.Tensor) -> torch.Tensor:
        B, T = idx.shape
        cfg = self.config

        cos, sin = self._get_rope(T, self.wte.weight.dtype, self.wte.weight.device)

        x = self.wte(idx)
        x = rms_norm(x)

        gate = self.smear_lambda * torch.sigmoid(self.smear_gate(x[:, 1:, :24]))
        x_smeared = x[:, 1:] + gate * x[:, :-1]
        x = torch.cat([x[:, :1], x_smeared], dim=1)

        x0 = x
        n_layer = cfg.n_layer
        backout_layer = n_layer // 2
        x_backout = None

        for i, block in enumerate(self.blocks):
            x = self.resid_lambdas[i] * x + self.x0_lambdas[i] * x0
            x = block(x, None, cos, sin, self.window_sizes[i])
            if i == backout_layer:
                x_backout = x

        if x_backout is not None:
            x = x - self.backout_lambda * x_backout
        x = rms_norm(x)

        softcap = 15.0
        logits = x @ self.wte.weight.t() if self.tie_embeddings else self.lm_head(x)
        logits = logits[..., : cfg.vocab_size].to(torch.float32)
        return softcap * torch.tanh(logits / softcap)

    @classmethod
    def from_pretrained(cls, path, map_location="cpu"):
        payload = torch.load(path, map_location=map_location, weights_only=False)
        config = YatFullGPTConfig(**payload["config"])
        model = cls(config)
        missing, unexpected = model.load_state_dict(payload["state_dict"], strict=False)
        real_missing = [k for k in missing if not k.startswith("rope_") and "_alpha_const" not in k]
        if real_missing:
            raise RuntimeError(f"Missing keys: {real_missing}")
        model.eval()
        return model


class _ConfigShim:
    def __init__(self, cfg: YatFullGPTConfig):
        self.sequence_len = cfg.sequence_len
        self.window_pattern = cfg.window_pattern
        self.n_layer = cfg.n_layer


__all__ = ["YatFullGPTConfig", "YatFull_GPT", "YatNMNAttention"]