#!/usr/bin/env python3
"""Generate all-sky (l, b) Mollweide maps for each Figure 3 study.

Each map shows the study's native sky coverage (or per-sightline sample),
the M31 position, and the fourteen XMM M31 fields, plus the |b| > 30 deg
boundary that Henley & Shelton (2013) actually imposed.
"""

from __future__ import annotations

from pathlib import Path

import astropy.units as u
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import scienceplots  # noqa: F401
from astropy.coordinates import Longitude

PROJECT_DIR = Path(__file__).resolve().parent
PAPER_DIR = PROJECT_DIR.parent / "paper_apj_v19_cgmsum_draft"

# M31 Galactic coordinates from explorer_summary.json
M31_L_DEG = 121.17432906153702
M31_B_DEG = -21.573308798437054

COLORS = {
    "blue": "#2676b8",
    "orange": "#d47a2c",
    "green": "#398b5f",
    "purple": "#7553a6",
    "cyan": "#2a95a5",
    "gray": "#6f7782",
    "red": "#b74842",
    "pink": "#c47ba8",
}


def _galactic_longitude_to_mollweide_rad(
    longitude_deg: np.ndarray | float,
) -> np.ndarray | float:
    """Project 0--360 deg Galactic longitude using the astro convention.

    Astropy wraps longitude at 180 deg. Matplotlib's Mollweide x coordinate
    increases to the right, so negating the wrapped value matches the standard
    astronomical convention: longitude increases to the left.
    """
    values = np.asarray(longitude_deg, dtype=float)
    wrapped_rad = Longitude(
        values * u.deg, wrap_angle=180 * u.deg
    ).to_value(u.rad)
    projected = -np.asarray(wrapped_rad, dtype=float)
    if values.ndim == 0:
        return float(projected)
    return projected


def _galactic_rectangle_vertices(
    l_lo_deg: float,
    l_hi_deg: float,
    b_lo_deg: float,
    b_hi_deg: float,
) -> list[tuple[np.ndarray, np.ndarray]]:
    """Return seam-safe Mollweide polygons for a Galactic rectangle."""
    if not 0 <= l_lo_deg < l_hi_deg <= 360:
        raise ValueError(
            "longitude interval must satisfy 0 <= l_lo < l_hi <= 360"
        )
    if not -90 <= b_lo_deg < b_hi_deg <= 90:
        raise ValueError(
            "latitude interval must satisfy -90 <= b_lo < b_hi <= 90"
        )

    cuts = [l_lo_deg]
    if l_lo_deg < 180 < l_hi_deg:
        cuts.append(180.0)
    cuts.append(l_hi_deg)

    polygons: list[tuple[np.ndarray, np.ndarray]] = []
    for lower, upper in zip(cuts[:-1], cuts[1:], strict=True):
        if upper <= 180:
            x_lower = -np.deg2rad(lower)
            x_upper = -np.deg2rad(upper)
        else:
            x_lower = np.deg2rad(360.0 - lower)
            x_upper = np.deg2rad(360.0 - upper)
        longitude_rad = np.array(
            [x_lower, x_upper, x_upper, x_lower, x_lower], dtype=float
        )
        latitude_rad = np.deg2rad(
            np.array([b_lo_deg, b_lo_deg, b_hi_deg, b_hi_deg, b_lo_deg])
        )
        polygons.append((longitude_rad, latitude_rad))
    return polygons


def _fill_galactic_rectangle(
    ax: plt.Axes,
    l_lo_deg: float,
    l_hi_deg: float,
    b_lo_deg: float,
    b_hi_deg: float,
    **kwargs: object,
) -> list:
    """Fill a Galactic rectangle, splitting it safely at l=180 deg."""
    artists = []
    label = kwargs.pop("label", None)
    for index, (longitude_rad, latitude_rad) in enumerate(
        _galactic_rectangle_vertices(l_lo_deg, l_hi_deg, b_lo_deg, b_hi_deg)
    ):
        artists.extend(
            ax.fill(
                longitude_rad,
                latitude_rad,
                label=label if index == 0 else None,
                **kwargs,
            )
        )
    return artists


def _load_halosat_2020_fields() -> pd.DataFrame:
    """Load and validate the 73-field Kaaret et al. (2020) fit catalog."""
    fields = pd.read_csv(
        PAPER_DIR / "m31_cgmsum_halosat_kaaret2020_southern_fields.csv"
    )
    required = {"obsid", "galactic_l_deg", "galactic_b_deg"}
    missing = required - set(fields.columns)
    if missing:
        raise ValueError(f"HaloSat field catalog missing columns: {sorted(missing)}")
    if len(fields) != 73 or fields["obsid"].nunique() != 73:
        raise ValueError("HaloSat 2020 fit catalog must contain 73 unique fields")
    if not fields["galactic_l_deg"].between(0, 360, inclusive="left").all():
        raise ValueError("HaloSat field longitudes must lie in [0, 360) deg")
    if not (fields["galactic_b_deg"] < -30).all():
        raise ValueError("HaloSat 2020 fit fields must all satisfy b < -30 deg")
    return fields


def _save(fig: plt.Figure, stem: str, column: str) -> None:
    for suffix in ("png", "pdf"):
        fig.savefig(
            PAPER_DIR / f"{stem}_{column}.{suffix}",
            dpi=300 if suffix == "png" else None,
            bbox_inches="tight",
            pad_inches=0.03,
        )
    plt.close(fig)


def _draw_mollweide_axes(ax: plt.Axes) -> None:
    """Configure a Mollweide projection with Galactic center at center."""
    ax.set_facecolor("#fafbfc")
    ax.grid(alpha=0.30, linewidth=0.4, color="#888888")
    ax.set_xticklabels([])  # l labels added separately
    ax.set_yticklabels([])
    ax.set_xlabel("")
    ax.set_ylabel("")


def _add_m31_marker(ax: plt.Axes, marker_size: float) -> list:
    """Plot M31 with a star marker and label."""
    l_rad = _galactic_longitude_to_mollweide_rad(M31_L_DEG)
    b_rad = np.deg2rad(M31_B_DEG)
    line = ax.scatter(
        l_rad,
        b_rad,
        marker="*",
        s=marker_size * 8,
        color=COLORS["red"],
        edgecolor="black",
        linewidth=0.6,
        zorder=10,
        label="M31 (l=121.17°, b=-21.57°)",
    )
    ax.annotate(
        "M31",
        xy=(l_rad, b_rad),
        xytext=(8, -8),
        textcoords="offset points",
        fontsize=8,
        weight="bold",
        color=COLORS["red"],
    )
    return [line]


def _add_galactic_plane(ax: plt.Axes, color: str = "#444444") -> None:
    """Draw the Galactic plane (b=0) as a dashed line."""
    l_line = np.deg2rad(np.linspace(-180, 180, 361))
    b_line = np.zeros_like(l_line)
    ax.plot(l_line, b_line, color=color, linestyle="--", linewidth=0.8, alpha=0.7, zorder=1)
    ax.text(
        np.deg2rad(0),
        np.deg2rad(2),
        "Galactic plane (b=0)",
        ha="center",
        fontsize=6,
        color=color,
        alpha=0.8,
    )


def _add_b30_boundary(ax: plt.Axes) -> None:
    """Shade the |b| < 30 deg forbidden zone (Henley & Shelton criterion)."""
    # Draw horizontal lines at b = +30 and b = -30
    l_line = np.deg2rad(np.linspace(-180, 180, 361))
    b_upper = np.full_like(l_line, np.deg2rad(30))
    b_lower = np.full_like(l_line, np.deg2rad(-30))
    ax.plot(l_line, b_upper, color=COLORS["gray"], linestyle=":", linewidth=0.9, alpha=0.7, zorder=2)
    ax.plot(l_line, b_lower, color=COLORS["gray"], linestyle=":", linewidth=0.9, alpha=0.7, zorder=2)


def _set_lon_labels(ax: plt.Axes, fontsize: float) -> None:
    """Add Galactic longitude labels around the Mollweide projection.

    Use standard 0--360 deg Galactic notation and the astronomical
    convention: l increases to the left, with l=0 deg at the centre and
    l=180 deg at both horizontal edges.
    """
    for x_deg, label_deg, ha in (
        (-180, 180, "right"),
        (-90, 90, "center"),
        (0, 0, "center"),
        (90, 270, "center"),
        (180, 180, "left"),
    ):
        ax.text(
            np.deg2rad(x_deg),
            np.deg2rad(-3),
            f"l={label_deg}°",
            ha=ha,
            va="top",
            fontsize=fontsize,
            color="#333333",
        )
    # Latitude tick labels along the right edge
    for b_deg in (-60, -30, 30, 60):
        ax.text(
            np.deg2rad(180),
            np.deg2rad(b_deg),
            f"b={b_deg}°",
            ha="left",
            va="center",
            fontsize=fontsize - 0.5,
            color="#333333",
        )


# ---------------------------------------------------------------------------
# Per-study maps
# ---------------------------------------------------------------------------


def plot_henley_sky(column: str, width: float) -> None:
    """H&S13: 86 high-|b| sightlines on the all-sky map."""
    data = pd.read_csv(
        PAPER_DIR / "m31_cgmsum_henley_shelton2013_absorbed_0p5_2p0.csv"
    )
    assert len(data) == 86, f"expected 86, got {len(data)}"
    assert (data["glat_deg"].abs() > 30).all(), "H&S13 sample must satisfy |b|>30"

    fig = plt.figure(figsize=(width, width * 0.55), constrained_layout=True)
    ax = fig.add_subplot(111, projection="mollweide")
    _draw_mollweide_axes(ax)
    _add_galactic_plane(ax)
    _add_b30_boundary(ax)

    # Plot 86 sightlines
    l_rad = _galactic_longitude_to_mollweide_rad(
        data["glon_deg"].to_numpy()
    )
    b_rad = np.deg2rad(data["glat_deg"].to_numpy())
    ax.scatter(
        l_rad,
        b_rad,
        s=14 if column == "1col" else 22,
        color=COLORS["blue"],
        alpha=0.65,
        edgecolor="none",
        zorder=3,
        label="H&S13 (N=86, |b|>30°)",
    )

    _add_m31_marker(ax, marker_size=14 if column == "1col" else 22)

    ax.legend(
        loc="lower left",
        fontsize=5 if column == "1col" else 8.5,
        framealpha=0.85,
    )
    _set_lon_labels(ax, fontsize=6 if column == "1col" else 9)

    fig.suptitle(
        "H&S13 sightlines on the all-sky (l, b) map",
        fontsize="medium",
        weight="bold",
    )
    _save(fig, "henley_sky_coverage", column)


def plot_efeds_sky(column: str, width: float) -> None:
    """Ponti+23 eFEDS: ~107.5 deg² field at l~220-235, b~20-40."""
    fig = plt.figure(figsize=(width, width * 0.55), constrained_layout=True)
    ax = fig.add_subplot(111, projection="mollweide")
    _draw_mollweide_axes(ax)
    _add_galactic_plane(ax)

    # eFEDS field: l=220-235, b=20-40 (rectangular approximation)
    _fill_galactic_rectangle(
        ax,
        220,
        235,
        20,
        40,
        color=COLORS["purple"],
        alpha=0.35,
        edgecolor=COLORS["purple"],
        linewidth=1.2,
        zorder=3,
        label="eFEDS (107.5 deg², l=220-235°, b=20-40°)",
    )

    _add_m31_marker(ax, marker_size=14 if column == "1col" else 22)

    ax.legend(
        loc="lower left",
        fontsize=5 if column == "1col" else 8.5,
        framealpha=0.85,
    )
    _set_lon_labels(ax, fontsize=6 if column == "1col" else 9)

    fig.suptitle(
        "Ponti+23 eFEDS footprint (l, b)",
        fontsize="medium",
        weight="bold",
    )
    _save(fig, "efeds_sky_coverage", column)


def plot_locatelli_sky(column: str, width: float) -> None:
    """Locatelli+24: eRASS1 western Galactic half (180 < l < 360 deg)."""
    fig = plt.figure(figsize=(width, width * 0.55), constrained_layout=True)
    ax = fig.add_subplot(111, projection="mollweide")
    _draw_mollweide_axes(ax)
    _add_galactic_plane(ax)

    # Western Galactic half: 180 < l < 360
    _fill_galactic_rectangle(
        ax,
        180,
        360,
        -90,
        90,
        color=COLORS["cyan"],
        alpha=0.18,
        edgecolor=COLORS["cyan"],
        linewidth=1.2,
        zorder=2,
        label="eRASS1 western half (180<l<360°)",
    )

    # Add hatching for the eastern excluded half
    _fill_galactic_rectangle(
        ax,
        0,
        180,
        -90,
        90,
        color="#dddddd",
        alpha=0.45,
        edgecolor=COLORS["gray"],
        linewidth=0.6,
        hatch="///",
        zorder=2,
        label="Excluded eastern half (0<l<180°)",
    )

    _add_m31_marker(ax, marker_size=14 if column == "1col" else 22)

    ax.legend(
        loc="lower left",
        fontsize=5 if column == "1col" else 8.5,
        framealpha=0.85,
    )
    _set_lon_labels(ax, fontsize=6 if column == "1col" else 9)

    fig.suptitle(
        "Locatelli+24 eRASS1 western-half map domain (l, b)",
        fontsize="medium",
        weight="bold",
    )
    _save(fig, "locatelli_sky_coverage", column)


def plot_ueda_sky(column: str, width: float) -> None:
    """Ueda+22: 130-field parent sample at 75<l<285, |b|>15."""
    fig = plt.figure(figsize=(width, width * 0.55), constrained_layout=True)
    ax = fig.add_subplot(111, projection="mollweide")
    _draw_mollweide_axes(ax)
    _add_galactic_plane(ax)

    # Parent sample 75<l<285, |b|>15. Split at l=180 map seam.
    _fill_galactic_rectangle(
        ax,
        75,
        285,
        15,
        90,
        color=COLORS["green"],
        alpha=0.18,
        edgecolor=COLORS["green"],
        linewidth=1.2,
        zorder=2,
        label="Ueda+22 parent sample (75<l<285°, |b|>15°)",
    )
    # Southern half, also split at the l=180 map seam.
    _fill_galactic_rectangle(
        ax,
        75,
        285,
        -90,
        -15,
        color=COLORS["green"],
        alpha=0.18,
        edgecolor=COLORS["green"],
        linewidth=1.2,
        zorder=2,
    )

    # Plot the 14 M31 XMM fields on top
    fields = pd.read_csv(
        PAPER_DIR / "m31_cgmsum_ueda2022_disk_m31_footprint_absorbed_0p5_2p0.csv"
    )
    ax.scatter(
        _galactic_longitude_to_mollweide_rad(
            fields["galactic_l_deg"].to_numpy()
        ),
        np.deg2rad(fields["galactic_b_deg"].to_numpy()),
        s=18 if column == "1col" else 28,
        marker="o",
        color=COLORS["orange"],
        edgecolor="black",
        linewidth=0.5,
        zorder=6,
        label="14 M31 XMM fields (l≈119.5-122.3°)",
    )

    _add_m31_marker(ax, marker_size=14 if column == "1col" else 22)

    ax.legend(
        loc="lower left",
        fontsize=5 if column == "1col" else 8.5,
        framealpha=0.85,
    )
    _set_lon_labels(ax, fontsize=6 if column == "1col" else 9)

    fig.suptitle(
        "Ueda+22 sample domain and 14 M31 XMM fields (l, b)",
        fontsize="medium",
        weight="bold",
    )
    _save(fig, "ueda_sky_coverage", column)


def plot_halosat_sky(column: str, width: float) -> None:
    """Kaaret+20: selection envelope and 73 actual southern fit centers."""
    fig = plt.figure(figsize=(width, width * 0.55), constrained_layout=True)
    ax = fig.add_subplot(111, projection="mollweide")
    _draw_mollweide_axes(ax)
    _add_galactic_plane(ax)
    _add_b30_boundary(ax)

    # The latitude cap is the selection envelope, not continuous field coverage.
    l_line = np.deg2rad(np.linspace(-180, 180, 361))
    b_line = np.full_like(l_line, np.deg2rad(-30))
    ax.fill_between(
        l_line,
        -np.pi / 2,
        b_line,
        color=COLORS["cyan"],
        alpha=0.10,
        edgecolor=COLORS["cyan"],
        linewidth=1.2,
        zorder=2,
        label="2020 selection envelope (b<-30°)",
    )

    halosat_fields = _load_halosat_2020_fields()
    ax.scatter(
        _galactic_longitude_to_mollweide_rad(
            halosat_fields["galactic_l_deg"].to_numpy()
        ),
        np.deg2rad(halosat_fields["galactic_b_deg"].to_numpy()),
        s=8 if column == "1col" else 18,
        marker="o",
        color=COLORS["cyan"],
        edgecolor="black",
        linewidth=0.35,
        zorder=4,
        label="Kaaret+20 fitted field centers (N=73)",
    )

    # Plot the 14 M31 XMM fields on top
    m31_fields = pd.read_csv(
        PAPER_DIR / "m31_cgmsum_halosat_kaaret2020_disk_adiabatic_halo_m31_extrapolation.csv"
    )
    ax.scatter(
        _galactic_longitude_to_mollweide_rad(
            m31_fields["galactic_l_deg"].to_numpy()
        ),
        np.deg2rad(m31_fields["galactic_b_deg"].to_numpy()),
        s=18 if column == "1col" else 28,
        marker="o",
        color=COLORS["orange"],
        edgecolor="black",
        linewidth=0.5,
        zorder=6,
        label="14 M31 XMM fields — out of domain",
    )

    _add_m31_marker(ax, marker_size=14 if column == "1col" else 22)

    ax.legend(
        loc="upper left",
        bbox_to_anchor=(0.02, 0.96),
        fontsize=5 if column == "1col" else 8.5,
        framealpha=0.85,
    )
    _set_lon_labels(ax, fontsize=6 if column == "1col" else 9)

    fig.suptitle(
        "HaloSat 2020 southern morphology sample vs M31 direction (l, b)",
        fontsize="medium",
        weight="bold",
    )
    _save(fig, "halosat_sky_coverage", column)


def plot_zhang_sky(column: str, width: float) -> None:
    """Zhang+24: external-galaxy stack (not an MW sightline)."""
    fig = plt.figure(figsize=(width, width * 0.55), constrained_layout=True)
    ax = fig.add_subplot(111, projection="mollweide")
    _draw_mollweide_axes(ax)
    _add_galactic_plane(ax)

    # Zhang is an external-galaxy stack: not a real MW LOS. Show the
    # M31 XMM fields as positional reference only.
    fields = pd.read_csv(
        PAPER_DIR / "m31_cgmsum_locatelli2024_reference_beta0p5_m31_footprint_predictions.csv"
    )
    ax.scatter(
        _galactic_longitude_to_mollweide_rad(
            fields["galactic_l_deg"].to_numpy()
        ),
        np.deg2rad(fields["galactic_b_deg"].to_numpy()),
        s=18 if column == "1col" else 28,
        marker="o",
        color=COLORS["orange"],
        edgecolor="black",
        linewidth=0.5,
        zorder=6,
        label="14 M31 XMM fields (positional reference only)",
    )

    # Add a note about the eROSITA all-sky coverage
    l_line = np.deg2rad(np.linspace(-180, 180, 361))
    ax.fill_between(
        l_line,
        np.deg2rad(-90),
        np.deg2rad(90),
        color=COLORS["orange"],
        alpha=0.06,
        zorder=1,
    )

    _add_m31_marker(ax, marker_size=14 if column == "1col" else 22)

    ax.text(
        0.5,
        0.96,
        "External-galaxy stack — no real Milky-Way LOS",
        transform=ax.transAxes,
        ha="center",
        va="top",
        weight="bold",
        fontsize=5 if column == "1col" else 8,
        color=COLORS["gray"],
    )

    ax.legend(
        loc="lower left",
        fontsize=5 if column == "1col" else 8.5,
        framealpha=0.85,
    )
    _set_lon_labels(ax, fontsize=6 if column == "1col" else 9)

    fig.suptitle(
        "Zhang+24 external-galaxy stack (no real M31 LOS)",
        fontsize="medium",
        weight="bold",
    )
    _save(fig, "zhang_sky_coverage", column)


def plot_grayson_sky(column: str, width: float) -> None:
    """Grayson+25: simulated X-ray stacks (EAGLE/SIMBA), not real LOS."""
    fig = plt.figure(figsize=(width, width * 0.55), constrained_layout=True)
    ax = fig.add_subplot(111, projection="mollweide")
    _draw_mollweide_axes(ax)
    _add_galactic_plane(ax)

    # Grayson is a simulation stack: not a real MW LOS
    fields = pd.read_csv(
        PAPER_DIR / "m31_cgmsum_locatelli2024_reference_beta0p5_m31_footprint_predictions.csv"
    )
    ax.scatter(
        _galactic_longitude_to_mollweide_rad(
            fields["galactic_l_deg"].to_numpy()
        ),
        np.deg2rad(fields["galactic_b_deg"].to_numpy()),
        s=18 if column == "1col" else 28,
        marker="o",

        color=COLORS["orange"],
        edgecolor="black",
        linewidth=0.5,
        zorder=6,
        label="14 M31 XMM fields (positional reference only)",
    )

    l_line = np.deg2rad(np.linspace(-180, 180, 361))
    ax.fill_between(
        l_line,
        np.deg2rad(-90),
        np.deg2rad(90),
        color=COLORS["purple"],
        alpha=0.06,
        zorder=1,
    )

    _add_m31_marker(ax, marker_size=14 if column == "1col" else 22)

    ax.text(
        0.5,
        0.96,
        "EAGLE/SIMBA synthetic stacks — no real Milky-Way LOS",
        transform=ax.transAxes,
        ha="center",
        va="top",
        weight="bold",
        fontsize=5 if column == "1col" else 8,
        color=COLORS["gray"],
    )

    ax.legend(
        loc="lower left",
        fontsize=5 if column == "1col" else 8.5,
        framealpha=0.85,
    )
    _set_lon_labels(ax, fontsize=6 if column == "1col" else 9)

    fig.suptitle(
        "Grayson+25 simulation stacks (no real M31 LOS)",
        fontsize="medium",
        weight="bold",
    )
    _save(fig, "grayson_sky_coverage", column)


def main() -> None:
    PAPER_DIR.mkdir(parents=True, exist_ok=True)
    with plt.style.context(["science", "no-latex"]):
        for column, width in (("1col", 3.35), ("2col", 7.0)):
            scale = 5.2 if column == "1col" else 8.5
            with plt.rc_context(
                {
                    "font.size": scale,
                    "axes.labelsize": scale,
                    "axes.titlesize": scale + 1,
                    "xtick.labelsize": scale - 0.3,
                    "ytick.labelsize": scale - 0.3,
                    "legend.fontsize": scale - 0.5,
                    "figure.titlesize": scale + 1.5,
                }
            ):
                plot_henley_sky(column, width)
                plot_efeds_sky(column, width)
                plot_locatelli_sky(column, width)
                plot_ueda_sky(column, width)
                plot_halosat_sky(column, width)
                plot_zhang_sky(column, width)
                plot_grayson_sky(column, width)
    print(f"Generated 28 sky-coverage assets in {PAPER_DIR}")


if __name__ == "__main__":
    main()