tests/test_height_comp.py

"""Geometry and fallback unit tests for the ball-centre height compensation
-- BG31 step 1.5 (13/08/2026).

A SYNTHETIC camera is built explicitly (a K with a known f and principal
point, a pose looking down at the table, a known height h), the 4 table
corners (z=0) are projected into pixels, fed to TableHomography (the REAL
production code), then the BALL CENTRES (z=R) are projected too: mapping plus
compensation has to recover the true centre foot on the cloth to better than
1mm (BRIEF 31 step 1.5). The uncompensated control has to be CLEARLY off
(two-digit millimetres far from the camera), so the test is not vacuous.

BG32 (13/08/2026) tightened the sanity bounds (h in [1.5, 6]m, f in
[800, 8000]px) and added ``cx_m``/``cy_m`` to the flag -- hence the extra test
that reproduces EXACTLY the junk case of shot 11 in the first smoke run
(h=0.686 / f=12767, which slipped through the wider BG31 bounds) and must
come out off.

Pure numpy -- no cv2 and no torch (camera.py keeps the poolcoach_cv rule).
"""

from __future__ import annotations

import sys
from pathlib import Path

import numpy as np
import pytest

from poolcoach_cv.camera import (CAM_F_MAX_PX, CAM_F_MIN_PX, CAM_H_MAX_M,
                                 CAM_H_MIN_M, HeightCompensation,
                                 decompose_homography)
from poolcoach_cv.homography import TableHomography

# the broadcast table frame (poolcoach_cv.broadcast -- not imported whole,
# to keep this light)
W, L = 1.27, 2.54
BALL_R = 0.028575
IMG_WH = (1920.0, 1080.0)


def make_projector(cam_xy=(0.6, -2.5), h=2.4, f=2200.0, wh=IMG_WH):
    """A synthetic pinhole camera: returns a function projecting 3D points in
    the table frame into pixels.

    The pose looks at the centre of the table (the real broadcast stance:
    outside a short rail, h above the cloth). K has no skew, square pixels and
    the principal point at the image centre -- exactly the assumptions the
    BG30 decomposition makes.
    """
    c = np.array([cam_xy[0], cam_xy[1], h], dtype=np.float64)
    z = np.array([W / 2, L / 2, 0.0]) - c
    z /= np.linalg.norm(z)
    x = np.cross(z, np.array([0.0, 0.0, 1.0]))
    x /= np.linalg.norm(x)
    y = np.cross(z, x)
    rot = np.stack([x, y, z])            # rows = camera axes in table frame

    def project(pts3):
        p = (np.atleast_2d(np.asarray(pts3, dtype=np.float64)) - c) @ rot.T
        return np.stack([f * p[:, 0] / p[:, 2] + wh[0] / 2,
                         f * p[:, 1] / p[:, 2] + wh[1] / 2], axis=1)

    return project


def table_h(project) -> TableHomography:
    corners_px = project([[0, 0, 0], [W, 0, 0], [W, L, 0], [0, L, 0]])
    return TableHomography(corners_px, table_w=W, table_l=L)


def ball_grid():
    xs = np.linspace(0.05, W - 0.05, 5)
    ys = np.linspace(0.05, L - 0.05, 7)
    return np.array([[x, y, BALL_R] for x in xs for y in ys])


# --------------------------------------------------------------- decompose

def test_decompose_recovers_the_synthetic_camera():
    """f/h/C are known in advance, so the decomposition has to return them
    almost exactly (closed-form maths with no noise -- what is left is float
    arithmetic)."""
    project = make_projector(cam_xy=(0.6, -2.5), h=2.4, f=2200.0)
    th = table_h(project)
    dec = decompose_homography(th.H, pp=(IMG_WH[0] / 2, IMG_WH[1] / 2))
    assert dec is not None
    assert abs(dec["h"] - 2.4) < 1e-3
    assert abs(dec["f_px"] - 2200.0) < 0.5
    assert abs(dec["cx"] - 0.6) < 1e-3
    assert abs(dec["cy"] - (-2.5)) < 1e-3


@pytest.mark.parametrize("cam_xy,h,f", [
    ((0.6, -2.5), 2.4, 2200.0),      # the BG30 stance (shot frame, camera
                                     # at y<0)
    ((0.635, 5.05), 2.41, 2236.0),   # the pilot decomposition numbers
                                     # (standard frame)
    ((-0.8, 3.6), 1.6, 1400.0),      # camera off to the left, and lower
])
def test_mapping_plus_compensation_recovers_positions_under_1mm(cam_xy, h, f):
    """Gate of BRIEF 31 step 1.5: known 3D points projected through the
    synthetic camera, then mapping plus compensation recovers the cloth-plane
    position to better than 1mm; without compensation it is clearly off."""
    project = make_projector(cam_xy=cam_xy, h=h, f=f)
    th = table_h(project)
    comp = HeightCompensation.from_homography(th.H, IMG_WH, BALL_R)
    assert comp.on, comp.reason

    balls = ball_grid()
    mapped = th.px_to_table(project(balls))     # the OLD convention (no
                                                # compensation)
    err_raw = np.linalg.norm(mapped - balls[:, :2], axis=1)
    err_comp = np.linalg.norm(comp.apply(mapped) - balls[:, :2], axis=1)
    # uncompensated: off by roughly k * (distance to the camera foot) -- it
    # has to be plainly visible
    assert err_raw.max() > 0.005
    assert err_comp.max() < 0.001


def test_meta_carries_h_f_cx_cy():
    project = make_projector(cam_xy=(0.6, -2.5), h=2.4, f=2200.0)
    comp = HeightCompensation.from_homography(table_h(project).H, IMG_WH,
                                              BALL_R)
    m = comp.meta()
    assert m["on"] is True
    assert abs(m["h_m"] - 2.4) < 1e-3
    assert abs(m["f_px"] - 2200.0) < 0.5
    # BG32: the camera foot in the table frame -- optional, same rule as
    # h_m/f_px
    assert abs(m["cx_m"] - 0.6) < 1e-3
    assert abs(m["cy_m"] - (-2.5)) < 1e-3
    assert "reason" not in m


# ---------------------------------------------------------------- fallback

def test_an_affine_homography_is_not_compensated_and_apply_is_identity():
    """Fronto-parallel (looking straight down): there is no perspective
    signal to estimate f from -> the decomposition returns None -> NO
    compensation, and apply is the identity."""
    corners_px = np.array([[500.0, 300.0], [627.0, 300.0],
                           [627.0, 554.0], [500.0, 554.0]])
    th = TableHomography(corners_px, table_w=W, table_l=L)
    assert decompose_homography(th.H, pp=(960.0, 540.0)) is None
    comp = HeightCompensation.from_homography(th.H, IMG_WH, BALL_R)
    assert not comp.on
    assert "f^2" in comp.reason
    pts = np.array([[0.3, 0.4], [1.0, 2.0]])
    assert np.array_equal(comp.apply(pts), pts)
    m = comp.meta()
    assert m["on"] is False and "reason" in m


def test_an_implausible_camera_is_not_compensated_and_says_why():
    """The decomposition runs but h is outside the sanity bounds -> no
    compensation, while meta still reports the measured h/f (and, since BG32,
    cx/cy) so it can be traced."""
    project = make_projector(cam_xy=(0.6, -20.0), h=CAM_H_MAX_M * 2, f=7000.0)
    comp = HeightCompensation.from_homography(table_h(project).H, IMG_WH,
                                              BALL_R)
    assert not comp.on
    assert "h =" in comp.reason
    m = comp.meta()
    assert m["on"] is False and m["h_m"] > CAM_H_MAX_M
    assert abs(m["cx_m"] - 0.6) < 1e-3
    assert abs(m["cy_m"] - (-20.0)) < 1e-2


def test_the_bg32_bounds_block_the_shot_11_junk_case():
    """Reproduces EXACTLY the junk case that slipped through the wider BG31
    bounds (shot 11, first smoke run: the decomposition gave h=0.686m,
    f=12767px, and it compensated with a junk camera). The BG32 bounds
    [1.5, 6]m have to block it: off, a reason mentioning h, meta carrying
    h/f/cx/cy for tracing, and apply the identity (better no compensation
    than compensation from a junk camera)."""
    project = make_projector(cam_xy=(0.6, -2.5), h=0.686, f=12766.8)
    comp = HeightCompensation.from_homography(table_h(project).H, IMG_WH,
                                              BALL_R)
    assert not comp.on
    assert "h =" in comp.reason
    m = comp.meta()
    assert m["on"] is False
    assert abs(m["h_m"] - 0.686) < 1e-3 and m["h_m"] < CAM_H_MIN_M
    assert abs(m["f_px"] - 12766.8) < 0.5
    assert abs(m["cx_m"] - 0.6) < 1e-3
    assert abs(m["cy_m"] - (-2.5)) < 1e-3
    pts = np.array([[0.3, 0.4], [1.0, 2.0]])
    assert np.array_equal(comp.apply(pts), pts)


@pytest.mark.parametrize("f", [700.0, 9000.0])
def test_an_f_outside_the_bounds_blocks_it_even_with_a_healthy_h(f):
    """The BG32 bound is on h AND f: h inside [1.5, 6]m but f outside
    [800, 8000]px (on either side) still has to come out off, with a reason
    about f."""
    project = make_projector(cam_xy=(0.6, -2.5), h=2.4, f=f)
    comp = HeightCompensation.from_homography(table_h(project).H, IMG_WH,
                                              BALL_R)
    assert not comp.on
    assert "f =" in comp.reason
    m = comp.meta()
    assert m["on"] is False
    assert abs(m["h_m"] - 2.4) < 1e-3
    assert not (CAM_F_MIN_PX <= m["f_px"] <= CAM_F_MAX_PX)
    assert abs(m["cx_m"] - 0.6) < 1e-3       # the f branch reports the
    assert abs(m["cy_m"] - (-2.5)) < 1e-3    # camera foot too


def test_a_missing_frame_size_blocks_compensation():
    project = make_projector()
    comp = HeightCompensation.from_homography(table_h(project).H, (0, 0),
                                              BALL_R)
    assert not comp.on
    assert "frame" in comp.reason


# ----------------------------------------- the worker-side adapter (step 1.2)

def _cv_worker():
    sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "scripts"))
    import cv_worker

    return cv_worker


def test_comp_homography_cls_applies_the_compensation_and_records_meta():
    """The subclass patched into bc.TableHomography: px_to_table returns
    COMPENSATED coordinates (<1mm) and the sink receives the height_comp meta
    -- exactly what analyze_once writes into the job JSON."""
    cw = _cv_worker()
    project = make_projector(cam_xy=(0.6, -2.5), h=2.4, f=2200.0)
    corners_px = project([[0, 0, 0], [W, 0, 0], [W, L, 0], [0, L, 0]])

    sink: dict = {}
    cls = cw.comp_homography_cls(IMG_WH, BALL_R, sink)
    th = cls(corners_px, table_w=W, table_l=L)
    assert sink["on"] is True
    assert abs(sink["h_m"] - 2.4) < 1e-3
    assert abs(sink["cx_m"] - 0.6) < 1e-3      # BG32: cx/cy into the job JSON
    assert abs(sink["cy_m"] - (-2.5)) < 1e-3

    balls = ball_grid()
    got = np.array([th.px_to_table(px) for px in project(balls)])
    err = np.linalg.norm(got - balls[:, :2], axis=1)
    assert err.max() < 0.001


def test_comp_homography_cls_falls_back_to_the_plain_mapping():
    """When the decomposition fails the subclass maps EXACTLY like the plain
    TableHomography (no compensation) and the sink reports on=False -- better
    no compensation than compensation from a junk camera."""
    cw = _cv_worker()
    corners_px = np.array([[500.0, 300.0], [627.0, 300.0],
                           [627.0, 554.0], [500.0, 554.0]])
    sink: dict = {}
    cls = cw.comp_homography_cls(IMG_WH, BALL_R, sink)
    th = cls(corners_px, table_w=W, table_l=L)
    assert sink["on"] is False and "reason" in sink
    base = TableHomography(corners_px, table_w=W, table_l=L)
    px = np.array([560.0, 400.0])
    assert np.allclose(th.px_to_table(px), base.px_to_table(px))