#!/usr/bin/env python3
"""Build a tiny, deterministic deepseek_v4 (DeepSeek-V4-Flash-Vision-Exp) fixture.

Purpose: a byte-reproducible random-init TEXT checkpoint plus a reduced config that
keeps the real deepseek_v4 field names (MLA q_lora_rank/o_lora_rank/o_groups, MoE
routed+shared experts, DSA index_*, num_hash_layers, MTP num_nextn_predict_layers,
swiglu_limit, sliding_window, yarn rope) so loaders, quant planners, and CI jobs can
exercise the new architecture's config parsing, expert-table sizing, and safetensors
load path without the real weights. No third-party deps (stdlib only).
"""
import hashlib
import json
import math
import os
import struct

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

# ---- fixture geometry (tiny on purpose; real names, small numbers) ----
VOCAB = 256
HIDDEN = 64
LAYERS = 4
HEADS = 4
KV_HEADS = 1                 # faithful: base num_key_value_heads == 1
HEAD_DIM = 16
INTER = 128                  # dense mlp intermediate (unused; all layers MoE here)
MOE_INTER = 32
N_ROUTED = 8
N_SHARED = 1
TOPK = 2
Q_LORA_RANK = 32
O_LORA_RANK = 32
O_GROUPS = 4
QK_NOPE_HEAD_DIM = 16
QK_ROPE_HEAD_DIM = 8
V_HEAD_DIM = 16


class SplitMix64:
    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():
    """Return {name: (shape, dtype, bytes)} in sorted-name consumption order.

    Reduced standard-attention + MoE convention (documented in README): the real
    MLA compressed tensors, DSA index tensors, hash-layer tensors, MTP head, and
    vision encoder are NOT included.
    """
    names = []
    for i in range(LAYERS):
        p = "model.layers.%d." % i
        names += [
            p + "input_layernorm.weight",
            p + "post_attention_layernorm.weight",
            p + "self_attn.q_a_proj.weight",
            p + "self_attn.q_b_proj.weight",
            p + "self_attn.kv_a_proj.weight",
            p + "self_attn.kv_b_proj.weight",
            p + "self_attn.q_a_layernorm.weight",
            p + "self_attn.kv_a_layernorm.weight",
            p + "self_attn.o_proj.weight",
            p + "mlp.gate.weight",
        ]
        for e in range(N_ROUTED):
            names += [
                p + "mlp.experts.%d.gate_proj.weight" % e,
                p + "mlp.experts.%d.up_proj.weight" % e,
                p + "mlp.experts.%d.down_proj.weight" % e,
            ]
        for s in range(N_SHARED):
            names += [
                p + "mlp.shared_experts.%d.gate_proj.weight" % s,
                p + "mlp.shared_experts.%d.up_proj.weight" % s,
                p + "mlp.shared_experts.%d.down_proj.weight" % s,
            ]
    names += ["model.embed_tokens.weight", "model.norm.weight"]

    shapes = {
        "model.embed_tokens.weight": (VOCAB, HIDDEN),
        "model.norm.weight": (HIDDEN,),
    }
    for i in range(LAYERS):
        p = "model.layers.%d." % i
        shapes[p + "input_layernorm.weight"] = (HIDDEN,)
        shapes[p + "post_attention_layernorm.weight"] = (HIDDEN,)
        shapes[p + "self_attn.q_a_proj.weight"] = (Q_LORA_RANK, HIDDEN)
        shapes[p + "self_attn.q_b_proj.weight"] = (HEADS * (QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM), Q_LORA_RANK)
        shapes[p + "self_attn.kv_a_proj.weight"] = (Q_LORA_RANK, HIDDEN)
        shapes[p + "self_attn.kv_b_proj.weight"] = (QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM + V_HEAD_DIM, Q_LORA_RANK)
        shapes[p + "self_attn.q_a_layernorm.weight"] = (Q_LORA_RANK,)
        shapes[p + "self_attn.kv_a_layernorm.weight"] = (Q_LORA_RANK,)
        shapes[p + "self_attn.o_proj.weight"] = (HIDDEN, HEADS * V_HEAD_DIM)
        shapes[p + "mlp.gate.weight"] = (N_ROUTED, HIDDEN)
        for e in range(N_ROUTED):
            shapes[p + "mlp.experts.%d.gate_proj.weight" % e] = (MOE_INTER, HIDDEN)
            shapes[p + "mlp.experts.%d.up_proj.weight" % e] = (MOE_INTER, HIDDEN)
            shapes[p + "mlp.experts.%d.down_proj.weight" % e] = (HIDDEN, MOE_INTER)
        for s in range(N_SHARED):
            shapes[p + "mlp.shared_experts.%d.gate_proj.weight" % s] = (MOE_INTER, HIDDEN)
            shapes[p + "mlp.shared_experts.%d.up_proj.weight" % s] = (MOE_INTER, HIDDEN)
            shapes[p + "mlp.shared_experts.%d.down_proj.weight" % s] = (HIDDEN, MOE_INTER)

    ones = {"model.norm.weight"}
    for i in range(LAYERS):
        p = "model.layers.%d." % i
        ones.add(p + "input_layernorm.weight")
        ones.add(p + "post_attention_layernorm.weight")
        ones.add(p + "self_attn.q_a_layernorm.weight")
        ones.add(p + "self_attn.kv_a_layernorm.weight")

    rng = SplitMix64(SEED)
    out = {}
    for name in sorted(names):
        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")
    pad = (-len(raw)) % 8
    raw += b" " * pad
    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__)), "deepseek_v4_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": ["DeepseekV4ForCausalLM"],
        "model_type": "deepseek_v4",
        "attention_bias": False,
        "attention_dropout": 0.0,
        "bos_token_id": 0,
        "eos_token_id": 1,
        "expert_dtype": "fp4",
        "head_dim": HEAD_DIM,
        "hidden_act": "silu",
        "hidden_size": HIDDEN,
        "initializer_range": 0.02,
        "max_position_embeddings": 256,
        "moe_intermediate_size": MOE_INTER,
        "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_hash_layers": 0,
        "num_key_value_heads": KV_HEADS,
        "num_nextn_predict_layers": 0,
        "o_groups": O_GROUPS,
        "o_lora_rank": O_LORA_RANK,
        "q_lora_rank": Q_LORA_RANK,
        "qk_rope_head_dim": QK_ROPE_HEAD_DIM,
        "rms_norm_eps": 1e-20,
        "rope_scaling": {
            "beta_fast": 32,
            "beta_slow": 1,
            "factor": 16,
            "original_max_position_embeddings": 65536,
            "type": "yarn",
        },
        "rope_theta": 10000,
        "routed_scaling_factor": 1.5,
        "scoring_func": "sqrtsoftplus",
        "sliding_window": 128,
        "swiglu_limit": 10.0,
        "tie_word_embeddings": False,
        "topk_method": "noaux_tc",
        "torch_dtype": "float32",
        "use_cache": True,
        "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": 0, "eos_token_id": 1, "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")

    total_params = sum(1 for t in tensors.values() for _ in range(t[0][0] * (t[0][1] if len(t[0]) > 1 else 1)))
    print("header_len=%d data_len=%d tensors=%d" % (header_len, data_len, len(tensors)))
    print("total_params=%d" % total_params)


if __name__ == "__main__":
    main()