# ruff: noqa: T201
"""Test methods to compare various convolution methods. Not for production use."""

# NOTE: Test functions ... this should not be used. Use convolve instead.

from typing import TYPE_CHECKING

import numpy as np
from gri_ell import Ell
from gri_pos import Pos
from gri_utils import conversion
from scipy.linalg import block_diag

if TYPE_CHECKING:
    from collections.abc import Generator, Sequence


def atwa_convolve(ells: Sequence[Ell] | Generator[Ell]) -> Ell:
    """Compare various convolve types.

    Args:
        ells(Sequence[Ell]): A set of ell objects. Uses xyz and info.

    Returns:
        Ell
    """
    np.set_printoptions(suppress=True, precision=5)
    ell = _position(ells)
    print("\nATWAI_ENU")
    # print(np.array(ell.cov.semi_axis_vec))
    print(ell.cov)
    print(ell.ellipse)

    ell2 = _std_cov(ell, ells)
    print("\nSTD COV")
    # print(np.array(ell2.cov.semi_axis_vec))
    print(ell2.cov)
    print(ell2.ellipse)

    ell3 = _bart_cov(ell, ells)
    print("\nBART COV")
    # print(np.array(ell2.cov.semi_axis_vec))
    print(ell3.cov)
    print(ell3.ellipse)

    ell4 = _bart_cov2(ell, ells)
    print("\nBART COV 2")
    # print(np.array(ell2.cov.semi_axis_vec))
    print(ell4.cov)
    print(ell4.ellipse)

    return ell3


def _position(ells: Sequence[Ell] | Generator[Ell]) -> Ell:
    """Get the position of the convolved point with direct covariance."""
    # 3Nx3 identity stack
    a = np.vstack([np.identity(3) for _ in ells])
    # 3Nx3N diagonalized information matrices
    w = block_diag(*(ell.info for ell in ells))
    # 3Nx1 xyz positions
    y = np.concatenate([e.xyz for e in ells])
    atw = a.T @ w  # 3x3N
    atwa = atw @ a  # 3x3
    atwy = atw @ y  # 3x1
    atwai = np.linalg.inv(atwa)  # 3x3
    xyz = atwai @ atwy  # 3x1
    atwai = conversion.xyz_to_enu_cov(xyz, atwai)
    return Ell.COV(Pos.XYZ(xyz), atwai)


def _std_cov(pos: Pos, ellipsoids: Sequence[Ell] | Generator[Ell]) -> Ell:
    """Get the covariance the standard way."""
    norm = 0
    s = np.zeros((3, 3), dtype=float)
    s_model = np.zeros((3, 3), dtype=float)
    n = 0
    for e in ellipsoids:
        n += 1
        w = e.info  # xyz, 1/m^2, 1 sigma
        s_model += w
        del_xyz = pos.delta_xyz(e)
        norm += float(del_xyz @ w @ del_xyz)
        s += np.outer(del_xyz, del_xyz)
    s_model = np.linalg.inv(s_model)
    x_2 = norm
    s_mod_inflated = (x_2 + 1) / n * s_model
    s_sample = s / n**2
    cov = s_mod_inflated + s_sample
    cov = conversion.xyz_to_enu_cov(pos.xyz, cov)
    return Ell.COV(pos, cov)


def _bart_cov(pos: Pos, ellipsoids: Sequence[Ell] | Generator[Ell]) -> Ell:
    """Get the sample_smi variation of covariance."""
    norm = 0
    s = np.zeros((3, 3), dtype=float)
    s_model = np.zeros((3, 3), dtype=float)
    n = 0
    for e in ellipsoids:
        n += 1
        w = e.info  # xyz, 1/m^2, 1 sigma
        s_model += w
        del_xyz = pos.delta_xyz(e)
        norm += float(del_xyz @ w @ del_xyz)
        # Bart's correction is to s
        theta = e.ellipse.ori_rad
        del_n1 = (pos.lla.lat_deg - e.lla.lat_deg) * 60 * 1.852 * 1000
        avg_lat = (pos.lla.lat_deg + e.lla.lat_deg) / 2
        del_e1 = (
            (pos.lla.lon_deg - e.lla.lon_deg)
            * 60
            * 1.852
            * np.cos(np.radians(avg_lat))
            * 1000
        )
        del_smi = del_e1 * np.cos(theta) - del_n1 * np.sin(theta)
        del_e = del_smi * np.cos(theta)
        del_n = -del_smi * np.sin(theta)
        enu_vec = np.array((del_e, del_n, pos.lla.alt_m - e.lla.alt_m))
        s += np.outer(enu_vec, enu_vec)
        # End bart's correction
    s_model = np.linalg.inv(s_model)
    x_2 = norm
    s_mod_inflated = (x_2 + 1) / n * s_model
    s_mod_inflated = conversion.xyz_to_enu_cov(pos.xyz, s_mod_inflated)
    s_sample = s / n**2
    cov = s_mod_inflated + s_sample
    return Ell.COV(pos, cov)


def _bart_cov2(pos: Pos, ellipsoids: Sequence[Ell] | Generator[Ell]) -> Ell:
    """Original implementations."""
    # s_model_inflated section
    n = len(list(ellipsoids))
    # chi squared
    chi_squared = 0
    for e in ellipsoids:
        w = e.info  # xyz, 1/m^2, 95%
        del_xyz = pos.delta_xyz(e)
        chi_squared += float(del_xyz @ w @ del_xyz)

    # inflation factor
    inflation_factor = (chi_squared + 1) / n

    # s model
    s_model = np.zeros((3, 3))
    for e in ellipsoids:
        s_model += e.info
    s_model = np.linalg.inv(s_model)

    # inflated model
    s_model_inflated = inflation_factor * s_model

    # rotate from ecef to enu
    s_model_inflated = conversion.xyz_to_enu_cov(pos.xyz, s_model_inflated)

    # s_sample_minor
    delta_enus = []
    for e in ellipsoids:
        orient_rad = e.ellipse.ori_rad

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

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

        delta_e = delta_smi * np.cos(orient_rad)
        delta_n = -delta_smi * np.sin(orient_rad)
        delta_up = pos.lla.alt_m - e.lla.alt_m

        delta_enus.append(np.array([[delta_e], [delta_n], [delta_up]]))

    s_sample_minor = 0
    for enu in delta_enus:
        s_sample_minor += enu @ enu.T
    s_sample_minor /= n**2

    cov = s_model_inflated + s_sample_minor
    return Ell.COV(pos, cov)
