#!/usr/bin/env python3
from __future__ import annotations

import argparse
import math
from pathlib import Path

import torch
import torch.distributed as dist
import torch.distributed.checkpoint as dcp
import torch.nn.functional as F
from tokenizers import Tokenizer


CKPT = Path(
    "runs/pretrain_8192_8gpu_dp8_mbs14_full_recompute_weighted_heldoutval_resume10000/"
    "checkpoints/iter_0107500"
)
TOKENIZER = Path("tokenizer/glm5.2/tokenizer.json")

HIDDEN = 1536
HEADS = 24
KV_HEADS = 4
HEAD_DIM = 64
DENSE_FFN = 4608
MOE_FFN = 6144
MOE_LAYERS = {2, 4, 6, 8, 10}
TOP_K = 4
N_LAYERS = 12
ROPE_BASE = 10000.0


def init_dist() -> None:
    if not dist.is_initialized():
        dist.init_process_group("gloo", init_method="tcp://127.0.0.1:29641", rank=0, world_size=1)


def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
    y = x.float() * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + eps)
    return (y * weight.float()).to(torch.bfloat16)


def rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1 = x[..., ::2]
    x2 = x[..., 1::2]
    out = torch.stack((-x2, x1), dim=-1)
    return out.flatten(-2)


def apply_rope(q: torch.Tensor, k: torch.Tensor, seq_len: int) -> tuple[torch.Tensor, torch.Tensor]:
    inv_freq = 1.0 / (ROPE_BASE ** (torch.arange(0, HEAD_DIM, 2, dtype=torch.float32) / HEAD_DIM))
    pos = torch.arange(seq_len, dtype=torch.float32)
    freqs = torch.einsum("t,d->td", pos, inv_freq)
    emb = torch.repeat_interleave(freqs, 2, dim=-1)
    cos = emb.cos()[None, :, None, :]
    sin = emb.sin()[None, :, None, :]
    qf = q.float()
    kf = k.float()
    return (qf * cos + rotate_half(qf) * sin).to(torch.bfloat16), (
        kf * cos + rotate_half(kf) * sin
    ).to(torch.bfloat16)


def swiglu(x: torch.Tensor) -> torch.Tensor:
    a, b = x.chunk(2, dim=-1)
    return F.silu(a.float()) * b.float()


def load_model_tensors(ckpt: Path) -> dict[str, torch.Tensor]:
    init_dist()
    metadata = dcp.FileSystemReader(str(ckpt)).read_metadata()
    state: dict[str, torch.Tensor] = {}
    for key, meta in metadata.state_dict_metadata.items():
        if key.startswith("optimizer.") or not hasattr(meta, "size"):
            continue
        state[key] = torch.empty(tuple(meta.size), dtype=meta.properties.dtype)
    print(f"loading {len(state)} model tensors from {ckpt}", flush=True)
    dcp.load(state, checkpoint_id=str(ckpt))
    print("loaded model tensors", flush=True)
    return state


def dense_mlp(x: torch.Tensor, state: dict[str, torch.Tensor], layer: int) -> torch.Tensor:
    prefix = f"decoder.layers.{layer}.mlp"
    h = rms_norm(x, state[f"{prefix}.linear_fc1.layer_norm_weight"])
    fc1 = F.linear(h.float(), state[f"{prefix}.linear_fc1.weight"].float())
    hidden = swiglu(fc1)
    return F.linear(hidden, state[f"{prefix}.linear_fc2.weight"].float()).to(torch.bfloat16)


def moe_mlp(x: torch.Tensor, state: dict[str, torch.Tensor], layer: int) -> torch.Tensor:
    prefix = f"decoder.layers.{layer}.mlp"
    h = rms_norm(x, state[f"decoder.layers.{layer}.pre_mlp_layernorm.weight"])
    router = F.linear(h.float(), state[f"{prefix}.router.weight"].float())
    topv, topi = torch.topk(router, TOP_K, dim=-1)
    probs = torch.softmax(topv, dim=-1)

    fc1_w = state[f"{prefix}.experts.experts.linear_fc1.weight"].float()
    fc2_w = state[f"{prefix}.experts.experts.linear_fc2.weight"].float()
    out = torch.zeros_like(h.float())
    flat_h = h.float().reshape(-1, HIDDEN)
    flat_topi = topi.reshape(-1, TOP_K)
    flat_probs = probs.reshape(-1, TOP_K)
    flat_out = out.reshape(-1, HIDDEN)
    for expert_id in range(fc1_w.shape[0]):
        mask = flat_topi == expert_id
        if not mask.any():
            continue
        token_idx, choice_idx = mask.nonzero(as_tuple=True)
        expert_in = flat_h[token_idx]
        expert_hidden = swiglu(F.linear(expert_in, fc1_w[expert_id]))
        expert_out = F.linear(expert_hidden, fc2_w[expert_id])
        flat_out[token_idx] += expert_out * flat_probs[token_idx, choice_idx].unsqueeze(-1)
    return out.to(torch.bfloat16)


def layer_forward(x: torch.Tensor, state: dict[str, torch.Tensor], layer: int) -> torch.Tensor:
    prefix = f"decoder.layers.{layer}"
    attn_in = rms_norm(x, state[f"{prefix}.self_attention.linear_qkv.layer_norm_weight"])
    qkv = F.linear(attn_in.float(), state[f"{prefix}.self_attention.linear_qkv.weight"].float())
    q, k, v = torch.split(qkv, [HEADS * HEAD_DIM, KV_HEADS * HEAD_DIM, KV_HEADS * HEAD_DIM], dim=-1)
    bsz, seq_len, _ = q.shape
    q = q.view(bsz, seq_len, HEADS, HEAD_DIM)
    k = k.view(bsz, seq_len, KV_HEADS, HEAD_DIM)
    v = v.view(bsz, seq_len, KV_HEADS, HEAD_DIM)
    q, k = apply_rope(q, k, seq_len)
    repeat = HEADS // KV_HEADS
    k = k.repeat_interleave(repeat, dim=2)
    v = v.repeat_interleave(repeat, dim=2)
    q = q.transpose(1, 2).float()
    k = k.transpose(1, 2).float()
    v = v.transpose(1, 2).float()
    scores = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(HEAD_DIM)
    mask = torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool), diagonal=1)
    scores = scores.masked_fill(mask[None, None, :, :], float("-inf"))
    ctx = torch.matmul(torch.softmax(scores, dim=-1), v)
    ctx = ctx.transpose(1, 2).contiguous().view(bsz, seq_len, HIDDEN)
    attn_out = F.linear(ctx, state[f"{prefix}.self_attention.linear_proj.weight"].float()).to(torch.bfloat16)
    x = x + attn_out
    x = x + (moe_mlp(x, state, layer) if layer in MOE_LAYERS else dense_mlp(x, state, layer))
    return x


@torch.inference_mode()
def logits_for_ids(ids: list[int], state: dict[str, torch.Tensor]) -> torch.Tensor:
    input_ids = torch.tensor(ids, dtype=torch.long).unsqueeze(0)
    x = F.embedding(input_ids, state["embedding.word_embeddings.weight"])
    for layer in range(N_LAYERS):
        x = layer_forward(x, state, layer)
    x = rms_norm(x, state["decoder.final_layernorm.weight"])
    return torch.matmul(x[:, -1, :].float(), state["embedding.word_embeddings.weight"].float().T)


def generate(prompt: str, state: dict[str, torch.Tensor], tokenizer: Tokenizer, max_new_tokens: int) -> str:
    ids = tokenizer.encode(prompt).ids
    print(f"prompt={prompt!r} input_tokens={len(ids)}", flush=True)
    for step in range(max_new_tokens):
        logits = logits_for_ids(ids, state)
        next_id = int(torch.argmax(logits, dim=-1)[0])
        ids.append(next_id)
        piece = tokenizer.decode([next_id], skip_special_tokens=False)
        print(f"step={step} id={next_id} piece={piece!r}", flush=True)
    return tokenizer.decode(ids, skip_special_tokens=False)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--max-new-tokens", type=int, default=8)
    parser.add_argument("--prompt", action="append", default=None)
    args = parser.parse_args()

    tokenizer = Tokenizer.from_file(str(TOKENIZER))
    prompts = args.prompt or [
        "The capital of France is",
        "中国的首都是",
    ]
    state = load_model_tensors(CKPT)
    for prompt in prompts:
        text = generate(prompt, state, tokenizer, args.max_new_tokens)
        print("=== GENERATED ===")
        print(text)


if __name__ == "__main__":
    main()