# ruff: noqa: T201
"""Plot the normalized space of the covariance matrix.

Run with:
python -m gri_sandbox.mahalanobis_curvature.main

Show how distances change when using covariance as a way both to visualize Mahalnobis
distance and covariances

Can be called with "--animate" to create a video animation of normalized space, or
without to display a single plot.

Pass in SMA, SMI, and ORI or by default it will use 3, 1, 15. If --animate is called,
ORI is ignored.

Animation is saved as "mahalanobis.mp4"
"""

from __future__ import annotations

import argparse
import math
import sys
from typing import TYPE_CHECKING

import matplotlib.animation as mpl_animation
import matplotlib.pyplot as plt

from gri_sandbox.mahalanobis_curvature import data, plots

if TYPE_CHECKING:
    from collections.abc import Generator

    import numpy as np


def _main() -> None:
    _get_args()


def make_image(sma: float, smi: float, ori: float) -> None:
    """Create an image with the covariance, normalized space, and windage plot.

    Args:
        sma(float): Semi Major Axis Length
        smi(float): Semi Minor Axis Length
        ori(float): Orientation in degrees clockwise (sma on y axis)
    """
    print(f"Making Image: SMA:{sma:.1f},  SMI:{smi:.1f}, ORI:{ori:.1f}")
    ori_rad = math.radians(ori)
    cov = data.covariance(sma, smi, ori_rad)

    plt.rcParams["figure.figsize"] = 18, 5
    _, axs = plt.subplots(1, 3)

    grids = data.get_grids(10)
    plots.plot_covariance(sma, smi, ori_rad, grids, axs[0])
    grids_norm = data.normalize_grid(cov, grids)
    plots.plot_wind(grids, grids_norm, axs[2])
    grids = data.get_grids(10, 500)
    grids_norm = data.normalize_grid(cov, grids)
    plots.plot_normalized(grids_norm, axs[1])

    plt.show()


def make_animation(sma: float, smi: float) -> None:
    """Create an animation with the covariance, normalized space, and windage plots.

    Args:
        sma(float): Semi Major Axis Length
        smi(float): Semi Minor Axis Length
    """
    print(f"Making Animation: SMA:{sma:.1f},  SMI:{smi:.1f}")
    plt.rcParams["figure.figsize"] = 18, 5
    fig, axs = plt.subplots(1, 3)

    def data_gen() -> Generator[
        tuple[
            float,
            list[tuple[np.ndarray, str]],
            list[tuple[np.ndarray, str]],
            list[tuple[np.ndarray, str]],
            list[tuple[np.ndarray, str]],
        ]
    ]:
        ori = 0.0
        while ori < 180:  # noqa: PLR2004
            print(f"Calculating orientation {ori}/180")
            ori_rad = math.radians(ori)
            cov = data.covariance(sma, smi, ori_rad)
            basic_grids = data.get_grids(10)
            basic_grids_norm = data.normalize_grid(cov, basic_grids)
            dense_grids = data.get_grids(10, 500)
            dense_grids_norm = data.normalize_grid(cov, dense_grids)
            ori += 0.5
            yield ori_rad, basic_grids, basic_grids_norm, dense_grids, dense_grids_norm

    titles = ("Covariance", "Normalized", "Diff")

    def draw(
        inputs: tuple[
            float,
            list[tuple[np.ndarray, str]],
            list[tuple[np.ndarray, str]],
            list[tuple[np.ndarray, str]],
            list[tuple[np.ndarray, str]],
        ],
    ) -> list:
        ori_rad, basic_grids, basic_grids_norm, dense_grids, dense_grids_norm = inputs
        for ax, title in zip(axs, titles, strict=False):
            ax.clear()
            ax.set_xlim(-10, 10)
            ax.set_ylim(-10, 10)
            ax.set_title(title)
        artists = []
        cov = data.covariance(sma, smi, ori_rad)
        basic_grids = data.get_grids(10)
        basic_grids_norm = data.normalize_grid(cov, basic_grids)
        dense_grids = data.get_grids(10, 500)
        dense_grids_norm = data.normalize_grid(cov, dense_grids)
        artists.extend(plots.plot_covariance(sma, smi, ori_rad, basic_grids, axs[0]))
        artists.extend(plots.plot_normalized(dense_grids_norm, axs[1]))
        artists.extend(plots.plot_wind(basic_grids, basic_grids_norm, axs[2]))
        return artists

    ani = mpl_animation.FuncAnimation(fig, draw, data_gen, blit=True, repeat=False)
    ani.save("mahalanobis.mp4", writer="ffmpeg", fps=20)


def _get_args() -> None:
    parser = argparse.ArgumentParser(
        description="Plot the normalized space of mahalanobis distance",
    )
    parser.add_argument(
        "sma",
        nargs="?",
        type=float,
        default=3,
        help="Semi Major Axis Length 0 to 10. Default(3.0)",
    )
    parser.add_argument(
        "smi",
        nargs="?",
        type=float,
        default=1,
        help="Semi Minor Axis Length 0 to 10. Default(1.0)",
    )
    parser.add_argument(
        "ori",
        nargs="?",
        type=float,
        default=15,
        help="Orientation in degrees. Default(15)",
    )
    parser.add_argument(
        "--animate",
        help="Create a video animation (mahalanobis.mp4). Ignores other options.",
        action="store_true",
        default=False,
    )
    args = parser.parse_args()
    # Do what we intend even if user gets it backwards
    sma = max(args.sma, args.smi)
    smi = min(args.sma, args.smi)

    if sma > 10 or smi < 0:  # noqa: PLR2004
        print("ERROR: SMA and SMI must be between 0 and 10")
        sys.exit(1)

    if args.animate:
        make_animation(sma, smi)
    else:
        make_image(sma, smi, args.ori)


if __name__ == "__main__":
    _main()
