from gri_ell import Ell
from gri_pos import Pos
from gri_utils import constants

from gri_convolve import smart_convolve


def test_simple():
    e0 = Ell.from_2d(Pos.XYZ(constants.ER_M, 0, 0), 10, 10, 0, alt_95_m=1)
    e1 = Ell.from_2d(Pos.XYZ(constants.ER_M, 1, 0), 10, 10, 0, alt_95_m=1)
    e2 = Ell.from_2d(Pos.XYZ(constants.ER_M, 2, 0), 10, 10, 0, alt_95_m=1)
    e3 = Ell.from_2d(Pos.XYZ(constants.ER_M, 100, 0), 10, 10, 0, alt_95_m=1)
    e4 = Ell.from_2d(Pos.XYZ(constants.ER_M, 101, 0), 10, 10, 0, alt_95_m=1)
    ans = smart_convolve((e0, e1, e2, e3, e4))
    assert ans is not None
    ell, ells, discard = ans
    assert len(ells) == 3
    assert 0 in ells
    assert 1 in ells
    assert 2 in ells
    assert len(discard) == 2
    assert 3 in discard
    assert 4 in discard
    xyz = Pos.XYZ(constants.ER_M, 1, 0).xyz
    assert ell.xyz == xyz, ell.xyz.arr() - xyz.arr()
    # TODO cov


def test_discards():
    e0 = Ell.from_2d(Pos.XYZ(constants.ER_M, 0, 202), 10, 10, 0, alt_95_m=1)
    e1 = Ell.from_2d(Pos.XYZ(constants.ER_M, 0, 201), 10, 10, 0, alt_95_m=1)
    e2 = Ell.from_2d(Pos.XYZ(constants.ER_M, 0, 0), 10, 10, 0, alt_95_m=1)
    e3 = Ell.from_2d(Pos.XYZ(constants.ER_M, 1, 0), 10, 10, 0, alt_95_m=1)
    e4 = Ell.from_2d(Pos.XYZ(constants.ER_M, 2, 0), 10, 10, 0, alt_95_m=1)
    e5 = Ell.from_2d(Pos.XYZ(constants.ER_M, 100, 0), 10, 10, 0, alt_95_m=1)
    e6 = Ell.from_2d(Pos.XYZ(constants.ER_M, 101, 0), 10, 10, 0, alt_95_m=1)
    ans = smart_convolve((e0, e1, e2, e3, e4, e5, e6))
    assert ans is not None
    ell, ells, discard = ans
    assert ell.xyz == Pos.XYZ(constants.ER_M, 1, 0).xyz
    # TODO cov
    assert len(ells) == 3
    assert 2 in ells
    assert 3 in ells
    assert 4 in ells
    assert len(discard) == 4
    assert 0 in discard
    assert 1 in discard
    assert 5 in discard
    assert 6 in discard


def test_min():
    e0 = Ell.from_2d(Pos.XYZ(constants.ER_M, 0, 0), 10, 10, 0, alt_95_m=1)
    e1 = Ell.from_2d(Pos.XYZ(constants.ER_M, 1, 0), 10, 10, 0, alt_95_m=1)
    ans = smart_convolve((e0, e1))
    assert ans is None


def test_min2():
    e0 = Ell.from_2d(Pos.XYZ(constants.ER_M, 0, 0), 10, 10, 0, alt_95_m=1)
    e1 = Ell.from_2d(Pos.XYZ(constants.ER_M, 1, 0), 10, 10, 0, alt_95_m=1)
    e2 = Ell.from_2d(Pos.XYZ(constants.ER_M, 100, 0), 10, 10, 0, alt_95_m=1)
    e3 = Ell.from_2d(Pos.XYZ(constants.ER_M, 101, 0), 10, 10, 0, alt_95_m=1)
    ans = smart_convolve((e0, e1, e2, e3))
    assert ans is None


def test_bounds():
    # Check that if we give it three points outside distance, it drops them all
    e0 = Ell.from_2d(Pos.XYZ(constants.ER_M, 0, 0), 10, 10, 0, alt_95_m=1)
    e1 = Ell.from_2d(Pos.XYZ(constants.ER_M, 100, 0), 10, 10, 0, alt_95_m=1)
    e2 = Ell.from_2d(Pos.XYZ(constants.ER_M, 200, 0), 10, 10, 0, alt_95_m=1)
    e3 = Ell.from_2d(Pos.XYZ(constants.ER_M, 400, 0), 10, 10, 0, alt_95_m=1)
    ans = smart_convolve((e0, e1, e2, e3))
    assert ans is None


DATA = """Lat,      Lon,      Alt,   SMA,   SMI, ORI
19.99690, 24.99855,   0,  0.315, 0.138,  23.727
19.99615, 25.00442, 450,  0.546, 0.507, 132.829
20.00023, 25.00200, 460,  0.515, 0.063,  83.023
19.99784, 25.00409, 470,  0.586, 0.329, 119.336
20.00195, 24.99712, 480,  4.166, 3.002, 125.774
19.99949, 25.00333, 500,  0.507, 0.083,  99.257
19.99603, 25.00319, 500,  0.448, 0.339, 142.945
20.00402, 24.99663, 500, 11.378, 8.963, 141.771
19.99659, 25.00266,  50,  0.701, 0.514, 143.758
"""


def read_data() -> list[list[float]]:
    data = [
        [float(s.strip()) for s in line.split(",")]
        for line in DATA.split("\n")[1:]
        if line.strip()
    ]
    return data  # noqa: RET504


def test_data():
    ells: list[Ell] = [
        Ell.from_2d(
            Pos.LLA(t[0], t[1], t[2]),
            t[3] * 1000,
            t[4] * 1000,
            t[5],
            alt_95_m=50 * constants.SIG_TO_95_1D**0.5,
        )
        for t in read_data()
    ]
    ans = smart_convolve(ells)
    assert ans is not None
    ell, used, discard = ans
    assert ell is not None
    assert len(discard) == 2
    assert 0 in discard
    assert 8 in discard
    assert len(used) == 7
    ell_data = Pos.LLA(
        19.99991540108,
        25.002065574735,
        480.01480168477,
        tolerance=(1e-6, 1e-6, 1e-3),
    )
    assert ell.lla == ell_data.lla, (
        f"\n{ell.lla.arr()}\n{ell_data.lla.arr()}\n{ell.lla.arr() - ell_data.lla.arr()}"
    )
