parking_solution/apps/vision/tests/test_vehicle.py

"""Vehicle stage (phase A) — pure post-processing on synthetic tensors, and the
recognizer composition over the stub with a fake detector. No weights needed."""

from __future__ import annotations

import os

import numpy as np
import pytest
from fastapi.testclient import TestClient

from vision_service.schemas import BBox, VehicleResult
from vision_service.vehicle import (
    COCO_VEHICLE_CLASSES,
    Detection,
    decode,
    letterbox,
    nms,
    pick_vehicle,
    vehicles_from_output,
)

SIZE = 64  # tiny "model" input: grids 8x8 + 4x4 + 2x2 = 84 rows
ROWS = (SIZE // 8) ** 2 + (SIZE // 16) ** 2 + (SIZE // 32) ** 2


def raw_output(hits: list[tuple[int, int, int, float, float, float, float]]) -> np.ndarray:
    """Build a YOLOX-style raw tensor [ROWS, 85] with the given (row, coco_class, _, obj,
    cls_score, log_w, log_h) hits; everything else is background."""
    raw = np.zeros((ROWS, 85), dtype=np.float32)
    raw[:, 2:4] = -10.0  # exp → ~0 size for background rows
    for row, cls, _, obj, score, lw, lh in hits:
        raw[row, 0:2] = 0.5  # centre of its grid cell
        raw[row, 2] = lw
        raw[row, 3] = lh
        raw[row, 4] = obj
        raw[row, 5 + cls] = score
    return raw


def test_decode_maps_grid_offsets_and_log_sizes_to_pixels() -> None:
    raw = raw_output([(0, 2, 0, 1.0, 1.0, np.log(2.0), np.log(3.0))])
    dec = decode(raw, SIZE)
    # Row 0 = stride-8 grid cell (0,0): centre (0.5+0)*8 = 4, size exp(log 2)*8 = 16 / 24.
    assert dec[0, :4].tolist() == [4.0, 4.0, 16.0, 24.0]
    # Last row = stride-32 cell (1,1): centre (0.5+1)*32 = 48.
    raw2 = raw_output([(ROWS - 1, 7, 0, 1.0, 1.0, 0.0, 0.0)])
    dec2 = decode(raw2, SIZE)
    assert dec2[ROWS - 1, :4].tolist() == [48.0, 48.0, 32.0, 32.0]


def test_vehicles_only_above_floor_mapped_to_vocabulary_and_scaled_back() -> None:
    raw = raw_output(
        [
            (0, 2, 0, 0.9, 0.9, np.log(2.0), np.log(2.0)),  # car, score .81
            (1, 0, 0, 0.99, 0.99, np.log(2.0), np.log(2.0)),  # person → ignored
            (2, 7, 0, 0.5, 0.5, np.log(2.0), np.log(2.0)),  # truck, score .25 → below floor
        ]
    )
    found = vehicles_from_output(raw, SIZE, scale=0.5, min_confidence=0.4)
    assert [d.body_type for d in found] == ["car"]
    assert round(found[0].confidence, 2) == 0.81
    # Box 16px wide in the letterboxed input → 32px in the original (scale 0.5).
    assert round(found[0].x2 - found[0].x1) == 32
    assert set(COCO_VEHICLE_CLASSES.values()) == {"car", "motorcycle", "bus", "truck"}


def test_nms_keeps_the_best_of_overlapping_boxes() -> None:
    boxes = np.array([[0, 0, 10, 10], [1, 1, 11, 11], [50, 50, 60, 60]], dtype=np.float32)
    scores = np.array([0.5, 0.9, 0.7], dtype=np.float32)
    assert sorted(nms(boxes, scores, 0.45)) == [1, 2]


def test_pick_prefers_the_box_holding_the_plate_else_the_largest() -> None:
    near = Detection("car", 0.9, 0, 0, 100, 100)
    far = Detection("truck", 0.8, 200, 200, 400, 400)  # larger
    inside = Detection("car", 0.7, 10, 10, 60, 60)  # tighter box also holding the plate
    assert pick_vehicle([near, far], None) is far
    assert pick_vehicle([near, far], BBox(x1=20, y1=20, x2=30, y2=30)) is near
    assert pick_vehicle([near, far, inside], BBox(x1=20, y1=20, x2=30, y2=30)) is inside
    assert pick_vehicle([near, far], BBox(x1=900, y1=900, x2=910, y2=910)) is far  # plate outside every box
    assert pick_vehicle([], None) is None


def test_letterbox_keeps_aspect_and_pads_with_114() -> None:
    pytest.importorskip("cv2")  # letterbox resizes with OpenCV — only present with the alpr extra
    frame = np.zeros((30, 60, 3), dtype=np.uint8)
    tensor, scale = letterbox(frame, 64)
    assert tensor.shape == (1, 3, 64, 64) and tensor.dtype == np.float32
    assert abs(scale - 64 / 60) < 1e-9
    assert tensor[0, 0, 63, 63] == 114.0  # padding
    assert tensor[0, 0, 0, 0] == 0.0  # image


class FakeDetector:
    model_version = "fake-vehicle"

    def __init__(self, result: VehicleResult | None) -> None:
        self.result = result
        self.calls: list[BBox | None] = []

    def detect(self, image_bytes: bytes, plate: BBox | None) -> VehicleResult | None:
        self.calls.append(plate)
        return self.result


def test_composition_fills_vehicle_over_the_stub_and_survives_a_failing_stage() -> None:
    from vision_service.recognizer import StubRecognizer, WithVehicle
    from vision_service.settings import Settings

    det = FakeDetector(VehicleResult(body_type="truck", confidence=0.77))
    rec = WithVehicle(StubRecognizer(Settings()), det)
    res = rec.analyze(b"jpeg-bytes")
    assert res.plate is None
    assert res.vehicle == VehicleResult(body_type="truck", confidence=0.77)
    assert res.model_version == "stub-0+fake-vehicle"
    assert det.calls == [None]

    class Boom:
        model_version = "boom"

        def detect(self, image_bytes: bytes, plate: BBox | None) -> VehicleResult | None:
            raise RuntimeError("no model")

    rec2 = WithVehicle(StubRecognizer(Settings()), Boom())
    res2 = rec2.analyze(b"jpeg-bytes")
    assert res2.vehicle is None
    assert rec2.ready is True
    assert "vehicle: RuntimeError: no model" in (rec2.error or "")


def test_app_reports_a_missing_model_file_and_keeps_serving() -> None:
    from vision_service.app import app

    os.environ["VISION_VEHICLE_MODEL_PATH"] = "/nonexistent/yolox.onnx"
    try:
        with TestClient(app) as client:
            health = client.get("/health").json()
            assert health["ready"] is True  # the plate stage (stub) is fine
            assert "vehicle:" in (health["detail"] or "")
            res = client.post("/analyze", content=b"x", headers={"content-type": "application/octet-stream"})
            assert res.status_code == 200
            assert res.json()["vehicle"] is None
    finally:
        os.environ.pop("VISION_VEHICLE_MODEL_PATH", None)


# ----------------------------------------------------------------------------------
# Phase B: the classifier stage over the detector
# ----------------------------------------------------------------------------------


def test_crop_vehicle_adds_the_margin_clamps_and_blurs_the_plate() -> None:
    cv2 = pytest.importorskip("cv2")
    from vision_service.vehicle import crop_vehicle

    frame = np.zeros((100, 200, 3), dtype=np.uint8)
    frame[40:50, 90:110] = (0, 255, 0)  # a green "plate"
    box = BBox(x1=50, y1=20, x2=150, y2=80)  # 100×60 → 8 % margin = 8 / 5 px
    crop = crop_vehicle(frame, box, None, 0.08)
    assert crop.shape == (70, 116, 3)
    edge = crop_vehicle(frame, BBox(x1=0, y1=0, x2=100, y2=60), None, 0.08)
    assert edge.shape == (65, 108, 3)  # clamped at the frame's top-left
    assert crop_vehicle(frame, BBox(x1=10, y1=10, x2=12, y2=12), None, 0.08) is None
    blurred = crop_vehicle(frame, box, BBox(x1=90, y1=40, x2=110, y2=50), 0.08)
    strip = blurred[40 - 20 + 5 : 50 - 20 + 5, 90 - 50 + 8 : 110 - 50 + 8, 1]  # plate strip, green channel
    assert strip.mean() < 200 and crop[20 + 5 : 30 + 5, 40 + 8 : 60 + 8, 1].mean() == 255
    assert cv2 is not None


class FakeClassifier:
    ready = True
    error = None
    min_confidence = 0.6
    model_version = "bodytype:vfake"

    def __init__(self, classes: list[str], answer: tuple[str, float] | None) -> None:
        self.classes = classes
        self.answer = answer
        self.calls = 0

    def classify(self, frame, box, plate):  # type: ignore[no-untyped-def]
        self.calls += 1
        if isinstance(self.answer, Exception):
            raise self.answer
        return self.answer


class FrameDetector:
    """A detector that answers on decoded frames (like YOLOX) with a fixed result."""

    model_version = "det"
    ready = True
    error = None

    def __init__(self, result: VehicleResult | None) -> None:
        self.result = result

    def detect_frame(self, frame, plate):  # type: ignore[no-untyped-def]
        return self.result

    def detect(self, image_bytes: bytes, plate: BBox | None) -> VehicleResult | None:
        raise AssertionError("the refined stage should share the decoded frame")


def _jpeg() -> bytes:
    cv2 = pytest.importorskip("cv2")
    ok, buf = cv2.imencode(".jpg", np.zeros((60, 80, 3), dtype=np.uint8))
    assert ok
    return bytes(buf)


def test_refined_detector_replaces_car_when_confident_else_keeps_the_detector() -> None:
    from vision_service.vehicle import RefinedVehicleDetector

    car = VehicleResult(body_type="car", confidence=0.85, bbox=BBox(x1=10, y1=10, x2=70, y2=50))
    sure = FakeClassifier(["sedan", "suv"], ("suv", 0.91))
    res = RefinedVehicleDetector(FrameDetector(car), sure).detect(_jpeg(), None)
    assert (
        res is not None and res.body_type == "suv" and res.confidence == 0.91 and res.detector_class == "car"
    )
    assert res.bbox == car.bbox

    unsure = FakeClassifier(["sedan", "suv"], ("suv", 0.4))
    res2 = RefinedVehicleDetector(FrameDetector(car), unsure).detect(_jpeg(), None)
    assert (
        res2 is not None
        and res2.body_type == "car"
        and res2.confidence == 0.85
        and res2.detector_class == "car"
    )

    # A class the classifier never trained on is left alone (its softmax means nothing there).
    bus = VehicleResult(body_type="bus", confidence=0.9, bbox=car.bbox)
    skip = FakeClassifier(["sedan", "suv"], ("suv", 0.99))
    res3 = RefinedVehicleDetector(FrameDetector(bus), skip).detect(_jpeg(), None)
    assert res3 == bus and skip.calls == 0
    # …unless it was: a classifier that knows trucks may override a truck.
    knows = FakeClassifier(["sedan", "truck", "van"], ("van", 0.8))
    truck = VehicleResult(body_type="truck", confidence=0.7, bbox=car.bbox)
    res4 = RefinedVehicleDetector(FrameDetector(truck), knows).detect(_jpeg(), None)
    assert res4 is not None and res4.body_type == "van" and res4.detector_class == "truck"


def test_refined_detector_survives_a_broken_classifier_and_reports_it() -> None:
    from vision_service.vehicle import RefinedVehicleDetector

    car = VehicleResult(body_type="car", confidence=0.85, bbox=BBox(x1=10, y1=10, x2=70, y2=50))
    boom = FakeClassifier(["sedan"], RuntimeError("bad graph"))  # type: ignore[arg-type]
    ref = RefinedVehicleDetector(FrameDetector(car), boom)
    assert ref.detect(_jpeg(), None) == car
    assert ref.error == "classifier: RuntimeError: bad graph"
    assert ref.ready is True and ref.model_version == "det+bodytype:vfake"
    # No box, or a detector that found nothing → nothing to classify.
    assert RefinedVehicleDetector(FrameDetector(None), boom).detect(_jpeg(), None) is None
    boxless = VehicleResult(body_type="car", confidence=0.85)
    assert RefinedVehicleDetector(FrameDetector(boxless), boom).detect(_jpeg(), None) == boxless


def test_classifier_without_files_is_not_ready_and_the_factory_skips_a_missing_model(tmp_path) -> None:  # type: ignore[no-untyped-def]
    from vision_service.recognizer import WithVehicle, build_recognizer
    from vision_service.settings import Settings
    from vision_service.vehicle import BodyTypeClassifier, RefinedVehicleDetector

    clf = BodyTypeClassifier(str(tmp_path / "bodytype.onnx"))
    assert clf.ready is False and "FileNotFoundError" in (clf.error or "")
    (tmp_path / "bodytype.json").write_text('{"format": "other"}')
    assert "unknown sidecar format" in (BodyTypeClassifier(str(tmp_path / "bodytype.onnx")).error or "")

    # A path with no file = the normal pre-model state: phase A only, no error in health.
    s = Settings(
        vehicle_model_path="/nonexistent/yolox.onnx", vehicle_classifier_path=str(tmp_path / "none.onnx")
    )
    rec = build_recognizer(s)
    assert isinstance(rec, WithVehicle)
    assert not isinstance(rec._detector, RefinedVehicleDetector)  # noqa: SLF001
    assert "classifier" not in (rec.error or "")