#!/usr/bin/env python3
"""Build a tiny, deterministic random-init glm_moe_dsa fixture (stdlib only).

Purpose: zai-org/GLM-5.3 (released 2026-08-25) is a 753B MoE, so nobody can load it in
CI or on a laptop. This fixture ships a ~0.3 MB random-init checkpoint that uses
the same reduced config schema (model_type glm_moe_dsa, routed + shared experts,
dense->MoE layer schedule) so loader, quant-pipeline, and CI tests can exercise
the new architecture without the real weights.

Random-init: NOT a trained model and not a quality claim. Naming and geometry are
documented in the README of the output folder.
"""
import hashlib
import json
import math
import os
import struct

M64 = (1 << 64) - 1
SEED = 20260901
SCALE = 0.02

# ---- tiny geometry (reduced from the real config, documented in README) ----
VOCAB = 256
HIDDEN = 64
LAYERS = 4
HEADS = 4
KV_HEADS = 4
HEAD_DIM = 16
DENSE_INTER = 128
MOE_INTER = 32
N_ROUTED = 8
TOPK = 2
N_SHARED = 1
FIRST_DENSE = 1
N_GROUP = 1


class SplitMix64:
    """SplitMix64 + Box-Muller, identical to the llama/t5 fixtures."""

    def __init__(self, seed):
        self.state = seed & M64
        self._spare = None

    def next_u64(self):
        self.state = (self.state + 0x9E3779B97F4A7C15) & M64
        z = self.state
        z = ((z ^ (z >> 30)) * 0xBF584A7F17C119E3) & M64
        z = ((z ^ (z >> 27)) * 0x94D049BB133111EB) & M64
        return z ^ (z >> 31)

    def uniform(self):
        return (self.next_u64() >> 11) / float(1 << 53)

    def gauss(self):
        if self._spare is not None:
            value, self._spare = self._spare, None
            return value
        u1 = 1.0 - self.uniform()
        u2 = self.uniform()
        radius = math.sqrt(-2.0 * math.log(u1))
        theta = 2.0 * math.pi * u2
        self._spare = radius * math.sin(theta)
        return radius * math.cos(theta)


def build_tensors():
    shapes = {
        "model.embed_tokens.weight": (VOCAB, HIDDEN),
        "model.norm.weight": (HIDDEN,),
    }
    ones = {"model.norm.weight"}
    for layer in range(LAYERS):
        p = "model.layers.%d." % layer
        shapes[p + "input_layernorm.weight"] = (HIDDEN,)
        shapes[p + "post_attention_layernorm.weight"] = (HIDDEN,)
        ones.add(p + "input_layernorm.weight")
        ones.add(p + "post_attention_layernorm.weight")
        shapes[p + "self_attn.q_proj.weight"] = (HEADS * HEAD_DIM, HIDDEN)
        shapes[p + "self_attn.k_proj.weight"] = (KV_HEADS * HEAD_DIM, HIDDEN)
        shapes[p + "self_attn.v_proj.weight"] = (KV_HEADS * HEAD_DIM, HIDDEN)
        shapes[p + "self_attn.o_proj.weight"] = (HIDDEN, HEADS * HEAD_DIM)
        if layer < FIRST_DENSE:
            shapes[p + "mlp.gate_proj.weight"] = (DENSE_INTER, HIDDEN)
            shapes[p + "mlp.up_proj.weight"] = (DENSE_INTER, HIDDEN)
            shapes[p + "mlp.down_proj.weight"] = (HIDDEN, DENSE_INTER)
        else:
            shapes[p + "mlp.gate.weight"] = (N_ROUTED, HIDDEN)
            for expert in range(N_ROUTED):
                e = p + "mlp.experts.%d." % expert
                shapes[e + "gate_proj.weight"] = (MOE_INTER, HIDDEN)
                shapes[e + "up_proj.weight"] = (MOE_INTER, HIDDEN)
                shapes[e + "down_proj.weight"] = (HIDDEN, MOE_INTER)
            shapes[p + "mlp.shared_experts.gate_proj.weight"] = (MOE_INTER, HIDDEN)
            shapes[p + "mlp.shared_experts.up_proj.weight"] = (MOE_INTER, HIDDEN)
            shapes[p + "mlp.shared_experts.down_proj.weight"] = (HIDDEN, MOE_INTER)

    rng = SplitMix64(SEED)
    out = {}
    for name in sorted(shapes):
        shape = shapes[name]
        count = 1
        for dim in shape:
            count *= dim
        if name in ones:
            values = [1.0] * count
        else:
            values = [rng.gauss() * SCALE for _ in range(count)]
        blob = b"".join(
            struct.pack("<f", struct.unpack("<f", struct.pack("<f", v))[0]) for v in values
        )
        out[name] = (list(shape), "F32", blob)
    return out


def write_safetensors(path, tensors, metadata):
    header = {"__metadata__": metadata}
    offset = 0
    blobs = []
    for name in sorted(tensors):
        shape, dtype, blob = tensors[name]
        header[name] = {"dtype": dtype, "shape": shape,
                        "data_offsets": [offset, offset + len(blob)]}
        offset += len(blob)
        blobs.append(blob)
    raw = json.dumps(header, separators=(",", ":")).encode("utf-8")
    raw += b" " * ((-len(raw)) % 8)
    with open(path, "wb") as handle:
        handle.write(struct.pack("<Q", len(raw)))
        handle.write(raw)
        for blob in blobs:
            handle.write(blob)
    return len(raw), offset


def main():
    out_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                           "glm_moe_dsa_tiny_fixture")
    os.makedirs(out_dir, exist_ok=True)

    tensors = build_tensors()
    metadata = {
        "format": "pt",
        "source": "usefulHuggingface",
        "generator": "SplitMix64 seed=%d Box-Muller scale=%s float32 row-major" % (SEED, SCALE),
    }
    header_len, data_len = write_safetensors(
        os.path.join(out_dir, "model.safetensors"), tensors, metadata)

    config = {
        "architectures": ["GlmMoeDsaForCausalLM"],
        "attention_bias": False,
        "attention_dropout": 0.0,
        "dtype": "float32",
        "ep_size": 1,
        "first_k_dense_replace": FIRST_DENSE,
        "head_dim": HEAD_DIM,
        "hidden_act": "silu",
        "hidden_size": HIDDEN,
        "intermediate_size": DENSE_INTER,
        "max_position_embeddings": 256,
        "model_type": "glm_moe_dsa",
        "moe_intermediate_size": MOE_INTER,
        "moe_layer_freq": 1,
        "moe_router_dtype": "float32",
        "n_group": N_GROUP,
        "n_routed_experts": N_ROUTED,
        "n_shared_experts": N_SHARED,
        "norm_topk_prob": True,
        "num_attention_heads": HEADS,
        "num_experts_per_tok": TOPK,
        "num_hidden_layers": LAYERS,
        "num_key_value_heads": KV_HEADS,
        "num_nextn_predict_layers": 0,
        "pad_token_id": 0,
        "rms_norm_eps": 1e-5,
        "rope_theta": 10000.0,
        "tie_word_embeddings": False,
        "vocab_size": VOCAB,
    }
    with open(os.path.join(out_dir, "config.json"), "w") as handle:
        json.dump(config, handle, indent=2, sort_keys=True)
        handle.write("\n")

    with open(os.path.join(out_dir, "generation_config.json"), "w") as handle:
        json.dump({"bos_token_id": 1, "eos_token_id": 2, "pad_token_id": 0,
                   "no_repeat_ngram_size": 4, "seed": SEED},
                  handle, indent=2, sort_keys=True)
        handle.write("\n")

    with open(os.path.join(out_dir, "tokenizer_config.json"), "w") as handle:
        json.dump({"model_max_length": 256, "bos_token": "<s>", "eos_token": "</s>",
                   "unk_token": "<unk>", "pad_token": "<pad>",
                   "model_input_names": ["input_ids"]},
                  handle, indent=2, sort_keys=True)
        handle.write("\n")

    with open(os.path.join(out_dir, "special_tokens_map.json"), "w") as handle:
        json.dump({"additional_special_tokens": ["<pad>", "<unk>"],
                   "bos_token": "<s>", "eos_token": "</s>",
                   "pad_token": "<pad>", "unk_token": "<unk>"},
                  handle, indent=2, sort_keys=True)
        handle.write("\n")

    lines = []
    for name in sorted(tensors):
        shape, dtype, blob = tensors[name]
        lines.append("%s  %s  %s  %d  %s" % (name, dtype, "x".join(map(str, shape)),
                                             len(blob), hashlib.sha256(blob).hexdigest()))
    with open(os.path.join(out_dir, "checksums.txt"), "w") as handle:
        handle.write("\n".join(lines) + "\n")

    print("header_len=%d data_len=%d tensors=%d" % (header_len, data_len, len(tensors)))
    print("total_params=%d" % (data_len // 4))


if __name__ == "__main__":
    main()