"""Figures for the dim_stars study (static PNG, matplotlib)."""
import os
import warnings

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

import dimlib as L
from common import DATA_MAIN as DATA, DATA_ORIG, FIGS, WORK, dpkl

warnings.filterwarnings("ignore")
SURVEYS = ("apogee", "lamost", "gspspec")
NAME = dict(apogee="APOGEE DR17", lamost="LAMOST DR11", gspspec="Gaia GSP-Spec", combined="Combined")
# first three categorical slots of the validated default palette (all-pairs safe), plus neutral inks
COL = dict(apogee="#2a78d6", lamost="#eb6834", gspspec="#1baf7a", combined="#0b0b0b")
INK, INK2, GRID = "#0b0b0b", "#52514e", "#dddcd6"
SURF = "#fcfcfb"
plt.rcParams.update({"font.size": 9.5, "axes.edgecolor": INK2, "axes.labelcolor": INK, "xtick.color": INK2,
                     "ytick.color": INK2, "axes.spines.top": False, "axes.spines.right": False,
                     "figure.facecolor": SURF, "axes.facecolor": SURF, "savefig.facecolor": SURF,
                     "axes.grid": True, "grid.color": GRID, "grid.linewidth": 0.6, "axes.axisbelow": True,
                     "legend.frameon": False})
FLAB = {0.25: "25%", 0.5: "50%", 0.75: "75%", 0.9: "90%"}


def fig1():
    """Distribution of the luminosity deficit in each survey."""
    fig, axes = plt.subplots(3, 1, figsize=(7.4, 7.8), sharex=True)
    for ax, sv in zip(axes, SURVEYS):
        h = pd.read_csv(os.path.join(DATA, "hist_D_%s.csv" % sv))
        x = 0.5 * (h.D_lo + h.D_hi)
        ax.step(x, np.maximum(h.n_before_colour_cut, 0.3), where="mid", color="#b5b4ad", lw=1.0,
                label="before colour-consistency cut")
        ax.step(x, np.maximum(h.n_clean, 0.3), where="mid", color=COL[sv], lw=1.6, label="clean sample")
        ax.set_yscale("log")
        top = h.n_clean.max()
        ax.set_ylim(0.5, top * 3)
        ax.set_xlim(-1.5, 0.8)
        for f in L.F_GRID:
            ax.axvline(np.log10(1 - f), color=INK2, lw=0.8, ls=(0, (4, 3)))
            if sv == SURVEYS[0]:
                ax.text(np.log10(1 - f), top * 3.6, FLAB[f], ha="center", va="bottom", fontsize=8.5, color=INK2)
        if sv == SURVEYS[0]:
            ax.text(-1.12, top * 3.6, "grey dimming f =", ha="right", va="bottom", fontsize=8.5, color=INK2)
        n = int(pd.read_csv(os.path.join(DATA, "cuts_%s.csv" % sv)).n_remaining.iloc[-1])
        ct = pd.read_csv(os.path.join(DATA, "core_tail_%s.csv" % sv)).iloc[0]
        ax.text(0.015, 0.93, "%s\n%s stars\ncore scatter\n%.3f dex" % (NAME[sv], format(n, ","), ct.sigma_core),
                transform=ax.transAxes, ha="left", va="top", color=INK, fontsize=9.5)
        if sv == SURVEYS[0]:
            ax.annotate("unresolved binaries\n(up to 2x the light, +0.30 dex)", xy=(0.30, h.n_clean[(x > 0.29) & (x < 0.31)].max() * 1.3),
                        xytext=(0.36, top * 0.35), fontsize=8.5, color=INK2, ha="left", va="center",
                        arrowprops=dict(arrowstyle="-", color=INK2, lw=0.7))
        ax.set_ylabel("stars per 0.01 dex")
        if sv == SURVEYS[0]:
            leg = ax.legend(loc="lower left", fontsize=8.5, bbox_to_anchor=(0.0, 0.10), frameon=True, framealpha=1.0,
                            facecolor=SURF, edgecolor="none")
            leg.set_zorder(5)
    axes[-1].set_xlabel(r"luminosity deficit  $D=\log_{10}(L_{\rm parallax}/L_{\rm spectroscopic})$   (less deficient of G and K$_s$)")
    fig.suptitle("Stars fainter than their spectra predict: a narrow core and a thin faint tail",
                 x=0.01, ha="left", fontsize=10.5, color=INK, y=0.995)
    fig.tight_layout(rect=(0, 0, 1, 0.975))
    fig.savefig(os.path.join(FIGS, "fig1_deficit_distributions.png"), dpi=170)
    plt.close(fig)


def fig2():
    """Observed absolute magnitude versus Teff with flagged stars."""
    cand = pd.read_csv(os.path.join(DATA, "candidates_all.csv"))
    fig, axes = plt.subplots(1, 3, figsize=(11.0, 4.5), sharey=True)
    for ax, sv in zip(axes, SURVEYS):
        s = pd.read_pickle(dpkl(sv, "final"))
        c = s[s.clean]
        ax.hexbin(c.teff, c.MK, gridsize=(70, 55), bins="log", cmap="Blues", mincnt=1, linewidths=0,
                  extent=(4000, 6500, -0.5, 7.5))
        cc = cand[cand.survey == sv]
        spec = np.zeros(len(cc), bool)
        three = np.zeros(len(cc), bool)
        unres = np.zeros(len(cc), bool)
        for f in (0.5, 0.75, 0.9):
            spec |= cc["flagS_%g" % f].values
            three |= cc["flag_%g" % f].values
            unres |= cc["flag_%g" % f].values & (cc["status_%g" % f] == "unresolved").values
        only = spec & ~three
        ax.scatter(cc.teff[only], cc.MK[only], s=7, facecolors="none", edgecolors=INK2, linewidths=0.5, alpha=0.8,
                   label="spectroscopic deficit only (%d)" % only.sum())
        ax.scatter(cc.teff[three & ~unres], cc.MK[three & ~unres], s=34, marker="s", facecolors=SURF, edgecolors=INK,
                   linewidths=1.0, label="three-way, resolved (%d)" % (three & ~unres).sum(), zorder=3)
        ax.scatter(cc.teff[unres], cc.MK[unres], s=34, color=INK, edgecolors=SURF, linewidths=0.6,
                   label="three-way, unresolved (%d)" % unres.sum(), zorder=4)
        ax.set_xlim(6500, 4000)
        ax.set_ylim(7.5, -0.5)
        ax.set_title(NAME[sv], fontsize=10, color=INK, loc="left")
        ax.set_xlabel(r"spectroscopic $T_{\rm eff}$ (K)")
        ax.legend(loc="upper left", fontsize=8, handletextpad=0.3)
    axes[0].set_ylabel(r"$M_{K_s}$ from Gaia parallax")
    fig.suptitle("Most stars with a >=50% spectroscopic deficit sit on the main sequence (wrong log g or Teff); "
                 "few are also below it",
                 x=0.01, ha="left", fontsize=10.5, color=INK)
    fig.tight_layout(rect=(0, 0, 1, 0.95))
    fig.savefig(os.path.join(FIGS, "fig2_hr_flagged.png"), dpi=170)
    plt.close(fig)


def fig3():
    """Selection efficiency of dimmed stars."""
    lim = pd.read_csv(os.path.join(DATA, "limits.csv"))
    fig, axes = plt.subplots(1, 2, figsize=(10.4, 4.4))
    ax = axes[0]
    for sv in SURVEYS:
        sf = pd.read_csv(os.path.join(DATA, "selection_function_%s.csv" % sv))
        g = sf.groupby("m_lo").agg(ns=("n_survey", "sum"), npar=("n_parent", "sum")).reset_index()
        ok = (g.npar >= 10) & (g.ns > 0)
        # the parent sample is a 2% random subset of Gaia DR3, hence the factor 50
        ax.plot(g.m_lo[ok] + 0.125, g.ns[ok] / (g.npar[ok] * 50.0), color=COL[sv], lw=1.8,
                label=NAME[sv] + (" (versus H)" if sv == "apogee" else " (versus G)"))
    ax.set_yscale("log")
    ax.set_ylim(1e-4, 2.0)
    ax.set_xlabel("apparent magnitude (2MASS H for APOGEE, Gaia G otherwise)")
    ax.set_ylabel("fraction of parent FGK stars in clean sample")
    ax.set_title("Survey selection functions inside each footprint", loc="left", fontsize=10, color=INK)
    ax.legend(fontsize=8.5, loc="lower left")
    ax = axes[1]
    for sv in SURVEYS:
        l = lim[(lim["sample"] == sv) & (lim.tier == "three-way")]
        ax.plot(l.f, l.mean_eps, "-o", color=COL[sv], lw=1.8, ms=5, label=NAME[sv])
        r = l.iloc[-1]
        ax.text(r.f + 0.03, r.mean_eps, "%.2f" % r.mean_eps, color=INK, fontsize=8.5, va="center")
    tw = pd.read_csv(os.path.join(DATA, "selection_eff_twin.csv"))
    tw = tw[tw.cls == "all"]
    for i, sv in enumerate(("apogee", "lamost")):
        t = tw[tw.survey == sv].groupby("f").eps_twin.agg(["min", "max"]).reset_index()
        t0 = tw[(tw.survey == sv) & (tw.model == tw.model.iloc[0])]
        ax.errorbar(t0.f + 0.012, t0.eps_twin, yerr=[t0.eps_twin.values - t["min"].values, t["max"].values - t0.eps_twin.values],
                    fmt="D", ms=5.5, mfc=SURF, mec=COL[sv], mew=1.4, ecolor=COL[sv], elinewidth=1.0, capsize=0,
                    zorder=5)
    ax.plot([], [], "D", ms=5.5, mfc=SURF, mec=INK2, mew=1.4, ls="none",
            label="twin cross-check, no selection model (bar: disc models)")
    fgrid = np.linspace(0.2, 0.92, 50)
    ax.plot(fgrid, (1 - fgrid) ** 1.5, color=INK2, lw=1.0, ls=(0, (4, 3)), label=r"uniform density, pure flux limit: $(1-f)^{3/2}$")
    ax.set_yscale("log")
    ax.set_ylim(0.02, 1.2)
    ax.set_yticks([0.03, 0.1, 0.3, 1.0])
    ax.set_yticklabels(["0.03", "0.1", "0.3", "1"])
    ax.yaxis.set_minor_formatter(matplotlib.ticker.NullFormatter())
    ax.set_xlim(0.18, 0.99)
    ax.set_xticks(list(L.F_GRID))
    ax.set_xlabel("fraction of starlight removed, f")
    ax.set_ylabel(r"modelled relative selection efficiency $\epsilon(f)$")
    ax.set_title("Dimmed stars are under-represented in flux-limited samples", loc="left", fontsize=10, color=INK)
    ax.legend(fontsize=8.5, loc="lower left")
    fig.tight_layout()
    fig.savefig(os.path.join(FIGS, "fig3_selection_efficiency.png"), dpi=170)
    plt.close(fig)


def fig4():
    """Upper limits versus dimming fraction, for both sets of astrometric cuts."""
    lim = pd.read_csv(os.path.join(DATA, "limits.csv"))
    lim = lim[lim.tier == "three-way"]
    lo = pd.read_csv(os.path.join(DATA_ORIG, "limits.csv"))
    lo = lo[(lo.tier == "three-way") & (lo["sample"] == "combined")]
    fig, ax = plt.subplots(figsize=(8.6, 5.4))
    for sv in SURVEYS:
        l = lim[lim["sample"] == sv]
        ax.plot(l.f, l.UL95_pop_vetted, "-o", color=COL[sv], lw=1.4, ms=4.5, label=NAME[sv] + " alone, tightened cuts",
                zorder=2, markeredgecolor=SURF, markeredgewidth=0.8)
    l = lim[lim["sample"] == "combined"]
    ax.fill_between(l.f.values, l.UL95_pop_vetted.values, lo.UL95_pop_vetted.values, color=INK, alpha=0.13, lw=0,
                    label="combined: range between the two cut sets", zorder=1)
    ax.plot(l.f, l.UL95_pop_vetted, "-o", color=INK, lw=2.6, ms=7.5, zorder=4, markeredgecolor=SURF, markeredgewidth=0.8,
            label="combined, tightened cuts (post hoc): RUWE < 1.2, parallax S/N >= 50")
    ax.plot(lo.f, lo.UL95_pop_vetted, "--s", color=INK, lw=1.8, ms=6.5, zorder=4, markeredgecolor=SURF, markeredgewidth=0.8,
            label="combined, pre-declared cuts: RUWE < 1.4, parallax S/N >= 20")
    for (_, r), (_, ro) in zip(l.iterrows(), lo.iterrows()):
        hi_, lo_ = max(r.UL95_pop_vetted, ro.UL95_pop_vetted), min(r.UL95_pop_vetted, ro.UL95_pop_vetted)
        ax.annotate("%.1e" % hi_, (r.f, hi_), xytext=(0, 9), textcoords="offset points", ha="center", va="bottom",
                    fontsize=8.5, color=INK, zorder=5, bbox=dict(boxstyle="round,pad=0.15", fc=SURF, ec="none", alpha=0.85))
        ax.annotate("%.1e" % lo_, (r.f, lo_), xytext=(0, -9), textcoords="offset points", ha="center", va="top",
                    fontsize=8.5, color=INK, zorder=5, bbox=dict(boxstyle="round,pad=0.15", fc=SURF, ec="none", alpha=0.85))
    # published infrared limits (Suazo et al. 2022 abstract: 300 K swarms at 90% completion)
    for y, lab in ((2e-5, "Suazo+22 infrared limit, < 100 pc"), (8e-4, "Suazo+22 infrared limit, < 5 kpc")):
        ax.plot([0.925, 0.965], [y, y], color=INK2, lw=1.6)
        ax.text(0.975, y, lab + "\n(300 K waste heat only)", ha="left", va="center", fontsize=8, color=INK2)
    ax.set_yscale("log")
    ax.set_xlim(0.2, 1.32)
    ax.set_ylim(3e-6, 1.5e-1)
    ax.set_xticks(list(L.F_GRID))
    ax.set_xlabel("fraction of starlight removed by a grey absorber, f", loc="left")
    ax.set_ylabel("95% upper limit on the fraction of stars")
    ax.set_title("Limits on grey dimming of FGK stars within 1 kpc under two sets of astrometric cuts\n"
                 "(f <= 0.5: main-sequence stars only; the limit quoted at each f is the weaker of the two)",
                 loc="left", fontsize=10.5, color=INK)
    ax.legend(fontsize=8.3, loc="upper right", title="population limits after vetting", title_fontsize=8.5)
    fig.tight_layout()
    fig.savefig(os.path.join(FIGS, "fig4_limits.png"), dpi=170)
    plt.close(fig)


def fig5():
    """Candidate rate versus astrometric quality under the pre-declared cuts (LAMOST)."""
    d = pd.read_csv(os.path.join(DATA, "astrometry_diagnostic.csv"))
    d = d[(d.survey == "lamost") & d.quantity.isin(["poe", "ruwe"])]
    fig, axes = plt.subplots(2, 2, figsize=(9.6, 6.4), sharey="row")
    for i, f in enumerate((0.25, 0.5)):
        for j, (qty, xlab, cut) in enumerate((("poe", "parallax / error", 50.0), ("ruwe", "RUWE", 1.2))):
            ax = axes[i, j]
            t = d[(d.f == f) & (d.quantity == qty)].reset_index(drop=True)
            x = np.arange(len(t))
            n = t.n_flagged.values
            ns = t.n_sensitive.values.astype(float)
            from scipy.stats import chi2
            lo_ = np.where(n > 0, 0.5 * chi2.ppf(0.16, 2 * np.maximum(n, 1)), 0.0) / ns
            hi_ = 0.5 * chi2.ppf(0.84, 2 * (n + 1)) / ns
            rate = n / ns
            ax.errorbar(x, rate, yerr=[rate - lo_, hi_ - rate], fmt="o", ms=7, color=COL["lamost"], mec=SURF, mew=0.8,
                        ecolor=COL["lamost"], elinewidth=1.4, capsize=0, zorder=3)
            for xx, r_, nn in zip(x, hi_, n):
                ax.text(xx, r_ * 1.25, "%d" % nn, ha="center", va="bottom", fontsize=8.5, color=INK2)
            labs = ["%g-%g" % (a, b) if np.isfinite(b) else ">= %g" % a for a, b in zip(t.lo, t.hi)]
            if qty == "ruwe":
                labs = ["< 1.0"] + labs[1:]
            ax.set_xticks(x)
            ax.set_xticklabels(labs)
            removed = (t.hi.values <= cut) if qty == "poe" else (t.lo.values >= cut)
            if removed.any():
                xr = x[removed]
                ax.axvspan(xr.min() - 0.5, xr.max() + 0.5, color=INK, alpha=0.07, lw=0, zorder=0)
                ax.text(0.5 * (xr.min() + xr.max()), 0.965, "removed by the\ntightened cut", transform=ax.get_xaxis_transform(),
                        ha="center", va="top", fontsize=8.5, color=INK2)
            ax.set_yscale("log")
            ax.set_xlim(-0.5, len(t) - 0.5)
            ax.grid(axis="x", visible=False)
            if j == 0:
                ax.set_ylabel("flagged stars per sensitive star\n(f = %g level)" % f)
            if i == 1:
                ax.set_xlabel(xlab)
        lo_y, hi_y = axes[i, 0].get_ylim()
        axes[i, 0].set_ylim(lo_y * 0.7, hi_y * 4)
    fig.suptitle("LAMOST, pre-declared cuts: the faint-tail rate rises steeply where the astrometry is poorest\n"
                 "(numbers: flagged stars per bin; bars: 68% Poisson intervals)",
                 x=0.01, ha="left", fontsize=10.5, color=INK, y=0.995)
    fig.tight_layout()
    fig.subplots_adjust(top=0.9)
    fig.savefig(os.path.join(FIGS, "fig5_tail_vs_astrometry.png"), dpi=170)
    plt.close(fig)
    old = os.path.join(FIGS, "fig5_cross_survey.png")        # figure of the first version, superseded
    if os.path.exists(old):
        os.remove(old)


if __name__ == "__main__":
    for fn in (fig1, fig2, fig3, fig4, fig5):
        fn()
        print("wrote", fn.__name__)