#!/usr/bin/env python3
"""
convert_transnetv2.py — verify and optimise the TransNetV2 scene-detector model.

The scene detector (src/nodes/scene_detector_node.hpp → ISceneDetector) consumes
the elya5/transnetv2 ONNX export with a fixed input contract:

    input  "input" : float32 [1, 100, 27, 48, 3]   RGB, channels-last, 0-255
    output "534"   : float32 [1, 100, 1]           per-frame boundary logits (used)
    output "535"   : float32 [1, 100, 1]           many-hot head (ignored)

This script provides the ORT/TRT conversion infra for that model:

  --verify   (default) assert the .onnx matches the contract above and run one
             dummy inference through onnxruntime, reporting output ranges.
  --ort-opt  write a graph-optimised .ort next to the model (ORT loads this
             faster; the runtime also caches its own under ./ort_cache).
  --trt      build a TensorRT engine via trtexec with the pinned 1x100x27x48x3
             profile (delegates to scripts/build_trt_engines.sh SCENE_MODEL).

Usage:
    python scripts/convert_transnetv2.py                    # verify default model
    python scripts/convert_transnetv2.py --ort-opt
    python scripts/convert_transnetv2.py --trt
    python scripts/convert_transnetv2.py --model path/to/transnetv2.onnx --verify
"""
import argparse
import os
import subprocess
import sys

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DEFAULT_MODEL = os.path.join(ROOT, "models", "transnetv2.onnx")

# The fixed contract the C++ scene detector depends on.
EXPECTED_INPUT_SHAPE = [1, 100, 27, 48, 3]
EXPECTED_OUTPUT_SHAPE = [1, 100, 1]


def verify(model_path: str) -> int:
    import numpy as np
    import onnx
    import onnxruntime as ort

    print(f"[verify] loading {model_path}")
    m = onnx.load(model_path, load_external_data=False)
    g = m.graph

    def shape(t):
        return [d.dim_value if d.HasField("dim_value") else d.dim_param
                for d in t.type.tensor_type.shape.dim]

    in_shape = shape(g.input[0])
    print(f"[verify] input '{g.input[0].name}': {in_shape}")
    if in_shape != EXPECTED_INPUT_SHAPE:
        print(f"[verify] ERROR: input shape {in_shape} != {EXPECTED_INPUT_SHAPE}")
        return 1

    out_names = [o.name for o in g.output]
    out0_shape = shape(g.output[0])
    print(f"[verify] outputs: {out_names}  primary '{out_names[0]}': {out0_shape}")
    if out0_shape != EXPECTED_OUTPUT_SHAPE:
        print(f"[verify] ERROR: primary output {out0_shape} != {EXPECTED_OUTPUT_SHAPE}")
        return 1

    # One dummy inference: a mid-grey clip should produce low boundary scores.
    sess = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
    dummy = np.full(EXPECTED_INPUT_SHAPE, 128.0, dtype=np.float32)
    logits = sess.run([out_names[0]], {g.input[0].name: dummy})[0]
    probs = 1.0 / (1.0 + np.exp(-logits))
    print(f"[verify] dummy inference OK — boundary prob "
          f"min={probs.min():.4f} max={probs.max():.4f} mean={probs.mean():.4f}")
    print("[verify] contract matches the C++ ISceneDetector. ✓")
    return 0


def ort_opt(model_path: str) -> int:
    import onnxruntime as ort

    out_path = os.path.splitext(model_path)[0] + ".ort"
    print(f"[ort-opt] writing graph-optimised model → {out_path}")
    so = ort.SessionOptions()
    so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    so.optimized_model_filepath = out_path
    # Constructing the session triggers optimisation + serialisation.
    ort.InferenceSession(model_path, so, providers=["CPUExecutionProvider"])
    print(f"[ort-opt] done: {out_path}")
    return 0


def trt(model_path: str) -> int:
    script = os.path.join(ROOT, "scripts", "build_trt_engines.sh")
    print(f"[trt] delegating to {script} (SCENE_MODEL={model_path})")
    env = dict(os.environ, SCENE_MODEL=model_path)
    # build_trt_engines.sh also builds ArcFace/SCRFD; that's harmless (and a
    # useful sanity check), but if you only want the scene engine, run trtexec
    # directly with the profile printed below.
    print("[trt] equivalent standalone command:")
    print(f"  trtexec --onnx={model_path} --fp16 "
          f"--minShapes=input:1x100x27x48x3 "
          f"--optShapes=input:1x100x27x48x3 "
          f"--maxShapes=input:1x100x27x48x3 "
          f"--saveEngine=trt_cache/transnetv2.100x27x48.fp16.engine --useCudaGraph")
    return subprocess.call(["bash", script], env=env)


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--model", default=DEFAULT_MODEL,
                    help=f"path to transnetv2.onnx (default: {DEFAULT_MODEL})")
    ap.add_argument("--verify", action="store_true", help="verify I/O contract (default)")
    ap.add_argument("--ort-opt", action="store_true", help="write optimised .ort")
    ap.add_argument("--trt", action="store_true", help="build a TensorRT engine")
    args = ap.parse_args()

    if not os.path.isfile(args.model):
        print(f"ERROR: model not found: {args.model}\n"
              f"Fetch it with: bash scripts/download_models.sh", file=sys.stderr)
        return 2

    # Default action is verify when nothing else is requested.
    if not (args.ort_opt or args.trt):
        args.verify = True

    rc = 0
    if args.verify:
        rc |= verify(args.model)
    if rc == 0 and args.ort_opt:
        rc |= ort_opt(args.model)
    if rc == 0 and args.trt:
        rc |= trt(args.model)
    return rc


if __name__ == "__main__":
    sys.exit(main())