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

import hashlib
import json
import math
import re
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
CONFIG = ROOT / "model/nemo_megatron/laoyao_2b_moe_nemo_megatron.yaml"
TOKENIZER = ROOT / "tokenizer/glm5.2/tokenizer.json"
EXPECTED_TOKENIZER_SHA256 = "19e773648cb4e65de8660ea6365e10acca112d42a854923df93db4a6f333a82d"


def read_scalar(name: str, text: str) -> int:
    match = re.search(rf"^\s*{re.escape(name)}:\s*([0-9]+)\s*$", text, re.M)
    if not match:
        raise SystemExit(f"missing integer field: {name}")
    return int(match.group(1))


def read_list(name: str, text: str) -> list[int]:
    match = re.search(rf"^\s*{re.escape(name)}:\s*\[([^\]]*)\]", text, re.M)
    if not match:
        raise SystemExit(f"missing list field: {name}")
    return [int(x.strip()) for x in match.group(1).split(",") if x.strip()]


def read_bool(name: str, text: str) -> bool:
    match = re.search(rf"^\s*{re.escape(name)}:\s*(true|false)\s*$", text, re.M)
    if not match:
        raise SystemExit(f"missing boolean field: {name}")
    return match.group(1) == "true"


def main() -> None:
    text = CONFIG.read_text(encoding="utf-8")
    tok_bytes = TOKENIZER.read_bytes()
    tok_sha = hashlib.sha256(tok_bytes).hexdigest()
    tok = json.loads(tok_bytes)
    vocab_size = len(tok.get("model", {}).get("vocab", {}))

    hidden = read_scalar("hidden_size", text)
    layers = read_scalar("num_layers", text)
    heads = read_scalar("num_attention_heads", text)
    groups = read_scalar("num_query_groups", text)
    dense_ffn = read_scalar("ffn_hidden_size", text)
    seq_len = read_scalar("seq_length", text)
    cfg_vocab = read_scalar("vocab_size", text)
    num_experts = read_scalar("num_moe_experts", text)
    topk = read_scalar("moe_router_topk", text)
    expert_ffn = read_scalar("moe_ffn_hidden_size", text)
    moe_layers = read_list("layer_indices", text)
    share_embeddings = read_bool("share_embeddings_and_output_weights", text)

    assert hidden % heads == 0, "hidden_size must divide num_attention_heads"
    assert heads % groups == 0, "num_attention_heads must divide num_query_groups for GQA"
    assert cfg_vocab == vocab_size, f"config vocab_size={cfg_vocab} != tokenizer vocab_size={vocab_size}"
    assert tok_sha == EXPECTED_TOKENIZER_SHA256, f"unexpected tokenizer sha256: {tok_sha}"
    assert all(0 <= i < layers for i in moe_layers), "MoE layer index out of range"
    assert 0 < topk <= num_experts, "moe_router_topk must be in [1, num_experts]"

    dense_layers = layers - len(moe_layers)
    # Rough parameter accounting for audit only. SwiGLU FFN uses 3 matrices.
    embed = cfg_vocab * hidden
    embedding_total = embed if share_embeddings else embed * 2
    attn_per_layer = hidden * hidden * (2 + 2 * groups / heads)  # q,o full; k,v grouped
    dense_ffn_per_layer = 3 * hidden * dense_ffn
    moe_ffn_per_layer_total = num_experts * 3 * hidden * expert_ffn
    moe_ffn_per_layer_active = topk * 3 * hidden * expert_ffn
    shared_total = embedding_total + layers * attn_per_layer + dense_layers * dense_ffn_per_layer
    active_expert_total = len(moe_layers) * moe_ffn_per_layer_active
    total_expert_total = len(moe_layers) * moe_ffn_per_layer_total
    dense_total = shared_total + total_expert_total
    active_total = shared_total + active_expert_total

    print("architecture_ok")
    print(f"tokenizer_sha256={tok_sha}")
    print(f"vocab_size={vocab_size}")
    print(f"share_embeddings_and_output_weights={str(share_embeddings).lower()}")
    print(f"layers={layers} hidden={hidden} heads={heads} query_groups={groups} seq_len={seq_len}")
    print(f"moe_layers={moe_layers} experts={num_experts} topk={topk} expert_ffn={expert_ffn}")
    print(f"rough_embedding_params={embedding_total/1e9:.3f}B")
    print(f"rough_shared_params={shared_total/1e9:.3f}B")
    print(f"rough_active_expert_params={active_expert_total/1e9:.3f}B")
    print(f"rough_total_expert_params={total_expert_total/1e9:.3f}B")
    print(f"rough_total_params={dense_total/1e9:.3f}B")
    print(f"rough_active_params={active_total/1e9:.3f}B")


if __name__ == "__main__":
    main()