"""Plot the covariance, normalized plot, and windage plots."""

from __future__ import annotations

import math
from typing import TYPE_CHECKING

import matplotlib.axes as mpl_axes
import matplotlib.patches as mpl_patches

if TYPE_CHECKING:
    import numpy as np


def plot_covariance(
    r1: float,
    r2: float,
    ori_rad: float,
    grids: list[tuple[np.ndarray, str]],
    ax: mpl_axes.Axes,
) -> list:
    """Plot the covariance matrix.

    Plots the matrix as an ellipse on the flat gridlines

    Return the artists (plots)
    """
    artists: list = []
    ax.set_xlim(-10, 10)
    ax.set_ylim(-10, 10)
    ax.set_title("Covariance")
    # plot gridlines
    for grid, c in grids:
        x, y = zip(*grid, strict=False)
        artists.extend(ax.plot(x, y, c=c, lw=1))
    # add ellipse
    ell = mpl_patches.Ellipse(
        (0, 0),
        width=r1 * 2,
        height=r2 * 2,
        angle=math.degrees(ori_rad),
        ec="k",
        fc="none",
        lw=3,
    )
    artists.append(ax.add_artist(ell))
    return artists


def plot_normalized(
    grids_norm: list[tuple[np.ndarray, str]],
    ax: mpl_axes.Axes,
) -> list:
    """Plot the normalized covariance and lines.

    Plot the normalized gridlines

    Return the artists (plots)
    """
    artists: list = []
    ax.set_xlim(-10, 10)
    ax.set_ylim(-10, 10)
    ax.set_title("Normalized")
    # plot gridlines
    for grid, c in grids_norm:
        x, y = zip(*grid, strict=False)
        artists.extend(ax.plot(x, y, c=c, lw=1))
    # add ellipse
    ell = mpl_patches.Ellipse(
        (0, 0),
        width=2,
        height=2,
        angle=0,
        ec="k",
        fc="none",
        lw=3,
    )
    artists.append(ax.add_artist(ell))
    return artists


def plot_wind(
    grids: list[tuple[np.ndarray, str]],
    grids_norm: list[tuple[np.ndarray, str]],
    ax: mpl_axes.Axes,
) -> list:
    """Plot the direction vectors of the movement of each point on the grid.

    Return the artists (plots)
    """
    artists: list = []
    ax.set_xlim(-10, 10)
    ax.set_ylim(-10, 10)
    ax.set_title("Wind")
    # plot gridlines
    for (grid, _), (grid_norm, _) in zip(grids, grids_norm, strict=False):
        for (x1, y1), (x2, y2) in zip(grid, grid_norm, strict=False):
            artists.extend(ax.plot((x1, x2), (y1, y2), c="k", lw=1))
    # add ellipse
    ell = mpl_patches.Ellipse(
        (0, 0),
        width=2,
        height=2,
        angle=0,
        ec="k",
        fc="none",
        lw=3,
    )
    artists.append(ax.add_artist(ell))
    return artists
