"""Quick test for CameraPlumbBob: project→unproject roundtrip and known intrinsics."""
import torch
import numpy as np

# Minimal stubs so we can run without the full anydata install
import sys, types

# Stub anydata.geometry.cameras.base
base_mod = types.ModuleType("anydata.geometry.cameras.base")
class CameraBaseStub(torch.nn.Module):
    def __init__(self, hw=None, Twc=None, Tcw=None):
        super().__init__()
        if Twc is None and Tcw is None:
            self._Tcw = torch.eye(4, dtype=self._K.dtype, device=self._K.device).unsqueeze(0).repeat(self._K.shape[0], 1, 1)
        else:
            self._Tcw = Tcw if Twc is None else torch.inverse(Twc)
        self._hw = hw
    @property
    def hw(self): return self._hw
    @property
    def K(self): return self._K
    @property
    def Tcw(self): return self._Tcw
    @property
    def Twc(self): return None if self._Tcw is None else self._Tcw.inverse()
    @property
    def device(self): return self._K.device
    @property
    def dtype(self): return self._K.dtype
    @property
    def b(self): return self._K.shape[0]
base_mod.CameraBase = CameraBaseStub
sys.modules.setdefault("anydata", types.ModuleType("anydata"))
sys.modules.setdefault("anydata.geometry", types.ModuleType("anydata.geometry"))
sys.modules.setdefault("anydata.geometry.cameras", types.ModuleType("anydata.geometry.cameras"))
sys.modules.setdefault("anydata.geometry.cameras.base", base_mod)
sys.modules.setdefault("anydata.geometry.camera_utils", types.ModuleType("anydata.geometry.camera_utils"))

# Stub anydata.utils.data — only need cat_channel_ones
utils_data = types.ModuleType("anydata.utils.data")
def cat_channel_ones(t, dim):
    ones = torch.ones_like(t[:, :1])
    return torch.cat([t, ones], dim=dim)
def unnorm_pixel_grid(uv, hw):
    return uv
utils_data.cat_channel_ones = cat_channel_ones
utils_data.unnorm_pixel_grid = unnorm_pixel_grid
sys.modules.setdefault("anydata.utils", types.ModuleType("anydata.utils"))
sys.modules.setdefault("anydata.utils.data", utils_data)
sys.modules.setdefault("anydata.utils.types", types.ModuleType("anydata.utils.types"))

from plumb_bob import CameraPlumbBob


def make_camera(fx, fy, cx, cy, dist_coeffs, hw):
    """Helper to build a CameraPlumbBob from explicit params."""
    K = torch.eye(3).unsqueeze(0)
    K[0, 0, 0] = fx
    K[0, 1, 1] = fy
    K[0, 0, 2] = cx
    K[0, 1, 2] = cy
    D = torch.tensor([dist_coeffs])
    return CameraPlumbBob(K=K, D=D, hw=hw)


def realsense_camera(hw=(720, 1280)):
    """Build camera with the colleague's RealSense plumb bob intrinsics."""
    return make_camera(
        fx=644.523681640625,
        fy=643.7030029296875,
        cx=647.5728759765625,
        cy=365.7279052734375,
        dist_coeffs=[
            -0.05548327788710594,   # k1
             0.06299213320016861,   # k2
             0.0003023890021722764, # p1
             0.001013805391266942,  # p2
            -0.01973005384206772,   # k3
        ],
        hw=hw,
    )


def test_roundtrip():
    """pixelToRay → rayToPixel should recover original pixels."""
    cam = realsense_camera()

    pixels = torch.tensor([
        [640, 360], [700, 400], [580, 320],
        [647.57, 365.73], [100, 100], [1200, 600],
    ], dtype=torch.float32)
    grid = pixels.T.unsqueeze(0)  # (1, 2, N)

    rays = cam.pixelToRay(grid)        # (1, 3, N)
    reprojected = cam.rayToPixel(rays)  # (1, 2, N)

    err = (reprojected[0] - grid[0]).abs().max().item()
    print(f"Roundtrip max error: {err:.6f} px")
    assert err < 0.01, f"Roundtrip error too large: {err}"


def test_principal_point_maps_to_optical_axis():
    """The principal point should map to a ray along the z-axis."""
    cam = realsense_camera()

    grid = torch.tensor([[[647.5728759765625], [365.7279052734375]]])  # (1, 2, 1)
    rays = cam.pixelToRay(grid)  # (1, 3, 1)

    ray = rays[0, :, 0]
    print(f"Principal point ray: [{ray[0]:.6f}, {ray[1]:.6f}, {ray[2]:.6f}]")
    assert ray[0].abs() < 1e-4 and ray[1].abs() < 1e-4, "Principal point should map to optical axis"


def test_distortion_param_count():
    """Plumb bob should have exactly 5 distortion params."""
    cam = realsense_camera()
    assert cam._D.shape[-1] == 5, f"Expected 5 distortion params, got {cam._D.shape[-1]}"
    print(f"Distortion params shape: {cam._D.shape} ✓")


if __name__ == "__main__":
    test_distortion_param_count()
    test_principal_point_maps_to_optical_axis()
    test_roundtrip()
    print("\nAll tests passed!")