#!/usr/bin/env python3
"""
infer_torch.py — ZeroGPU detection backend for the HF Space (v10 .pt).

Why a raw forward instead of ultralytics' predict():
  spatial_logic.postprocess_onnx does conf-filter -> class-agnostic NMS ->
  **top-3** extraction, and those top-3 slots are what sphinx_corrector's
  Viterbi consumes. model.predict() returns top-1 only, after its own NMS,
  which would gut the entire NLP correction layer. DetectionModel's raw
  forward returns the same [1, 4+nc, N] tensor the ONNX export (nms=False)
  produces — pixel cxcywh in rows 0-3, independent sigmoid class scores in
  rows 4.. — so postprocess_onnx is reused UNCHANGED.

Verified against artifacts/best_modelv10.pt: ultralytics 8.4.102, yolo11l,
imgsz 1024, nc=150, class order identical to class_map50_v9.json (0
mismatches, transmutations m4/f34/o29 still at idx 30/61/70). So the tensor
is [1, 154, 21504], exactly as v9.

ZeroGPU rules honoured here:
  * `import spaces` before torch.
  * The model is built on CPU at module scope and only moved to CUDA INSIDE
    the @spaces.GPU function — ZeroGPU patches CUDA init, so a module-scope
    .to('cuda') fails.
  * postprocess_onnx runs OUTSIDE the GPU function: quota is duration-based,
    so we hold the GPU only for the forward pass, and we avoid pickling
    Detection dataclasses back across the fork boundary.
"""

from __future__ import annotations

import os
from pathlib import Path
from typing import Callable

import numpy as np

import spaces                      # must precede torch (ZeroGPU)
import torch
from ultralytics import YOLO

import spatial_logic as SL

ROOT       = Path(__file__).parent
WEIGHTS    = Path(os.getenv('SPHINX_WEIGHTS',
                            ROOT / 'artifacts' / 'best_modelv10.pt'))
GPU_SECONDS = int(os.getenv('SPHINX_GPU_SECONDS', '60'))

# Built once, on CPU.
#   .eval()  — the detect head only emits the concatenated inference tensor in
#              eval mode; training mode returns the per-stride feature maps.
#   .float() — best_modelv10.pt stores a MIX of HalfStorage and FloatStorage
#              tensors (verified by inspecting the checkpoint), so feeding a
#              float32 input would raise "expected scalar type Half but found
#              Float". Upcasting everything to fp32 removes the dtype mismatch
#              and keeps this backend numerically comparable to the fp32 ONNX
#              path the pipeline was validated against.
_model = YOLO(str(WEIGHTS)).model.float().eval()
for _p in _model.parameters():
    _p.requires_grad_(False)


def _canvas_to_tensor(canvas: np.ndarray) -> torch.Tensor:
    """BGR uint8 HWC -> RGB float32 NCHW in [0,1] (matches make_onnx_infer_fn)."""
    x = canvas[:, :, ::-1].astype(np.float32) / 255.0     # BGR -> RGB
    x = np.ascontiguousarray(x.transpose(2, 0, 1))[None]  # HWC -> NCHW
    return torch.from_numpy(x)


@spaces.GPU(duration=GPU_SECONDS)
def _forward(canvas: np.ndarray) -> np.ndarray:
    """Raw forward on the letterboxed canvas. Returns [1, 4+nc, N] as numpy."""
    dev = 'cuda' if torch.cuda.is_available() else 'cpu'
    model = _model.to(dev)
    x = _canvas_to_tensor(canvas).to(dev)
    with torch.inference_mode():
        out = model(x)
    # DetectionModel in eval mode returns either the tensor or a
    # (tensor, feature_maps) tuple depending on version — normalise.
    if isinstance(out, (list, tuple)):
        out = out[0]
    return out.float().cpu().numpy()


def make_torch_infer_fn(
    class_names : list[str],
    conf_thresh : float = SL.CONF_THRESHOLD,
    iou_thresh  : float = SL.NMS_IOU,
    imgsz       : int   = 1024,
) -> Callable[[np.ndarray], list]:
    """
    Drop-in replacement for SL.make_onnx_infer_fn — same signature, same
    Callable[[bgr], list[Detection]] contract, bboxes in original-image pixels.

    Letterbox (never stretch): fragile breakpoint #9. Reuses SL.letterbox so
    preprocessing is provably identical to the ONNX path.
    """
    def infer(bgr: np.ndarray) -> list:
        canvas, scale, dx, dy = SL.letterbox(bgr, imgsz)
        raw  = _forward(canvas)
        dets = SL.postprocess_onnx(raw, class_names, conf_thresh, iou_thresh)
        # undo letterbox: original = (model - d) / scale
        for d in dets:
            x1, y1, x2, y2 = d.bbox
            d.bbox = ((x1 - dx) / scale, (y1 - dy) / scale,
                      (x2 - dx) / scale, (y2 - dy) / scale)
        return dets

    return infer