#!/usr/bin/env python3
"""Export real expert matrices for a bounded group32 NPU compatibility probe."""
import argparse
import json
from pathlib import Path
import numpy as np
import torch
from safetensors import safe_open
from official_int4_reference import unpack4

p = argparse.ArgumentParser(description=__doc__)
p.add_argument("--official", type=Path, required=True)
p.add_argument("--source", type=Path, required=True)
p.add_argument("--output", type=Path, required=True)
a = p.parse_args()
a.output.mkdir(parents=True, exist_ok=True)
torch.set_num_threads(4)
indices = {str(root): json.loads((root / "model.safetensors.index.json").read_text())["weight_map"]
           for root in (a.official, a.source)}
def get(root, name):
    with safe_open(root / indices[str(root)][name], framework="pt") as f:
        return f.get_tensor(name)
records = []
for proj in ("gate", "up", "down"):
    name = f"model.layers.1.mlp.experts.0.{proj}_proj.weight"
    original = get(a.source, name).float()
    n, k = original.shape
    q = unpack4(get(a.official, name+"_packed"), (n,k))
    s = get(a.official, name+"_scale").float()
    for kind in ("official_group32", "old_channel"):
        if kind == "old_channel":
            s = torch.maximum(original.amax(1).clamp_min(0)/7, (-original.amin(1)).clamp_min(0)/8).unsqueeze(1)
            q = torch.round(original/s).clamp(-8,7).to(torch.int8)
        group = 32 if kind == "official_group32" else k
        case = a.output / (proj+"_"+kind)
        case.mkdir(exist_ok=True)
        np.array([k,n,group], dtype=np.int32).tofile(case / "shape.bin")
        q.numpy().tofile(case / "codes.bin")
        s.numpy().tofile(case / "scales.bin")
        weight = (q.reshape(n,-1,group).float()*s[:,:,None]).reshape(n,k)
        for rows in (1,128):
            rng = torch.Generator().manual_seed(314159+rows)
            x = torch.randn(rows,k,generator=rng)
            # Explicit independent integer oracle for the added groupwise A8.
            xg = x.reshape(rows,-1,group)
            xs = xg.abs().amax(-1)/127
            xq = torch.round(xg/xs[:,:,None]).clamp(-127,127)
            oracle = torch.zeros(rows,n)
            for g in range(k//group):
                oracle += (xq[:,g,:] @ q[:,g*group:(g+1)*group].float().T)*xs[:,g,None]*s[:,g]
            x.numpy().tofile(case / f"input-{rows}.bin")
            oracle.numpy().tofile(case / f"oracle-{rows}.bin")
            official_out = x @ weight.T
            original_out = x @ original.T
            records.append({"case":case.name,"rows":rows,"group":group,"k":k,"n":n,
                "weight_only_output_relative_rmse":float(torch.linalg.vector_norm(official_out-original_out)/torch.linalg.vector_norm(original_out)),
                "extra_a8_output_relative_rmse":float(torch.linalg.vector_norm(oracle-official_out)/torch.linalg.vector_norm(official_out))})
(a.output / "manifest.json").write_text(json.dumps({"description":"one actual layer1 expert, synthetic Gaussian activations; not full model or task accuracy", "cases": records},indent=2)+"\n")
print(json.dumps(records,indent=2))