"""Oblique / probe's-eye multiplanar reformation.

Surgeons plan trajectories in views aligned to the *probe*, not to the scanner
axes: the **in-line view** (the plane containing the trajectory, showing the
whole path from entry to target) and the **probe's-eye view** (the cross-section
perpendicular to the trajectory at a chosen depth — what the advancing tip
"sees"). Both are arbitrary-plane reformats sampled from the volume by trilinear
interpolation, so they work for any orientation, unlike the axis-aligned MPR.

Pure NumPy on top of ``Volume.sample_world``; unit-tested headlessly.
"""
from __future__ import annotations
import numpy as np
from .volume import Volume
from .planning import Trajectory


def orthonormal_frame(normal, up=None):
    """Right-handed (right, up, normal) frame with `normal` as the plane normal."""
    n = np.asarray(normal, float)
    L = np.linalg.norm(n)
    n = n / L if L > 0 else np.array([0.0, 0.0, 1.0])
    if up is None:
        up = np.array([0.0, 0.0, 1.0]) if abs(n[2]) < 0.9 else np.array([0.0, 1.0, 0.0])
    up = np.asarray(up, float)
    right = np.cross(up, n)
    rl = np.linalg.norm(right)
    right = right / rl if rl > 0 else np.array([1.0, 0.0, 0.0])
    up2 = np.cross(n, right)                       # re-orthogonalize
    return right, up2, n


def _sample_plane(volume: Volume, center, right, up, us, vs, order=1):
    """Sample the volume on the grid center + u*right + v*up (image[v_index, u_index])."""
    center = np.asarray(center, float)
    out = np.empty((len(vs), len(us)), np.float32)
    for j, v in enumerate(vs):
        base = center + v * up
        for i, u in enumerate(us):
            out[j, i] = volume.sample_world(base + u * right, order=order)
    return out


def oblique_slice(volume: Volume, center, normal, up=None, fov_mm=80.0,
                  spacing_mm=1.0, order=1) -> np.ndarray:
    """Reformat a square patch on an arbitrary plane through `center`."""
    right, up2, _ = orthonormal_frame(normal, up)
    half = fov_mm / 2.0
    ax = np.arange(-half, half + 1e-9, spacing_mm)
    return _sample_plane(volume, center, right, up2, ax, ax, order)


def probes_eye(volume: Volume, traj: Trajectory, depth_mm=None, fov_mm=40.0,
               spacing_mm=0.5, up=None, order=1) -> np.ndarray:
    """Cross-section perpendicular to the trajectory at `depth_mm` from entry
    (default: at the target). This is the probe's-eye view."""
    L = traj.length
    depth = L if depth_mm is None else float(np.clip(depth_mm, 0.0, L))
    center = traj.entry + traj.direction * depth
    return oblique_slice(volume, center, traj.direction, up=up,
                         fov_mm=fov_mm, spacing_mm=spacing_mm, order=order)


def along_trajectory(volume: Volume, traj: Trajectory, width_mm=40.0, pad_mm=10.0,
                     spacing_mm=0.5, side=None, order=1):
    """In-line view: the plane containing the trajectory, from entry−pad to
    target+pad. Returns (image, us) where `us` is distance from entry (mm),
    so us=0 is the entry and us=length is the target; the path is the row v=0."""
    E = traj.entry
    d = traj.direction
    if side is None:
        ref = np.array([1.0, 0, 0]) if abs(d[0]) < 0.9 else np.array([0, 1.0, 0])
        side = np.cross(d, ref)
    side = np.asarray(side, float)
    sl = np.linalg.norm(side)
    side = side / sl if sl > 0 else np.array([1.0, 0.0, 0.0])
    us = np.arange(-pad_mm, traj.length + pad_mm + 1e-9, spacing_mm)
    vs = np.arange(-width_mm / 2, width_mm / 2 + 1e-9, spacing_mm)
    img = np.empty((len(vs), len(us)), np.float32)
    for j, v in enumerate(vs):
        for i, u in enumerate(us):
            img[j, i] = volume.sample_world(E + d * u + side * v, order=order)
    return img, us
