import numpy as np
from astropy import units as u
import ligo.skymap.plot
from matplotlib import pyplot as plt
from m4opt import skygrid

areas = np.asarray([1000, 500, 100, 50]) * u.deg**2
methods = [
    skygrid.geodesic,
    skygrid.golden_angle_spiral,
    skygrid.healpix,
    skygrid.sinusoidal,
]

fig, axs = plt.subplots(
    len(methods),
    len(areas),
    figsize=(8, 6),
    tight_layout=True,
    gridspec_kw=dict(wspace=0.1, hspace=0.1),
    subplot_kw=dict(projection="astro globe", center="0d 25d"),
)

for method, ax in zip(methods, axs[:, 0]):
    ax.text(
        -0.2,
        0.5,
        method.__name__,
        rotation=90,
        ha="center",
        va="center",
        transform=ax.transAxes,
    )

for area, ax in zip(areas, axs[0, :]):
    ax.text(
        0.5,
        1.2,
        area.to_string(format="latex"),
        ha="center",
        va="center",
        transform=ax.transAxes,
    )

for method, axrow in zip(methods, axs):
    for area, ax in zip(areas, axrow):
        for coord in ax.coords:
            coord.set_ticklabel_visible(False)
            coord.set_ticks_visible(False)
        ax.plot_coord(method(area), ".")
        ax.grid()