import numpy as np
from gri_ell import Ell
from gri_pos import Pos
from gri_utils import constants

from gri_convolve import cluster_convolve
from gri_convolve.convolve.cluster_convolve import _max_spread_ori


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)
    positions, pos_ells, discard = cluster_convolve((e0, e1, e2, e3, e4))
    assert len(positions) == 1
    assert len(pos_ells) == 1
    assert len(pos_ells[0]) == 3
    assert positions[0] == Pos.XYZ(constants.ER_M, 1, 0)
    assert len(discard) == 2
    assert 3 in discard
    assert 4 in discard


def test_simple_multi():
    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)
    e5 = Ell.from_2d(Pos.XYZ(constants.ER_M, 102, 0), 10, 10, 0, alt_95_m=1)
    e6 = Ell.from_2d(Pos.XYZ(constants.ER_M, 201, 0), 10, 10, 0, alt_95_m=1)
    e7 = Ell.from_2d(Pos.XYZ(constants.ER_M, 202, 0), 10, 10, 0, alt_95_m=1)
    positions, pos_ells, discard = cluster_convolve(
        (e0, e1, e2, e3, e4, e5, e6, e7),
        max_ori_spread=False,
    )
    assert len(positions) == 2
    assert len(pos_ells) == 2
    assert len(pos_ells[0]) == 3
    assert len(pos_ells[1]) == 3
    answers = [Pos.XYZ(constants.ER_M, 1, 0), Pos.XYZ(constants.ER_M, 101, 0)]
    assert positions[0] != positions[1]
    assert positions[0] in answers
    assert positions[1] in answers
    assert len(discard) == 2
    assert 6 in discard
    assert 7 in discard


def test_breakpoint():
    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, 300, 0), 10, 10, 0, alt_95_m=1)
    e4 = Ell.from_2d(Pos.XYZ(constants.ER_M, 401, 0), 10, 10, 0, alt_95_m=1)
    e5 = Ell.from_2d(Pos.XYZ(constants.ER_M, 502, 0), 10, 10, 0, alt_95_m=1)
    e6 = Ell.from_2d(Pos.XYZ(constants.ER_M, 201, 0), 10, 10, 0, alt_95_m=1)
    e7 = Ell.from_2d(Pos.XYZ(constants.ER_M, 202, 0), 10, 10, 0, alt_95_m=1)
    e8 = Ell.from_2d(Pos.XYZ(constants.ER_M, 203, 0), 10, 10, 0, alt_95_m=1)
    positions, pos_ells, discard = cluster_convolve(
        (e0, e1, e2, e3, e4, e5, e6, e7, e8),
        max_ori_spread=False,
    )
    assert len(positions) == 2
    assert len(pos_ells) == 2
    assert len(pos_ells[0]) == 3
    assert len(pos_ells[1]) == 3
    answers = [Pos.XYZ(constants.ER_M, 1, 0), Pos.XYZ(constants.ER_M, 202, 0)]
    assert positions[0] != positions[1]
    assert positions[0] in answers, f"{positions[0]}"
    assert positions[1] in answers, f"{positions[1]}"
    assert len(discard) == 3
    assert 3 in discard
    assert 4 in discard
    assert 5 in discard


def test_sorting():
    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, 200, 0), 10, 10, 0, alt_95_m=1)
    e4 = Ell.from_2d(Pos.XYZ(constants.ER_M, 201, 0), 10, 10, 0, alt_95_m=1)
    e5 = Ell.from_2d(Pos.XYZ(constants.ER_M, 202, 0), 10, 10, 0, alt_95_m=1)
    e6 = Ell.from_2d(Pos.XYZ(constants.ER_M, 201, 0.1), 10, 10, 0, alt_95_m=1)
    e7 = Ell.from_2d(Pos.XYZ(constants.ER_M, 202, 0.1), 10, 10, 0, alt_95_m=1)
    e8 = Ell.from_2d(Pos.XYZ(constants.ER_M, 203, 0.1), 10, 10, 0, alt_95_m=1)
    positions, pos_ells, discard = cluster_convolve(
        (e0, e1, e2, e3, e4, e5, e6, e7, e8),
    )
    assert len(positions) == 2
    assert len(pos_ells) == 2
    assert len(discard) == 0
    assert positions[0].ellipse.sma_95 < positions[1].ellipse.sma_95
    assert len(pos_ells[0]) > len(pos_ells[1])
    assert all(e in pos_ells[0] for e in [3, 4, 5, 6, 7, 8])
    assert all(e in pos_ells[1] for e in [0, 1, 2])


def test_max_spread():
    ells: list[tuple[int, Ell]] = []
    for idx, ori in enumerate(np.linspace(0, 179, 12)):
        ells.append((idx, Ell.from_2d(Pos.LLA(0, 0, 0), 10, 1, float(ori))))
    idxs = [e[0] for e in ells]
    assert idxs == sorted(idxs)
    oris = [float(e[1].ellipse.ori_deg) for e in ells]
    assert oris == sorted(oris)

    e2 = _max_spread_ori(ells)
    idxs = [e[0] for e in e2]
    assert idxs == [0, 3, 6, 9, 1, 4, 7, 10, 2, 5, 8, 11]

    e2 = _max_spread_ori(ells[:-1])
    idxs = [e[0] for e in e2]
    assert idxs == [0, 3, 6, 9, 1, 4, 7, 10, 2, 5, 8]

    e2 = _max_spread_ori(ells[:-2])
    idxs = [e[0] for e in e2]
    assert idxs == [0, 3, 6, 9, 1, 4, 7, 2, 5, 8]

    e2 = _max_spread_ori(ells[:-3])
    idxs = [e[0] for e in e2]
    assert idxs == [0, 3, 6, 1, 4, 7, 2, 5, 8]

    e2 = _max_spread_ori(ells[:-4])
    idxs = [e[0] for e in e2]
    assert idxs == [0, 2, 4, 6, 1, 3, 5, 7]

    # for 1 to 4 points, just ensure it doesn't error
    for ii in range(1, 4):
        e2 = _max_spread_ori(ells[:ii])
