"""The simulated 5-axis solve: residuals, the stage envelope, and the frames.

The solver is the part of this feature that is easiest to get wrong and hardest
to notice: a wrong derivation still returns two plausible angles. So the central
test here does not check a formula — it brute-forces the stage envelope and
requires the solver to be no worse than the best pose that exists.

That check is what caught the four wrong versions before it: rotating the
sensor-frame normal instead of the platform-frame one, composing the stages in
the wrong order, a seeded search that skipped the right basin, and a sweep that
ranked yaws by the wrong quantity and so reported a quarter of the patch as
unreachable when all of it aligns.
"""

from __future__ import annotations

import numpy as np
import pytest

from orisys_autocal.motion.kinematics import (
    ALIGNED_TOL_DEG,
    PlatformStack,
    Pose,
    residual_for,
    rotation_about,
)
from orisys_autocal.sensor.frames import (
    AXIS_MAP,
    describe,
    mount_rotation,
    to_platform,
    to_sensor,
)

#: Normals to exercise, including the awkward ones.
NORMALS = (
    np.array([0.0, 1.0, 0.0]),  # the patch's nominal normal
    np.array([0.0, -1.0, 0.0]),  # straight into the tool
    np.array([1.0, 0.0, 0.0]),
    np.array([0.0, 0.0, 1.0]),
    np.array([0.3, 0.5, -0.81]),
    np.array([-0.6, 0.2, 0.77]),
    np.array([0.0, 0.3, 0.954]),
)


def _brute_force_best(
    stack: PlatformStack, unit_platform: np.ndarray, *, step: float = 0.25
) -> float:
    """The lowest residual any pose in the envelope achieves — by search."""
    tool = stack.tool_direction
    yaw_limit = min(180.0, stack.yaw_limit_deg)
    pitch_limit = min(360.0, stack.pitch_limit_deg)
    best = 1e9
    for yaw in np.arange(-yaw_limit, yaw_limit + step, step):
        turned = rotation_about(stack.yaw_axis, float(yaw)) @ unit_platform
        for pitch in np.arange(-pitch_limit, pitch_limit + step, step):
            rotated = rotation_about(stack.pitch_axis, float(pitch)) @ turned
            residual = float(np.degrees(np.arccos(np.clip(rotated @ tool, -1.0, 1.0))))
            best = min(best, residual)
    return best


# ------------------------------------------------------------------- the mount


def test_the_axis_map_is_a_rotation():
    """A mapping that mirrors the part is worse than none: it looks fine."""
    matrix = mount_rotation()
    assert float(np.linalg.det(matrix)) == pytest.approx(1.0)
    assert np.allclose(matrix @ matrix.T, np.eye(3))


def test_the_sensor_y_maps_to_the_platform_z():
    """The bench's frame relation, stated once so it can be corrected once."""
    assert np.allclose(to_platform(np.array([0.0, 1.0, 0.0])), [0.0, 0.0, 1.0])


def test_the_sensor_x_maps_to_the_platform_x():
    assert np.allclose(to_platform(np.array([1.0, 0.0, 0.0])), [1.0, 0.0, 0.0])


def test_the_mount_round_trips():
    vector = np.array([0.3, -0.5, 0.81])
    assert np.allclose(to_sensor(to_platform(vector)), vector)


def test_a_mirroring_map_is_refused():
    """A map that reflects the part is worse than none: the numbers look sane."""
    with pytest.raises(ValueError) as excinfo:
        mount_rotation({"x": "x", "y": "y", "z": "-z"})  # det = -1
    assert "not a rotation" in str(excinfo.value)


def test_an_incomplete_map_is_refused():
    with pytest.raises(ValueError) as excinfo:
        mount_rotation({"x": "x", "y": "z"})
    assert "no entry" in str(excinfo.value)


def test_the_mapping_is_described_for_the_ui():
    assert describe() == "sensor x->x, y->z, z->-y"


def test_swapping_the_mapping_changes_the_answer():
    """The one-line edit the bench may need, proven to actually do something."""
    # Asymmetric on purpose: a point on a coordinate axis maps to the same place
    # under both rotations, which would make this test pass without the mapping
    # doing anything at all.
    point = np.array([1.0, 5.0, 2.0])
    default = PlatformStack()
    # A different, still right-handed mapping: sensor x->y, y->z, z->x.
    swapped = PlatformStack(axis_map={"x": "y", "y": "z", "z": "x"})
    assert float(np.linalg.det(swapped.mount)) == pytest.approx(1.0)
    assert not np.allclose(
        default.point_to_platform(point), swapped.point_to_platform(point)
    )


# --------------------------------------------------------------- the solving


def test_the_bench_envelope_is_the_default():
    stack = PlatformStack()
    assert stack.yaw_limit_deg == 90.0
    assert stack.pitch_limit_deg == 360.0


@pytest.mark.parametrize("normal", NORMALS)
def test_the_solver_is_never_worse_than_a_brute_force_search(normal):
    """The test that pins the derivation down."""
    stack = PlatformStack()
    unit = stack.normal_to_platform(normal)
    unit = unit / np.linalg.norm(unit)
    pose = stack.solve(np.array([0.0, 5.0, 0.0]), normal)
    best = _brute_force_best(stack, unit, step=0.5)
    assert pose.residual_deg <= best + 0.05, (
        f"the solver returned {pose.residual_deg:.3f} deg where "
        f"{best:.3f} deg is achievable"
    )


# --------------------------------------------- the fast sweep vs the full grid


def _random_units(count: int, seed: int) -> list[np.ndarray]:
    generator = np.random.default_rng(seed)
    vectors = generator.normal(size=(count, 3))
    return [vector / np.linalg.norm(vector) for vector in vectors]


def test_the_fast_sweep_agrees_with_the_full_grid_sweep():
    """The closed-form sweep must find what the exhaustive grid finds.

    The shipped search solves the best pitch at each yaw in closed form, which is
    what made a full patch 14x faster. The exhaustive grid is kept as the slow
    reference precisely so this comparison exists: an optimisation in this file has
    silently cost 68 degrees before, while looking entirely reasonable.
    """
    stack = PlatformStack()
    for unit in _random_units(12, seed=20240517)[:6]:
        fast_yaw, fast_pitch, fast_residual = stack._search_alignment(  # noqa: SLF001
            unit, stack.tool_direction
        )
        ref_yaw, ref_pitch, ref_residual = stack.search_alignment_reference(unit, step=5.0)
        # The reference is quantised to the sweep step and is not polished, so it
        # may be slightly worse; it must never be better than the shipped path.
        assert fast_residual <= ref_residual + 1e-9, (
            f"fast sweep left {fast_residual:.6f} deg where the full grid found "
            f"{ref_residual:.6f} deg (fast {fast_yaw:.3f}/{fast_pitch:.3f}, "
            f"reference {ref_yaw:.3f}/{ref_pitch:.3f})"
        )
        # And the polished answer should be far better than the grid's own step.
        assert fast_residual <= ref_residual + 0.05


def test_the_closed_form_pitch_beats_a_direct_scan():
    """Guard the sign convention the closed form depends on.

    There are two cancellation-free coefficients and picking the wrong one returns
    the antipode: a pitch 180 degrees away that still looks like a number. This is
    the check that caught exactly that, so it is stated directly against a numeric
    scan of the pitch rather than against the rest of the solver.
    """
    stack = PlatformStack()
    for unit in _random_units(6, seed=99):
        yaws = np.array([-80.0, -35.0, -10.0, 0.0, 15.0, 47.0, 90.0])
        closed_form = stack._best_pitch_per_yaw(unit, yaws)  # noqa: SLF001
        for yaw, pitch in zip(yaws, closed_form):
            scan = np.arange(-360.0, 360.0, 0.1)
            yawed = stack._turn(stack.yaw_axis, np.array([float(yaw)]), unit)[0]  # noqa: SLF001
            # The residual is monotone in the dot product, so comparing the dot is
            # enough and avoids a million arccos calls.
            tool = stack.tool_direction
            cosine, sine = np.cos(np.radians(scan)), np.sin(np.radians(scan))
            stationary, first, second, sign_first, sign_second = (
                stack._rotation_coefficients(stack.pitch_axis)  # noqa: SLF001
            )
            dot = (
                yawed[stationary] * tool[stationary]
                + (yawed[first] * cosine + sign_first * yawed[second] * sine) * tool[first]
                + (sign_second * yawed[first] * sine + yawed[second] * cosine) * tool[second]
            )
            achieved = stack._cell(unit, float(yaw), float(pitch))  # noqa: SLF001
            achieved_dot = float(np.dot(achieved, tool))
            assert achieved_dot >= float(dot.max()) - 1e-9, (
                f"closed-form pitch {pitch:.3f} deg at yaw {yaw:.1f} reaches dot "
                f"{achieved_dot:.6f}, but a direct scan reaches {dot.max():.6f}"
            )


def _independent_residual(
    stack: PlatformStack, normal_sensor: np.ndarray, yaw_deg: float, pitch_deg: float
) -> float:
    """The residual worked out from raw matrices built here, not by the solver.

    ``residual_for`` shares ``rotation_about`` with the solver, so it cannot catch
    an error in that shared helper. This builds the two 3x3 matrices from the
    rotation formulas directly, so the two paths agree only if both are right.
    """
    normal = stack.normal_to_platform(np.asarray(normal_sensor, dtype=np.float64))
    unit = normal / np.linalg.norm(normal)

    yaw, pitch = np.radians(yaw_deg), np.radians(pitch_deg)
    about_yaw = np.array(
        [
            [np.cos(yaw), -np.sin(yaw), 0.0],
            [np.sin(yaw), np.cos(yaw), 0.0],
            [0.0, 0.0, 1.0],
        ]
    )
    about_pitch = np.array(
        [
            [np.cos(pitch), 0.0, np.sin(pitch)],
            [0.0, 1.0, 0.0],
            [-np.sin(pitch), 0.0, np.cos(pitch)],
        ]
    )
    turned = about_pitch @ about_yaw @ unit
    return float(
        np.degrees(np.arccos(np.clip(float(turned @ stack.tool_direction), -1.0, 1.0)))
    )


@pytest.mark.parametrize("normal", NORMALS)
def test_a_returned_pose_actually_achieves_its_residual(normal):
    """Independent of how it was found: apply the angles and measure."""
    stack = PlatformStack()
    pose = stack.solve(np.array([0.0, 5.0, 0.0]), normal)
    measured = residual_for(stack, normal, pose.yaw_deg, pose.pitch_deg)
    assert measured == pytest.approx(pose.residual_deg, abs=1e-6)


@pytest.mark.parametrize("normal", NORMALS)
def test_the_residual_survives_an_independently_built_rotation(normal):
    """The shared rotation helper cannot be the only thing agreeing with itself."""
    stack = PlatformStack()
    pose = stack.solve(np.array([0.0, 5.0, 0.0]), normal)
    measured = _independent_residual(stack, normal, pose.yaw_deg, pose.pitch_deg)
    assert measured == pytest.approx(pose.residual_deg, abs=1e-6)


def test_every_sampled_patch_normal_is_reachable_within_the_yaw_limit():
    """The whole patch aligns, and the yaw limit is not what stops it.

    An earlier revision of this solver claimed a quarter of the patch was
    unreachable at 65 deg. That was the solver's bug. The real answer is that the
    360 deg pitch stage does the work and the +-90 deg yaw limit costs at most a
    quarter of a degree anywhere on the patch, so no re-mount is needed.
    """
    stack = PlatformStack()
    worst = 0.0
    for normal in NORMALS:
        pose = stack.solve(np.array([0.0, 5.0, 0.0]), normal)
        assert pose.ok
        worst = max(worst, pose.residual_deg)
    assert worst < 0.5, f"worst residual across the sampled patch is {worst:.3f} deg"


def test_a_generic_normal_aligns_almost_exactly():
    """Away from the yaw clamp the solver should reach the polish floor."""
    stack = PlatformStack()
    for normal in NORMALS[:6]:
        pose = stack.solve(np.array([0.0, 5.0, 0.0]), normal)
        assert pose.residual_deg < 0.01, (
            f"normal {normal} left {pose.residual_deg:.4f} deg behind"
        )


def test_the_nominal_normal_needs_no_rotation():
    """Sensor +Y is the platform +Z, so the patch's centre is already aligned."""
    stack = PlatformStack()
    pose = stack.solve(np.array([0.0, 5.0, 0.0]), np.array([0.0, 1.0, 0.0]))
    assert pose.ok
    assert pose.residual_deg == pytest.approx(0.0, abs=1e-6)


def test_a_normal_straight_into_the_tool_is_flagged_indeterminate():
    stack = PlatformStack()
    pose = stack.solve(np.array([0.0, 5.0, 0.0]), np.array([0.0, -1.0, 0.0]))
    assert pose.ok
    assert not pose.yaw_determined
    assert "indeterminate" in pose.reason


def test_a_degenerate_normal_is_refused():
    stack = PlatformStack()
    pose = stack.solve(np.array([0.0, 5.0, 0.0]), np.zeros(3))
    assert not pose.ok
    assert "degenerate" in pose.reason


def test_the_yaw_limit_is_respected():
    stack = PlatformStack(yaw_limit_deg=90.0)
    for normal in NORMALS:
        pose = stack.solve(np.array([0.0, 5.0, 0.0]), normal)
        assert abs(pose.yaw_deg) <= 90.0 + 1e-6


def test_the_pitch_limit_is_respected():
    stack = PlatformStack(pitch_limit_deg=45.0)
    for normal in NORMALS:
        pose = stack.solve(np.array([0.0, 5.0, 0.0]), normal)
        assert abs(pose.pitch_deg) <= 45.0 + 1e-6


def test_tightening_the_envelope_can_only_leave_more_residual():
    """A narrower stage cannot align better than a wider one."""
    normal = np.array([-0.6, 0.2, 0.77])
    point = np.array([0.0, 5.0, 0.0])
    wide = PlatformStack(yaw_limit_deg=180.0).solve(point, normal).residual_deg
    narrow = PlatformStack(yaw_limit_deg=30.0).solve(point, normal).residual_deg
    assert narrow >= wide - 1e-6


def test_an_unbounded_yaw_can_align_arbitrary_normals():
    """The envelope, not the solver, is what leaves residual."""
    stack = PlatformStack(yaw_limit_deg=180.0, pitch_limit_deg=180.0)
    worst = 0.0
    for index in range(40):
        normal = np.random.default_rng(index).normal(size=3)
        normal = normal / np.linalg.norm(normal)
        worst = max(worst, stack.solve(np.array([0.0, 5.0, 0.0]), normal).residual_deg)
    assert worst < 0.5, f"worst residual {worst:.3f} deg with a free yaw"


def test_the_reported_xy_is_the_contact_point_not_the_origin():
    stack = PlatformStack()
    point = np.array([2.0, 5.0, 3.0])
    pose = stack.solve(point, np.array([0.0, 1.0, 0.0]))
    platform = stack.point_to_platform(point)
    assert pose.target_xy_mm[0] == pytest.approx(platform[0])
    assert pose.target_xy_mm[1] == pytest.approx(platform[1])


def test_the_plate_lifts_the_point_along_the_tool_axis():
    stack = PlatformStack()
    point = np.array([1.0, 5.0, 2.0])
    plain = stack.point_to_platform(point)
    lifted = stack.point_to_platform(point, plate_thickness_mm=2.5)
    assert np.allclose(lifted - plain, [0.0, 0.0, 2.5])


def test_the_tool_offset_shifts_the_reported_point():
    stack = PlatformStack(tool_offset_mm=(1.0, 2.0, 3.0))
    point = np.array([0.0, 5.0, 0.0])
    pose = stack.solve(point, np.array([0.0, 1.0, 0.0]))
    plain = PlatformStack().solve(point, np.array([0.0, 1.0, 0.0]))
    assert pose.target_xy_mm[0] == pytest.approx(plain.target_xy_mm[0] + 1.0)
    assert pose.target_xy_mm[1] == pytest.approx(plain.target_xy_mm[1] + 2.0)


def test_the_approach_sign_is_validated():
    with pytest.raises(ValueError):
        PlatformStack(approach_sign=0)


def test_swapping_the_approach_sign_changes_the_target():
    """The sign is the difference between pressing in and pulling away."""
    normal = np.array([0.0, 1.0, 0.0])
    point = np.array([0.0, 5.0, 0.0])
    inward = PlatformStack(approach_sign=-1).solve(point, normal)
    outward = PlatformStack(approach_sign=1).solve(point, normal)
    assert inward.residual_deg == pytest.approx(0.0, abs=1e-6)
    assert outward.pitch_deg != pytest.approx(inward.pitch_deg)


def test_a_pose_describes_itself():
    pose = Pose(ok=True, target_xy_mm=(1.0, 2.0), yaw_deg=3.0, pitch_deg=4.0, residual_deg=0.01)
    text = pose.describe()
    assert "yaw" in text and "pitch" in text and "residual" in text
    assert Pose(ok=False, reason="nope").describe() == "no pose: nope"


def test_an_aligned_pose_reports_that_it_reached_tolerance():
    aligned = Pose(ok=True, residual_deg=0.0)
    assert aligned.reached_tolerance and aligned.reachable
    # ``reachable`` is the coarser statement: within half a degree counts as
    # "the machine can put the tool on this normal".
    loose = Pose(ok=True, residual_deg=ALIGNED_TOL_DEG * 10)
    assert not loose.reached_tolerance
    assert loose.reachable
    far = Pose(ok=True, residual_deg=30.0)
    assert not far.reached_tolerance and not far.reachable