#!/usr/bin/env python3
"""Build a small, auditable snapshot from Kimi K3's public model artifacts.

The script intentionally does not download a checkpoint. It consumes:

1. the public config and safetensors index;
2. safetensors JSON headers fetched with HTTP Range;
3. two small byte ranges containing one KDA parameter prefix and one MoE
   router prefix;
4. a local checkout of the official FlashKDA repository.

Raw model bytes stay local. The generated JSON contains only aggregate
statistics, public shapes, revisions, checksums, and a clearly labelled
synthetic-input router stress probe.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import platform
import re
import subprocess
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import torch


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=Path, required=True)
    parser.add_argument("--index", type=Path, required=True)
    parser.add_argument("--hf-model", type=Path, required=True)
    parser.add_argument("--kda-slice", type=Path, required=True)
    parser.add_argument("--router-prefix", type=Path, required=True)
    parser.add_argument("--mla-header", type=Path, required=True)
    parser.add_argument("--vision-header", type=Path, required=True)
    parser.add_argument("--flashkda-dir", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--synthetic-tokens", type=int, default=2048)
    parser.add_argument("--seed", type=int, default=20260729)
    parser.add_argument("--captured-at", default=None)
    return parser.parse_args()


def read_json(path: Path) -> dict[str, Any]:
    return json.loads(path.read_text())


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def summarize(values: torch.Tensor) -> dict[str, Any]:
    flat = values.detach().float().flatten().cpu()
    points = torch.tensor([0, 0.01, 0.1, 0.25, 0.5, 0.75, 0.9, 0.99, 1])
    quantiles = torch.quantile(flat, points).tolist()
    labels = ["min", "p01", "p10", "p25", "p50", "p75", "p90", "p99", "max"]
    return {
        "count": flat.numel(),
        "mean": flat.mean().item(),
        "std": flat.std().item(),
        "quantiles": dict(zip(labels, quantiles, strict=True)),
    }


def tensor_fact(header: dict[str, Any], name: str) -> dict[str, Any]:
    entry = header[name]
    return {
        "name": name,
        "dtype": entry["dtype"],
        "shape": entry["shape"],
        "bytes": entry["data_offsets"][1] - entry["data_offsets"][0],
    }


def parse_benchmark(path: Path, heads: int = 96) -> dict[str, Any]:
    text = path.read_text()
    section = text.split(f"### `T=8192`, `H={heads}`, `D=128`", 1)[1].split("###", 1)[0]
    rows = {}
    for label, flash, chunk, speedup, gdn, gdn_speedup in re.findall(
        r"\| ([^|]+?) \| ([0-9.]+) \| ([0-9.]+) \| ([0-9.]+)× \| ([0-9.]+) \| ([0-9.]+)× \|",
        section,
    ):
        rows[label.strip()] = {
            "flash_kda_ms": float(flash),
            "fla_chunk_kda_ms": float(chunk),
            "speedup_vs_chunk_kda": float(speedup),
            "fla_chunk_gdn_ms": float(gdn),
            "speedup_vs_gdn": float(gdn_speedup),
        }
    return {"sequence": 8192, "heads": heads, "dimension": 128, "rows": rows}


def load_metrics(load: torch.Tensor) -> dict[str, Any]:
    mean = load.mean()
    ordered = load.sort().values
    count = load.numel()
    indices = torch.arange(1, count + 1, device=load.device, dtype=torch.float32)
    gini = ((2 * indices - count - 1) * ordered).sum() / (count * ordered.sum())
    quantiles = torch.quantile(
        load,
        torch.tensor([0, 0.1, 0.25, 0.5, 0.75, 0.9, 0.99, 1], device=load.device),
    ).tolist()
    labels = ["min", "p10", "p25", "p50", "p75", "p90", "p99", "max"]
    return {
        "mean": mean.item(),
        "std": load.std().item(),
        "cv": (load.std() / mean).item(),
        "gini": gini.item(),
        "zero_experts": int((load == 0).sum()),
        "quantiles": dict(zip(labels, quantiles, strict=True)),
    }


def main() -> None:
    args = parse_args()
    config = read_json(args.config)
    text = config["text_config"]
    index = read_json(args.index)
    hf_model = read_json(args.hf_model)
    mla_header = read_json(args.mla_header)
    vision_header = read_json(args.vision_header)

    names = list(index["weight_map"])
    shard_files = [
        item
        for item in hf_model["siblings"]
        if re.fullmatch(r"model-\d+-of-\d+\.safetensors", item["rfilename"])
    ]
    shard_sizes = [item["size"] for item in shard_files]

    kda_raw = args.kda_slice.read_bytes()
    if len(kda_raw) != 49_664:
        raise ValueError(f"unexpected KDA slice length: {len(kda_raw)}")
    a_log = torch.frombuffer(bytearray(kda_raw[:512]), dtype=torch.float32).clone()
    dt_bias = torch.frombuffer(bytearray(kda_raw[512:]), dtype=torch.float32).clone().view(96, 128)

    router_raw = args.router_prefix.read_bytes()
    if len(router_raw) != 13_488_640:
        raise ValueError(f"unexpected router prefix length: {len(router_raw)}")
    correction_bias = torch.frombuffer(
        bytearray(router_raw[:3584]), dtype=torch.float32
    ).clone()
    router_weight = torch.frombuffer(
        bytearray(router_raw[643_584:13_488_640]), dtype=torch.bfloat16
    ).clone().view(896, 7168)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    torch.manual_seed(args.seed)
    if device.type == "cuda":
        torch.cuda.manual_seed_all(args.seed)
    router_device = router_weight.to(device)
    bias_device = correction_bias.to(device)
    synthetic = torch.randn(
        args.synthetic_tokens, 7168, device=device, dtype=torch.float32
    )
    synthetic *= torch.rsqrt(synthetic.square().mean(-1, keepdim=True) + 1e-6)
    scores = torch.sigmoid(synthetic.to(torch.bfloat16) @ router_device.T).float()
    unbiased_ids = scores.topk(16, dim=-1).indices
    biased_ids = (scores + bias_device).topk(16, dim=-1).indices

    def loads(ids: torch.Tensor) -> torch.Tensor:
        return torch.bincount(ids.flatten(), minlength=896).float()

    unbiased_load = loads(unbiased_ids)
    biased_load = loads(biased_ids)
    overlap = torch.tensor(
        [
            len(set(unbiased_ids[row].tolist()) & set(biased_ids[row].tolist()))
            for row in range(args.synthetic_tokens)
        ],
        dtype=torch.float32,
    )
    router_norms = router_weight.float().norm(dim=1)

    # This is deliberately a hypothesis probe, not a canonical forward pass:
    # the checkpoint stores A_log[128], while public code/API expect A_log[H=96].
    channelwise_log_decay = -5.0 * torch.sigmoid(torch.exp(a_log).view(1, 128) * dt_bias)
    channelwise_retention = torch.exp(channelwise_log_decay)

    per_expert_bytes = 3 * (5_505_024 + 344_064)
    all_routed_expert_bytes = per_expert_bytes * 896 * 92
    kda_layers = text["linear_attn_config"]["kda_layers"]
    mla_layers = text["linear_attn_config"]["full_attn_layers"]

    flash_revision = subprocess.check_output(
        ["git", "-C", str(args.flashkda_dir), "rev-parse", "HEAD"],
        text=True,
    ).strip()

    captured_at = args.captured_at or datetime.now(timezone.utc).isoformat()
    result = {
        "schema_version": 1,
        "captured_at": captured_at,
        "evidence_boundary": {
            "checkpoint_forward_run": False,
            "raw_weights_committed": False,
            "router_inputs": "deterministic synthetic RMS-normalized vectors, not token hidden states",
            "kda_retention_probe": "noncanonical channel-wise interpretation used only to expose the A_log shape ambiguity",
        },
        "provenance": {
            "huggingface_model": "moonshotai/Kimi-K3",
            "huggingface_revision": hf_model["sha"],
            "flashkda_revision": flash_revision,
            "sha256": {
                "config": sha256(args.config),
                "index": sha256(args.index),
                "kda_slice": sha256(args.kda_slice),
                "router_prefix": sha256(args.router_prefix),
            },
        },
        "checkpoint": {
            "tensor_data_bytes": index["metadata"]["total_size"],
            "tensor_data_tb": index["metadata"]["total_size"] / 1e12,
            "tensor_data_tib": index["metadata"]["total_size"] / 2**40,
            "shards": len(shard_files),
            "shard_file_bytes": {
                "sum": sum(shard_sizes),
                "min": min(shard_sizes),
                "max": max(shard_sizes),
                "mean": sum(shard_sizes) / len(shard_sizes),
            },
            "tensor_entries": len(names),
            "tensor_counts": {
                "expert_packed": sum(
                    bool(re.search(r"experts\.\d+\.w[123]\.weight_packed$", name))
                    for name in names
                ),
                "expert_scales": sum(
                    bool(re.search(r"experts\.\d+\.w[123]\.weight_scale$", name))
                    for name in names
                ),
                "router_weight": sum(
                    name.endswith("block_sparse_moe.gate.weight") for name in names
                ),
                "router_correction_bias": sum(
                    name.endswith("gate.e_score_correction_bias") for name in names
                ),
                "attnres_proj": sum(
                    bool(re.search(r"(_res_proj|output_attn_res_proj)\.weight$", name))
                    for name in names
                ),
                "attnres_norm": sum(
                    bool(re.search(r"(_res_norm|output_attn_res_norm)\.weight$", name))
                    for name in names
                ),
                "kda_a_log": sum(name.endswith("self_attn.A_log") for name in names),
                "kda_dt_bias": sum(name.endswith("self_attn.dt_bias") for name in names),
                "vision": sum(name.startswith("vision_tower.") for name in names),
                "projector": sum(name.startswith("mm_projector.") for name in names),
            },
            "derived_routed_expert_bytes": all_routed_expert_bytes,
            "derived_routed_expert_share": all_routed_expert_bytes
            / index["metadata"]["total_size"],
        },
        "configuration": {
            "layers": text["num_hidden_layers"],
            "dense_layers": text["first_k_dense_replace"],
            "hidden": text["hidden_size"],
            "vocabulary": text["vocab_size"],
            "context": text["max_position_embeddings"],
            "kda_layers": kda_layers,
            "mla_layers": mla_layers,
            "heads": text["num_attention_heads"],
            "head_dim": text["linear_attn_config"]["head_dim"],
            "attnres_block": text["attn_res_block_size"],
            "experts": text["num_experts"],
            "active_experts": text["num_experts_per_token"],
            "shared_experts": text["num_shared_experts"],
            "latent_width": text["routed_expert_hidden_size"],
            "expert_intermediate": text["moe_intermediate_size"],
            "situ_beta": text["activation_situ_beta"],
            "situ_linear_beta": text["activation_situ_linear_beta"],
            "mla_nope": text["mla_use_nope"],
            "mla_output_gate": text["mla_use_output_gate"],
        },
        "tensor_examples": {
            "mla_layer_4": [
                tensor_fact(
                    mla_header,
                    "language_model.model.layers.3.self_attn.kv_a_proj_with_mqa.weight",
                ),
                tensor_fact(
                    mla_header,
                    "language_model.model.layers.3.self_attn.kv_b_proj.weight",
                ),
                tensor_fact(
                    mla_header,
                    "language_model.model.layers.3.self_attn.q_a_proj.weight",
                ),
                tensor_fact(
                    mla_header,
                    "language_model.model.layers.3.self_attn.q_b_proj.weight",
                ),
                tensor_fact(
                    mla_header,
                    "language_model.model.layers.3.self_attn.g_proj.weight",
                ),
            ],
            "vision": [
                tensor_fact(vision_header, "vision_tower.patch_embed.proj.weight"),
                tensor_fact(vision_header, "vision_tower.patch_embed.pos_emb.weight"),
                tensor_fact(vision_header, "vision_tower.encoder.blocks.0.wqkv.weight"),
                tensor_fact(vision_header, "vision_tower.encoder.blocks.26.wqkv.weight"),
                tensor_fact(vision_header, "vision_tower.encoder.final_layernorm.weight"),
            ],
            "routed_expert_0": [
                {"name": "w1.weight_packed", "dtype": "U8", "shape": [3072, 1792], "bytes": 5_505_024},
                {"name": "w1.weight_scale", "dtype": "U8", "shape": [3072, 112], "bytes": 344_064},
                {"name": "w2.weight_packed", "dtype": "U8", "shape": [3584, 1536], "bytes": 5_505_024},
                {"name": "w2.weight_scale", "dtype": "U8", "shape": [3584, 96], "bytes": 344_064},
                {"name": "w3.weight_packed", "dtype": "U8", "shape": [3072, 1792], "bytes": 5_505_024},
                {"name": "w3.weight_scale", "dtype": "U8", "shape": [3072, 112], "bytes": 344_064},
            ],
        },
        "parameter_audit": {
            "a_log_checkpoint_shape": [128],
            "a_log_public_code_shape": [96],
            "dt_bias_shape": [96, 128],
            "beta_projection_shape": [96, 7168],
            "status": "observed shape inconsistency; runtime meaning unresolved",
            "a_log": summarize(a_log),
            "a_rate_exp": summarize(torch.exp(a_log)),
            "dt_bias": summarize(dt_bias),
            "channelwise_hypothesis": {
                "log_decay": summarize(channelwise_log_decay),
                "one_step_retention": summarize(channelwise_retention),
                "retention_after_64_steps": summarize(channelwise_retention.pow(64)),
            },
            "router_correction_bias": summarize(correction_bias),
            "router_row_l2": summarize(router_norms),
            "router_bias_norm_correlation": torch.corrcoef(
                torch.stack([correction_bias, router_norms])
            )[0, 1].item(),
        },
        "router_stress_probe": {
            "seed": args.seed,
            "synthetic_tokens": args.synthetic_tokens,
            "hidden_rms": synthetic.square().mean().sqrt().item(),
            "without_correction_bias": load_metrics(unbiased_load),
            "with_correction_bias": load_metrics(biased_load),
            "membership_overlap_mean": overlap.mean().item(),
            "tokens_changed": int((overlap < 16).sum()),
            "changed_fraction": (overlap < 16).float().mean().item(),
            "mean_replacements_per_token": (16 - overlap).mean().item(),
        },
        "flashkda": {
            "supported_architectures": ["90a", "100a", "103a", "120a"],
            "requirements": {"cuda": ">=12.9", "pytorch": ">=2.4", "gpu": "SM90+"},
            "official_benchmarks": {
                "h20": parse_benchmark(args.flashkda_dir / "BENCHMARK_H20.md"),
                "gb200": parse_benchmark(args.flashkda_dir / "BENCHMARK_GB200.md"),
            },
            "local_environment": {
                "python": platform.python_version(),
                "torch": torch.__version__,
                "torch_cuda": torch.version.cuda,
                "gpu": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None,
                "capability": list(torch.cuda.get_device_capability(0))
                if torch.cuda.is_available()
                else None,
                "libc": list(platform.libc_ver()),
            },
            "baseline_host_build": {
                "status": "blocked in the default CUDA 12.8 environment; superseded by the separate CUDA 13 runtime probe",
                "attempt_1": "system g++ 15 exceeds CUDA 12.8 host compiler range",
                "attempt_2": "temporary g++ 13 reaches nvcc, then CUDA 12.8 headers conflict with current glibc math declarations",
                "interpretation": "this snapshot records the first host path only; see src/data/k3-flashkda-runtime.json for the successful isolated build and RTX 5090 execution",
            },
        },
    }

    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(result, ensure_ascii=False, indent=2) + "\n")
    print(json.dumps({"output": str(args.output), "bytes": args.output.stat().st_size}))


if __name__ == "__main__":
    main()