"""Convolution of all points.

Find the center of the points via convolution and find the error region via different
inflation methods.
"""

from typing import TYPE_CHECKING, Literal

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

if TYPE_CHECKING:
    from collections.abc import Sequence


def convolve(
    ellipsoids: Sequence[Ell],
    *,
    inflation: Literal["none", "std", "bart"] = "bart",
) -> Ell:
    """Find the convolved point location from all the points.

    Does not throw away any points. Choose an inflation technique to estimate the error
    region.

    Args:
        ellipsoids: A set of ell objects. Uses xyz and info.
        inflation: "none" performs strict covariance calculations by combining
            information matrices. "std" inflates covariance based on the sample
            distribution, "bart" inflates covariances based on SMI reduction rules.

    Returns:
        Ell
    """
    # Find the position and basic uninflated covariance
    s_model = np.add.reduce([e.info for e in ellipsoids])  # xyz, 1/m^2, 1-sigma
    s_var = np.add.reduce([e.info @ e.xyz for e in ellipsoids])
    cov = np.linalg.inv(s_model)
    pos = Pos.XYZ(cov @ s_var)
    ell = Ell.COV(pos, cov)
    if inflation == "none":
        return ell

    # Begin normalized and smi-normalized inflation of covariance
    s_model = cov
    n = len(ellipsoids)
    x_2 = 0
    s = np.zeros((3, 3), dtype=float)
    for e in ellipsoids:
        del_xyz = pos.delta_xyz(e)
        x_2 += float(del_xyz @ e.info @ del_xyz)
        if inflation == "std":
            s += np.outer(del_xyz, del_xyz)
        elif inflation == "bart":
            theta = e.ellipse.ori_rad

            delta_north = 1.852 * 60 * 1000 * (pos.lla.lat_deg - e.lla.lat_deg)
            delta_east = (
                1.852
                * 60
                * 1000
                * (pos.lla.lon_deg - e.lla.lon_deg)
                * np.cos(np.radians((pos.lla.lat_deg + e.lla.lat_deg) / 2))
            )

            delta_smi = delta_east * np.cos(theta) - delta_north * np.sin(theta)

            del_e = delta_smi * np.cos(theta)
            del_n = -delta_smi * np.sin(theta)
            del_u = pos.lla.alt_m - e.lla.alt_m

            enu_vec = np.array((del_e, del_n, del_u))

            s += np.outer(enu_vec, enu_vec)
        else:
            raise ValueError("Bad inflation setting")
    s_sample = s / n**2
    s_mod_inflated = (x_2 + 1) / n * s_model
    # s_mod_inflated is in xyz
    # std s_sample is in xyz, bart s_sample is in enu
    if inflation == "std":
        cov = conversion.xyz_to_enu_cov(pos.xyz, s_mod_inflated + s_sample)
    else:
        s_mod_inflated = conversion.xyz_to_enu_cov(pos.xyz, s_mod_inflated)
        cov = s_mod_inflated + s_sample
    return Ell.COV(pos, cov)
