#!/usr/bin/env python3
"""Simplify exported graphs (static shapes) and write one Pulsar2 5.1 build config per
graph, reproducing the minus-v0.1 recipe (docs/LFM2_VL_AX650_OPTIMIZATION.md, D.5):
  - NPU3 (all 3 cores), U16 everywhere,
  - FP32 on every RMSNorm chain op (Mul/ReduceMean/Add/Sqrt/Div/Mul): 2 chains per conv
    layer, 4 per attention layer (operator, q_layernorm, k_layernorm, ffn) + post's norm,
  - FP32 runtime I/O (indices S32), Numpy calibration from gen_calib.py,
  - vision: onnxsim off, smooth-quant + FP32 conv bias (AXERA vision-encoder practice).
  python make_configs.py <build_dir>
"""
import json, os, sys
import onnx
from onnxsim import simplify

B = sys.argv[1]
meta = json.load(open(f"{B}/export_meta.json"))
graphs = ["vision_encoder_512"] + [f"l{i}_{t}_fused" for i, t in enumerate(meta["layer_types"])] + ["post_d"]
os.makedirs(f"{B}/sim", exist_ok=True)
os.makedirs(f"{B}/configs", exist_ok=True)
summary = {}
for g in graphs:
    m = onnx.load(f"{B}/onnx/{g}.onnx")
    if g != "vision_encoder_512":
        m, ok = simplify(m)
        assert ok, g
    onnx.save(m, f"{B}/sim/{g}.onnx", save_as_external_data=m.ByteSize() > 1.9e9)
    norm_ops = [n.name for n in m.graph.node
                if "norm" in n.name.lower() and n.op_type in ("Mul", "ReduceMean", "Add", "Sqrt", "Div", "Pow")]
    inputs = [i.name for i in m.graph.input]
    is_vision = g == "vision_encoder_512"
    cfg = {
        "model_type": "ONNX",
        "npu_mode": "NPU3",
        "target_hardware": "AX650",
        "onnx_opt": {"enable_onnxsim": False},
        "quant": {
            "input_configs": [{"tensor_name": n, "calibration_dataset": f"/data/calib/{g}/{n}.tar",
                               "calibration_format": "Numpy", "calibration_size": -1} for n in inputs],
            "layer_configs": [{"start_tensor_names": ["DEFAULT"], "end_tensor_names": ["DEFAULT"],
                               "data_type": "U16"}],
            "calibration_method": "MinMax",
            "precision_analysis": True,
            "precision_analysis_method": "EndToEnd",
            "conv_bias_data_type": "FP32",
            "enable_smooth_quant": is_vision,
        },
        "input_processors": [{"tensor_name": n, "src_dtype": "S32" if n == "indices" else "FP32"}
                             for n in inputs],
        "compiler": {"check": 0},
    }
    if norm_ops and not is_vision:
        cfg["quant"]["layer_configs"].append({"layer_names": norm_ops, "data_type": "FP32"})
    json.dump(cfg, open(f"{B}/configs/{g}.json", "w"), indent=1)
    summary[g] = {"nodes": len(m.graph.node), "fp32_norm_ops": len(norm_ops),
                  "op_types": sorted({n.op_type for n in m.graph.node})}
    print(g, summary[g]["nodes"], "nodes,", len(norm_ops), "FP32 norm ops", flush=True)
json.dump(summary, open(f"{B}/configs/summary.json", "w"), indent=1)