"""Find multiple convolved points with outlier detection."""

from typing import TYPE_CHECKING

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

from .smart_convolve import smart_convolve

if TYPE_CHECKING:
    from collections.abc import Callable, Generator, Sequence

    from gri_pos import LLA


def cluster_convolve(  # noqa: PLR0913
    ellipsoids: Sequence[Ell] | Generator[Ell],
    *,
    max_norm: float = 2,
    min_pts: int = 3,
    max_pts: int | None = None,
    min_sma_m: float = 0,
    max_ori_spread: bool = True,
    alt_post_process: Callable[[Sequence[tuple[LLA, list[LLA]]]], Sequence[float]]
    | None = None,
    alt_post_process_fudge_m: float = 0,
    # Can give a function that does the spline in utils, can take nearest, nearest+int,
    # spline, spline+int # spline would need scipy in utils ... do we want that? Probs.
    # alt can be handled by:
    #   dted: https://pypi.org/project/dted/
    #   elevation:  https://pypi.org/project/elevation/
    #   google: https://developers.google.com/maps/documentation/javascript/elevation
    #           https://airsci.com/google-maps-elevation-api-python-example/
) -> tuple[list[Ell], list[list[int]], list[int]]:
    """Find the convolved point location from all the points.

    Returns a list of Ell convolved locations, the list of Ells per position that were
    used by that location (such that return[1][1] is the list of Ells in return[1]), and
    a list of remaining / discarded Ells.

    Args:
        ellipsoids(Sequence[Ell]|Generator[Ell]): A set of ell objects. Uses xyz and
            info.
        max_norm(float, optional): Maximum mahalanobis distance from convolved point to
            allow. If an input it outside that range (input ell to point distance), add
            it to the discard list and return it. Defaults to 2.0
        min_pts(int, optional): Minimum allowed points in a convolved answer. Defaults
            to 3
        max_pts(int, optional): Maximum allowed points in a convolved answer. If there
            are leftover points (eg: 5 points passed, 3 max and 3 min, 2 left over that
            cannot be used), they will be added to discard. This takes more processing.
        min_sma_m(float, optional): Do not allow final convolved sma to go below min
            sma unless min_pts is reached. This will split up convolutions. If there are
            leftover points that cannot be used, they will be added to discard. This
            is iterative and takes a lot more processing.
        max_ori_spread(bool, optional): If set to true, for any clusters that need to be
            split up for min_sma or max_pts, first sort the data by orientation, then
            split into quarters and shuffle so that maximal orientation differences are
            seen by the convolver.
        alt_post_process(Callable[[Sequence[tuple[LLA,Sequence[LLA]]]], Sequence[float]], optional):
            A function that takes a list of (LLA,list[LLA]) where LLA is an np.ndarray
            vector that is latitude degrees, longitude degrees, altitude meters. The
            first element is the convolved location and the second element is the list
            of positions (Ells) as LLAs that went into that location.  Sequences are
            required to enable batch processing by the function.

            It should then output a list of altitudes in the same order as the input
            list.

            This package provides altitude.nearest, which replaces the Ell altitude with
            the nearest Ell altitude (often a decent estimate).
        alt_post_process_fudge_m(float, optional): Adds a set amount of meters to the
            altitude after any other alt_post_process function.

    Returns:
        (list[Ell], list[list[used Ells indexes]], list[discarded Ells indexes])
    """  # noqa: E501
    # To keep track of the index, make it a list of tuple[idx,ell]
    ells = [(idx, ell) for idx, ell in enumerate(ellipsoids)]
    if max_ori_spread:
        # sort by ori, split into quarters and shuffle
        ells = _max_spread_ori(ells)

    locations, loc_ells, discards = _first_pass(ells, max_norm, min_pts)

    # Now we have the collections that should go together (roughly), let's post process
    reprocess: list[int] = []
    for idx, (loc, loc_ell) in enumerate(zip(locations, loc_ells, strict=False)):
        if (
            max_pts is not None and len(loc_ell) > max_pts
        ) or loc.ellipse.sma_95 < min_sma_m:
            reprocess.append(idx)
    for idx in reprocess:
        new_locs, new_ells, new_discards = _reprocess(
            [ells[i] for i in loc_ells[idx]],
            max_norm,
            min_pts,
            max_pts,
            min_sma_m,
            max_ori_spread,
        )
        locations.extend(new_locs)
        loc_ells.extend(new_ells)
        discards.extend(new_discards)
    for idx in reprocess[::-1]:
        del locations[idx]
        del loc_ells[idx]

    # Adjust the altitude with a post process function if supplied and add any
    # fudge altitude distances
    if alt_post_process is not None or alt_post_process_fudge_m != 0:
        locations = _fix_alts(
            ells,
            locations,
            loc_ells,
            alt_post_process,
            alt_post_process_fudge_m,
        )

    # Sort the index lists
    discards.sort()
    loc_ells = [sorted(idxs) for idxs in loc_ells]
    # Now sort by min ell sma (smallest to largest)
    if len(locations) > 0:
        tups = sorted(
            zip(locations, loc_ells, strict=True),
            key=lambda x: x[0].ellipse.sma_95,
        )
        locations, loc_ells = zip(*tups, strict=True)
    return list(locations), list(loc_ells), discards


def _max_spread_ori(ells: list[tuple[int, Ell]]) -> list[tuple[int, Ell]]:
    """Spread the data into quarters by ori."""
    ells = sorted(ells, key=lambda x: x[1].ellipse.ori_deg)
    a = np.array([e[0] for e in ells], dtype=int)
    add = np.array([-1] * int(np.ceil(len(a) / 4) * 4 - len(a)), dtype=int)
    a = np.append(a, add)
    a = a.reshape((4, -1))
    a = a.T.flatten()
    a = a[a >= 0]
    return [ells[idx] for idx in a]


def _reprocess(  # noqa: PLR0913
    loc_ells: list[tuple[int, Ell]],
    max_norm: float = 2,
    min_pts: int = 3,
    max_pts: int | None = None,
    min_sma_m: float = 0,
    max_ori_spread: bool = False,  # noqa: FBT001, FBT002
) -> tuple[list[Ell], list[list[int]], list[int]]:
    """Find any indexes (loc,loc_ells) that we need to reprocess for max_pts or min_sma.

    Reprocess them in place. Remove from lists, add to tail end with new answers.
    """
    if max_ori_spread:
        # respread
        loc_ells = _max_spread_ori(loc_ells)

    if max_pts is None:
        max_pts = len(loc_ells)

    start = 0
    stop = max_pts
    ells = [ell for _, ell in loc_ells]
    new_locs: list[Ell] = []
    new_loc_ells: list[list[int]] = []
    new_discarded: list[int] = []
    while stop - start >= min_pts:
        ans = smart_convolve(
            ells[start:stop],
            max_norm=max_norm,
            min_pts=min_pts,
        )
        if ans is None:
            # odd case where we can't find an answer within max_pts but could with
            # the whole number of ells. Skip this batch
            start = stop
            stop = min(start + max_pts + 1, len(ells))
            new_discarded.extend(list(range(start, stop)))
            continue
        ell, used, discard = ans
        if min_sma_m > 0 and ell.ellipse.sma_95 < min_sma_m and stop - start > min_pts:
            stop -= 1
            continue
        # discards are unlikely in this case, just skip handling them
        new_discarded.extend([loc_ells[d + start][0] for d in discard])
        new_locs.append(ell)
        new_loc_ells.append([loc_ells[u + start][0] for u in used])
        start = stop
        stop = min(start + max_pts, len(ells))
    # Add any unused at the tail
    new_discarded.extend(range(start, stop))
    return new_locs, new_loc_ells, new_discarded


def _fix_alts(
    ells: list[tuple[int, Ell]],
    locations: list[Ell],
    loc_ells: list[list[int]],
    alt_post_process: Callable[[Sequence[tuple[LLA, list[LLA]]]], Sequence[float]]
    | None = None,
    alt_post_process_fudge_m: float = 0,
) -> list[Ell]:
    """Use the provided function to adjust altitudes, and add any fudge factor."""
    # Get base alts
    if alt_post_process is not None:
        # Convert to LLAs
        pos = [e.lla for e in locations]
        pos_lists = [[ells[idx][1].lla for idx in idx_list] for idx_list in loc_ells]
        # Get alts
        alts = alt_post_process(list(zip(pos, pos_lists, strict=True)))
    else:
        # Get existing alts
        alts = [e.lla.alt_m for e in locations]
    # Replace alts
    new_ell: list[Ell] = []
    for alt, ell in zip(alts, locations, strict=True):
        new_ell.append(
            Ell.COV(
                Pos.LLA(
                    ell.lla.lat_deg,
                    ell.lla.lon_deg,
                    alt + alt_post_process_fudge_m,
                ),
                ell.cov,
            ),
        )
    return new_ell


def _first_pass(
    ells: list[tuple[int, Ell]],
    max_norm: float = 2,
    min_pts: int = 3,
) -> tuple[list[Ell], list[list[int]], list[int]]:
    """First smart convolve pass.

    Get all convolved locations to cluster data regardless of max size or min sma.
    """
    locations: list[Ell] = []
    loc_ells: list[list[int]] = []
    # To keep track of the index, make it a list of tuple[idx,ell]
    # First pass
    while ans := smart_convolve(
        [ell for _, ell in ells],
        max_norm=max_norm,
        min_pts=min_pts,
    ):
        ell, used, discard = ans
        # Add location
        locations.append(ell)
        # Add the original indexes to the location ell
        loc_ells.append([e[0] for idx, e in enumerate(ells) if idx in used])
        # redo ells list to be the remainint tuples of unused ells
        ells = [e for idx, e in enumerate(ells) if idx in discard]

    return locations, loc_ells, [e[0] for e in ells]
