# ruff: noqa: S311, T201
"""Compare surface distance measurements.

Run with:
python -m gri_sandbox.distance_shootout.main

Try the gri_utils.distance surface measurements. Compare them for cpu time and accuracy.
Memory should be roughly equivalent already (no major objects are used).

Compare at low latitudes, mid latitudes, and high, with distances that are close, mid,
and far.

vincenty's formula is reported to be accurate (to the WGS84 ellipsoid) to about 0.5mm,
so we'll use that as our benchmark for accuracy.
"""

from __future__ import annotations

import math
import random
import statistics
import time
from dataclasses import dataclass, field
from typing import TYPE_CHECKING

import matplotlib.pyplot as plt
from gri_utils import constants, distance
from matplotlib import collections as mc

from gri_sandbox.distance_shootout import dist_fcc, dist_lambert

if TYPE_CHECKING:
    from collections.abc import Callable, Sequence

    import numpy as np
    from numpy.typing import NDArray

print(distance.surface_dist_simple((0, 0), (0, 1)))


@dataclass
class Algo:
    """Class for tracking algorith information."""

    name: str
    func: Callable[[np.ndarray | Sequence, np.ndarray | Sequence], NDArray | float]
    time: float = 0
    deltas: list[float] = field(default_factory=list)


# def straightline(lla1: np.ndarray | Sequence, lla2: np.ndarray | Sequence) -> float:
#     """Calculate the straight line distance."""
#     xyz1 = conversion.lla_to_xyz((lla1[0], lla1[1], 0))
#     xyz2 = conversion.lla_to_xyz((lla2[0], lla2[1], 0))
#     return distance.dist(xyz1, xyz2)


ALGORITHMS: dict[str, Algo] = {
    "fcc": Algo("fcc", dist_fcc.surface_dist_fcc_m),
    "haversine": Algo("haversine", distance.surface_dist_haversine),
    "lambert": Algo("lambert", dist_lambert.surface_dist_lambert_m),
    # "straight": Algo("straight", straightline),
    "simple": Algo("simple", distance.surface_dist_simple),
    "vincenty": Algo("vincenty", distance.surface_dist_vincenty),
}
"""Dictionary of algorithm to amount of time in that algorithm for total computation"""


def _main() -> None:
    smoketest()
    trials = 10000
    start_m = 1000
    stop_m = 50000

    true_dists: list[float] = []
    for lat in (5, 40, 75):
        positions = generate(lat, start_m, stop_m, trials)
        # _plot_positions(positions)
        # Get vincenty distances
        # TODO convert to ns_epoch timer?
        true_algo = ALGORITHMS["vincenty"]
        while true_algo.time == 0:
            try:
                start = time.time()
                true_dists: list[float] = [
                    float(true_algo.func(tup[0], tup[1])) for tup in positions
                ]
                true_algo.time += time.time() - start
                true_algo.deltas = [0.0] * len(true_dists)
            except ValueError:
                # vincenty can fail to converge sometimes, let's make sure we have all
                # valid data
                true_algo.time = 0
                true_algo.deltas = []

        for algo in ALGORITHMS.values():
            if algo.name == "vincenty":
                continue
            start = time.time()
            distances: list[float] = [
                float(algo.func(tup[0], tup[1])) for tup in positions
            ]
            algo.time = time.time() - start
            # now compare accuracy
            algo.deltas = [
                (td - d) * 100 / td for td, d in zip(true_dists, distances, strict=True)
            ]

        print("-------------------------------------")
        print(f"{lat} degrees Latitude")
        print(f"   {' ' * 12} {'ns/calc':>7s} {'Pct StDev':>12s} {'Pct Med':>12s}")
        for _algo_key, algo in sorted(ALGORITHMS.items()):
            median_delta = statistics.median(algo.deltas)
            stdev_delta = statistics.stdev(algo.deltas)
            ns = algo.time * 1e9 / trials
            print(
                "   "
                f"{algo.name:12s} "
                f"{ns:7.0f} "
                f"{stdev_delta:12.6f} "
                f"{median_delta:12.6f}",
            )

        plot_algos(lat, true_dists)


def plot_algos(core_lat: float, true_dists: list[float]) -> None:
    """Plot the data in ALGORITHMS."""
    # Get min/max
    axis_min = 0
    axis_max = 0
    for algo in ALGORITHMS.values():
        axis_min = min(axis_min, *algo.deltas)
        axis_max = max(axis_max, *algo.deltas)
    _fig, _axes = plt.subplots(len(ALGORITHMS) - 1, 1, figsize=(10, 12))
    idx = 0
    for _algo_key, algo in sorted(ALGORITHMS.items()):
        if algo.name == "vincenty":
            continue
        plt.subplot(len(ALGORITHMS) - 1, 1, idx + 1)
        plt.scatter(true_dists, algo.deltas, label=algo.name, s=1)
        plt.title(f"{algo.name}")
        plt.xlabel("Distance(m)")
        plt.ylabel("Error(pct)")
        plt.ylim(math.floor(axis_min * 10) / 10, math.ceil(axis_max * 10) / 10)
        idx += 1
    plt.suptitle(f"{core_lat} Degrees Latitude")
    plt.tight_layout()
    plt.show()


def smoketest() -> None:
    """Verify our vincenty calculation still equals the looked up answers."""
    # Looked up vincenty positions in online calculators
    positions: list[tuple[tuple, tuple, float]] = [
        ((0, 0, 0), (0, 0, 0), 0),
        ((0, 0, 0), (0, 1, 0), 111319.491),
        ((0, 0, 0), (1, 0, 0), 110574.389),
        ((10.2345, 15.6789, 0), (11.5789, 17.9876, 0), 292937.392),
        ((0, 0, 0), (40, 120, 0), 12521126.888),
        ((0, 0, 0), (40, 60, 0), 7500166.649),
        ((0, 0, 0), (40, -60, 0), 7500166.649),
        ((40, 0, 0), (40, -60, 0), 5020978.634),
        ((-5.123, 1.234, 0), (5.123, -2.345, 0), 1200809.23),
        ((0.1, 0.1, 0), (0.101, 0.1, 0), 110.574),
        ((0.1, 0.1, 0), (0.1, 0.101, 0), 111.319),
        ((40.1, 40.1, 0), (40.101, 40.1, 0), 111.037),
        ((40.1, 40.1, 0), (40.1, 40.101, 0), 85.269),
        ((80.1, 80.1, 0), (80.101, 80.1, 0), 111.661),
        ((80.1, 80.1, 0), (80.1, 80.101, 0), 19.202),
    ]

    for idx, pos in enumerate(positions):
        dist = distance.surface_dist_vincenty(pos[0], pos[1], converge_m=0.001)
        abs_m = abs(dist - pos[2])
        assert abs_m <= 0.001, (  # noqa: PLR2004, S101
            f"   smoketest #{idx:02d} : "
            f"{pos[0][0]:>7.3f},{pos[0][1]:>7.3f} : "
            f"{pos[1][0]:>7.3f},{pos[1][1]:>7.3f} : "
            f"Abs: {abs_m:>6.3f}",
        )


def generate(
    core_lat_deg: float,
    min_dist_m: float,
    max_dist_m: float,
    count: int,
) -> list[tuple[tuple[float, float], tuple[float, float]]]:
    """Generate a list of (position1,position2) centered around this lat and dist."""
    positions: list[tuple[tuple[float, float], tuple[float, float]]] = []
    core_lat = math.radians(core_lat_deg)
    deg_2_m = math.radians(constants.ER_M)
    dist_diff = (max_dist_m - min_dist_m) / count
    while len(positions) < count:
        dist_m = min_dist_m + len(positions) * dist_diff
        theta = random.random() * 2 * math.pi
        rad = dist_m / 2
        lat1 = core_lat_deg + math.cos(theta) * rad / deg_2_m
        lon1 = math.sin(theta) * rad / math.cos(core_lat) / deg_2_m
        theta = 2 * math.pi - theta
        lat2 = core_lat_deg - math.cos(theta) * rad / deg_2_m
        lon2 = math.sin(theta) * rad / math.cos(core_lat) / deg_2_m
        positions.append(((lat1, lon1), (lat2, lon2)))
    return positions


def _plot_positions(
    positions: list[tuple[tuple[float, float], tuple[float, float]]],
) -> None:
    """Plot the position line segments that were generated."""
    # Need to reverse x and y for lat/lon
    pos2: list[tuple[tuple[float, float], tuple[float, float]]] = []
    for p1, p2 in positions:
        pos2.append(((p1[1], p1[0]), (p2[1], p2[0])))
    lc = mc.LineCollection(pos2)
    _fix, ax = plt.subplots()
    ax.add_collection(lc)
    ax.autoscale()
    ax.margins(0.1)
    plt.show()


if __name__ == "__main__":
    _main()
