"""Tests for the automatic trajectory optimizer (pure scoring/ranking)."""
import numpy as np
import pytest

from sulcora.volume import Volume
from sulcora.planning import Trajectory
from sulcora import optimizer as O


def _blob(voxels, shape=(21, 21, 41)):
    d = np.zeros(shape, np.uint8)
    for v in voxels:
        d[v] = 1
    return Volume(d, np.eye(4))


def _struct(name, voxels, min_clearance_mm=1.0, weight=1.0, shape=(21, 21, 41)):
    return O.RiskStructure(name, _blob(voxels, shape), min_clearance_mm, weight)


# --------------------------------------------------------------------------- #
# Candidate entry generation
# --------------------------------------------------------------------------- #
def test_candidate_entries_on_cap_within_cone():
    target = np.array([0, 0, 0.0])
    ents = O.candidate_entries(target, radius=50, axis=(0, 0, 1),
                               max_angle_deg=45, n_theta=4, n_phi=8)
    # every entry is `radius` from the target
    assert np.allclose(np.linalg.norm(ents - target, axis=1), 50.0, atol=1e-6)
    # and within the cone (angle to +Z <= 45deg)
    cosines = (ents - target)[:, 2] / 50.0
    assert np.all(cosines >= np.cos(np.radians(45)) - 1e-9)


# --------------------------------------------------------------------------- #
# Ranking
# --------------------------------------------------------------------------- #
def test_optimizer_prefers_safer_corridor():
    struct = _struct("blob", [(5, 10, 20)])
    target = np.array([10, 10, 0.0])
    safe = np.array([10, 10, 40.0])        # vertical line, ~5 mm from blob
    risky = np.array([5, 10, 40.0])        # angles toward the blob, ~2.5 mm
    plans = O.optimize_trajectory(target, [risky, safe], [struct])
    assert np.allclose(plans[0].entry, safe)                 # safer ranked first
    assert plans[0].min_clearance_mm > plans[1].min_clearance_mm
    assert plans[0].score >= plans[1].score


def test_hard_min_clearance_makes_infeasible_and_sorts_last():
    struct = _struct("blob", [(5, 10, 20)], min_clearance_mm=4.0)
    target = np.array([10, 10, 0.0])
    safe = np.array([10, 10, 40.0])        # ~5 mm -> feasible
    risky = np.array([5, 10, 40.0])        # ~2.5 mm -> violates 4 mm margin
    plans = O.optimize_trajectory(target, [risky, safe], [struct])
    assert plans[0].feasible and np.allclose(plans[0].entry, safe)
    assert not plans[-1].feasible


def test_min_clearance_and_nearest_across_structures():
    s1 = _struct("far", [(5, 10, 20)])     # ~5 mm
    s2 = _struct("near", [(8, 10, 20)])    # ~2 mm
    plan = O.score_trajectory((10, 10, 40), (10, 10, 0), [s1, s2],
                              _trees={s.name: O._tree(O._mask_world_points(s.mask))
                                      for s in (s1, s2)})
    assert plan.nearest == "near"
    assert abs(plan.min_clearance_mm - 2.0) < 0.2
    assert abs(plan.clearances["far"] - 5.0) < 0.2


def test_angle_and_depth_constraints():
    struct = _struct("blob", [(0, 0, 0)])   # far from the paths below
    target = np.array([10, 10, 0.0])
    steep = np.array([40, 10, 5.0])         # nearly horizontal -> large angle from +Z
    plans = O.optimize_trajectory(target, [steep], [struct], max_angle_deg=30)
    assert not plans[0].feasible            # angle exceeds 30 deg
    deep = np.array([10, 10, 200.0])        # very long
    plans2 = O.optimize_trajectory(target, [deep], [struct], max_depth_mm=50)
    assert not plans2[0].feasible


def test_deformation_factory_makes_ranking_deformation_aware():
    from sulcora.deformation import ProbeDeformation
    struct = _struct("blob", [(5, 15, 15)], shape=(31, 31, 31))
    target = np.array([9, 30, 15.0]); entry = np.array([9, 0, 15.0])
    static = O.score_trajectory(entry, target, [struct],
                                _trees={"blob": O._tree(O._mask_world_points(struct.mask))})
    fac = lambda tr: ProbeDeformation(tr, probe_radius_mm=1.5, influence_radius_mm=18.0)
    deformed = O.score_trajectory(entry, target, [struct], deformation_factory=fac)
    assert deformed.clearances["blob"] != static.clearances["blob"]


def test_plan_safest_generates_and_ranks():
    struct = _struct("blob", [(5, 10, 20)])
    plans = O.plan_safest((10, 10, 5), [struct], radius=35, max_angle_deg=45,
                          n_theta=4, n_phi=8)
    assert len(plans) > 1
    assert plans[0].score >= plans[-1].score          # sorted best-first
    assert plans[0].band in ("HIGH", "CAUTION", "LOW")
