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

from gri_convolve import altitude, cluster_convolve


def test_max_pts():
    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, 3, 0), 10, 10, 0, alt_95_m=1)
    e4 = Ell.from_2d(Pos.XYZ(constants.ER_M, 4, 0), 10, 10, 0, alt_95_m=1)
    e5 = Ell.from_2d(Pos.XYZ(constants.ER_M, 5, 0), 10, 10, 0, alt_95_m=1)
    e6 = Ell.from_2d(Pos.XYZ(constants.ER_M, 6, 0), 10, 10, 0, alt_95_m=1)
    # Evenly divisible
    positions, pos_ells, discard = cluster_convolve(
        (e0, e1, e2, e3, e4, e5),
        min_pts=3,
        max_pts=3,
    )
    assert len(positions) == 2
    assert len(pos_ells) == 2
    assert len(pos_ells[0]) == 3
    assert len(pos_ells[1]) == 3
    assert len(discard) == 0
    # 1 left over
    positions, pos_ells, discard = cluster_convolve(
        (e0, e1, e2, e3, e4, e5, e6),
        min_pts=3,
        max_pts=3,
    )
    assert len(positions) == 2
    assert len(pos_ells) == 2
    assert len(pos_ells[0]) == 3
    assert len(pos_ells[1]) == 3
    assert len(discard) == 1


def test_min_sma():
    e0 = Ell.from_2d(Pos.XYZ(constants.ER_M, 0, 0), 100, 10, 0, alt_95_m=1)
    e1 = Ell.from_2d(Pos.XYZ(constants.ER_M, 1, 0), 100, 10, 0, alt_95_m=1)
    e2 = Ell.from_2d(Pos.XYZ(constants.ER_M, 2, 0), 100, 10, 0, alt_95_m=1)
    e3 = Ell.from_2d(Pos.XYZ(constants.ER_M, 3, 0), 100, 10, 0, alt_95_m=1)
    e4 = Ell.from_2d(Pos.XYZ(constants.ER_M, 4, 0), 100, 10, 0, alt_95_m=1)
    e5 = Ell.from_2d(Pos.XYZ(constants.ER_M, 5, 0), 100, 10, 0, alt_95_m=1)
    positions, pos_ells, discard = cluster_convolve(
        (e0, e1, e2, e3, e4, e5),
        min_pts=1,
        min_sma_m=90,
        max_ori_spread=False,
    )
    assert len(positions) == 6
    assert len(pos_ells) == 6
    assert len(discard) == 0
    positions, pos_ells, discard = cluster_convolve(
        (e0, e1, e2, e3, e4, e5),
        min_pts=1,
        min_sma_m=0.00001,
    )
    assert len(positions) == 1
    assert len(pos_ells) == 1
    assert len(pos_ells[0]) == 6
    assert len(discard) == 0


def test_spread():
    e0 = Ell.from_2d(Pos.XYZ(constants.ER_M, 0, 0), 100, 10, 10, alt_95_m=1)
    e1 = Ell.from_2d(Pos.XYZ(constants.ER_M, 1, 0), 100, 10, 20, alt_95_m=1)
    e2 = Ell.from_2d(Pos.XYZ(constants.ER_M, 2, 0), 100, 10, 30, alt_95_m=1)
    e3 = Ell.from_2d(Pos.XYZ(constants.ER_M, 3, 0), 100, 10, 40, alt_95_m=1)
    e4 = Ell.from_2d(Pos.XYZ(constants.ER_M, 2, 0), 100, 10, 50, alt_95_m=1)
    e5 = Ell.from_2d(Pos.XYZ(constants.ER_M, 3, 0), 100, 10, 60, alt_95_m=1)
    e6 = Ell.from_2d(Pos.XYZ(constants.ER_M, 2, 0), 100, 10, 70, alt_95_m=1)
    e7 = Ell.from_2d(Pos.XYZ(constants.ER_M, 3, 0), 100, 10, 80, alt_95_m=1)
    positions, pos_ells, _discard = cluster_convolve(
        (e0, e1, e2, e3, e4, e5, e6, e7),
        min_pts=2,
        max_pts=2,
        max_ori_spread=True,
    )
    assert len(positions) == 4
    assert abs(pos_ells[0][1] - pos_ells[0][0]) == 2
    assert (pos_ells[1][1] - pos_ells[1][0]) == 2


def test_alt_methods():
    e0 = Ell.from_2d(Pos.LLA(0, 0.00001, 0), 10, 10, 0, alt_95_m=10)
    e1 = Ell.from_2d(Pos.LLA(0, 0.00002, 1), 10, 10, 0, alt_95_m=10)
    e2 = Ell.from_2d(Pos.LLA(0, 0.00003, 3), 10, 10, 0, alt_95_m=10)
    positions, pos_ells, discard = cluster_convolve((e0, e1, e2))
    assert len(positions) == 1
    assert len(pos_ells) == 1
    assert len(pos_ells[0]) == 3
    assert len(discard) == 0
    pos = positions[0]
    assert pos.lla.alt_m > 1.1

    positions, pos_ells, discard = cluster_convolve(
        (e0, e1, e2),
        alt_post_process=altitude.nearest,
    )
    assert len(positions) == 1
    pos = positions[0]
    assert len(pos_ells) == 1
    assert len(pos_ells[0]) == 3
    assert len(discard) == 0
    assert pos.lla.alt_m == 1

    positions, pos_ells, discard = cluster_convolve(
        (e0, e1, e2),
        alt_post_process=altitude.nearest,
        alt_post_process_fudge_m=-5,
    )
    assert len(positions) == 1
    pos = positions[0]
    assert len(pos_ells) == 1
    assert len(pos_ells[0]) == 3
    assert len(discard) == 0
    assert pos.lla.alt_m == -4

    positions, pos_ells, discard = cluster_convolve(
        (e0, e1, e2),
        alt_post_process_fudge_m=-5,
    )
    assert len(positions) == 1
    pos = positions[0]
    assert len(pos_ells) == 1
    assert len(pos_ells[0]) == 3
    assert len(discard) == 0
    assert pos.lla.alt_m < 0


def test_no_answer():
    e0 = Ell.from_2d(Pos.LLA(0, 0.00001, 0), 10, 10, 0, alt_95_m=10)
    e1 = Ell.from_2d(Pos.LLA(0, 0.00002, 1), 10, 10, 0, alt_95_m=10)
    e2 = Ell.from_2d(Pos.LLA(0, 0.00003, 3), 10, 10, 0, alt_95_m=10)
    positions, pos_ells, discard = cluster_convolve((e0, e1, e2), min_pts=4)
    assert len(positions) == 0
    assert len(pos_ells) == 0
    assert len(discard) == 3
