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

Purpose: Qwen/Qwen3.8-Flash-Next (released 2026-08-24) is a large multimodal
Qwen4ExpForConditionalGeneration MoE with a hybrid linear+full attention stack,
ngram vocab, an indexer, PLE, hyper-connections, and an MTP head, so nobody can
instantiate it in CI or on a laptop. This fixture ships a ~0.3 MB random-init
TEXT checkpoint plus a reduced config that keeps the real qwen4_exp / qwen4_exp_text
field names (layer_types, linear_*, ngram_*, indexer_*, ple_*, hc_*, mtp,
rope_parameters, vision_config) 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.

Random-init: NOT trained and NOT a quality claim. The tensor set is a REDUCED
standard-attention + MoE convention; the real model's linear-attention (conv/ssm),
ngram, indexer, PLE, hyper-connection, MTP, and vision/projector tensors are NOT
included (see README omissions). Geometry is documented in the README.
"""
import hashlib
import json
import math
import os
import struct

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

# ---- tiny geometry (reduced from the real text_config, documented in README) ----
VOCAB = 256
HIDDEN = 64
LAYERS = 4
HEADS = 4
KV_HEADS = 2
HEAD_DIM = 16
MOE_INTER = 32
SHARED_INTER = 32
NUM_EXPERTS = 8
TOPK = 2


class SplitMix64:
    """SplitMix64 + Box-Muller, identical to the llama/t5/glm 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)
        shapes[p + "mlp.gate.weight"] = (NUM_EXPERTS, HIDDEN)
        for expert in range(NUM_EXPERTS):
            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_expert.gate_proj.weight"] = (SHARED_INTER, HIDDEN)
        shapes[p + "mlp.shared_expert.up_proj.weight"] = (SHARED_INTER, HIDDEN)
        shapes[p + "mlp.shared_expert.down_proj.weight"] = (HIDDEN, SHARED_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__)),
                           "qwen4_exp_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)

    layer_types = ["linear_attention", "linear_attention", "linear_attention",
                   "full_attention"]
    text_config = {
        "attention_bias": False,
        "attention_dropout": 0.0,
        "bos_token_id": 1,
        "dtype": "float32",
        "eos_token_id": 2,
        "full_attention_interval": 4,
        "hc_count": 2,
        "hc_lowrank": 32,
        "head_dim": HEAD_DIM,
        "heads_per_ngram": 2,
        "hidden_act": "silu",
        "hidden_size": HIDDEN,
        "indexer_budget": 32,
        "indexer_compress_ratio": 4,
        "indexer_head_dim": 16,
        "indexer_kv_heads": 1,
        "indexer_n_heads": 2,
        "initializer_range": 0.02,
        "layer_types": layer_types,
        "linear_conv_kernel_dim": 4,
        "linear_key_head_dim": 16,
        "linear_num_key_heads": 2,
        "linear_num_value_heads": 4,
        "linear_value_head_dim": 16,
        "make_ngram_vocab_size_divisible_by": 128,
        "mamba_ssm_dtype": "float32",
        "max_position_embeddings": 256,
        "model_type": "qwen4_exp_text",
        "moe_intermediate_size": MOE_INTER,
        "mtp": {
            "hybrid": True,
            "layer_types": ["full_attention"],
            "mtp_use_hidden_state_from_layer": None,
            "num_hidden_layers": 1,
            "rope_theta": 10000000,
        },
        "mtp_num_hidden_layers": 1,
        "mtp_use_dedicated_embeddings": False,
        "ngram_size": 3,
        "ngram_vocab_size_base": 4096,
        "num_attention_heads": HEADS,
        "num_experts": NUM_EXPERTS,
        "num_experts_per_tok": TOPK,
        "num_hidden_layers": LAYERS,
        "num_key_value_heads": KV_HEADS,
        "output_gate_type": "sigmoid",
        "output_router_logits": False,
        "pad_token_id": 0,
        "partial_rotary_factor": 0.25,
        "ple_conv_kernel_size": 4,
        "ple_embed_dim": HIDDEN,
        "ple_layer_ids": [2],
        "rms_norm_eps": 1e-6,
        "rope_parameters": {
            "mrope_interleaved": True,
            "mrope_section": [4, 4, 8],
            "partial_rotary_factor": 0.25,
            "rope_theta": 10000000,
            "rope_type": "default",
        },
        "router_aux_loss_coef": 0.001,
        "shared_expert_intermediate_size": SHARED_INTER,
        "split_ngram_parts": 128,
        "tie_word_embeddings": False,
        "use_cache": True,
        "vocab_size": VOCAB,
    }
    config = {
        "architectures": ["Qwen4ExpForConditionalGeneration"],
        "image_token_id": 248056,
        "language_model_only": True,
        "model_type": "qwen4_exp",
        "text_config": text_config,
        "tie_word_embeddings": False,
        "video_token_id": 248057,
        "vision_config": {
            "model_type": "qwen4_exp",
            "_note": "placeholder; this fixture ships NO vision/projector tensors",
        },
        "vision_end_token_id": 248054,
        "vision_start_token_id": 248053,
    }
    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()