# /// script
# requires-python = ">=3.11,<3.13"
# dependencies = [
#   "coreai-core==1.0.0b2",
#   "coreai-torch==0.4.1",
#   "coreai-opt==0.2.1",
#   "scikit-learn>=1.7.2",
#   "torch",
#   "transformers>=5.9,<6",
#   "huggingface_hub",
#   "safetensors",
#   "numpy",
#   "pillow",
#   "ml-dtypes<=0.6.0",
#   "sentencepiece",
#   "protobuf",
#   "datasets",
#   "tokenizers",
# ]
# ///
"""North Star: export LiquidAI/LFM2.5-8B-A1B (MoE) to a CoreAI bundle on HF Jobs.

Uses OUR lfm2_moe export def (routed SwitchGLU/GatherMM = reads only top-4 experts,
short-conv as native conv1d). VETS bit-exact vs the HF reference on a truncated model
FIRST — aborts the (expensive) full export if the def is wrong — then exports the full
8B bounded to 4096 ctx and pushes the bundle to the Hub.
"""
import os, sys, shutil, subprocess, glob, json

ROOT = "/tmp/work"
os.makedirs(ROOT, exist_ok=True)
os.chdir(ROOT)
BASE = "LiquidAI/LFM2.5-8B-A1B"
OUT_REPO = "AbdallaHolmes/LFM2.5-8B-A1B-CoreAI-NorthStar"
CTX = 4096

def run(cmd, **kw):
    print(f"\n$ {' '.join(cmd) if isinstance(cmd,list) else cmd}", flush=True)
    return subprocess.run(cmd, check=True, **kw)

# ── 1. clone Apple's coreai-models + drop in OUR patch ────────────────────────
run(["git", "clone", "--depth", "1", "https://github.com/john-rocky/coreai-models"])
from huggingface_hub import snapshot_download
patch = snapshot_download("AbdallaHolmes/coreai-models-lfm-patch", repo_type="model")
pkg = f"{ROOT}/coreai-models/python/src/coreai_models"
shutil.copy(f"{patch}/lfm2.py",     f"{pkg}/models/macos/lfm2.py")
shutil.copy(f"{patch}/lfm2_moe.py", f"{pkg}/models/macos/lfm2_moe.py")
shutil.copy(f"{patch}/registry.py", f"{pkg}/models/registry.py")
shutil.copy(f"{patch}/macos.py",    f"{pkg}/export/macos.py")
shutil.copy(f"{patch}/pipeline.py", f"{pkg}/export/pipeline.py")
print("patch applied", flush=True)

# make the package importable + CLI available without pulling its old pins
sys.path.insert(0, f"{ROOT}/coreai-models/python/src")

# ── 2. VET the MoE def bit-exact vs HF reference (truncated) ──────────────────
import torch
from transformers import AutoConfig
from transformers.models.lfm2_moe.modeling_lfm2_moe import Lfm2MoeForCausalLM as HFRef
from safetensors import safe_open
from coreai_models.models.macos.lfm2_moe import Lfm2MoeForCausalLM
from coreai_models.models.macos.lfm2 import num_attention_layers, num_conv_layers

N = 6  # covers dense(0,1)+MoE + conv + attention layers
DT = torch.float32
cfg = AutoConfig.from_pretrained(BASE)
cfg.num_hidden_layers = N
cfg.layer_types = list(cfg.layer_types[:N])
repo = snapshot_download(BASE, allow_patterns=["*.safetensors", "config.json"])

tensors = {}
for f in glob.glob(os.path.join(repo, "*.safetensors")):
    with safe_open(f, framework="pt", device="cpu") as sf:
        for k in sf.keys():
            if k.startswith("model.layers.") and int(k.split(".")[2]) >= N:
                continue
            tensors[k] = sf.get_tensor(k).to(DT)

hf = HFRef(cfg).to(DT)
hf_miss, hf_unexp = hf.load_state_dict(dict(tensors), strict=False)
hf.tie_weights(); hf.eval()
# DIAGNOSTIC: if HF's own reference did not load the experts (per-expert checkpoint vs
# fused gate_up_proj in this transformers version), the reference is garbage and the vet
# would false-negative. Report it, and if so, convert the checkpoint to fused for HF.
_hf_real_miss = [m for m in hf_miss if "rotary" not in m and "lm_head" not in m]
print(f"[VET] HF ref load: {len(_hf_real_miss)} missing, {len(hf_unexp)} unexpected", flush=True)
if _hf_real_miss:
    print(f"[VET] HF missing sample: {_hf_real_miss[:6]}", flush=True)
    # HF uses fused experts.gate_up_proj / experts.down_proj — build them from per-expert w1/w3/w2
    fused = dict(tensors)
    import re as _re
    for i in range(N):
        pfx = f"model.layers.{i}.feed_forward.experts"
        if f"{pfx}.0.w1.weight" not in fused:
            continue
        ne = 0
        while f"{pfx}.{ne}.w1.weight" in fused:
            ne += 1
        gate_up = torch.stack([
            torch.cat([fused.pop(f"{pfx}.{e}.w1.weight"), fused.pop(f"{pfx}.{e}.w3.weight")], dim=0)
            for e in range(ne)
        ])  # (E, 2*inter, hidden)
        down = torch.stack([fused.pop(f"{pfx}.{e}.w2.weight") for e in range(ne)])  # (E, hidden, inter)
        fused[f"model.layers.{i}.feed_forward.experts.gate_up_proj"] = gate_up
        fused[f"model.layers.{i}.feed_forward.experts.down_proj"] = down
    hf = HFRef(cfg).to(DT)
    hf_miss2, _ = hf.load_state_dict(fused, strict=False)
    hf.tie_weights(); hf.eval()
    still = [m for m in hf_miss2 if "rotary" not in m and "lm_head" not in m]
    print(f"[VET] HF ref after fused-convert: {len(still)} missing (want 0)", flush=True)

ours = Lfm2MoeForCausalLM(cfg)
sd = dict(tensors); ours._mutate_state_dict(sd)
miss, unexp = ours.load_state_dict(sd, strict=False)
real_miss = [m for m in miss if "lm_head" not in m]
assert not real_miss, f"ours missing keys: {real_miss[:8]}"
assert not unexp, f"ours unexpected keys: {unexp[:8]}"
ours = ours.to(DT).eval()

torch.manual_seed(0)
prompt = torch.randint(1, cfg.vocab_size, (1, 7), dtype=torch.int32)
step = torch.randint(1, cfg.vocab_size, (1, 1), dtype=torch.int32)
with torch.no_grad():
    ho = hf(prompt.long(), use_cache=True)
    hs = hf(step.long(), past_key_values=ho.past_key_values, use_cache=True)
    n_attn = num_attention_layers(cfg); n_conv = num_conv_layers(cfg)
    hd = getattr(cfg, "head_dim", None) or cfg.hidden_size // cfg.num_attention_heads
    kc = torch.zeros(n_attn, 1, cfg.num_key_value_heads, 64, hd, dtype=DT)
    vc = torch.zeros_like(kc)
    cd = getattr(cfg, "conv_dim", cfg.hidden_size)
    cv = torch.zeros(n_conv, 1, cd, cfg.conv_L_cache - 1, dtype=DT)
    op = ours(prompt, torch.arange(7, dtype=torch.int32).unsqueeze(0), kc, vc, cv)
    os_ = ours(step, torch.arange(8, dtype=torch.int32).unsqueeze(0), kc, vc, cv)

d1 = (ho.logits - op).abs().max().item()
d2 = (hs.logits - os_).abs().max().item()
scale = ho.logits.abs().max().item()
tol = 3e-3 * max(scale, 1.0)
print(f"\n[VET] prefill Δ={d1:.6f} decode Δ={d2:.6f} scale={scale:.1f} tol={tol:.4f}", flush=True)
if not (d1 < tol and d2 < tol):
    print("[VET] ❌ MISMATCH — our MoE def diverges from HF. ABORTING export (credit saved).", flush=True)
    sys.exit(1)
print("[VET] ✅ MoE def is BIT-EXACT vs HF reference. Proceeding to full export.", flush=True)
del hf, ours, tensors, sd

# ── 3. full 8B export (routed MoE + conv1d, 4-bit, bounded 4096, mmap path) ────
env = dict(os.environ)
env["PYTHONPATH"] = f"{ROOT}/coreai-models/python/src:" + env.get("PYTHONPATH", "")
run([sys.executable, "-m", "coreai_models.llm.export", BASE,
     "--experimental", "--compute-precision", "float16",
     "--max-context-length", str(CTX),
     "--output-dir", f"{ROOT}/out"], env=env)

# inject chat_template into tokenizer_config (FIX4) so the bundle chats correctly
bdir = glob.glob(f"{ROOT}/out/*")[0]
tc_path = os.path.join(bdir, "tokenizer", "tokenizer_config.json")
jinja = os.path.join(bdir, "tokenizer", "chat_template.jinja")
if os.path.exists(tc_path) and os.path.exists(jinja):
    tc = json.load(open(tc_path))
    if "chat_template" not in tc:
        tc["chat_template"] = open(jinja).read()
        json.dump(tc, open(tc_path, "w"), indent=2, ensure_ascii=False)
        print("chat_template injected", flush=True)

# ── 4. push bundle to the Hub ─────────────────────────────────────────────────
from huggingface_hub import HfApi
api = HfApi()
api.create_repo(OUT_REPO, repo_type="model", exist_ok=True)
api.upload_folder(folder_path=f"{ROOT}/out", repo_id=OUT_REPO, repo_type="model")
print(f"\n✅ DONE — bundle pushed to https://huggingface.co/{OUT_REPO}", flush=True)
print("Download to the Mac and compile+bench with Tools/CoreAIBench.", flush=True)