import numpy as np
import pytest
from aarecommon.math.coordinate import (
    AerotechCoordinate,
    Coordinate,
    SmargonCoordinate,
    positive_coords,
)


def test_coordinate_addition():
    c1 = Coordinate(x=1, y=2, z=3)
    c2 = Coordinate(x=4, y=5, z=6)
    c3 = c1 + c2
    assert c3.x == 5
    assert c3.y == 7
    assert c3.z == 9


def test_coordinate_subtraction():
    c1 = Coordinate(x=5, y=7, z=9)
    c2 = Coordinate(x=1, y=2, z=3)
    c3 = c1 - c2
    assert c3.x == 4
    assert c3.y == 5
    assert c3.z == 6


def test_coordinate_multiplication_scalar():
    c1 = Coordinate(x=1, y=2, z=3)
    c2 = c1 * 2.0
    assert c2.x == 2.0
    assert c2.y == 4.0
    assert c2.z == 6.0


def test_coordinate_dot_product():
    c1 = Coordinate(x=1, y=2, z=3)
    c2 = Coordinate(x=4, y=5, z=6)
    dot = c1 * c2
    assert dot == (1 * 4 + 2 * 5 + 3 * 6)


def test_coordinate_division():
    c1 = Coordinate(x=2, y=4, z=6)
    c2 = c1 / 2.0
    assert c2.x == 1.0
    assert c2.y == 2.0
    assert c2.z == 3.0


def test_coordinate_division_by_zero():
    c1 = Coordinate(x=2, y=4, z=6)
    with pytest.raises(ValueError, match="Cannot divide by zero"):
        _ = c1 / 0.0


def test_coordinate_normalize():
    c1 = Coordinate(x=3, y=0, z=4)
    c2 = c1.normalize()
    assert c2.x == 0.6
    assert c2.y == 0.0
    assert c2.z == 0.8


def test_coordinate_normalize_zero():
    c1 = Coordinate(x=0, y=0, z=0)
    with pytest.raises(ValueError, match="Cannot normalize a zero-magnitude vector."):
        c1.normalize()


def test_coordinate_rotate_x():
    c1 = Coordinate(x=1, y=1, z=0)
    # Rotate 90 degrees around X. (1, 1, 0) -> (1, 0, 1)
    c2 = c1.rotate(90, "x")
    assert np.isclose(c2.x, 1)
    assert np.isclose(c2.y, 0)
    assert np.isclose(c2.z, 1)


def test_coordinate_rotate_y():
    c1 = Coordinate(x=1, y=0, z=1)
    # Rotate 90 degrees around Y. (1, 0, 1) -> (1, 0, -1)
    c2 = c1.rotate(90, "y")
    assert np.isclose(c2.x, 1)
    assert np.isclose(c2.y, 0)
    assert np.isclose(c2.z, -1)


def test_coordinate_rotate_z():
    c1 = Coordinate(x=1, y=0, z=0)
    # Rotate 90 degrees around Z. (1, 0, 0) -> (0, 1, 0)
    c2 = c1.rotate(90, "z")
    assert np.isclose(c2.x, 0)
    assert np.isclose(c2.y, 1)
    assert np.isclose(c2.z, 0)


def test_coordinate_rotate_invalid_axis():
    c1 = Coordinate(x=1, y=1, z=1)
    with pytest.raises(ValueError, match="Invalid axis"):
        c1.rotate(90, "w")


def test_smargon_coordinate_equality():
    s1 = SmargonCoordinate(sh_mm=Coordinate(x=1, y=1, z=1), phi_deg=10, chi_deg=20)
    s2 = SmargonCoordinate(sh_mm=Coordinate(x=1.05, y=0.95, z=1.01), phi_deg=10.05, chi_deg=19.95)
    assert s1 == s2

    s3 = SmargonCoordinate(sh_mm=Coordinate(x=2, y=1, z=1), phi_deg=10, chi_deg=20)
    assert s1 != s3


def test_aerotech_coordinate_equality():
    a1 = AerotechCoordinate(at_mm=Coordinate(x=1, y=1, z=1), omega_deg=10)
    a2 = AerotechCoordinate(at_mm=Coordinate(x=1.005, y=0.995, z=1.001), omega_deg=10.005)
    assert a1 == a2

    a3 = AerotechCoordinate(at_mm=Coordinate(x=1.1, y=1, z=1), omega_deg=10)
    assert a1 != a3


def test_positive_coords():
    c1 = Coordinate(x=1, y=1, z=1)
    assert positive_coords(c1) == c1

    with pytest.raises(ValueError, match="Coordinates must be positive"):
        positive_coords(Coordinate(x=-1, y=1, z=1))

    with pytest.raises(ValueError, match="Coordinates must be positive"):
        positive_coords(Coordinate(x=1, y=-1, z=1))


def test_coordinate_unsupported_ops():
    c1 = Coordinate(x=1, y=1, z=1)
    assert c1.__add__(1) == NotImplemented
    assert c1.__sub__(1) == NotImplemented
    assert c1.__mul__("string") == NotImplemented
    assert c1.__truediv__("string") == NotImplemented


def test_smargon_coordinate_eq_not_implemented():
    s1 = SmargonCoordinate(sh_mm=Coordinate(x=1, y=1, z=1), phi_deg=10, chi_deg=20)
    assert s1.__eq__(1) == NotImplemented


def test_aerotech_coordinate_eq_not_implemented():
    a1 = AerotechCoordinate(at_mm=Coordinate(x=1, y=1, z=1), omega_deg=10)
    assert a1.__eq__(1) == NotImplemented