"""Find convolved point with outlier detection."""

from typing import TYPE_CHECKING

import numpy as np
from gri_utils import constants

from .convolve import convolve

if TYPE_CHECKING:
    from collections.abc import Generator, Sequence

    from gri_ell import Ell


def smart_convolve(
    ellipsoids: Sequence[Ell] | Generator[Ell],
    *,
    max_norm: float = 2,
    min_pts: int = 3,
) -> None | tuple[Ell, list[int], list[int]]:
    """Find the convolved point location from all the points.

    Returns an Ell (or None), the list of Ells (by index) used to make that
    ellipsoid, and the remaining Ells (by index) that were excluded.

    Useful to tell all the points that should go together (by norm) and remove the
    points that do not belong.

    Note that outlier convolve can (incorrectly) return no answer when an answer does
    exist if the data is not pre-clustered decently. EG: a valid cluster of five points
    and a invalid group of twenty points further away that does not convolve will
    exclude the five points first by norm, then exclude each other. Roughly pre-cluster
    your data before using this function.

    Args:
        ellipsoids(Sequence[Ell]|Generator[Ell]): A set of ell objects. Uses xyz and
            info.
        max_norm(float, optional): Maximum ellipse norm 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

    Returns:
        tuple[Ell, list[used Ells idx], list[discarded Ells idx]] or None
    """
    # To keep track of the index, make it a list of tuple[idx,ell]
    ells = [(idx, ell) for idx, ell in enumerate(ellipsoids)]
    # Verify length is minimal
    if len(ells) < min_pts:
        return None
    unused: list[int] = []
    # loop as long as we have enough ells until none are discarded or all are discarded.
    while len(ells) >= min_pts:
        ell, discard_idx = _test_norm([ell for _, ell in ells], max_norm)
        if discard_idx is None:
            # Answer found!
            return ell, [idx for idx, _ in ells], unused
        # Remove index from ells and add to unused
        unused.append(ells[discard_idx][0])
        del ells[discard_idx]
    # Not enough left and no answer found
    return None


def _test_norm(ells: list[Ell], max_norm: float) -> tuple[Ell, int | None]:
    """Get a discarded index at max norm distance or None if all are contained."""
    ell = convolve(ells, inflation="none")
    norms = [e.dist_mahalanobis(ell) / constants.SIG_TO_95_3D**0.5 for e in ells]
    idx = int(np.argmax(norms))
    return ell, idx if norms[idx] > max_norm else None
