import numpy as np
import pytest

from gri_utils.constants import stats
from gri_utils.ellipsoids import (
    cov_projection,
    cov_slice,
    cov_to_alt_param,
    cov_to_params,
    cov_to_params_95,
    params_95_to_cov,
    params_to_cov,
)
from gri_utils.matrices import rotate_cw

ARR_2D = np.array(((4, 0), (0, 9)))
ARR_3D = np.array(((4, 0, 0), (0, 9, 0), (0, 0, 25)))
SEQ_2D = ((4, 0), (0, 9))
SEQ_3D = ((4, 0, 0), (0, 9, 0), (0, 0, 25))


def test_cov_to_params():
    ans = cov_to_params(ARR_2D)
    assert np.allclose(ans, (3, 2, 0))
    ans = cov_to_params(SEQ_2D)
    assert np.allclose(ans, (3, 2, 0))
    for ori in range(0, 180, 10):
        arr2 = rotate_cw(np.radians(ori), ARR_2D)
        ans = cov_to_params(arr2)
        assert np.allclose(ans, (3, 2, ori))
        arr2 = rotate_cw(np.radians(ori), SEQ_2D)
        ans = cov_to_params(arr2)
        assert np.allclose(ans, (3, 2, ori))

    with pytest.raises(ValueError):
        cov_to_params(((1, 2, 3), (4, 5, 6)))


def test_params_to_cov_2d():
    arr = params_to_cov(3, 2, 0)
    assert np.allclose(arr, ARR_2D)
    for ori in range(0, 180, 10):
        arr = params_to_cov(3, 2, ori)
        assert np.allclose(arr, rotate_cw(np.radians(ori), ARR_2D))


def test_params_to_cov_3d():
    arr = params_to_cov(3, 2, 0, alt=5)
    assert np.allclose(arr, ARR_3D)
    for ori in range(0, 180, 10):
        arr = params_to_cov(3, 2, ori, alt=5)
        assert np.allclose(arr, rotate_cw(np.radians(ori), ARR_3D))


def test_cov_to_alt_param():
    alt = cov_to_alt_param(ARR_3D)
    assert alt == pytest.approx(5)
    alt = cov_to_alt_param(SEQ_3D)
    assert alt == pytest.approx(5)

    with pytest.raises(ValueError):
        cov_to_alt_param(ARR_2D)


def test_cov_projection():
    ans = cov_projection(ARR_3D)
    assert np.allclose(ans, ARR_2D)


def test_cov_slice_basic():
    arr, center = cov_slice(ARR_3D)
    assert np.allclose(arr, ARR_2D)
    assert np.allclose(center, (0, 0, 0))
    arr, center = cov_slice(SEQ_3D)
    assert np.allclose(arr, ARR_2D)
    assert np.allclose(center, (0, 0, 0))
    arr, center = cov_slice(ARR_3D, 0)
    assert np.allclose(arr, ARR_2D)
    assert np.allclose(center, (0, 0, 0))
    arr, center = cov_slice(ARR_3D, 0, (1, 2, 3))
    assert np.allclose(arr, ARR_2D)
    assert np.allclose(center, (1, 2, 3))

    with pytest.raises(ValueError):
        cov_slice(ARR_2D)

    with pytest.raises(ValueError):
        cov_slice(ARR_3D, 5)
    cov_slice(ARR_3D, 3, (0, 1, 6))
    with pytest.raises(ValueError):
        cov_slice(ARR_3D, 6, (0, 1, 6))


def test_cov_slice_heights():
    # moves vertically
    ell = ((4, 0, 0), (0, 4, 0), (0, 0, 100))
    arr, center = cov_slice(ell, 9)
    assert np.allclose(center, (0, 0, 9))
    assert arr[0, 0] == pytest.approx(arr[1, 1])
    assert arr[0, 0] < 2

    # rotate a tall ellipse around the y axis
    ell = rotate_cw(np.radians(-45), ((4, 0, 0), (0, 4, 0), (0, 0, 100)), 1)
    arr, center = cov_slice(ell)
    assert np.allclose(center, (0, 0, 0))
    assert arr[0, 0] > 2
    assert arr[1, 1] == pytest.approx(4)

    with pytest.raises(ValueError):
        cov_slice(ell, 8)


def test_cov_slice2():
    ell = rotate_cw(np.radians(-45), ((4, 0, 0), (0, 4, 0), (0, 0, 100)), 1)
    # top inside
    arr, center = cov_slice(ell, 7)
    assert center[2] == pytest.approx(7)
    assert center[1] == pytest.approx(0)
    assert 6 < center[0] < 7
    assert arr[0, 0] > arr[1, 1]
    assert arr[1, 1] < 1


def test_params_95_to_cov():
    cov = params_95_to_cov(5, 4, 0, 3)
    sma = 5 / stats.SIG_TO_95_2D**0.5
    smi = 4 / stats.SIG_TO_95_2D**0.5
    alt = 3 / stats.SIG_TO_95_1D**0.5
    assert np.allclose(cov, [[smi**2, 0, 0], [0, sma**2, 0], [0, 0, alt**2]])


def test_cov_to_params_95():
    sma = 5 / stats.SIG_TO_95_2D**0.5
    smi = 4 / stats.SIG_TO_95_2D**0.5
    alt = 3 / stats.SIG_TO_95_1D**0.5
    sma_95, smi_95, ori_deg, alt_95 = cov_to_params_95(
        [[smi**2, 0, 0], [0, sma**2, 0], [0, 0, alt**2]],
    )
    assert sma_95 == pytest.approx(5)
    assert smi_95 == pytest.approx(4)
    assert alt_95 == pytest.approx(3)
    assert ori_deg == pytest.approx(0)
    # 2D
    sma_95, smi_95, ori_deg, alt_95 = cov_to_params_95(
        [[smi**2, 0], [0, sma**2]],
    )
    assert sma_95 == pytest.approx(5)
    assert smi_95 == pytest.approx(4)
    assert alt_95 == pytest.approx(0)
    assert ori_deg == pytest.approx(0)
    # Err
    with pytest.raises(ValueError):
        sma_95, smi_95, ori_deg, alt_95 = cov_to_params_95([[1], [1]])
