"""As smoke_tiny.py, with the tiny Qwen3.5-MoE quantized to block FP8 on load (transformers
FP8Experts, as Qwen3.8-Flash-Next-FP8 loads): observer, prune, save, reload, saved names."""

import json
import tempfile
from pathlib import Path
from types import SimpleNamespace

import torch
from safetensors import safe_open
from transformers import AutoModelForCausalLM, FineGrainedFP8Config
from transformers.models.qwen3_5_moe.configuration_qwen3_5_moe import Qwen3_5MoeTextConfig

from reap.main import _setup_observer
from reap.prune import prune

torch.manual_seed(0)
cfg = Qwen3_5MoeTextConfig(
    vocab_size=256, hidden_size=256, num_hidden_layers=2, num_attention_heads=4, num_key_value_heads=2,
    head_dim=64, moe_intermediate_size=128, shared_expert_intermediate_size=128, num_experts=8,
    num_experts_per_tok=2, layer_types=["linear_attention", "full_attention"], linear_num_key_heads=2,
    linear_num_value_heads=2, linear_key_head_dim=64, linear_value_head_dim=64,
)
with tempfile.TemporaryDirectory() as d:
    AutoModelForCausalLM.from_config(cfg).save_pretrained(Path(d) / "bf16")
    model = AutoModelForCausalLM.from_pretrained(
        Path(d) / "bf16", quantization_config=FineGrainedFP8Config(weight_block_size=[128, 128]), device_map="cuda"
    ).eval()
    print(type(model.model.layers[0].mlp.experts).__name__, model.model.layers[0].mlp.experts.gate_up_proj.dtype,
          tuple(model.model.layers[0].mlp.experts.gate_up_proj_scale_inv.shape))
    observer = _setup_observer(model, SimpleNamespace(renormalize_router_weights=True, record_pruning_metrics_only=True))
    ids = torch.randint(0, 256, (1, 64)).cuda()
    with torch.no_grad(), observer.set_attention_mask(torch.ones_like(ids)):
        ref = model(input_ids=ids).logits
    state = {k: {m: (v.mean if hasattr(v, "mean") and not torch.is_tensor(v) else v) for m, v in s.items()} for k, s in observer.state.items()}
    print("reap[0]", [round(float(x), 4) for x in state[0]["reap"]])
    observer.close_hooks()
    out = prune(state, model, SimpleNamespace(prune_method="reap", perserve_super_experts=False, perserve_outliers=False), 4, Path(d) / "pruned")
    keys = {}
    for f in Path(out).glob("*.safetensors"):
        with safe_open(f, "pt") as fh:
            keys |= {k: fh.get_slice(k).get_dtype() for k in fh.keys()}
    cfg_saved = json.load(open(Path(out) / "config.json"))
    print("saved num_experts", cfg_saved["num_experts"], "| quantization_config", cfg_saved.get("quantization_config", {}).get("quant_method"))
    print("saved expert keys", sorted((k, v) for k, v in keys.items() if "layers.0.mlp.experts.0." in k))
    re = AutoModelForCausalLM.from_pretrained(out, device_map="cuda").eval()
    print("reloaded", type(re.model.layers[0].mlp.experts).__name__, tuple(re.model.layers[0].mlp.experts.gate_up_proj.shape))
    with torch.no_grad():
        print("reloaded forward ok", tuple(re(input_ids=ids).logits.shape))