#!/usr/bin/env python3
"""Mistral-Small-3.1-24B-Base-2503 -> thrasher base: strip vision, claim ChatML.

    python prep_base.py --src <base snapshot dir> --out <dir> [--init copy|mean|none]
    python prep_base.py --src <dir with tokenizer jsons> --out <dir> --tokenizer-only

One streaming pass over the shards (CPU-only, ~one tensor in memory at a time;
fine on the box or locally).
"""
from __future__ import annotations

import argparse
import json
import shutil
import sys
from pathlib import Path

CLAIMS = {  # id -> (old, new, donor_id)
    20: ("<SPECIAL_20>", "<|im_start|>", 1),   # donor <s>
    21: ("<SPECIAL_21>", "<|im_end|>", 2),     # donor </s>
}


BROKEN_REGEX = (r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+"
                r"|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+")
FIXED_REGEX = (r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*"
               r"[\p{Ll}\p{Lm}\p{Lo}\p{M}]+|[^\r\n\p{L}\p{N}]?"
               r"[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+[\p{Ll}\p{Lm}\p{Lo}\p{M}]*"
               r"|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n/]*|\s*[\r\n]+|\s+(?!\S)|\s+")
DROP_PREFIXES = ("vision_tower.", "multi_modal_projector.")
LM_PREFIX = "language_model."
ROW_KEYS = ("model.embed_tokens.weight", "lm_head.weight")


def replace_deep(obj, mapping: dict[str, str]):
    if isinstance(obj, str):
        return mapping.get(obj, obj)
    if isinstance(obj, list):
        return [replace_deep(x, mapping) for x in obj]
    if isinstance(obj, dict):
        return {k: replace_deep(v, mapping) for k, v in obj.items()}
    return obj


def jload(p: Path):
    with open(p) as f:
        return json.load(f)


def jdump(obj, p: Path):
    with open(p, "w") as f:
        json.dump(obj, f, indent=2, ensure_ascii=False)
        f.write("\n")


def prep_tokenizer(src: Path, out: Path, template_text: str) -> None:
    strmap = {old: new for old, new, _ in CLAIMS.values()}

    tj = jload(src / "tokenizer.json")
    renamed = 0
    for tok in tj["added_tokens"]:
        if tok["content"] in strmap:
            tok["content"] = strmap[tok["content"]]
            renamed += 1
    vocab = tj["model"]["vocab"]
    for old, new, _ in CLAIMS.values():
        assert old in vocab, f"{old} not in vocab — wrong base?"
        assert new not in vocab, f"{new} already in vocab"
        vocab[new] = vocab.pop(old)
    assert renamed == len(CLAIMS), f"renamed {renamed} added_tokens, expected {len(CLAIMS)}"

    split = tj["pre_tokenizer"]["pretokenizers"][0]["pattern"]
    assert split["Regex"] == BROKEN_REGEX, "pre_tokenizer not the known-broken pattern — re-diff before baking"
    split["Regex"] = FIXED_REGEX
    jdump(tj, out / "tokenizer.json")

    tc = replace_deep(jload(src / "tokenizer_config.json"), strmap)
    tc["eos_token"] = "<|im_end|>"
    tc["chat_template"] = template_text
    jdump(tc, out / "tokenizer_config.json")

    sm = replace_deep(jload(src / "special_tokens_map.json"), strmap)
    eos = sm.get("eos_token")
    if isinstance(eos, dict):
        eos["content"] = "<|im_end|>"
    else:
        sm["eos_token"] = "<|im_end|>"
    jdump(sm, out / "special_tokens_map.json")

    (out / "chat_template.jinja").write_text(template_text)

    # round-trip proof, not guess
    from tokenizers import Tokenizer
    tok = Tokenizer.from_file(str(out / "tokenizer.json"))
    ids = tok.encode("<|im_start|>user\nhi<|im_end|>\n").ids
    assert ids[0] == 1 and 20 in ids and 21 in ids, f"claim round-trip failed: {ids}"
    assert tok.decode([20, 21], skip_special_tokens=False) == "<|im_start|><|im_end|>"
    print(f"[tokenizer] claimed: " + ", ".join(
        f"{new}={i}" for i, (_, new, _) in CLAIMS.items()))
    print(f"[tokenizer] round-trip ids for ChatML probe: {ids}")


def prep_configs(src: Path, out: Path) -> None:
    cfg = jload(src / "config.json")
    text = cfg["text_config"]
    text.update({
        "architectures": ["MistralForCausalLM"],
        "model_type": "mistral",
        "torch_dtype": cfg.get("torch_dtype", "bfloat16"),
        "tie_word_embeddings": False,
        "bos_token_id": 1,
        "eos_token_id": 21,
    })
    jdump(text, out / "config.json")

    gen = {"bos_token_id": 1, "eos_token_id": [21]}
    if (src / "generation_config.json").exists():
        g = jload(src / "generation_config.json")
        g.update(gen)
        g.pop("pad_token_id", None)
        gen = g
    gen["transformers_version"] = None
    gen = {k: v for k, v in gen.items() if v is not None}
    jdump(gen, out / "generation_config.json")
    print("[config] MistralForCausalLM, untied, eos_token_id=[21]")


def prep_weights(src: Path, out: Path, init: str) -> None:
    import torch
    from safetensors import safe_open
    from safetensors.torch import save_file

    index = jload(src / "model.safetensors.index.json")
    wmap = index["weight_map"]
    shards: dict[str, list[str]] = {}
    for key, shard in wmap.items():
        shards.setdefault(shard, []).append(key)

    new_map: dict[str, str] = {}
    total = 0
    n_drop = n_keep = 0
    donor_rows: dict[str, dict[int, torch.Tensor]] = {}  # row_key -> {donor_id: row}

    shard_names = sorted(shards)
    for si, shard in enumerate(shard_names, 1):
        out_name = f"model-{si:05d}-of-{len(shard_names):05d}.safetensors"
        tensors: dict[str, torch.Tensor] = {}
        with safe_open(src / shard, framework="pt") as f:
            for key in sorted(shards[shard]):
                if key.startswith(DROP_PREFIXES):
                    n_drop += 1
                    continue
                assert key.startswith(LM_PREFIX), f"unexpected key {key}"
                nk = key[len(LM_PREFIX):]
                t = f.get_tensor(key)
                if nk in ROW_KEYS:
                    t = claim_rows(nk, t, init, donor_rows)
                tensors[nk] = t
                n_keep += 1
        if not tensors:
            continue
        save_file(tensors, str(out / out_name), metadata={"format": "pt"})
        for nk, t in tensors.items():
            new_map[nk] = out_name
            total += t.numel() * t.element_size()
        print(f"[weights] {shard} -> {out_name}  ({len(tensors)} tensors)")

    jdump({"metadata": {"total_size": total}, "weight_map": new_map},
          out / "model.safetensors.index.json")
    print(f"[weights] kept {n_keep}, dropped {n_drop}, total {total/1e9:.2f} GB")
    assert n_keep == 363 and n_drop == 222, "key census mismatch vs 2026-08-26 index"


def claim_rows(name: str, t, init: str, donor_rows) -> "torch.Tensor":
    import torch
    live = t[1000:]  # rows past the control block are all trained BPE tokens
    live_norm = live.float().norm(dim=1)
    print(f"[liveness] {name}: live rows norm mean {live_norm.mean():.4f} "
          f"(p5 {live_norm.quantile(0.05):.4f})")
    for tid, (_, new, donor) in CLAIMS.items():
        print(f"[liveness]   row {tid} ({new}): norm {t[tid].float().norm():.4f}, "
              f"donor row {donor}: {t[donor].float().norm():.4f}")
    if init == "none":
        return t
    t = t.clone()
    for tid, (_, _, donor) in CLAIMS.items():
        t[tid] = t[donor] if init == "copy" else live.float().mean(0).to(t.dtype)
    return t


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--src", required=True, type=Path)
    ap.add_argument("--out", required=True, type=Path)
    ap.add_argument("--template", type=Path,
                    default=Path(__file__).parent.parent / "configs" / "thrasher.jinja")
    ap.add_argument("--init", choices=("copy", "mean", "none"), default="copy")
    ap.add_argument("--tokenizer-only", action="store_true")
    args = ap.parse_args()

    args.out.mkdir(parents=True, exist_ok=True)
    template_text = args.template.read_text()

    prep_tokenizer(args.src, args.out, template_text)
    if args.tokenizer_only:
        print("[done] tokenizer-only")
        return
    prep_configs(args.src, args.out)
    prep_weights(args.src, args.out, args.init)
    print("[done] prepped base at", args.out)


if __name__ == "__main__":
    main()