#!/usr/bin/env python3
"""Export the authoritative BatDetect2 detector/classifier checkpoint to ONNX."""

from __future__ import annotations

import argparse
import hashlib
import json
import subprocess
import sys
from pathlib import Path

EXPECTED_UK_LABELS = [
    "MYOMYS",
    "MYOALC",
    "CNESER",
    "PIPNAT",
    "BARBAR",
    "MYONAT",
    "MYODAU",
    "MYOBRA",
    "PIPPIP",
    "MYOBEC",
    "PIPPYG",
    "RHIHIP",
    "NYCLEI",
    "RHIFER",
    "PLEAUR",
    "NYCNOC",
    "PLEAUS",
]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--source", type=Path, required=True, help="BatDetect2 repository checkout")
    parser.add_argument(
        "--checkpoint",
        type=Path,
        help="defaults to the checkout's bundled batdetect2_uk_same.ckpt",
    )
    parser.add_argument("--output", type=Path, default=Path("models/batdetect2-uk-same.onnx"))
    parser.add_argument("--opset", type=int, default=17)
    parser.add_argument(
        "--accept-noncommercial-license",
        action="store_true",
        help="acknowledge that the checkpoint is CC BY-NC 4.0",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if not args.accept_noncommercial_license:
        raise SystemExit(
            "BatDetect2's bundled model is CC BY-NC 4.0; rerun with "
            "--accept-noncommercial-license after reviewing those terms"
        )
    source = args.source.resolve()
    checkpoint = args.checkpoint or (
        source / "src/batdetect2/models/checkpoints/batdetect2_uk_same.ckpt"
    )
    if not checkpoint.is_file() or not (source / "src/batdetect2").is_dir():
        raise SystemExit("--source or --checkpoint does not contain the expected BatDetect2 files")
    sys.path.insert(0, str(source / "src"))
    try:
        import numpy as np
        import onnx
        import onnxruntime as ort
        import torch
        from batdetect2.train import load_model_from_checkpoint
    except ImportError as exc:
        raise SystemExit(
            "Install BatDetect2 and ONNX export dependencies in an isolated environment; "
            "see docs/MODELS.md"
        ) from exc

    model, _ = load_model_from_checkpoint(checkpoint)
    source_labels = list(model.class_names)
    labels = [label.upper() for label in source_labels]
    if labels != EXPECTED_UK_LABELS:
        raise SystemExit(
            "checkpoint class order differs from the locked Android contract:\n"
            f"expected={EXPECTED_UK_LABELS}\nactual={source_labels}"
        )

    class DetectorOnly(torch.nn.Module):
        def __init__(self, detector: torch.nn.Module) -> None:
            super().__init__()
            self.detector = detector

        def forward(self, input: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
            output = self.detector(input)
            return output.detection_probs, output.class_probs

    wrapper = DetectorOnly(model.detector).eval()
    dummy = torch.zeros((1, 1, 128, 256), dtype=torch.float32)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    with torch.inference_mode():
        torch.onnx.export(
            wrapper,
            (dummy,),
            args.output,
            input_names=["input"],
            output_names=["detection_probs", "class_probs"],
            opset_version=args.opset,
            do_constant_folding=True,
            dynamo=False,
        )
    onnx.checker.check_model(onnx.load(args.output))
    with torch.inference_mode():
        detection, classes = wrapper(dummy)
        torch.manual_seed(20260818)
        parity_input = torch.rand((1, 1, 128, 256), dtype=torch.float32)
        reference_detection, reference_classes = wrapper(parity_input)
    runtime = ort.InferenceSession(str(args.output), providers=["CPUExecutionProvider"])
    converted_detection, converted_classes = runtime.run(
        ["detection_probs", "class_probs"],
        {runtime.get_inputs()[0].name: parity_input.numpy()},
    )
    detection_error = float(np.max(np.abs(reference_detection.numpy() - converted_detection)))
    class_error = float(np.max(np.abs(reference_classes.numpy() - converted_classes)))
    maximum_absolute_error = max(detection_error, class_error)
    if maximum_absolute_error > 1e-4:
        raise SystemExit(
            f"PyTorch/ONNX parity failed: max absolute error {maximum_absolute_error:.8g}"
        )
    try:
        source_revision = subprocess.run(
            ["git", "-C", str(source), "rev-parse", "HEAD"],
            check=True,
            capture_output=True,
            text=True,
        ).stdout.strip()
    except (OSError, subprocess.CalledProcessError):
        source_revision = "unknown"
    try:
        checkpoint_identity = str(checkpoint.resolve().relative_to(source))
    except ValueError:
        checkpoint_identity = checkpoint.name
    metadata = {
        "backend": "batdetect2",
        "source": "https://github.com/macaodha/batdetect2",
        "source_revision": source_revision,
        "checkpoint": checkpoint_identity,
        "input_name": "input",
        "input_shape": [1, 1, 128, 256],
        "outputs": {
            "detection_probs": list(detection.shape),
            "class_probs": list(classes.shape),
        },
        "labels": labels,
        "source_labels": source_labels,
        "opset": args.opset,
        "license": "CC BY-NC 4.0",
        "commercial_use": False,
        "onnx_bytes": args.output.stat().st_size,
        "onnx_sha256": hashlib.sha256(args.output.read_bytes()).hexdigest(),
        "source_runtime_parity": {
            "fixture": "seeded random float32 tensor (seed 20260818)",
            "detection_maximum_absolute_error": detection_error,
            "class_maximum_absolute_error": class_error,
            "tolerance": 1e-4,
        },
    }
    args.output.with_suffix(".json").write_text(json.dumps(metadata, indent=2) + "\n")
    print(f"wrote {args.output}")


if __name__ == "__main__":
    main()