#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
LULUV2 inference-only runtime.

This file intentionally contains only the code needed to load and run a
standalone native-bf16 LULUV2 checkpoint. It contains only the
runtime loader, tokenizer bridge, decoder modules, and two-pass inference path
needed for local generation.

Runtime behavior:
- loads a local checkpoint supplied by the user/repo;
- uses local tokenizer files;
- does not download or load any external model weights;
- preserves the VWM/two-pass inference path when present in the checkpoint.
"""

from __future__ import annotations

import json
import math
import os
import time
from dataclasses import dataclass
from types import SimpleNamespace
from typing import Dict, Optional, Tuple

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

_TRANSFORMERS_IMPORT_ERROR = None
try:
    from transformers import AutoTokenizer as _HFAutoTokenizer
    try:
        from transformers import AutoConfig as _HFAutoConfig
    except Exception:
        _HFAutoConfig = None
except Exception as _e:
    _TRANSFORMERS_IMPORT_ERROR = _e
    _HFAutoTokenizer = None
    _HFAutoConfig = None

class _TokenOutput(dict):
    def __getattr__(self, name):
        try:
            return self[name]
        except KeyError as exc:
            raise AttributeError(name) from exc
    def to(self, device):
        out = _TokenOutput()
        for k, v in self.items():
            out[k] = v.to(device) if torch.is_tensor(v) else v
        return out

class _LocalTokenizer:
    def __init__(self, path: str, tokenizer_file: Optional[str] = None, **kwargs):
        import json as _json
        try:
            from tokenizers import Tokenizer as _TokenizerCore
        except Exception as exc:
            raise RuntimeError(
                "transformers import failed and tokenizers is unavailable. "
                "Install tokenizers or use a matching torch/transformers pair."
            ) from exc
        self.name_or_path = path or tokenizer_file or "<local-tokenizer>"
        if tokenizer_file:
            tok_file = tokenizer_file
            base_dir = os.path.dirname(os.path.abspath(tok_file))
        else:
            base_dir = os.path.abspath(path)
            tok_file = os.path.join(base_dir, "tokenizer.json")
        if not os.path.exists(tok_file):
            raise FileNotFoundError(f"Local tokenizer.json not found: {tok_file}")
        self._tok = _TokenizerCore.from_file(tok_file)
        self.vocab_size = int(self._tok.get_vocab_size())
        self.model_max_length = 10**9
        self.truncation_side = "left"
        self.chat_template = None
        self.eos_token = None
        self.pad_token = None
        cfg_path = os.path.join(base_dir, "tokenizer_config.json")
        sp_path = os.path.join(base_dir, "special_tokens_map.json")
        for p in (cfg_path, sp_path):
            if os.path.exists(p):
                try:
                    data = _json.load(open(p, "r", encoding="utf-8"))
                except Exception:
                    data = {}
                if self.chat_template is None and isinstance(data.get("chat_template"), str):
                    self.chat_template = data.get("chat_template")
                for key, attr in (("eos_token", "eos_token"), ("pad_token", "pad_token")):
                    val = data.get(key)
                    if isinstance(val, dict):
                        val = val.get("content")
                    if isinstance(val, str):
                        setattr(self, attr, val)
        if self.eos_token is None:
            for cand in ("<|im_end|>", "<|endoftext|>", "</s>"):
                if self._tok.token_to_id(cand) is not None:
                    self.eos_token = cand
                    break
        if self.pad_token is None:
            self.pad_token = self.eos_token
        self.eos_token_id = self._tok.token_to_id(self.eos_token) if self.eos_token else None
        self.pad_token_id = self._tok.token_to_id(self.pad_token) if self.pad_token else self.eos_token_id

    def __len__(self):
        return self.vocab_size

    def __call__(self, text, return_tensors=None, truncation=False, max_length=None, add_special_tokens=True, **kwargs):
        if isinstance(text, (list, tuple)):
            encoded = [self._encode_one(t, add_special_tokens, truncation, max_length) for t in text]
            maxlen = max(len(x) for x in encoded) if encoded else 0
            pad = self.pad_token_id if self.pad_token_id is not None else 0
            arr = [x + [pad] * (maxlen - len(x)) for x in encoded]
            if return_tensors == "pt":
                return _TokenOutput(input_ids=torch.tensor(arr, dtype=torch.long))
            return _TokenOutput(input_ids=arr)
        ids = self._encode_one(str(text), add_special_tokens, truncation, max_length)
        if return_tensors == "pt":
            return _TokenOutput(input_ids=torch.tensor([ids], dtype=torch.long))
        return _TokenOutput(input_ids=ids)

    def _encode_one(self, text, add_special_tokens=True, truncation=False, max_length=None):
        enc = self._tok.encode(text, add_special_tokens=bool(add_special_tokens))
        ids = list(enc.ids)
        if truncation and max_length is not None and len(ids) > int(max_length):
            if self.truncation_side == "left":
                ids = ids[-int(max_length):]
            else:
                ids = ids[:int(max_length)]
        return ids

    def decode(self, ids, skip_special_tokens=True, **kwargs):
        if torch.is_tensor(ids):
            ids = ids.detach().cpu().tolist()
        if ids and isinstance(ids[0], list):
            ids = ids[0]
        return self._tok.decode([int(x) for x in ids], skip_special_tokens=bool(skip_special_tokens))

    def apply_chat_template(self, messages, tokenize=False, add_generation_prompt=False, **kwargs):
        chunks = []
        for m in messages:
            role = str(m.get("role", "user"))
            content = str(m.get("content", ""))
            chunks.append(f"<|im_start|>{role}\n{content}<|im_end|>")
        if add_generation_prompt:
            chunks.append("<|im_start|>assistant\n")
        text = "\n".join(chunks)
        if tokenize:
            return self(text, add_special_tokens=False).input_ids
        return text

class _AutoTokenizerShim:
    @staticmethod
    def from_pretrained(path, *args, **kwargs):
        if _HFAutoTokenizer is not None:
            return _HFAutoTokenizer.from_pretrained(path, *args, **kwargs)
        return _LocalTokenizer(path)

class _AutoConfigShim:
    @staticmethod
    def from_pretrained(path, *args, **kwargs):
        if _HFAutoConfig is not None:
            return _HFAutoConfig.from_pretrained(path, *args, **kwargs)
        raise RuntimeError(
            "AutoConfig requested, but transformers failed to import. "
            "Use --no-config-download / embedded model_config for LULUV2."
        )
AutoTokenizer = _AutoTokenizerShim
AutoConfig = _AutoConfigShim

torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
if hasattr(torch, "set_float32_matmul_precision"):
    torch.set_float32_matmul_precision("high")
try:
    if torch.cuda.is_available():
        torch.backends.cuda.enable_flash_sdp(True)
        torch.backends.cuda.enable_mem_efficient_sdp(True)
        torch.backends.cuda.enable_math_sdp(False)
except Exception:
    pass



def parse_dtype(name: str):
    name = str(name).strip().lower()
    if name in {"bf16", "bfloat16"}:
        return torch.bfloat16
    if name in {"fp16", "float16", "half"}:
        return torch.float16
    if name in {"fp32", "float32"}:
        return torch.float32
    raise ValueError(f"Unknown dtype: {name}")


def human_bytes(n: float) -> str:
    units = ["B", "KB", "MB", "GB", "TB"]
    x = float(n)
    i = 0
    while x >= 1024.0 and i < len(units) - 1:
        x /= 1024.0
        i += 1
    return f"{x:.2f} {units[i]}"


def safe_torch_load(path: str, map_location="cpu"):
    # PyTorch 2.6+ defaults may warn around weights_only. This checkpoint stores
    # Python metadata plus tensors, so weights_only=False is intentional.
    try:
        return torch.load(path, map_location=map_location, weights_only=False)
    except TypeError:
        return torch.load(path, map_location=map_location)


def module_has_vwm(sd: Dict[str, torch.Tensor], prefix: str) -> bool:
    return f"{prefix}.A" in sd and f"{prefix}.B" in sd and f"{prefix}.c" in sd


def linear_shape_from_state(sd: Dict[str, torch.Tensor], prefix: str) -> Tuple[int, int, bool]:
    if module_has_vwm(sd, prefix):
        out_features = int(sd[f"{prefix}.A"].shape[0])
        in_features = int(sd[f"{prefix}.B"].shape[0])
        has_bias = f"{prefix}.bias" in sd
        return in_features, out_features, has_bias
    wkey = f"{prefix}.weight"
    if wkey not in sd:
        raise KeyError(f"Cannot infer Linear shape for {prefix}; missing {wkey} and VWM A/B/c")
    out_features, in_features = sd[wkey].shape
    has_bias = f"{prefix}.bias" in sd
    return int(in_features), int(out_features), has_bias


def make_linear_from_state(sd: Dict[str, torch.Tensor], prefix: str) -> nn.Module:
    in_features, out_features, has_bias = linear_shape_from_state(sd, prefix)
    if module_has_vwm(sd, prefix):
        rank = int(sd[f"{prefix}.c"].shape[0])
        return VWMFactorizedLinear(in_features, out_features, rank, bias=has_bias, name=prefix)
    return nn.Linear(in_features, out_features, bias=has_bias)


def module_has_vwm_embedding(sd: Dict[str, torch.Tensor], prefix: str) -> bool:
    return f"{prefix}.A" in sd and f"{prefix}.B" in sd and f"{prefix}.c" in sd


def embedding_shape_from_state(sd: Dict[str, torch.Tensor], prefix: str) -> Tuple[int, int]:
    if module_has_vwm_embedding(sd, prefix):
        return int(sd[f"{prefix}.A"].shape[0]), int(sd[f"{prefix}.B"].shape[0])
    wkey = f"{prefix}.weight"
    if wkey not in sd:
        raise KeyError(f"Cannot infer embedding shape for {prefix}; missing dense or VWM embedding tensors")
    return int(sd[wkey].shape[0]), int(sd[wkey].shape[1])


def make_embedding_from_state(sd: Dict[str, torch.Tensor], prefix: str) -> nn.Module:
    vocab_size, hidden_size = embedding_shape_from_state(sd, prefix)
    if module_has_vwm_embedding(sd, prefix):
        rank = int(sd[f"{prefix}.c"].shape[0])
        return VWMFactorizedEmbedding(vocab_size, hidden_size, rank, name=prefix)
    return nn.Embedding(vocab_size, hidden_size)


def expand_shared_banks_into_state(ckpt: Dict, sd: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
    """Expand experimental shared-bank storage into normal per-module A/B/c tensors."""
    banks = ckpt.get("shared_banks")
    if not banks:
        return sd
    out = dict(sd)
    n = 0
    for bank_id, bank in banks.items():
        A = bank["A"]
        B = bank["B"]
        modules = bank.get("modules", {})
        for prefix, m in modules.items():
            out[f"{prefix}.A"] = A
            out[f"{prefix}.B"] = B
            out[f"{prefix}.c"] = m["c"]
            if "bias" in m and m["bias"] is not None:
                out[f"{prefix}.bias"] = m["bias"]
            n += 1
    print(f"[shared-bank] expanded {len(banks)} banks into {n} VWM modules")
    return out


# -----------------------------
# VWM linear used by the exported checkpoint
# -----------------------------


class VWMFactorizedLinear(nn.Module):
    """
    W ~= A diag(c) B^T
    y = ((x @ B) * c) @ A^T + bias

    This matches LULU2 exporter's exported VWMFactorizedLinear
    state names: A, B, c, optional bias.
    """

    def __init__(self, in_features: int, out_features: int, rank: int, bias: bool = True, name: str = ""):
        super().__init__()
        self.in_features = int(in_features)
        self.out_features = int(out_features)
        self.rank = int(rank)
        self.name = name
        self.A = nn.Parameter(torch.empty(out_features, rank), requires_grad=False)
        self.B = nn.Parameter(torch.empty(in_features, rank), requires_grad=False)
        self.c = nn.Parameter(torch.empty(rank), requires_grad=False)
        self.bias = nn.Parameter(torch.zeros(out_features), requires_grad=False) if bias else None

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Compute in the activation dtype/device. Parameters are already moved by model.to(...).
        t = torch.matmul(x, self.B.to(dtype=x.dtype))
        t = t * self.c.to(dtype=x.dtype)
        y = torch.matmul(t, self.A.to(dtype=x.dtype).transpose(0, 1))
        if self.bias is not None:
            y = y + self.bias.to(dtype=x.dtype)
        return y




class VWMFactorizedEmbedding(nn.Module):
    """Runtime for exported VWM embedding: E ~= A diag(c) B^T."""
    def __init__(self, num_embeddings: int, embedding_dim: int, rank: int, name: str = "model.embed_tokens"):
        super().__init__()
        self.num_embeddings = int(num_embeddings)
        self.embedding_dim = int(embedding_dim)
        self.rank = int(rank)
        self.name = name
        self.A = nn.Parameter(torch.empty(num_embeddings, rank), requires_grad=False)
        self.B = nn.Parameter(torch.empty(embedding_dim, rank), requires_grad=False)
        self.c = nn.Parameter(torch.empty(rank), requires_grad=False)

    @property
    def weight(self):
        # Dense materialization only for compatibility/debug. Normal forward avoids this.
        return (self.A * self.c.view(1, -1)) @ self.B.T

    def forward(self, input_ids: torch.LongTensor) -> torch.Tensor:
        a = F.embedding(input_ids, self.A)
        t = a * self.c.to(dtype=a.dtype)
        return torch.matmul(t, self.B.to(dtype=a.dtype).transpose(0, 1))


class TiedEmbeddingLMHead(nn.Module):
    """LM head tied to the model embedding matrix, dense or VWM."""
    def __init__(self, embedding_module: nn.Module):
        super().__init__()
        self.embedding_module = embedding_module

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        emb = self.embedding_module
        if isinstance(emb, VWMFactorizedEmbedding):
            # logits = h @ E.T = (h @ B) * c @ A.T
            t = torch.matmul(hidden_states, emb.B.to(dtype=hidden_states.dtype))
            t = t * emb.c.to(dtype=hidden_states.dtype)
            return torch.matmul(t, emb.A.to(dtype=hidden_states.dtype).transpose(0, 1))
        return F.linear(hidden_states, emb.weight.to(dtype=hidden_states.dtype))

# -----------------------------
# LULU2 decoder architecture
# -----------------------------


class LuluRMSNorm(nn.Module):
    def __init__(self, hidden_size: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size), requires_grad=False)
        self.variance_epsilon = float(eps)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.float()
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        return self.weight.to(dtype=input_dtype) * hidden_states.to(input_dtype)


class LuluRotaryEmbedding(nn.Module):
    def __init__(self, dim: int, max_position_embeddings: int = 32768, base: float = 1000000.0):
        super().__init__()
        self.dim = int(dim)
        self.max_position_embeddings = int(max_position_embeddings)
        self.base = float(base)
        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float() / self.dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    @torch.no_grad()
    def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        # position_ids: [B, T]
        inv_freq = self.inv_freq.to(device=x.device)
        freqs = torch.einsum("bt,d->btd", position_ids.float(), inv_freq.float())
        emb = torch.cat((freqs, freqs), dim=-1)
        return emb.cos().to(dtype=x.dtype), emb.sin().to(dtype=x.dtype)


def rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)


def apply_rotary_pos_emb(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    # q/k: [B, H, T, D], cos/sin: [B, T, D]
    cos = cos.unsqueeze(1)
    sin = sin.unsqueeze(1)
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed


def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
    if n_rep == 1:
        return hidden_states
    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)


class LuluVWMMLP(nn.Module):
    def __init__(self, cfg, sd: Dict[str, torch.Tensor], layer_idx: int):
        super().__init__()
        p = f"model.layers.{layer_idx}.mlp"
        self.gate_proj = make_linear_from_state(sd, f"{p}.gate_proj")
        self.up_proj = make_linear_from_state(sd, f"{p}.up_proj")
        self.down_proj = make_linear_from_state(sd, f"{p}.down_proj")

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class LuluVWMAttention(nn.Module):
    def __init__(self, cfg, sd: Dict[str, torch.Tensor], layer_idx: int):
        super().__init__()
        self.layer_idx = int(layer_idx)
        self.hidden_size = int(cfg.hidden_size)
        self.num_heads = int(cfg.num_attention_heads)
        self.num_key_value_heads = int(getattr(cfg, "num_key_value_heads", self.num_heads))
        self.head_dim = int(getattr(cfg, "head_dim", self.hidden_size // self.num_heads))
        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
        self.scaling = self.head_dim ** -0.5
        self.attention_dropout = float(getattr(cfg, "attention_dropout", 0.0))

        p = f"model.layers.{layer_idx}.self_attn"
        self.q_proj = make_linear_from_state(sd, f"{p}.q_proj")
        self.k_proj = make_linear_from_state(sd, f"{p}.k_proj")
        self.v_proj = make_linear_from_state(sd, f"{p}.v_proj")
        self.o_proj = make_linear_from_state(sd, f"{p}.o_proj")

        rope_theta = float(getattr(cfg, "rope_theta", 1000000.0))
        max_pos = int(getattr(cfg, "max_position_embeddings", 32768))
        self.rotary_emb = LuluRotaryEmbedding(self.head_dim, max_position_embeddings=max_pos, base=rope_theta)

    def forward(self, hidden_states: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor:
        bsz, q_len, _ = hidden_states.size()

        query_states = self.q_proj(hidden_states)
        key_states = self.k_proj(hidden_states)
        value_states = self.v_proj(hidden_states)

        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)

        cos, sin = self.rotary_emb(value_states, position_ids)
        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

        key_states = repeat_kv(key_states, self.num_key_value_groups)
        value_states = repeat_kv(value_states, self.num_key_value_groups)

        # Full forward is causal. This generation script recomputes the full prefix each token.
        attn_output = F.scaled_dot_product_attention(
            query_states,
            key_states,
            value_states,
            attn_mask=None,
            dropout_p=0.0,
            is_causal=True,
            scale=self.scaling,
        )
        attn_output = attn_output.transpose(1, 2).contiguous().reshape(bsz, q_len, self.hidden_size)
        return self.o_proj(attn_output)


class LuluVWMDecoderLayer(nn.Module):
    def __init__(self, cfg, sd: Dict[str, torch.Tensor], layer_idx: int):
        super().__init__()
        self.self_attn = LuluVWMAttention(cfg, sd, layer_idx)
        self.mlp = LuluVWMMLP(cfg, sd, layer_idx)
        self.input_layernorm = LuluRMSNorm(cfg.hidden_size, eps=getattr(cfg, "rms_norm_eps", 1e-6))
        self.post_attention_layernorm = LuluRMSNorm(cfg.hidden_size, eps=getattr(cfg, "rms_norm_eps", 1e-6))

    def forward(self, hidden_states: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor:
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        hidden_states = self.self_attn(hidden_states, position_ids=position_ids)
        hidden_states = residual + hidden_states

        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states
        return hidden_states


class LuluVWMModel(nn.Module):
    def __init__(self, cfg, sd: Dict[str, torch.Tensor]):
        super().__init__()
        self.config = cfg
        vocab_size, hidden_size = embedding_shape_from_state(sd, "model.embed_tokens")
        self.embed_tokens = make_embedding_from_state(sd, "model.embed_tokens")
        self.layers = nn.ModuleList([LuluVWMDecoderLayer(cfg, sd, i) for i in range(int(cfg.num_hidden_layers))])
        self.norm = LuluRMSNorm(hidden_size, eps=getattr(cfg, "rms_norm_eps", 1e-6))

    def forward(self, input_ids: torch.LongTensor, position_ids: Optional[torch.LongTensor] = None) -> torch.Tensor:
        bsz, seq_len = input_ids.shape
        if position_ids is None:
            position_ids = torch.arange(seq_len, device=input_ids.device, dtype=torch.long).unsqueeze(0).expand(bsz, -1)
        hidden_states = self.embed_tokens(input_ids)
        for layer in self.layers:
            hidden_states = layer(hidden_states, position_ids=position_ids)
        return self.norm(hidden_states)


class LuluVWMForCausalLM(nn.Module):
    def __init__(self, cfg, sd: Dict[str, torch.Tensor]):
        super().__init__()
        self.config = cfg
        self.model = LuluVWMModel(cfg, sd)
        _, hidden_size = embedding_shape_from_state(sd, "model.embed_tokens")
        self.tie_word_embeddings = bool(getattr(cfg, "tie_word_embeddings", False))
        if module_has_vwm(sd, "lm_head") or "lm_head.weight" in sd:
            self.lm_head = make_linear_from_state(sd, "lm_head")
        else:
            self.tie_word_embeddings = True
            self.lm_head = TiedEmbeddingLMHead(self.model.embed_tokens)

    def forward(self, input_ids: torch.LongTensor, position_ids: Optional[torch.LongTensor] = None):
        hidden_states = self.model(input_ids=input_ids, position_ids=position_ids)
        logits = self.lm_head(hidden_states)
        return SimpleNamespace(logits=logits)


# -----------------------------
# config loading / inference
# -----------------------------


def infer_minimal_config_from_state(sd: Dict[str, torch.Tensor], model_id: str = "") -> SimpleNamespace:
    if "model.embed_tokens.weight" in sd:
        hidden_size = int(sd["model.embed_tokens.weight"].shape[1])
        vocab_size = int(sd["model.embed_tokens.weight"].shape[0])
    elif module_has_vwm_embedding(sd, "model.embed_tokens"):
        vocab_size = int(sd["model.embed_tokens.A"].shape[0])
        hidden_size = int(sd["model.embed_tokens.B"].shape[0])
    else:
        raise ValueError("Checkpoint is missing model.embed_tokens dense or VWM tensors. Use a full standalone checkpoint, not a delta checkpoint.")
    layer_ids = []
    for k in sd.keys():
        if k.startswith("model.layers."):
            try:
                layer_ids.append(int(k.split(".")[2]))
            except Exception:
                pass
    num_hidden_layers = max(layer_ids) + 1 if layer_ids else 0
    inter_key = "model.layers.0.mlp.gate_proj.weight"
    if inter_key in sd:
        intermediate_size = int(sd[inter_key].shape[0])
    else:
        intermediate_size = 4864

    # Best known defaults for LULU2. If you export
    # model_config into the checkpoint, these assumptions are not used.
    num_attention_heads = 14
    num_key_value_heads = 2
    head_dim = hidden_size // num_attention_heads
    if head_dim * num_attention_heads != hidden_size:
        # Fallback if a different decoder variant is used and no config is present.
        # This requires explicit command-line override in practice.
        num_attention_heads = 1
        num_key_value_heads = 1
        head_dim = hidden_size

    return SimpleNamespace(
        model_type="luluv2",
        model_id=model_id,
        vocab_size=vocab_size,
        hidden_size=hidden_size,
        intermediate_size=intermediate_size,
        num_hidden_layers=num_hidden_layers,
        num_attention_heads=num_attention_heads,
        num_key_value_heads=num_key_value_heads,
        head_dim=head_dim,
        rms_norm_eps=1e-6,
        rope_theta=1000000.0,
        max_position_embeddings=32768,
        attention_dropout=0.0,
        tie_word_embeddings=False,
    )


def namespace_from_dict(d: Dict) -> SimpleNamespace:
    return SimpleNamespace(**d)


def load_runtime_config(ckpt: Dict, sd: Dict[str, torch.Tensor], args) -> SimpleNamespace:
    if "model_config" in ckpt and isinstance(ckpt["model_config"], dict):
        print("[config] using model_config embedded in checkpoint")
        d = dict(ckpt["model_config"])
        if ckpt.get("tie_word_embeddings") is True:
            d["tie_word_embeddings"] = True
        return namespace_from_dict(d)

    model_id = args.model_id or ckpt.get("model_id") or ckpt.get("args", {}).get("model_id") or "LULU2"

    if args.no_config_download:
        print("[config] no embedded config and --no-config-download set; using LULU2 defaults")
        cfg = infer_minimal_config_from_state(sd, model_id=model_id)
        if ckpt.get("tie_word_embeddings") is True:
            cfg.tie_word_embeddings = True
        return cfg

    print(f"[config] loading config metadata only from {model_id}; no model weights are loaded")
    cfg = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
    return cfg


# -----------------------------
# generation
# -----------------------------


def build_chat_prompt(tokenizer, user_prompt: str, system_prompt: str = "You are a helpful assistant. Answer directly and naturally.") -> str:
    messages = []
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})
    messages.append({"role": "user", "content": user_prompt})
    try:
        return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    except Exception:
        return f"system\n{system_prompt}\nuser\n{user_prompt}\nassistant\n"


@torch.no_grad()
def sample_next(logits: torch.Tensor, temperature: float = 0.0, top_k: int = 0, top_p: float = 1.0) -> torch.Tensor:
    if temperature <= 0.0:
        return torch.argmax(logits, dim=-1, keepdim=True)

    logits = logits / max(temperature, 1e-6)
    if top_k and top_k > 0:
        k = min(int(top_k), logits.size(-1))
        thresh = torch.topk(logits, k, dim=-1).values[..., -1, None]
        logits = torch.where(logits >= thresh, logits, torch.full_like(logits, -float("inf")))
    if top_p < 1.0:
        sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
        probs = torch.softmax(sorted_logits, dim=-1)
        cumulative_probs = torch.cumsum(probs, dim=-1)
        sorted_indices_to_remove = cumulative_probs > top_p
        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
        sorted_indices_to_remove[..., 0] = False
        sorted_logits = sorted_logits.masked_fill(sorted_indices_to_remove, -float("inf"))
        logits = torch.full_like(logits, -float("inf")).scatter(1, sorted_indices, sorted_logits)
    probs = torch.softmax(logits, dim=-1)
    return torch.multinomial(probs, num_samples=1)


@torch.no_grad()
def generate_text(model, tokenizer, prompt: str, device, max_new_tokens: int = 120, temperature: float = 0.0, top_k: int = 0, top_p: float = 1.0, max_context: int = 2048) -> Tuple[str, float]:
    model.eval()
    enc = tokenizer(prompt, return_tensors="pt")
    input_ids = enc.input_ids.to(device)
    eos_id = tokenizer.eos_token_id
    t0 = time.time()
    start_len = int(input_ids.shape[1])

    for _ in range(max_new_tokens):
        ctx = input_ids[:, -max_context:]
        out = model(ctx)
        next_logits = out.logits[:, -1, :].float()
        next_id = sample_next(next_logits, temperature=temperature, top_k=top_k, top_p=top_p)
        input_ids = torch.cat([input_ids, next_id.to(input_ids.device)], dim=-1)
        if eos_id is not None and int(next_id.item()) == int(eos_id):
            break

    dt = time.time() - t0
    new_tokens = max(1, int(input_ids.shape[1]) - start_len)
    return tokenizer.decode(input_ids[0], skip_special_tokens=True), new_tokens / max(dt, 1e-9)


def load_tokenizer(args, ckpt):
    tok_path = args.tokenizer or ckpt.get("tokenizer_dir") or ckpt.get("model_id") or ckpt.get("args", {}).get("model_id") or args.model_id
    if not tok_path:
        raise ValueError("Tokenizer path/name is required. Pass --tokenizer <local-dir-or-model-id>.")
    # If checkpoint stores a relative tokenizer_dir like "tokenizer", resolve it
    # relative to the checkpoint location so no HF lookup is needed.
    ckpt_dir = os.path.dirname(os.path.abspath(args.checkpoint))
    if tok_path and not os.path.isabs(tok_path):
        maybe_local = os.path.join(ckpt_dir, tok_path)
        if os.path.isdir(maybe_local):
            tok_path = maybe_local
    print(f"[tokenizer] {tok_path}")
    tok = AutoTokenizer.from_pretrained(tok_path, trust_remote_code=True, local_files_only=bool(args.local_files_only))
    if tok.pad_token_id is None and tok.eos_token_id is not None:
        tok.pad_token = tok.eos_token
    return tok


# -----------------------------
# main

# Public model aliases used by the UI/runtime.
Lulu2RMSNorm = LuluRMSNorm
Lulu2RotaryEmbedding = LuluRotaryEmbedding
Lulu2VWMMLP = LuluVWMMLP
Lulu2VWMAttention = LuluVWMAttention
Lulu2VWMDecoderLayer = LuluVWMDecoderLayer
Lulu2VWMModel = LuluVWMModel
Lulu2ForCausalLM = LuluVWMForCausalLM


class Pass2RefinementAdapter(nn.Module):
    """Small gated residual adapter conditioned on pass-1 layer state."""

    def __init__(self, hidden_size: int, rank: int, gate_init: float = -5.0):
        super().__init__()
        self.hidden_size = int(hidden_size)
        self.rank = int(rank)
        self.x_norm = LuluRMSNorm(hidden_size)
        self.cond_norm = LuluRMSNorm(hidden_size)
        self.down = nn.Linear(2 * hidden_size, rank, bias=False)
        self.up = nn.Linear(rank, hidden_size, bias=False)
        self.gate = nn.Parameter(torch.tensor(float(gate_init)))

        nn.init.normal_(self.down.weight, mean=0.0, std=0.02 / math.sqrt(max(1, hidden_size)))
        # Zero init means the two-pass model starts exactly as pass 1.
        nn.init.zeros_(self.up.weight)

    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
        z = torch.cat([self.x_norm(x), self.cond_norm(cond)], dim=-1)
        delta = self.up(F.silu(self.down(z)))
        return torch.sigmoid(self.gate).to(dtype=x.dtype) * delta


@dataclass
class Pass2Config:
    adapter_rank: int = 64
    adapter_gate_init: float = -5.0
    layer_gate_init: float = -5.0
    pass_embed_scale: float = 0.0
    mode: str = "refine_pass1_residual"


class Lulu2TwoPassForCausalLM(nn.Module):
    """
    Wraps a loaded LULU2 base model.

    Pass 1: normal LULU2 decoder forward, producing the pass-1 residual stream.
    Pass 2: starts from pass-1 residual stream and adds small gated refinements.

    h2_i = h2_i + sigmoid(layer_gate_i) * (BaseLayer_i(h2_i) - h2_i)
                 + Adapter_i(h2_i, pass1_state_i)

    With zero-initialized adapter up-projections and negative gates, the model
    starts extremely close to the loaded LULU2 checkpoint and learns refinements.
    """

    def __init__(self, base: Lulu2ForCausalLM, cfg: Pass2Config):
        super().__init__()
        self.base = base
        self.pass2_config = cfg
        hidden = int(base.config.hidden_size)
        n_layers = int(base.config.num_hidden_layers)
        self.pass_embed = nn.Parameter(torch.randn(2, hidden) * float(cfg.pass_embed_scale))
        self.layer_gates = nn.Parameter(torch.full((n_layers,), float(cfg.layer_gate_init)))
        self.adapters = nn.ModuleList([
            Pass2RefinementAdapter(hidden, int(cfg.adapter_rank), gate_init=float(cfg.adapter_gate_init))
            for _ in range(n_layers)
        ])

    def _position_ids(self, input_ids: torch.LongTensor, position_ids: Optional[torch.LongTensor] = None):
        if position_ids is not None:
            return position_ids
        bsz, seq_len = input_ids.shape
        return torch.arange(seq_len, device=input_ids.device, dtype=torch.long).unsqueeze(0).expand(bsz, -1)

    def forward_pass1_features(self, input_ids: torch.LongTensor, position_ids: Optional[torch.LongTensor] = None):
        position_ids = self._position_ids(input_ids, position_ids)
        h = self.base.model.embed_tokens(input_ids)
        h = h + self.pass_embed[0].to(dtype=h.dtype).view(1, 1, -1)
        layer_states = []
        for layer in self.base.model.layers:
            h = layer(h, position_ids=position_ids)
            layer_states.append(h)
        return h, layer_states, position_ids

    def forward(self, input_ids: torch.LongTensor, position_ids: Optional[torch.LongTensor] = None, return_pass1_logits: bool = False):
        h1_resid, pass1_states, position_ids = self.forward_pass1_features(input_ids, position_ids=position_ids)
        h1 = self.base.model.norm(h1_resid)

        # Pass 2 refines pass 1; it does not discard pass 1.
        h2 = h1_resid + self.pass_embed[1].to(dtype=h1_resid.dtype).view(1, 1, -1)
        for i, layer in enumerate(self.base.model.layers):
            before = h2
            layer_out = layer(h2, position_ids=position_ids)
            layer_delta = layer_out - before
            layer_gate = torch.sigmoid(self.layer_gates[i]).to(dtype=h2.dtype)
            adapter_delta = self.adapters[i](h2, pass1_states[i])
            h2 = before + layer_gate * layer_delta + adapter_delta

        h2 = self.base.model.norm(h2)
        logits2 = self.base.lm_head(h2)

        if return_pass1_logits:
            with torch.no_grad():
                logits1 = self.base.lm_head(h1)
        else:
            logits1 = None
        return SimpleNamespace(logits=logits2, pass1_logits=logits1)


@torch.no_grad()

def load_lulu2_base(args, device, dtype):
    print("[guard] LULUV2 VWM runtime: no AutoModelForCausalLM.from_pretrained call and no external-model weights loaded.")
    print(f"[load] {args.checkpoint} ({human_bytes(os.path.getsize(args.checkpoint))})")
    ckpt = safe_torch_load(args.checkpoint, map_location="cpu")
    if "model" not in ckpt:
        raise ValueError("Checkpoint missing model state dict")
    sd = expand_shared_banks_into_state(ckpt, ckpt["model"])
    cfg = load_runtime_config(ckpt, sd, args)
    print(f"[config] hidden={cfg.hidden_size} layers={cfg.num_hidden_layers}")
    base = Lulu2ForCausalLM(cfg, sd)
    missing, unexpected = base.load_state_dict(sd, strict=False)
    print(f"[state:base] missing={len(missing)} unexpected={len(unexpected)}")
    if missing:
        print("[state:base] first missing:", missing[:10])
    if unexpected:
        print("[state:base] first unexpected:", unexpected[:10])
    base.to(device=device, dtype=dtype)
    return ckpt, base