"""Extrae los 300 expertos seleccionados por capa del 397B-FP8 -> modelo 300exp-FP8.
Prunea experts + gate (router 512->300), mantiene shared_expert/atencion/vision, quita MTP.
"""
import json, os, re, time
import numpy as np
import torch
from safetensors import safe_open
from safetensors.torch import save_file

MODEL="/model"; OUT="/model300"
os.makedirs(OUT, exist_ok=True)
sel=np.load("/expert_selection.npy")  # [L,300]
L,N=sel.shape
print(f"seleccion [{L},{N}]", flush=True)
idx=json.load(open(f"{MODEL}/model.safetensors.index.json"))
wmap=idx["weight_map"]
files=sorted(set(wmap.values()))
handles={f:safe_open(os.path.join(MODEL,f),framework="pt",device="cpu") for f in files}
def get(k): return handles[wmap[k]].get_tensor(k)

# mapa: experto original -> nuevo indice, por capa
remap={l:{int(o):n for n,o in enumerate(sel[l])} for l in range(L)}

new={}; t0=time.time()
exp_re=re.compile(r"^model\.language_model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(.+)$")
gate_re=re.compile(r"^model\.language_model\.layers\.(\d+)\.mlp\.gate\.weight$")
for ki,k in enumerate(wmap.keys()):
    if k.startswith("mtp."):
        continue  # quitar MTP
    me=exp_re.match(k)
    if me:
        l,e,rest=int(me.group(1)),int(me.group(2)),me.group(3)
        if e in remap[l]:
            ne=remap[l][e]
            new[f"model.language_model.layers.{l}.mlp.experts.{ne}.{rest}"]=get(k)
        continue
    mg=gate_re.match(k)
    if mg:
        l=int(mg.group(1))
        w=get(k)  # [512, hidden]
        new[k]=w[sel[l]].contiguous()  # [300, hidden]
        continue
    # resto (atencion, shared_expert, norms, embeds, vision, lm_head): copiar
    new[k]=get(k)
    if ki%5000==0: print(f"  {ki} tensores, {time.time()-t0:.0f}s", flush=True)

print(f"nuevo state: {len(new)} tensores. guardando shards...", flush=True)
# guardar en shards ~5GB
items=list(new.items()); shard=[]; sz=0; si=0; sidx={}
SHARD_MAX=5e9
def flush_shard(shard,si):
    fn=f"model-{si:05d}.safetensors"
    save_file(dict(shard), os.path.join(OUT,fn), metadata={"format":"pt"})
    for kk,_ in shard: sidx[kk]=fn
    return fn
for k,v in items:
    shard.append((k,v)); sz+=v.numel()*v.element_size()
    if sz>=SHARD_MAX:
        flush_shard(shard,si); print(f"  shard {si} ({sz/1e9:.1f}GB)",flush=True); shard=[]; sz=0; si+=1
if shard: flush_shard(shard,si)
# index + config
json.dump({"metadata":{},"weight_map":sidx}, open(os.path.join(OUT,"model.safetensors.index.json"),"w"))
cfg=json.load(open(f"{MODEL}/config.json"))
cfg["text_config"]["num_experts"]=N
cfg["text_config"]["mtp_num_hidden_layers"]=0
json.dump(cfg, open(os.path.join(OUT,"config.json"),"w"), indent=2)
# copiar tokenizer/processor configs
import shutil
for f in os.listdir(MODEL):
    if f.endswith(".json") and "index" not in f and f!="config.json" or "token" in f.lower() or "processor" in f.lower() or "merges" in f or "vocab" in f:
        try: shutil.copy(os.path.join(MODEL,f), os.path.join(OUT,f))
        except: pass
print(f"DONE: 300exp guardado en {OUT}, {si+1} shards, {time.time()-t0:.0f}s", flush=True)