"""Tests for oblique / probe's-eye reformation (pure geometry + sampling)."""
import numpy as np
from sulcora.volume import Volume
from sulcora.planning import Trajectory
from sulcora import reslice as R


def _ramp_x(shape=(40, 40, 40)):
    """Volume whose intensity equals the world-x coordinate (identity affine)."""
    d = np.zeros(shape, np.float32)
    d[:] = np.arange(shape[0])[:, None, None]
    return Volume(d, np.eye(4))


def test_orthonormal_frame_is_orthonormal():
    right, up, n = R.orthonormal_frame((0, 0, 1))
    for a in (right, up, n):
        assert abs(np.linalg.norm(a) - 1.0) < 1e-9
    assert abs(np.dot(right, up)) < 1e-9
    assert abs(np.dot(right, n)) < 1e-9
    assert np.allclose(np.cross(right, up), n, atol=1e-9)   # right-handed


def test_oblique_axial_matches_x_ramp():
    vol = _ramp_x()
    # a plane with normal +z through (20,20,20): the image samples world x along
    # the 'right' axis, so its centre pixel should read ~20.
    img = R.oblique_slice(vol, center=(20, 20, 20), normal=(0, 0, 1),
                          fov_mm=20.0, spacing_mm=1.0)
    ci = img.shape[0] // 2
    assert abs(img[ci, ci] - 20.0) < 0.5
    # values vary along exactly one image axis (the one aligned with world x)
    assert np.ptp(img, axis=0).max() > np.ptp(img, axis=1).max() or \
           np.ptp(img, axis=1).max() > np.ptp(img, axis=0).max()


def test_probes_eye_centered_on_trajectory():
    vol = _ramp_x()
    traj = Trajectory(entry=(10, 20, 20), target=(30, 20, 20))  # along +x
    # perpendicular cross-section at mid-depth -> centred at world x=20
    img = R.probes_eye(vol, traj, depth_mm=10.0, fov_mm=16.0, spacing_mm=1.0)
    ci = img.shape[0] // 2
    assert abs(img[ci, ci] - 20.0) < 0.6
    assert img.shape[0] == img.shape[1]                         # square FOV


def test_along_trajectory_spans_entry_to_target():
    vol = _ramp_x()
    traj = Trajectory(entry=(8, 20, 20), target=(32, 20, 20))   # along +x, len 24
    img, us = R.along_trajectory(vol, traj, width_mm=10.0, pad_mm=4.0, spacing_mm=1.0)
    assert us[0] < 0 and us[-1] > traj.length                    # padded past both ends
    # centre row runs along the path; intensity (=world x) increases with u
    mid = img[img.shape[0] // 2]
    assert mid[-1] > mid[0]
    # us=0 is the entry (world x=8), us=length is the target (world x=32)
    i0 = int(np.argmin(np.abs(us - 0.0)))
    iL = int(np.argmin(np.abs(us - traj.length)))
    assert abs(mid[i0] - 8.0) < 0.8 and abs(mid[iL] - 32.0) < 0.8
