# Copyright (c) Gensyn Ltd.
# Licensed under the Apache License, Version 2.0.
"""open-1b model for Hugging Face transformers.

Plain-PyTorch inference port of the training model (which runs on RepOps,
Gensyn's bitwise-reproducible kernel library). Module and parameter names
match the training checkpoint FQNs exactly, so the released
``model.safetensors`` holds the training tensors byte-for-byte with no
renaming — the same bytes the per-step audit hashes commit to.

Faithfulness notes (this port is for inference, not for audit replay):
  - Linear layers reproduce the training LSQ int8 W8A8 forward on the int8
    grid: per-tensor dynamic activation scale ``amax/127``, per-channel
    learned weight scale, round-clamp to [-127, 127], GEMM, dequantize.
    The GEMM here accumulates the integer-valued operands in fp32 rather
    than int32, and attention P@V runs in bf16 SDPA where training used a
    block-structured int8 flash kernel — so logits are numerically close
    to, but not bit-identical with, the training stack. Bit-exact replay
    is what https://github.com/gensyn-ai/pretraining-audit-cli is for.
  - ``config.quantized_forward=False`` skips the int8 emulation and runs
    plain GEMMs on the bf16 master weights.
"""

from __future__ import annotations

import math
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel
from transformers.cache_utils import Cache, DynamicCache
from transformers.generation import GenerationMixin
from transformers.masking_utils import (
    create_causal_mask,
    create_sliding_window_causal_mask,
)
from transformers.modeling_outputs import CausalLMOutputWithPast

from .configuration_open1b import Open1BConfig

_INT8_QMAX = 127.0
_FLASH_BLOCK = 32  # training flash-attention block size; windows align to it


# --------------------------------------------------------------------------- #
# Modules — names and semantics mirror pretrain.model (the training code)
# --------------------------------------------------------------------------- #


class Open1BRMSNorm(nn.Module):
    """RMSNorm computed in fp32, result cast back before the (optional) gain —
    the exact op order of the training ``RMSNorm``."""

    def __init__(self, dim: int, eps: float, elementwise_affine: bool = True) -> None:
        super().__init__()
        self.eps = eps
        if elementwise_affine:
            self.weight = nn.Parameter(torch.ones(dim))
        else:
            self.register_parameter("weight", None)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        in_dtype = x.dtype
        xf = x.float()
        rms = xf.pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
        out = (xf * rms).to(in_dtype)
        return out if self.weight is None else out * self.weight


class Open1BLinear(nn.Module):
    """Linear that optionally emulates the training LSQ int8 W8A8 forward.

    When ``quantized`` (an LSQ-trained layer with ``quantized_forward`` on),
    the forward reproduces ``repop.qat.lsq.LSQQuantizedMatmulFunction``:
      x_scale = amax(|x|)/127 (per tensor, dynamic);  x_q = rc(x / x_scale)
      s = |weight_scale|.clamp_min(1e-12) (per channel); w_q = rc(w / s)
      y = (x_q @ w_q^T) * (x_scale * s)
    with rc = round-clamp to [-127, 127]. Quantization math runs in fp32;
    the GEMM accumulates integer-valued fp32 operands (see module docstring).
    """

    def __init__(self, in_features: int, out_features: int, quantized: bool) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.empty(out_features, in_features))
        if quantized:
            self.weight_scale = nn.Parameter(torch.ones(out_features))
        else:
            self.register_parameter("weight_scale", None)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.weight_scale is None:
            return F.linear(x, self.weight)
        s = self.weight_scale.float().abs().clamp_min(1e-12)
        w_q = (self.weight.float() / s.unsqueeze(1)).clamp(-_INT8_QMAX, _INT8_QMAX).round()
        xf = x.float()
        x_scale = xf.abs().amax().clamp_min(1e-12) / _INT8_QMAX
        x_q = (xf / x_scale).round().clamp(-_INT8_QMAX, _INT8_QMAX)
        y = F.linear(x_q, w_q) * (x_scale * s)
        return y.to(x.dtype)


def _rope_cos_sin(
    position_ids: torch.Tensor, head_dim: int, theta: float
) -> tuple[torch.Tensor, torch.Tensor]:
    """fp32 RoPE angles for absolute ``position_ids`` [B, T] ->
    cos/sin [B, T, head_dim/2]. Computed on the fly (no buffers: non-persistent
    buffers do not survive transformers' meta-device loading)."""
    inv_freq = 1.0 / (
        theta
        ** (
            torch.arange(
                0, head_dim, 2, dtype=torch.float32, device=position_ids.device
            )
            / head_dim
        )
    )
    freqs = position_ids.to(torch.float32).unsqueeze(-1) * inv_freq
    return torch.cos(freqs), torch.sin(freqs)


def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    """Interleaved-pair RoPE, exactly as the training ``apply_rope``:
    pairs are (x[..., 0::2], x[..., 1::2]), NOT the transformers half-split.

    ``x``: [B, T, H, head_dim]; ``cos``/``sin``: [B or 1, T, head_dim/2].
    """
    x1 = x[..., 0::2]
    x2 = x[..., 1::2]
    cos = cos.unsqueeze(-2)  # broadcast over heads
    sin = sin.unsqueeze(-2)
    rx1 = x1 * cos - x2 * sin
    rx2 = x1 * sin + x2 * cos
    out = torch.empty_like(x)
    out[..., 0::2] = rx1
    out[..., 1::2] = rx2
    return out


class Open1BAttention(nn.Module):
    """GQA with gain-free QK-norm and per-layer full/sliding-window choice."""

    def __init__(self, config: Open1BConfig, layer_idx: int) -> None:
        super().__init__()
        self.layer_idx = layer_idx
        self.n_heads = config.num_attention_heads
        self.n_kv_heads = config.num_key_value_heads
        self.head_dim = config.head_dim
        self.n_rep = self.n_heads // self.n_kv_heads
        # Per-layer window from config.layer_types (which encodes the training
        # formula: last layer of each swa_full_every group and the final layer
        # are full-causal, the rest slide). 0 = full causal.
        self.window = (
            0 if config.layer_types[layer_idx] == "full_attention"
            else config.sliding_window
        )

        d = config.hidden_size
        quantized = config.quantized_forward and (
            layer_idx >= config.qat_exempt_first_n_blocks
        )
        self.wq = Open1BLinear(d, self.n_heads * self.head_dim, quantized)
        self.wk = Open1BLinear(d, self.n_kv_heads * self.head_dim, quantized)
        self.wv = Open1BLinear(d, self.n_kv_heads * self.head_dim, quantized)
        self.wo = Open1BLinear(d, d, quantized)
        if config.qk_norm:
            self.q_norm = Open1BRMSNorm(
                self.head_dim, config.rms_norm_eps,
                elementwise_affine=config.qk_norm_gain,
            )
            self.k_norm = Open1BRMSNorm(
                self.head_dim, config.rms_norm_eps,
                elementwise_affine=config.qk_norm_gain,
            )
        else:
            self.q_norm = self.k_norm = None

    def forward(
        self,
        x: torch.Tensor,
        cos: torch.Tensor,
        sin: torch.Tensor,
        mask_full: Optional[torch.Tensor],
        mask_windowed: Optional[torch.Tensor],
        past_key_values: Optional[Cache],
        cache_position: Optional[torch.Tensor],
    ) -> torch.Tensor:
        B, T, _ = x.shape
        q = self.wq(x).view(B, T, self.n_heads, self.head_dim)
        k = self.wk(x).view(B, T, self.n_kv_heads, self.head_dim)
        v = self.wv(x).view(B, T, self.n_kv_heads, self.head_dim)
        if self.q_norm is not None:
            q = self.q_norm(q)
            k = self.k_norm(k)
        q = _apply_rope(q, cos, sin)
        k = _apply_rope(k, cos, sin)

        q = q.transpose(1, 2)  # [B, H, T, D]
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)
        if past_key_values is not None:
            k, v = past_key_values.update(
                k, v, self.layer_idx, {"cache_position": cache_position}
            )
        if self.n_rep > 1:
            k = k.repeat_interleave(self.n_rep, dim=1)
            v = v.repeat_interleave(self.n_rep, dim=1)

        mask = mask_windowed if self.window > 0 else mask_full
        out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
        out = out.transpose(1, 2).contiguous().view(B, T, -1)
        return self.wo(out)


class Open1BFFN(nn.Module):
    """SwiGLU with the fused gate+up projection layout of the checkpoint."""

    def __init__(self, config: Open1BConfig, quantized: bool) -> None:
        super().__init__()
        self.w_gate_up = Open1BLinear(
            config.hidden_size, 2 * config.intermediate_size, quantized
        )
        self.w_down = Open1BLinear(
            config.intermediate_size, config.hidden_size, quantized
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate, up = self.w_gate_up(x).chunk(2, dim=-1)
        return self.w_down(F.silu(gate) * up)


class Open1BBlock(nn.Module):
    def __init__(self, config: Open1BConfig, layer_idx: int) -> None:
        super().__init__()
        self.norm1 = Open1BRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.attn = Open1BAttention(config, layer_idx)
        self.norm2 = Open1BRMSNorm(config.hidden_size, config.rms_norm_eps)
        quantized = config.quantized_forward and (
            layer_idx >= config.qat_exempt_first_n_blocks
        )
        self.ffn = Open1BFFN(config, quantized)

    def forward(self, x, cos, sin, mask_full, mask_windowed, past_key_values,
                cache_position):
        x = x + self.attn(self.norm1(x), cos, sin, mask_full, mask_windowed,
                          past_key_values, cache_position)
        x = x + self.ffn(self.norm2(x))
        return x


class Open1BEmbedding(nn.Module):
    """Untied embedding + LM head, matching the training FQNs
    (``embedding.tok_embeddings`` / ``embedding.output``)."""

    def __init__(self, config: Open1BConfig) -> None:
        super().__init__()
        self.tok_embeddings = nn.Embedding(config.vocab_size, config.hidden_size)
        self.output = Open1BLinear(config.hidden_size, config.vocab_size,
                                   quantized=False)  # head is full precision


class Open1BForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = Open1BConfig
    base_model_prefix = ""
    main_input_name = "input_ids"
    supports_gradient_checkpointing = False
    _no_split_modules = ["Open1BBlock"]

    def __init__(self, config: Open1BConfig) -> None:
        super().__init__(config)
        self.embedding = Open1BEmbedding(config)
        self.emb_norm = (
            Open1BRMSNorm(config.hidden_size, config.rms_norm_eps)
            if config.embedding_norm else None
        )
        self.blocks = nn.ModuleList(
            Open1BBlock(config, i) for i in range(config.num_hidden_layers)
        )
        self.norm_out = Open1BRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.post_init()

    # -- HF plumbing --------------------------------------------------------
    def get_input_embeddings(self):
        return self.embedding.tok_embeddings

    def set_input_embeddings(self, value):
        self.embedding.tok_embeddings = value

    def get_output_embeddings(self):
        return self.embedding.output

    def set_output_embeddings(self, value):
        self.embedding.output = value

    def _init_weights(self, module):
        # Training init: trunc normal(0, 0.02) at ±2σ; norms at 1. Only used
        # for from-config construction — released checkpoints load everything.
        if isinstance(module, (Open1BLinear, nn.Embedding)):
            nn.init.trunc_normal_(module.weight, std=0.02, a=-0.04, b=0.04)
            if getattr(module, "weight_scale", None) is not None:
                with torch.no_grad():
                    scale = module.weight.abs().mean(dim=1) * (
                        2.0 / math.sqrt(_INT8_QMAX)
                    )
                    module.weight_scale.copy_(scale.clamp_min(1e-6))
        elif isinstance(module, Open1BRMSNorm) and module.weight is not None:
            nn.init.ones_(module.weight)

    # -- forward -------------------------------------------------------------
    def forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        past_key_values: Optional[Cache] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        labels: Optional[torch.Tensor] = None,
        use_cache: Optional[bool] = None,
        cache_position: Optional[torch.Tensor] = None,
        **kwargs,
    ) -> CausalLMOutputWithPast:
        if inputs_embeds is None:
            inputs_embeds = self.embedding.tok_embeddings(input_ids)
        h = inputs_embeds
        B, T = h.shape[:2]

        use_cache = use_cache if use_cache is not None else self.config.use_cache
        if use_cache and past_key_values is None:
            # config-aware: sliding layers get rolling window-sized buffers
            past_key_values = DynamicCache(config=self.config)
        past_len = (
            past_key_values.get_seq_length() if past_key_values is not None else 0
        )
        if cache_position is None:
            cache_position = torch.arange(
                past_len, past_len + T, device=h.device
            )
        if position_ids is None:
            position_ids = cache_position.unsqueeze(0)

        cos, sin = _rope_cos_sin(
            position_ids, self.config.head_dim, self.config.rope_theta
        )  # [B or 1, T, head_dim/2]
        # transformers' mask helpers understand the hybrid layer structure
        # (config.layer_types) AND the cache layout — with a config-aware
        # DynamicCache the sliding layers keep rolling window-sized buffers,
        # so masks must be built against the physical cache, not absolute
        # positions.
        #
        # Window semantics — measured against the training kernel
        # (repop causal_flash_attention, flash block size 32): the window is
        # BLOCK-ALIGNED, not per-key. Query i attends key j iff
        #   j <= i  and  floor(j/32) >= floor(i/32) - (window/32 - 1)
        # i.e. the query's own 32-token block plus the previous window/32 - 1
        # blocks. This is a subset of the exact [i-window+1, i] sliding
        # window (so the rolling cache always retains every visible key);
        # the extra block constraint is AND-ed onto the standard mask.
        blocks_back = self.config.sliding_window // _FLASH_BLOCK - 1

        def _block_window_overlay(batch_idx, head_idx, q_idx, kv_idx):
            return (kv_idx // _FLASH_BLOCK) >= (q_idx // _FLASH_BLOCK) - blocks_back

        mask_kwargs = dict(
            config=self.config,
            inputs_embeds=h,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            position_ids=position_ids,
            allow_is_causal_skip=False,  # always materialise; we call SDPA raw
        )
        mask_full = create_causal_mask(**mask_kwargs)
        mask_windowed = create_sliding_window_causal_mask(
            **mask_kwargs, and_mask_function=_block_window_overlay
        )

        if self.emb_norm is not None:
            h = self.emb_norm(h)
        for block in self.blocks:
            h = block(h, cos, sin, mask_full, mask_windowed,
                      past_key_values, cache_position)
        h = self.norm_out(h)
        logits = self.embedding.output(h)

        loss = None
        if labels is not None:
            loss = self.loss_function(
                logits=logits, labels=labels, vocab_size=self.config.vocab_size,
                **kwargs,
            )
        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=past_key_values if use_cache else None,
        )