#!/usr/bin/env python3
"""Run YOLO26-depth ONNX models and save depth heatmap visualizations.

Example:
    python onnx_model/infer_depth_onnx.py \
        --model onnx_model/yolo26n-depth.onnx \
        --source wm9955/depth8/images/kv1_NYUdata_NYU0532.jpg
"""

from __future__ import annotations

import argparse
import time
from pathlib import Path

import cv2
import numpy as np
import onnxruntime as ort


def parse_args() -> argparse.Namespace:
    """Parse command-line arguments."""
    root = Path(__file__).resolve().parent
    parser = argparse.ArgumentParser(description="YOLO26-depth ONNX Runtime inference")
    parser.add_argument("--model", type=Path, default=root / "yolo26n-depth.onnx", help="Path to an ONNX model")
    parser.add_argument("--wavematrix", action="store_true", help="Use wavematrix engine for inference")
    parser.add_argument("--source", type=Path, required=True, help="Path to one input image")
    parser.add_argument("--output-dir", type=Path, default=root / "runs", help="Directory for output files")
    parser.add_argument("--alpha", type=float, default=0.6, help="Heatmap opacity in the overlay, in [0, 1]")
    parser.add_argument("--cmap", choices=("jet", "inferno"), default="jet", help="OpenCV heatmap color map")
    parser.add_argument(
        "--mode",
        choices=("disparity", "metric"),
        default="disparity",
        help="Color normalization: inverse depth or metric depth",
    )
    parser.add_argument(
        "--postprocess",
        choices=("default", "predictor"),
        default="default",
        help="Depth resize path: direct ONNX output or Ultralytics predictor-compatible geometry",
    )
    args = parser.parse_args()
    if not 0.0 <= args.alpha <= 1.0:
        parser.error("--alpha must be in [0, 1]")
    return args

def model_input_size(session, use_wavematrix: bool) -> tuple[str, tuple[int, int]]:
    """Return the image input name and static (height, width) required by the ONNX model."""
    inputs = session.get_inputs()
    if len(inputs) != 1:
        raise ValueError(f"Expected one model input, got {len(inputs)}")

    image_input = inputs[0]
    shape = image_input.shape
    if not isinstance(shape[2], int) or not isinstance(shape[3], int):
        raise ValueError(f"Dynamic ONNX input shapes are not supported by this script: {shape}")
    if use_wavematrix:
        return image_input.name, (shape[1], shape[2])
    else:
        return image_input.name, (shape[2], shape[3])


def letterbox(image: np.ndarray, new_shape: tuple[int, int]) -> tuple[np.ndarray, tuple[int, int, int, int]]:
    """Resize and center-pad an image exactly as Ultralytics LetterBox does for a fixed ONNX input."""
    height, width = image.shape[:2]
    target_height, target_width = new_shape
    gain = min(target_height / height, target_width / width)
    resized_width, resized_height = round(width * gain), round(height * gain)
    pad_width, pad_height = target_width - resized_width, target_height - resized_height
    left, right = round(pad_width / 2 - 0.1), round(pad_width / 2 + 0.1)
    top, bottom = round(pad_height / 2 - 0.1), round(pad_height / 2 + 0.1)

    if (width, height) != (resized_width, resized_height):
        image = cv2.resize(image, (resized_width, resized_height), interpolation=cv2.INTER_LINEAR)
    image = cv2.copyMakeBorder(image, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114))
    return image, (top, bottom, left, right)


def preprocess(image: np.ndarray, input_size: tuple[int, int], use_wavematrix: bool = False) -> tuple[np.ndarray, tuple[int, int, int, int]]:
    """Letterbox BGR input, convert it to RGB NCHW float32, and normalize to [0, 1]."""
    image, padding = letterbox(image, input_size)
    if use_wavematrix:
        tensor = image[..., ::-1][None]
        return np.ascontiguousarray(tensor, dtype=np.uint8), padding
    else:
        tensor = image[..., ::-1].transpose(2, 0, 1)[None]
        return np.ascontiguousarray(tensor, dtype=np.float32) / 255.0, padding


def output_depth_map(output: np.ndarray) -> np.ndarray:
    """Extract one 2D depth map from a standard ONNX depth output tensor."""
    output = np.asarray(output)
    if output.ndim == 4 and output.shape[:2] == (1, 1):
        return output[0, 0]
    elif output.ndim == 3 and output.shape[0] == 1:
        return output[0]
    elif output.ndim == 2:
        return output
    raise ValueError(f"Expected one depth map with shape (1, 1, H, W), got {output.shape}")


def restore_depth(
    output: np.ndarray,
    original_shape: tuple[int, int],
    padding: tuple[int, int, int, int],
) -> np.ndarray:
    """Remove LetterBox padding and resize the exported ONNX output to the original image size."""
    depth = output_depth_map(output)

    top, bottom, left, right = padding
    end_y = depth.shape[0] - bottom if bottom else depth.shape[0]
    end_x = depth.shape[1] - right if right else depth.shape[1]
    depth = depth[top:end_y, left:end_x]
    if depth.size == 0:
        raise ValueError("LetterBox crop is empty; model output and input geometry are inconsistent")
    return cv2.resize(depth.astype(np.float32), (original_shape[1], original_shape[0]), interpolation=cv2.INTER_LINEAR)


def restore_depth_predictor_aligned(output: np.ndarray, original_shape: tuple[int, int]) -> np.ndarray:
    """Match ``DepthPredictor.postprocess`` geometry from an exported 4x-upsampled ONNX output.

    The PyTorch model passes its H/4 head map directly to ``ops.scale_masks``. ONNX export upsamples that map by four
    first, so this function reconstructs the H/4 map and applies the same crop-rounding and bilinear resize sequence.
    It reduces the otherwise visible interpolation-only difference when comparing ONNX output with ``YOLO.predict()``.
    """
    depth = output_depth_map(output)
    height, width = depth.shape
    if height % 4 or width % 4:
        raise ValueError(f"Expected an exported depth output divisible by 4, got {depth.shape}")
    head = cv2.resize(depth, (width // 4, height // 4), interpolation=cv2.INTER_LINEAR)
    original_height, original_width = original_shape
    gain = min(head.shape[0] / original_height, head.shape[1] / original_width)
    pad_width = (head.shape[1] - round(original_width * gain)) / 2
    pad_height = (head.shape[0] - round(original_height * gain)) / 2
    top, left = round(pad_height - 0.1), round(pad_width - 0.1)
    bottom = head.shape[0] - round(pad_height + 0.1)
    right = head.shape[1] - round(pad_width + 0.1)
    head = head[top:bottom, left:right]
    if head.size == 0:
        raise ValueError("Predictor-compatible crop is empty; model output and input geometry are inconsistent")
    return cv2.resize(head.astype(np.float32), (original_width, original_height), interpolation=cv2.INTER_LINEAR)


def colorize_depth(depth: np.ndarray, cmap: str, mode: str) -> np.ndarray:
    """Create the Ultralytics-style BGR depth heatmap; invalid depths are black."""
    depth = np.asarray(depth, dtype=np.float32)
    valid = np.isfinite(depth) & (depth > 0)
    values = np.where(valid, 1.0 / np.where(valid, depth, 1.0), 0.0) if mode == "disparity" else depth
    pool = values[valid]
    if pool.size:
        low, high = np.percentile(pool, (2, 98)) if mode == "disparity" else (pool.min(), pool.max())
    else:
        low, high = 0.0, 1.0
    high = max(float(high), float(low) + 1e-6)
    normalized = np.clip((values - low) / (high - low), 0.0, 1.0)
    color_map = {"jet": cv2.COLORMAP_JET, "inferno": cv2.COLORMAP_INFERNO}[cmap]
    heatmap = cv2.applyColorMap((normalized * 255).astype(np.uint8), color_map)
    heatmap[~valid] = 0
    return heatmap


def save_results(
    source: Path,
    output_dir: Path,
    image: np.ndarray,
    depth: np.ndarray,
    heatmap: np.ndarray,
    alpha: float,
) -> tuple[Path, Path, Path]:
    """Save metric depth values, a standalone heatmap, and a heatmap overlay."""
    output_dir.mkdir(parents=True, exist_ok=True)
    stem = source.stem
    depth_path = output_dir / f"{stem}_depth.npy"
    heatmap_path = output_dir / f"{stem}_depth_heatmap.png"
    overlay_path = output_dir / f"{stem}_depth_overlay.jpg"
    image_heatmap = np.hstack((image, heatmap))
    np.save(depth_path, depth)
    if not cv2.imwrite(str(heatmap_path), image_heatmap):
        raise IOError(f"Failed to write {heatmap_path}")
    overlay = cv2.addWeighted(image, 1.0 - alpha, heatmap, alpha, 0)
    if not cv2.imwrite(str(overlay_path), overlay):
        raise IOError(f"Failed to write {overlay_path}")
    return depth_path, heatmap_path, overlay_path


def main() -> None:
    """Run one image through a YOLO26-depth ONNX model."""
    args = parse_args()
    if not args.model.is_file():
        raise FileNotFoundError(f"ONNX model not found: {args.model}")
    image = cv2.imread(str(args.source))
    if image is None:
        raise FileNotFoundError(f"Unable to read image: {args.source}")

    if args.wavematrix:
        import axengine as axe
        session = axe.InferenceSession(str(args.model))
    else:
        session = ort.InferenceSession(str(args.model))
    input_name, input_size = model_input_size(session, args.wavematrix)
    input_tensor, padding = preprocess(image, input_size, args.wavematrix)

    start = time.perf_counter()
    output = session.run(None, {input_name: input_tensor})[0]
    inference_ms = (time.perf_counter() - start) * 1000
    if args.postprocess == "predictor":
        depth = restore_depth_predictor_aligned(output, image.shape[:2])
    else:
        depth = restore_depth(output, image.shape[:2], padding)
    heatmap = colorize_depth(depth, args.cmap, args.mode)
    depth_path, heatmap_path, overlay_path = save_results(
        args.source, args.output_dir, image, depth, heatmap, args.alpha
    )

    valid = depth[np.isfinite(depth) & (depth > 0)]
    print(f"Providers: {session.get_providers()}")
    print(f"Input: {args.source} ({image.shape[1]}x{image.shape[0]}) -> {input_size[1]}x{input_size[0]}")
    print(f"Postprocess: {args.postprocess}")
    print(f"Inference: {inference_ms:.1f} ms")
    print(f"Depth (m): min={valid.min():.3f}, median={np.median(valid):.3f}, max={valid.max():.3f}")
    print(f"Saved depth: {depth_path}")
    print(f"Saved heatmap: {heatmap_path}")
    print(f"Saved overlay: {overlay_path}")


if __name__ == "__main__":
    main()