"""
Runs a single image through an ONNX model and saves the raw output
tensor of every output layer as a binary file.

Each file is named after the full layer name with a .bin extension.
"""

import argparse
import os
import numpy as np
import onnxruntime as ort

STR_LEN = 80

def main():
    parser = argparse.ArgumentParser(
        description="Dump all output layer tensors of an ONNX model to raw binary files"
    )
    parser.add_argument("model", type=str, help="Path to the ONNX model file")
    parser.add_argument(
        "-i", "--image_size", type=int, default=416,
        help="Model input size (default: 416)"
    )
    parser.add_argument(
        "-o", "--output_dir", type=str, default="layer_dumps",
        help="Directory to save the binary files (default: layer_dumps)"
    )
    args = parser.parse_args()

    # ---- Setup ----
    print("-" * STR_LEN)
    print("DUMP OUTPUT LAYERS")
    print("-" * STR_LEN)

    sess = ort.InferenceSession(args.model, providers=["CPUExecutionProvider"])

    input_meta = sess.get_inputs()[0]
    input_name = input_meta.name
    print(f"Model input : {input_name}  shape={input_meta.shape}")

    output_metas = sess.get_outputs()
    print(f"Model outputs: {len(output_metas)} layer(s)")
    for meta in output_metas:
        print(f"  - {meta.name}  shape={meta.shape}")
    print("-" * STR_LEN)

    # ---- Preprocess image ----
    img_arr = np.zeros(input_meta.shape, dtype=np.float32)

    # ---- Inference ----
    output_names = [m.name for m in output_metas]
    results = sess.run(output_names, {input_name: img_arr})

    # ---- Save tensors ----
    os.makedirs(args.output_dir, exist_ok=True)

    for name, tensor in zip(output_names, results):
        filename = f"{name}.bin".replace("/", "_")
        filepath = os.path.join(args.output_dir, filename)
        tensor = np.array(tensor)
        tensor.tofile(filepath)
        print(f"Saved {filepath}  dtype={tensor.dtype}  shape={tensor.shape}  "
              f"size={os.path.getsize(filepath)} bytes")

    print("-" * STR_LEN)
    print(f"Done. {len(results)} layer(s) dumped to '{args.output_dir}'")
    print("-" * STR_LEN)


if __name__ == "__main__":
    main()