"""Grid-based covariance functions."""

from __future__ import annotations

import math

import numpy as np


def get_grids(r: int, num: int | None = None) -> list[tuple[np.ndarray, str]]:
    """Get x,y points for plotting gridlines.

    Args:
        r (int): The max (and -r for min) of the axis, eg: r=10 means -10 to 10
        num (int): The number of points per axis line to plot, defaults to r*2+1

    Returns:
        a list of gridlines (as a list of x,y pairs, length num) coupled with a color
        string
    """
    if num is None:
        num = r * 2 + 1
    # List of 2xN of x,y pairs
    grids: list[tuple[np.ndarray, str]] = []
    # special case, center axis
    grids.append((np.linspace((-r, 0), (r, 0), num), "k"))
    grids.append((np.linspace((0, -r), (0, r), num), "k"))
    for p in range(-r, r + 1):
        if p == 0:
            continue
        # horiz
        grids.append((np.linspace((-r, p), (r, p), num), "b"))
        grids.append((np.linspace((p, -r), (p, r), num), "r"))
    return grids


# TODO check how covariance works with eigenvalue decomposition ... if the larger of the
# two (sma) is in the x2 position, then our sma/smi/ori may need to have another 90
# degrees added if the assumption is that ori is sma degrees (clockwise?anticlockwise?)
# from north
def covariance(r1: float, r2: float, ori_rad: float) -> np.ndarray:
    """Get the covariance matrix.

    Args:
        r1: Largest radius (aligned north / up)
        r2: Smallest radius
        ori_rad: Radians counterclockwise from north to rotate
    """
    cos = math.cos(ori_rad)
    sin = math.sin(ori_rad)
    x1 = r1 * r1 * cos * cos + r2 * r2 * sin * sin
    x2 = r1 * r1 * sin * sin + r2 * r2 * cos * cos
    c12 = (r1 * r1 - r2 * r2) * sin * cos
    return np.array([[x1, c12], [c12, x2]])
    # make sure it's correct ... eigenvals are the radii, eighenvecs are the rotation
    # (eigenvals, eigenvecs) = np.linalg.eig(cov)
    # assert abs(eigenvals[0] ** 0.5 - r1) < 0.000001
    # assert abs(eigenvals[1] ** 0.5 - r2) < 0.000001
    # assert abs(math.acos(eigenvecs[0][0]) - ori_rad) < 0.00001


def normalize_grid(
    cov: np.ndarray,
    grids: list[tuple[np.ndarray, str]],
) -> list[tuple[np.ndarray, str]]:
    """Normalize grids to the covariance.

    Args:
        cov (np.ndarray): The covariance matrix
        grids (list[tuple[np.ndarray, str]]): The gridlines in (list[[x,y]], color) form

    Returns:
        copy of grids with each point moved to the Malanobis distance of that point
    """
    invcov = np.linalg.inv(cov)
    grids_norm: list[tuple[np.ndarray, str]] = []
    for grid, c in grids:
        grid_copy = grid.copy()
        for xy in grid_copy:
            # TODO possible to do this outside loop with broadcasting?
            d = (xy @ invcov @ xy.T) ** 0.5
            norm = np.linalg.norm(xy)
            if norm == 0:
                continue
            scale = d / np.linalg.norm(xy)
            xy *= scale  # noqa: PLW2901
        grids_norm.append((grid_copy, c))
    return grids_norm
