import numpy as np
import pytest

from gri_utils import rotate


def test_rot_2d():
    em = np.array(((1, 0), (0, 4)))
    em2 = rotate.rotate_ccw(np.pi, em)
    assert np.allclose(em, em2)
    em2 = rotate.rotate_ccw(np.pi, ((1, 0), (0, 4)))
    assert np.allclose(em, em2)
    em2 = rotate.rotate_ccw(np.pi / 2, em)
    assert np.allclose(em2, np.array(((4, 0), (0, 1))))


def test_rot_2d_vec():
    vec = (1, 0)
    assert np.allclose(rotate.rotate_ccw(np.pi / 2, vec), (0, 1))
    assert np.allclose(rotate.rotate_ccw(np.pi, vec), np.array((-1, 0)))


def test_rot_x():
    em = np.array(((1, 0, 0), (0, 4, 0), (0, 0, 9)))
    em2 = rotate.rotate_ccw(np.pi, em, 0)
    assert np.allclose(em, em2)
    em2 = rotate.rotate_ccw(np.pi, ((1, 0, 0), (0, 4, 0), (0, 0, 9)), 0)
    assert np.allclose(em, em2)
    em2 = rotate.rotate_ccw(np.pi / 2, em, 0)
    assert np.allclose(em2, np.array(((1, 0, 0), (0, 9, 0), (0, 0, 4))))
    ang = np.pi / 10
    s = np.sin(ang)  # .309
    c = np.cos(ang)  # .951
    em2 = rotate.rotate_ccw(ang, em, 0)
    vals, vecs = np.linalg.eig(em2)
    # sort by size of abs(vec) into ((val,uvec),...)
    eigens = tuple(
        (abs(val), np.sign(val) * uvec) for val, uvec in zip(vals, vecs.T, strict=False)
    )
    eigens = tuple(sorted(eigens, key=lambda x: x[0], reverse=True))
    # in size order, so z,y,x
    x = np.array([1, 0, 0])
    y = np.array([0, c, s])
    z = np.array([0, -s, c])
    # fmt: off
    assert eigens[0][1] == pytest.approx(z) or eigens[0][1] == pytest.approx(-z)
    assert eigens[1][1] == pytest.approx(y) or eigens[1][1] == pytest.approx(-y)
    assert eigens[2][1] == pytest.approx(x) or eigens[2][1]== pytest.approx(-x)
    # fmt: on


def test_rot_y():
    em = np.array(((1, 0, 0), (0, 4, 0), (0, 0, 9)))
    em2 = rotate.rotate_ccw(np.pi, em, 1)
    assert np.allclose(em, em2)
    em2 = rotate.rotate_ccw(np.pi, ((1, 0, 0), (0, 4, 0), (0, 0, 9)), 1)
    assert np.allclose(em, em2)
    em2 = rotate.rotate_ccw(np.pi / 2, em, 1)
    assert np.allclose(em2, np.array(((9, 0, 0), (0, 4, 0), (0, 0, 1))))
    ang = np.pi / 10
    s = np.sin(ang)  # .309
    c = np.cos(ang)  # .951
    em2 = rotate.rotate_ccw(ang, em, 1)
    vals, vecs = np.linalg.eig(em2)
    # sort by size of abs(vec) into ((val,uvec),...)
    eigens = tuple(
        (abs(val), np.sign(val) * uvec) for val, uvec in zip(vals, vecs.T, strict=False)
    )
    eigens = tuple(sorted(eigens, key=lambda x: x[0], reverse=True))
    # in size order, so z,y,x
    x = np.array([c, 0, -s])
    y = np.array([0, 1, 0])
    z = np.array([s, 0, c])
    # fmt: off
    assert eigens[0][1] == pytest.approx(z) or eigens[0][1] == pytest.approx(-z)
    assert eigens[1][1] == pytest.approx(y) or eigens[1][1] == pytest.approx(-y)
    assert eigens[2][1] == pytest.approx(x) or eigens[2][1] == pytest.approx(-x)
    # fmt: on


def test_rot_z():
    em = np.array(((1, 0, 0), (0, 4, 0), (0, 0, 9)))
    em2 = rotate.rotate_ccw(np.pi, em)
    assert np.allclose(em, em2)
    em2 = rotate.rotate_ccw(np.pi, ((1, 0, 0), (0, 4, 0), (0, 0, 9)))
    assert np.allclose(em, em2)
    em2 = rotate.rotate_ccw(np.pi / 2, em)
    assert np.allclose(em2, np.array(((4, 0, 0), (0, 1, 0), (0, 0, 9))))
    ang = np.pi / 10
    s = np.sin(ang)  # .309
    c = np.cos(ang)  # .951
    em2 = rotate.rotate_ccw(ang, em)
    vals, vecs = np.linalg.eig(em2)
    # sort by size of abs(vec) into ((val,uvec),...)
    eigens = tuple(
        (abs(val), np.sign(val) * uvec) for val, uvec in zip(vals, vecs.T, strict=False)
    )
    eigens = tuple(sorted(eigens, key=lambda x: x[0], reverse=True))
    # in size order, so z,y,x
    x = np.array([c, s, 0])
    y = np.array([-s, c, 0])
    z = np.array([0, 0, 1])
    # fmt: off
    assert eigens[0][1] == pytest.approx(z) or eigens[0][1] == pytest.approx(-z)
    assert eigens[1][1] == pytest.approx(y) or eigens[1][1] == pytest.approx(-y)
    assert eigens[2][1] == pytest.approx(x) or eigens[2][1] == pytest.approx(-x)
    # fmt: on


def test_rot_x_vec():
    vec = (1, 1, 1)
    assert np.allclose(rotate.rotate_ccw(np.pi / 2, vec, 0), (1, -1, 1))


def test_rot_y_vec():
    vec = (1, 1, 1)
    assert np.allclose(rotate.rotate_ccw(np.pi / 2, vec, 1), (1, 1, -1))


def test_rot_z_vec():
    vec = (1, 1, 1)
    assert np.allclose(rotate.rotate_ccw(np.pi / 2, vec, 2), (-1, 1, 1))


def test_rot_cw():
    em = np.array(((1, 0, 0), (0, 4, 0), (0, 0, 9)))
    for t in np.linspace(0, np.pi * 2, 18):
        assert np.allclose(rotate.rotate_ccw(t, em), rotate.rotate_cw(-t, em))


def test_errors():
    with pytest.raises(ValueError):
        rotate.rotate_cw(0.123, ((1, 2), (3, 4)), axis=1)
    with pytest.raises(ValueError):
        rotate.rotate_cw(0.123, ((1, 2, 3), (4, 5, 6), (7, 8, 9)), axis=-1)
    with pytest.raises(ValueError):
        rotate.rotate_cw(0.123, ((1, 2, 3), (4, 5, 6), (7, 8, 9)), axis=3)
    with pytest.raises(ValueError):
        rotate.rotate_cw(0.123, np.array(1), axis=3)
    with pytest.raises(ValueError):
        rotate.rotate_cw(
            0.123,
            (
                (((1, 2), (3, 4)), ((5, 6), (7, 8))),
                (((1, 2), (3, 4)), ((5, 6), (7, 8))),
            ),  # type: ignore  # noqa: PGH003
        )
