"""Tests for rotated coordinate system transforms."""
from __future__ import annotations

import math

import numpy as np
import pytest

from ai_mouse.coord import encode_trajectory, decode_trajectory


class TestEncodeTrajectory:
    """Test pixel → rotated normalised frame."""

    def test_start_maps_to_origin(self):
        start = (100, 200)
        end = (400, 500)
        points = np.array([[100, 200]], dtype=float)
        result = encode_trajectory(points, start, end)
        np.testing.assert_allclose(result[0], [0.0, 0.0], atol=1e-10)

    def test_end_maps_to_one_zero(self):
        start = (100, 200)
        end = (400, 500)
        points = np.array([[400, 500]], dtype=float)
        result = encode_trajectory(points, start, end)
        np.testing.assert_allclose(result[0], [1.0, 0.0], atol=1e-10)

    def test_midpoint_maps_to_half_zero(self):
        start = (0, 0)
        end = (200, 0)
        points = np.array([[100, 0]], dtype=float)
        result = encode_trajectory(points, start, end)
        np.testing.assert_allclose(result[0], [0.5, 0.0], atol=1e-10)

    def test_lateral_offset_positive(self):
        """Point at (100, 50) with horizontal start→end has lateral = 50/200 = 0.25."""
        start = (0, 0)
        end = (200, 0)
        # For horizontal u=(1,0), v=(-0, 1)=(0,1).
        # Point (100, 50): forward = 100/200=0.5, lateral = 50/200=0.25
        points = np.array([[100, 50]], dtype=float)
        result = encode_trajectory(points, start, end)
        np.testing.assert_allclose(result[0], [0.5, 0.25], atol=1e-10)

    def test_various_angles(self):
        """Encode/decode round-trip works for various angles."""
        angles = [0, 45, 90, 135, 180, -45, -90, -135]
        for deg in angles:
            rad = math.radians(deg)
            start = (400, 300)
            dist = 200
            end = (int(400 + dist * math.cos(rad)), int(300 + dist * math.sin(rad)))
            # Create a curved path
            t = np.linspace(0, 1, 20)
            px = start[0] + t * (end[0] - start[0]) + 20 * np.sin(t * math.pi)
            py = start[1] + t * (end[1] - start[1]) + 20 * np.cos(t * math.pi)
            points = np.stack([px, py], axis=1)

            encoded = encode_trajectory(points, start, end)
            assert encoded[0, 0] == pytest.approx(0.0, abs=0.2)
            assert encoded[-1, 0] == pytest.approx(1.0, abs=0.2)


class TestDecodeTrajectory:
    """Test rotated normalised frame → pixel."""

    def test_origin_maps_to_start(self):
        start = (100, 200)
        end = (400, 500)
        normalised = np.array([[0.0, 0.0]], dtype=float)
        result = decode_trajectory(normalised, start, end)
        np.testing.assert_allclose(result[0], [100, 200], atol=1e-10)

    def test_one_zero_maps_to_end(self):
        start = (100, 200)
        end = (400, 500)
        normalised = np.array([[1.0, 0.0]], dtype=float)
        result = decode_trajectory(normalised, start, end)
        np.testing.assert_allclose(result[0], [400, 500], atol=1e-10)


class TestRoundTrip:
    """Encode then decode should return original points."""

    def test_round_trip_horizontal(self):
        start = (50, 100)
        end = (350, 100)
        points = np.array([[50, 100], [150, 130], [250, 90], [350, 100]], dtype=float)
        encoded = encode_trajectory(points, start, end)
        decoded = decode_trajectory(encoded, start, end)
        np.testing.assert_allclose(decoded, points, atol=1e-8)

    def test_round_trip_diagonal(self):
        start = (100, 100)
        end = (500, 400)
        rng = np.random.default_rng(42)
        points = np.column_stack([
            np.linspace(100, 500, 30) + rng.normal(0, 10, 30),
            np.linspace(100, 400, 30) + rng.normal(0, 10, 30),
        ])
        encoded = encode_trajectory(points, start, end)
        decoded = decode_trajectory(encoded, start, end)
        np.testing.assert_allclose(decoded, points, atol=1e-8)

    def test_round_trip_vertical(self):
        """Vertical movement (angle=90°) doesn't collapse."""
        start = (300, 50)
        end = (300, 450)
        points = np.array([[300, 50], [310, 200], [295, 350], [300, 450]], dtype=float)
        encoded = encode_trajectory(points, start, end)
        decoded = decode_trajectory(encoded, start, end)
        np.testing.assert_allclose(decoded, points, atol=1e-8)