#!/usr/bin/env python3
"""Compare original flat ONNX outputs with the 4D BPU-wrapper ONNX."""

import argparse
from pathlib import Path

import numpy as np
import onnxruntime as ort


def load_samples(calibration_dir, limit, flat_dim):
    paths = sorted(Path(calibration_dir).glob("*.bin"))[:limit]
    if not paths:
        raise FileNotFoundError(f"No calibration .bin files found in {calibration_dir}")
    samples = []
    for path in paths:
        sample = np.fromfile(path, dtype=np.float32)
        if sample.size != flat_dim:
            raise ValueError(f"{path} has {sample.size} float32 values, expected {flat_dim}")
        samples.append(sample.reshape(1, flat_dim))
    return samples


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--flat-onnx", type=Path, required=True)
    parser.add_argument("--bpu4d-onnx", type=Path, required=True)
    parser.add_argument("--calibration-dir", type=Path, default=Path("calibration_data"))
    parser.add_argument("--flat-dim", type=int, default=450)
    parser.add_argument("--limit", type=int, default=64)
    args = parser.parse_args()

    flat = ort.InferenceSession(str(args.flat_onnx), providers=["CPUExecutionProvider"])
    wrapped = ort.InferenceSession(str(args.bpu4d_onnx), providers=["CPUExecutionProvider"])
    flat_input = flat.get_inputs()[0].name
    wrapped_input = wrapped.get_inputs()[0].name

    max_abs = 0.0
    max_mean_abs = 0.0
    for sample in load_samples(args.calibration_dir, args.limit, args.flat_dim):
        out_flat = flat.run(None, {flat_input: sample})[0]
        out_wrapped = wrapped.run(None, {wrapped_input: sample.reshape(1, 1, 1, args.flat_dim)})[0]
        diff = np.abs(out_flat - out_wrapped)
        max_abs = max(max_abs, float(diff.max()))
        max_mean_abs = max(max_mean_abs, float(diff.mean()))

    print(f"checked={min(args.limit, len(list(args.calibration_dir.glob('*.bin'))))}")
    print(f"max_abs_diff={max_abs:.9g}")
    print(f"max_mean_abs_diff={max_mean_abs:.9g}")


if __name__ == "__main__":
    main()
