"""Figures (PNG, written to figures/).  Reads only the derived tables in data/.

fig1_completeness_matrix.png   completeness by location and size: in archive / catalogued / recognised
fig2_survey_calibration.png    NEO completeness, simulator fit, measured elongation weight, discovery-magnitude test
fig3_parked_orbits.png         detection probability for objects parked near Earth
fig4_surfaces_and_belts.png    lunar / martian image coverage; belt catalogue completeness
fig5_found_probability.png     probability a single artefact has been found; interstellar limits
"""
import json
import os

import matplotlib
matplotlib.use("Agg")
import matplotlib.ticker
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.colors import LogNorm

from common import DATA, FIGS, D_from_H, H_from_D

plt.rcParams.update({"font.size": 10, "axes.titlesize": 11, "axes.labelsize": 10, "figure.dpi": 130,
                     "axes.spines.top": False, "axes.spines.right": False})
INK, MUTED = "#222222", "#777777"
C_BLUE, C_ORANGE, C_GREEN, C_RED, C_PURPLE, C_GREY = "#2c6fbb", "#e07b00", "#2a9d5c", "#c0392b", "#7d4fb3", "#8a8a8a"


def rd(name):
    return pd.read_csv(os.path.join(DATA, name))


def size_ticks(ax, axis="x"):
    t = [1, 10, 100, 1000, 10000]
    lab = ["1 m", "10 m", "100 m", "1 km", "10 km"]
    if axis == "x":
        ax.set_xticks(t)
        ax.set_xticklabels(lab)
    else:
        ax.set_yticks(t)
        ax.set_yticklabels(lab)


def fig1():
    M = rd("completeness_matrix.csv")
    M = M[M.albedo == 0.2]
    locs = list(dict.fromkeys(M.location))
    prov = M.drop_duplicates("location").set_index("location").provenance
    sizes = np.sort(M.size_m.unique())
    fig, axes = plt.subplots(1, 3, figsize=(20.5, 6.8), sharey=True)
    norm = LogNorm(vmin=1e-4, vmax=1.0)
    for ax, col, title in zip(axes, ["C_arch_c", "C_cat_c", "C_rec_c"],
                              ["1. In archived data\n(surfaces: imaged at 2 pixels or more; orbits: as panel 2)",
                               "2. Catalogued\n(orbits: designated object; surfaces: imaged x examined [judgement])",
                               "3. Recognised as a non-human artefact\n(central scenario of the recognition model)"]):
        Z = np.array([[M[(M.location == l) & np.isclose(M.size_m, s)][col].iloc[0] for s in sizes] for l in locs])
        im = ax.imshow(np.clip(Z, 1e-4, 1), norm=norm, cmap="viridis", aspect="auto",
                       extent=[np.log10(sizes[0]) - 0.05, np.log10(sizes[-1]) + 0.05, len(locs), 0])
        ax.set_xticks([0, 1, 2, 3, 4])
        ax.set_xticklabels(["1 m", "10 m", "100 m", "1 km", "10 km"])
        ax.set_xlabel("object diameter (albedo 0.2)")
        ax.set_title(title, loc="left", fontsize=10.5)
        ax.set_yticks(np.arange(len(locs)) + 0.5)
        ax.set_yticklabels([f"{l}  [{prov[l]}]" for l in locs])
        for i in range(len(locs)):
            for D in [1, 10, 100, 1000, 10000]:
                v = Z[i, np.argmin(np.abs(sizes - D))]
                txt = "%.2f" % v if v >= 0.095 else ("0" if v < 5e-7 else ("%.0e" % v).replace("e-0", "e-"))
                ax.text(np.log10(D) + (0.14 if D == 1 else (-0.14 if D == 10000 else 0)), i + 0.5, txt, ha="center", va="center",
                        fontsize=8, color="white" if v < 0.3 else "black")
        ax.spines[:].set_visible(False)
    cb = fig.colorbar(im, ax=axes, fraction=0.02, pad=0.01)
    cb.set_label("fraction of the location's phase space or surface")
    fig.suptitle("How completely has each parking location been searched for an inert object of a given size?  "
                 "[in brackets: provenance of tier 1]", x=0.01, ha="left", fontsize=13)
    fig.savefig(os.path.join(FIGS, "fig1_completeness_matrix.png"), bbox_inches="tight")
    plt.close(fig)


def fig2():
    comp = rd("neo_completeness.csv")
    cur = rd("survey_fit_curves.csv")
    val = rd("survey_validation.csv")
    ew_all = rd("elong_weight_measured.csv")
    if "variant" not in ew_all.columns:
        ew_all["variant"] = "base"
    ew = ew_all[ew_all.variant == "base"].reset_index(drop=True)
    vh = rd("discovery_magnitude_hist.csv")
    vt = rd("discovery_magnitude_test.csv").set_index("variant")
    fig, axes = plt.subplots(2, 2, figsize=(13.5, 9.6))
    ax = axes[0, 0]
    Hc = 0.5 * (comp.H1 + comp.H2)
    ax.plot(Hc, comp.compl_diff_neomod2, "o", color=INK, ms=4, label="catalogue / NEOMOD2 population")
    ax.plot(Hc, comp.compl_diff_hc21, "s", color=C_GREY, ms=4, mfc="none", label="catalogue / Harris & Chodas 2021")
    b = cur[cur.variant == "base"]
    for col, c, lab in [("lt2005", C_BLUE, "2005"), ("lt2016", C_GREEN, "2016"), ("final", C_RED, "2026.8")]:
        o = "obs_" + col
        ax.plot(b.H, b[o], ".", color=c, ms=5)
        ax.plot(b.H, b["pred_" + col], "-", color=c, lw=1.6, label=f"simulator (base), discovered by {lab}")
    for v, ls in [("shallow", ":"), ("deep", "--"), ("vdeep", "-.")]:
        d = cur[cur.variant == v]
        ax.plot(d.H, d.pred_final, ls, color=C_RED, lw=1, alpha=0.8, label=f"{v} variant, 2026.8")
    ax.set_yscale("log")
    ax.set_ylim(1e-4, 1.5)
    ax.set_xlabel("absolute magnitude H")
    ax.set_ylabel("fraction of near-Earth asteroids catalogued")
    ax.set_title("(a) NEO completeness and simulator fit", loc="left")
    ax.legend(fontsize=7.5, frameon=False, loc="lower left")
    top = ax.secondary_xaxis("top", functions=(lambda H: H, lambda H: H))
    hh = np.array([17.5, 20, 22.5, 25, 27.5])
    top.set_xticks(hh)
    top.set_xticklabels([("%.0f m" % D_from_H(h, 0.2)) if D_from_H(h, 0.2) < 1000 else ("%.1f km" % (D_from_H(h, 0.2) / 1e3)) for h in hh], fontsize=8)
    ax = axes[0, 1]
    v = val[val.variant == "base"]
    subs = list(dict.fromkeys(v.subset))
    cols = [C_BLUE, C_ORANGE, C_GREEN, C_PURPLE, C_RED, INK]
    for s, c in zip(subs, cols):
        d = v[(v.subset == s) & (v.n_known >= 10)]
        ax.plot(0.5 * (d.H1 + d.H2), d.ratio_pred_over_known, "o-", color=c, ms=4, lw=1.3, label=s.split("(")[0])
    ax.axhline(1, color=MUTED, lw=0.8)
    ax.axhspan(0.5, 2.0, color="#dddddd", alpha=0.5, lw=0)
    ax.set_yscale("log")
    ax.set_ylim(0.2, 5)
    ax.set_yticks([0.25, 0.5, 1, 2, 4])
    ax.set_yticklabels(["0.25", "0.5", "1", "2", "4"])
    ax.yaxis.set_minor_formatter(matplotlib.ticker.NullFormatter())
    ax.set_xlabel("absolute magnitude H")
    ax.set_ylabel("predicted / catalogued")
    ax.set_title("(b) Orbit classes not used in the fit (base variant)", loc="left")
    ax.legend(fontsize=7.5, frameon=False, ncol=2, loc="upper left")
    ax = axes[1, 0]
    x = ew.elong_centre.values
    yc = np.clip(ew.w_c.values, 1e-5, None)
    ax.errorbar(x, yc, yerr=[yc - np.clip(ew.w_lo.values, 1e-5, None), np.clip(ew.w_hi.values, 1e-5, None) - yc], fmt="o", color=INK, ms=5,
                capsize=3, lw=1.2, label="measured (base): real first observations / simulator\nprediction for flat coverage (grey: real objects per bin)")
    for vname, c in [("shallow", C_GREEN), ("deep", C_RED), ("vdeep", C_PURPLE)]:
        d = ew_all[ew_all.variant == vname]
        if len(d):
            ax.plot(d.elong_centre, np.clip(d.w_c, 1e-5, None), "-", color=c, lw=0.9, alpha=0.8, label=f"measured, {vname} variant")
    xx = np.linspace(30, 180, 301)
    ax.plot(xx, np.clip(np.interp(xx, x, ew.w_c.values), 1e-5, None), "-", color=INK, lw=1)
    if "w_assumed_ramp" in ew.columns:
        t = 0.32
        ramp = np.clip((xx - 60) / 30, 0, 1) * (t + (1 - t) * np.clip((xx - 60) / 120, 0, 1))
        ax.plot(xx, np.clip(ramp, 1e-5, None), "--", color=C_BLUE, lw=1.3, label="assumed in the first version (edge at 60 deg)")
    ax.axvline(60, color=C_ORANGE, lw=0.8)
    ax.text(61, 2e-4, "Sun-Earth L4/L5", color=C_ORANGE, fontsize=8, rotation=90, va="bottom")
    for xi, n in zip(x, ew.n_objects):
        ax.text(xi, 2.2, str(int(n)), fontsize=6.5, ha="center", color=MUTED)
    ax.set_yscale("log")
    ax.set_ylim(1e-4, 4)
    ax.set_xlim(15, 182)
    ax.set_xlabel("solar elongation (deg)")
    ax.set_ylabel("survey coverage weight (1 = near opposition)")
    ax.set_title("(c) Measured survey coverage versus solar elongation", loc="left")
    ax.legend(fontsize=7.5, frameon=False, loc="lower right")
    ax = axes[1, 1]
    r = vh[vh.variant == "REAL"]
    xc = 0.5 * (r.V_lo.values + r.V_hi.values)
    ax.bar(xc, r.n / r.n.sum(), 0.46, color="#bbbbbb", label="real first observations (%d Earth-like / co-orbital asteroids)" % int(r.n.sum()))
    for vname, c, ls in [("shallow", C_GREEN, ":"), ("base", C_BLUE, "-"), ("deep", C_RED, "--"), ("vdeep", C_PURPLE, "-.")]:
        d = vh[vh.variant == vname]
        if len(d):
            ax.plot(xc, d.n / d.n.sum(), ls, color=c, lw=1.6, label=f"simulator, {vname} (PIT mean {vt.loc[vname, 'pit_mean']:.2f})")
    ax.set_xlim(17, 24.5)
    ax.set_xlabel("apparent V magnitude on the discovery night (with trailing term)")
    ax.set_ylabel("fraction of discoveries per 0.5 mag")
    ax.set_title("(d) Discovery-magnitude test (PIT mean 0.5 = calibrated)", loc="left")
    ax.legend(fontsize=7.5, frameon=False, loc="upper left")
    fig.tight_layout()
    fig.savefig(os.path.join(FIGS, "fig2_survey_calibration.png"), bbox_inches="tight")
    plt.close(fig)


def fig3():
    cur = rd("parked_pdet_curves.csv")
    fig, axes = plt.subplots(1, 2, figsize=(14, 5.4), sharey=True)
    groups = [
        ("(a) Heliocentric orbits near Earth's", [
            ("Earth quasi-satellites", C_BLUE), ("Earth-like heliocentric orbits (0.9-1.1 au)", C_GREEN),
            ("Sun-Earth L4/L5, parked near the point (libration < 10 deg)", C_ORANGE),
            ("Sun-Earth L4/L5, wide tadpole orbits", C_PURPLE), ("Earth horseshoe co-orbitals", C_RED)]),
        ("(b) Geocentric orbits beyond the geostationary belt", [
            ("High Earth orbit (60,000-300,000 km)", C_BLUE), ("Earth-Moon L4/L5 and lunar-distance orbits", C_ORANGE),
            ("Distant geocentric orbits (0.6-1.2 million km)", C_GREEN)]),
    ]
    alt = ["p_detect", "p_wlo", "p_whi", "p_klo", "p_khi", "p_theory", "p_theory_klo", "p_theory_khi"]
    for ax, (title, items) in zip(axes, groups):
        for name, c in items:
            d = cur[cur.location == name]
            use = [k for k in alt if k in d.columns and d[k].notna().all()]
            lo = d.groupby("H")[use].min().min(axis=1)
            hi = d.groupby("H")[use].max().max(axis=1)
            D = D_from_H(lo.index.values, 0.2)
            ax.fill_between(D, lo.values, hi.values, color=c, alpha=0.18, lw=0)
            u = d[d.variant == "base"].sort_values("H")
            ax.plot(D_from_H(u.H.values, 0.2), u.p_detect, color=c, lw=1.8,
                    label=name.replace(" (0.9-1.1 au)", "").replace(", parked near the point (libration < 10 deg)", ", parked at the point"))
            if "p_detect_uncalibrated" in u.columns and u.p_detect_uncalibrated.notna().all():
                ax.plot(D_from_H(u.H.values, 0.2), u.p_detect_uncalibrated, color=c, lw=1, ls=":")
        ax.set_xscale("log")
        ax.set_xlim(0.3 if "Geocentric" in title else 3, 1e4 if "Helio" in title else 100)
        ax.set_xlabel("object diameter (albedo 0.2)")
        ax.set_title(title, loc="left")
        ax.legend(fontsize=8, frameon=False, loc="upper left" if "Helio" in title else "lower right")
        ax.axhline(0.5, color=MUTED, lw=0.6, ls="--")
    axes[0].set_ylabel("probability of being a catalogued object by 2026.8")
    size_ticks(axes[0])
    axes[1].set_xticks([0.3, 1, 3, 10, 30, 100])
    axes[1].set_xticklabels(["0.3 m", "1 m", "3 m", "10 m", "30 m", "100 m"])
    # real objects for orientation
    for nm, H, y, ha in [("Kamo'oalewa", 24.31, 0.13, "center"), ("2023 FW13", 26.03, 0.07, "center"), ("2020 XL5", 20.28, 0.045, "left"),
                         ("2010 TK7", 20.78, 0.10, "right")]:
        axes[0].annotate(nm, (D_from_H(H, 0.2), 0.0), (D_from_H(H, 0.2), y), fontsize=7.5, ha=ha, color=MUTED,
                         arrowprops=dict(arrowstyle="-", color=MUTED, lw=0.6))
    for nm, H, y in [("2020 CD3", 31.74, 0.07), ("2006 RH120", 29.5, 0.07), ("2020 SO", 28.4, 0.17)]:
        axes[1].annotate(nm, (D_from_H(H, 0.2), 0.0), (D_from_H(H, 0.2), y), fontsize=7.5, ha="center", color=MUTED,
                         arrowprops=dict(arrowstyle="-", color=MUTED, lw=0.6))
    axes[1].text(0.98, 0.42, "dotted: before the minimoon correction\nbands: depth variants, minimoon-population\nrange, two trailing-loss laws", transform=axes[1].transAxes,
                 fontsize=8, color=MUTED, va="top", ha="right")
    axes[0].text(0.98, 0.15, "bands: depth variants and 95% bounds\nof the measured elongation weight", transform=axes[0].transAxes,
                 fontsize=8, color=MUTED, va="bottom", ha="right")
    fig.suptitle("Would existing surveys have catalogued an inert object parked near Earth?", x=0.01, ha="left", fontsize=12)
    fig.tight_layout()
    fig.savefig(os.path.join(FIGS, "fig3_parked_orbits.png"), bbox_inches="tight")
    plt.close(fig)


def fig4():
    moon, ctx = rd("moon_nac_samples.csv"), rd("mars_ctx_samples.csv")
    hy = rd("hirise_coverage_by_year.csv")
    bh, bs = rd("belt_H_histograms.csv"), rd("belt_completeness.csv").set_index("region")
    fig, axes = plt.subplots(1, 3, figsize=(16, 4.8))
    ax = axes[0]
    for d, c, lab in [(moon.best_res, INK, "Moon: LROC NAC, best image"),
                      (moon.best_res_goodlight, C_BLUE, "Moon: Sun 5-70 deg above horizon"),
                      (ctx.best_res, C_RED, "Mars: CTX, best image")]:
        x = np.sort(d.dropna().values)
        ax.step(x, (np.arange(len(x)) + 1) / len(d), where="post", color=c, lw=1.6, label=lab)
    ax.set_xscale("log")
    ax.set_xlim(0.35, 80)
    ax.set_xticks([0.5, 1, 2, 5, 10])
    ax.set_xticklabels(["0.5", "1", "2", "5", "10"])
    ax.set_xlabel("pixel scale (m)")
    ax.set_ylabel("fraction of surface imaged at this scale or finer")
    ax.set_title("(a) Orbital image coverage", loc="left")
    ax.legend(fontsize=7.5, frameon=False, loc="center right")
    ax = axes[1]
    ax.plot(hy.year, 100 * hy.union_fraction, "o-", color=C_ORANGE, ms=4)
    ax.set_xticks([2006, 2010, 2014, 2018, 2022, 2026])
    ax.set_xlabel("year")
    ax.set_ylabel("% of Mars covered by HiRISE (0.25-0.5 m per pixel)")
    ax.set_title("(b) HiRISE union coverage of Mars", loc="left")
    ax.set_ylim(0, None)
    ax = axes[2]
    for reg, c in [("Inner main belt (2.0-2.5 au)", C_BLUE), ("Middle main belt (2.5-2.82 au)", C_GREEN),
                   ("Outer main belt (2.82-3.27 au)", C_ORANGE), ("Jupiter Trojans (4.8-5.5 au)", C_RED),
                   ("Trans-Neptunian (30-100 au)", C_PURPLE)]:
        b = bh[bh.region == reg].sort_values("H")
        r = np.clip(b.ratio_to_powerlaw.values, 0, 1)
        r = np.where(b.H.values < bs.loc[reg, "H_mode"] - 1.0, 1.0, r)
        r = np.minimum.accumulate(r)
        ax.plot(D_from_H(b.H.values, 0.2), r, color=c, lw=1.6, label=reg)
    ax.set_xscale("log")
    ax.set_xlim(100, 1e6)
    ax.set_xticks([100, 1e3, 1e4, 1e5, 1e6])
    ax.set_xticklabels(["100 m", "1 km", "10 km", "100 km", "1000 km"])
    ax.set_xlabel("diameter (albedo 0.2)")
    ax.set_ylabel("catalogued / power-law extrapolation")
    ax.set_title("(c) Catalogue completeness beyond Mars", loc="left")
    ax.legend(fontsize=7.5, frameon=False, loc="lower right")
    fig.tight_layout()
    fig.savefig(os.path.join(FIGS, "fig4_surfaces_and_belts.png"), bbox_inches="tight")
    plt.close(fig)


def fig5():
    P = rd("found_probability.csv")
    iso = rd("iso_detection_volume.csv")
    fig, axes = plt.subplots(1, 3, figsize=(16.5, 5.0))
    schemes = list(dict.fromkeys(P.scheme))
    cols = [INK, C_BLUE, C_ORANGE]
    ax = axes[0]
    for s, c in zip(schemes, cols):
        d = P[(P.scheme == s) & (P.albedo == 0.2)].sort_values("size_m")
        ax.fill_between(d.size_m, d.P_cat_lo, d.P_cat_hi, color=c, alpha=0.15, lw=0)
        ax.plot(d.size_m, d.P_cat_c, color=c, lw=1.8, label=s)
        ax.plot(d.size_m, d.P_arch_c, color=c, lw=1.1, ls="--")
    ax.set_xscale("log")
    ax.set_xlabel("object diameter (albedo 0.2)")
    size_ticks(ax)
    ax.set_ylim(0, 1)
    ax.set_title("(a) In archived data (dashed) and catalogued (solid)", loc="left")
    ax.set_ylabel("probability for one artefact at a location drawn from the prior")
    ax.legend(fontsize=7.5, frameon=False, loc="upper left", title="prior over locations", title_fontsize=8)
    ax.text(0.98, 0.03, "bands: low-high scenario envelopes", transform=ax.transAxes, fontsize=7.5, color=MUTED, ha="right")
    ax = axes[1]
    for s, c in zip(schemes, cols):
        d = P[(P.scheme == s) & (P.albedo == 0.2)].sort_values("size_m")
        ax.fill_between(d.size_m, np.clip(d.P_rec_lo, 1e-7, 1), d.P_rec_hi, color=c, alpha=0.15, lw=0)
        ax.plot(d.size_m, d.P_rec_c, color=c, lw=1.8, label=s)
        d2 = P[(P.scheme == s) & (P.albedo == 0.01)].sort_values("size_m")
        ax.plot(d2.size_m, d2.P_rec_c, color=c, lw=1, ls=":")
    ax.set_xscale("log")
    ax.set_xlabel("object diameter")
    size_ticks(ax)
    ax.set_yscale("log")
    ax.set_ylim(1e-6, 1)
    ax.set_title("(b) Catalogued and recognised as non-human (scenario model)", loc="left")
    ax.text(0.98, 0.03, "solid: albedo 0.2; dotted: albedo 0.01\nbands: all judgement parameters low / all high", transform=ax.transAxes,
            fontsize=7.5, color=MUTED, ha="right")
    ax = axes[2]
    b = iso[(iso.kinematics == "disc")]
    for v, ls in [("base", "-"), ("shallow", ":"), ("vdeep", "--")]:
        d = b[b.variant == v].sort_values("H")
        ax.plot(d.D_m_p0p2, d.n_UL95_au3, ls, color=INK, lw=1.6 if v == "base" else 1, label=f"95% upper limit, 1998-2026 ({v})")
    d = b[b.variant == "base"].sort_values("H")
    ax.plot(d.D_m_p0p2, d.n_UL95_after_lsst_au3, color=C_BLUE, lw=1.6, label="after 10 yr at Rubin depth (scaling)")
    c = iso[(iso.kinematics == "cold") & (iso.variant == "base")].sort_values("H")
    ax.plot(c.D_m_p0p2, c.n_UL95_au3, color=C_GREY, lw=1, label="kinematically cold population")
    ax.plot([D_from_H(22.08, 0.2)], [0.2], "*", color=C_RED, ms=11, label="1I/'Oumuamua: 0.2 per au$^3$ (Do et al. 2018)")
    try:
        js = json.load(open(os.path.join(DATA, "iso_summary.json")))
        n1 = js["n_oumuamua_like_au3_disc_base"]
        lo, hi = js["n_oumuamua_like_au3_poisson95"]
        ax.errorbar([D_from_H(22.08, 0.2) * 1.25], [n1], yerr=[[n1 - lo], [hi - n1]], fmt="o", color=C_ORANGE, ms=5, capsize=3,
                    label="this work: one detection (95% Poisson interval)")
    except Exception:      # noqa
        pass
    ax.set_xscale("log")
    ax.set_yscale("log")
    ax.set_xlabel("object diameter (albedo 0.2)")
    ax.set_ylabel("number density (au$^{-3}$)")
    size_ticks(ax)
    ax.set_title("(c) Inert interstellar objects: density limits", loc="left")
    ax.legend(fontsize=7.5, frameon=False, loc="upper right")
    sec = ax.secondary_yaxis("right", functions=(lambda n: n * 8.775e16, lambda n: n / 8.775e16))
    sec.set_ylabel("objects per star (0.1 stars per pc$^3$)")
    fig.tight_layout()
    fig.savefig(os.path.join(FIGS, "fig5_found_probability.png"), bbox_inches="tight")
    plt.close(fig)


def main():
    for f in (fig1, fig2, fig3, fig4, fig5):
        f()
        print("[fig]", f.__name__, "done")


if __name__ == "__main__":
    main()