"""Step 15: figures (PNG) for the report."""
import os
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from scipy.stats import norm
from common import DATA, FIGS, mad_sigma

# palette (first three categorical slots validate for all-pairs use; grey for context)
BLUE, ORANGE, AQUA, GREY, INK, INK2, SURF = "#2a78d6", "#eb6834", "#1baf7a", "#b9b8b2", "#0b0b0b", "#52514e", "#fcfcfb"
plt.rcParams.update({
    "figure.facecolor": SURF, "axes.facecolor": SURF, "savefig.facecolor": SURF, "font.size": 10,
    "axes.edgecolor": "#8a8984", "axes.labelcolor": INK, "xtick.color": INK2, "ytick.color": INK2,
    "axes.grid": True, "grid.color": "#e6e5e0", "grid.linewidth": 0.6, "axes.spines.top": False,
    "axes.spines.right": False, "legend.frameon": False, "axes.titlesize": 11, "axes.titleweight": "bold"})
FS = (0.5, 0.75, 0.9)


def fig_tf():
    cf = pd.read_csv(os.path.join(DATA, "cf4_resid.csv"))
    al = pd.read_csv(os.path.join(DATA, "alfalfa_resid.csv"), low_memory=False)
    al = al[al.gold]
    fc = pd.read_csv(os.path.join(DATA, "cf4_tf_fits.csv")).set_index("relation")
    fa = pd.read_csv(os.path.join(DATA, "alfalfa_tf_fits.csv")).set_index("relation")
    fin = pd.read_csv(os.path.join(DATA, "candidates_final.csv"))
    fin = fin[fin.f == 0.75]
    fig, axes = plt.subplots(1, 2, figsize=(12, 5.2), sharey=True)
    xs = np.linspace(1.9, 2.75, 100)
    for ax, d, fit, title, src in ((axes[0], cf[cf.logLi.notna()], fc.loc["i"], "Cosmicflows-4 (i band; H0 = 75)", "cf4"),
                                   (axes[1], al, fa.loc["i"], "ALFALFA-SDSS gold (i band; H0 = 70)", "alfalfa")):
        ax.scatter(d.logV, d.logLi, s=3, color=GREY, alpha=0.5, lw=0, rasterized=True, label=f"sample (N = {len(d):,})")
        y = fit.c0 + fit.c1 * (xs - fit.x0) + fit.c2 * (xs - fit.x0) ** 2
        ax.plot(xs, y, color=INK, lw=2, label="robust fit")
        for f, ls in zip(FS, ("--", "-.", ":")):
            ax.plot(xs, y + np.log10(1 - f), color=INK2, lw=1.2, ls=ls, label=f"f = {f:g} dimming")
        key = "PGC" if src == "cf4" else "AGC"
        c = fin[fin.src == src].merge(d[[key, "logV", "logLi"]].rename(columns={"logV": "x", "logLi": "y"}), on=key, how="left")
        ex, un = c[c.status != "UNRESOLVED"], c[c.status == "UNRESOLVED"]
        ax.scatter(ex.x, ex.y, s=34, marker="x", color=BLUE, lw=1.4, label=f"candidate, explained ({len(ex)})")
        ax.scatter(un.x, un.y, s=46, marker="o", facecolor=ORANGE, edgecolor=SURF, lw=1.0, label=f"candidate, unresolved ({len(un)})")
        ax.set_title(title)
        ax.set_xlabel("log10 rotation velocity [km/s]")
        ax.set_xlim(1.88, 2.72)
        ax.set_ylim(7.6, 11.6)
        ax.legend(loc="lower right", fontsize=8.5)
    axes[0].set_ylabel("log10 i-band luminosity [Lsun]")
    fig.suptitle("Tully-Fisher relation and candidates more than 1.5 mag underluminous (f >= 0.75)", y=0.995, fontsize=12)
    fig.tight_layout()
    fig.savefig(os.path.join(FIGS, "fig1_tully_fisher.png"), dpi=150)
    plt.close(fig)


def fig_resid():
    u = pd.read_csv(os.path.join(DATA, "union_sample.csv"))
    fin = pd.read_csv(os.path.join(DATA, "candidates_final.csv"))
    fig, axes = plt.subplots(1, 2, figsize=(12, 4.8))
    ax = axes[0]
    bins = np.arange(-2.0, 1.61, 0.05)
    ax.hist(u.D, bins=bins, histtype="stepfilled", color=GREY, alpha=0.7, label=f"all sample galaxies ({len(u):,})")
    c = u[u.auto_clean]
    ax.hist(c.D, bins=bins, histtype="step", color=BLUE, lw=2, label=f"pass catalogue flags ({len(c):,})")
    sig = mad_sigma(c.D)
    x = np.linspace(-2, 1.6, 400)
    ax.plot(x, len(c) * 0.05 * norm.pdf(x, np.median(c.D), sig), color=INK, lw=1.2, ls="--", label=f"Gaussian, robust sigma = {sig:.2f} dex")
    for f in FS:
        t = np.log10(1 - f)
        ax.axvline(t, color=INK2, lw=1, ls=":")
        ax.text(t, 2.2e3, f" f={f:g}", rotation=90, va="top", ha="right", fontsize=8.5, color=INK2)
    ax.set_yscale("log")
    ax.set_ylim(0.6, 3e3)
    ax.set_xlabel("Tully-Fisher residual D [dex]  (negative = underluminous for its rotation speed)")
    ax.set_ylabel("galaxies per 0.05 dex")
    ax.set_title("Residual distribution, union sample")
    ax.legend(loc="upper right", fontsize=8.5)
    # vetting funnel
    ax = axes[1]
    fun = pd.read_csv(os.path.join(DATA, "vetting_funnel.csv"))
    stages = ["N_raw", "N_after_stage1", "N_after_stage2", "N_after_visual", "N_unresolved_baryon_deficient"]
    names = ["below\nthreshold", "pass catalogue\nflags", "pass object\nchecks =\nUNRESOLVED",
             "also survive\nvisual vetting\n(variant)*", "unresolved and\nbaryon-\ndeficient"]
    w = 0.26
    for k, (f, col) in enumerate(zip(FS, (BLUE, ORANGE, AQUA))):
        r = fun[fun.f == f].iloc[0]
        vals = [r[s] for s in stages]
        xpos = np.arange(len(stages)) + (k - 1) * (w + 0.02)
        ok = [np.isfinite(vv) for vv in vals]
        ax.bar(xpos[ok], np.array(vals)[ok], width=w, color=col, label=f"f >= {f:g}")
        for xx, vv in zip(xpos, vals):
            if np.isfinite(vv):
                ax.text(xx, max(vv, 0.8) * 1.12, f"{int(vv)}", ha="center", va="bottom", fontsize=8, color=INK)
    ax.set_yscale("log")
    ax.set_ylim(0.7, 2500)
    ax.set_xticks(np.arange(len(stages)))
    ax.set_xticklabels(names, fontsize=8.5)
    ax.set_ylabel("number of candidates")
    ax.set_title("Vetting funnel")
    ax.grid(axis="x", visible=False)
    ax.legend(loc="upper right")
    ax.text(0.0, -0.30, "*visual vetting (f >= 0.75 and 0.9 lists only) rejected a similar share of blind controls,\n"
            " so it is not used for the primary limits", transform=ax.transAxes, fontsize=8, color=INK2)
    fig.tight_layout()
    fig.savefig(os.path.join(FIGS, "fig2_residuals_and_vetting.png"), dpi=150)
    plt.close(fig)


def fig_limits():
    lim = pd.read_csv(os.path.join(DATA, "limits.csv"))
    var = pd.read_csv(os.path.join(DATA, "limit_variants.csv"))
    tt = pd.read_csv(os.path.join(DATA, "limits_vs_temperature.csv"))
    fig = plt.figure(figsize=(14.2, 5.6))
    gs = fig.add_gridspec(1, 3, width_ratios=[1.55, 1, 1], wspace=0.22)
    ax = fig.add_subplot(gs[0])
    prim = lim[lim.case.str.startswith("PRIMARY")].sort_values("f_min")
    gold = lim[lim.case == "stratum: ALFALFA gold subsample (uniform HI selection)"].sort_values("f_min")
    rng = var[var.in_range.astype(bool)].groupby("f_min").UL95_fraction.agg(["min", "max"])
    # range over defensible analysis variants (thin whisker), primary value (marker)
    for r in prim.itertuples():
        ax.plot([r.f_min, r.f_min], [rng.loc[r.f_min, "min"], rng.loc[r.f_min, "max"]], color=BLUE, lw=1.5, zorder=3)
    ax.plot(prim.f_min, prim.UL95_fraction, color=BLUE, lw=0, marker="o", ms=9, mec=SURF, mew=1.2, zorder=4,
            label="this work, primary (whisker: range over analysis variants)")
    ax.plot(gold.f_min + 0.016, gold.UL95_fraction, color=ORANGE, lw=0, marker="s", ms=7.5, mec=SURF, mew=1.2, zorder=4,
            label="this work, uniformly HI-selected ALFALFA subsample alone")
    for r in prim.itertuples():
        ax.annotate(f"{100 * r.UL95_fraction:.2g}%  ({r.N_unresolved} / {r.N_eff:.0f})", (r.f_min, r.UL95_fraction),
                    textcoords="offset points", xytext=(10, -15), ha="left", va="center", fontsize=8.5, color=INK)
    # literature (all at f = 0.75, the Annis criterion); drawn just left of our f = 0.75 point
    xl = 0.728
    ax.plot([xl, xl], [0.04, 0.06], color=INK2, lw=4, solid_capstyle="round", zorder=3)
    ax.annotate("Annis 1999 (57 discs), re-derived\nby Zackrisson+15: 4-6%", (xl, 0.049), textcoords="offset points",
                xytext=(-9, -2), ha="right", va="center", fontsize=8.5, color=INK2)
    ax.scatter([xl], [0.0275], marker="v", s=70, color=INK2, zorder=3)
    ax.annotate("Zackrisson+15, unvetted: 11 / ~400\n(quoted as 3%; 95% Poisson limit 4.5%)", (xl, 0.0275),
                textcoords="offset points", xytext=(-9, -4), ha="right", va="center", fontsize=8.5, color=INK2)
    ax.scatter([xl], [0.003], marker="v", s=70, facecolor="none", edgecolor=INK2, zorder=3)
    ax.annotate("Zackrisson+15 if all 11 are natural: 1 / 400", (xl, 0.003), textcoords="offset points",
                xytext=(-9, -3), ha="right", fontsize=8.5, color=INK2)
    ax.set_yscale("log")
    ax.set_ylim(2e-3, 0.5)
    ax.set_xlim(0.44, 1.06)
    ax.set_xticks([0.5, 0.75, 0.9])
    ax.set_xticklabels(["0.5\n(0.75 mag)", "0.75\n(1.5 mag)", "0.9\n(2.5 mag)"])
    ax.set_xlabel("minimum fraction f of starlight intercepted")
    ax.set_ylabel("95% upper limit on fraction of HI-rich disc galaxies")
    ax.set_title("Limit vs harvested fraction (waste heat < ~250 K)")
    ax.legend(loc="upper right", fontsize=8.5)
    # ---- limits as a function of waste-heat temperature (small multiples, shared y axis)
    axT = None
    for j, f in enumerate((0.75, 0.9)):
        axT = fig.add_subplot(gs[1 + j], sharey=axT) if axT is not None else fig.add_subplot(gs[1 + j])
        t = tt[tt.f_min == f]
        axT.plot(t.T_K, t.UL95_any_cause, color=BLUE, lw=2, marker="o", ms=6, mec=SURF, mew=1, label="any cause (all unresolved candidates)")
        tw = t[t.UL95_waste_heat_at_T.notna()]
        axT.plot(tw.T_K, tw.UL95_waste_heat_at_T, color=AQUA, lw=2, marker="D", ms=6, mec=SURF, mew=1,
                 label="only candidates whose W3/W4 flux allows it")
        for T, dx in ((100, 0), (600, -6)):
            r = t[t.T_K == T].iloc[0]
            axT.annotate(f"{100 * r.UL95_any_cause:.2g}%", (T, r.UL95_any_cause), textcoords="offset points", xytext=(dx, 8),
                         ha="center", fontsize=8.5, color=INK)
            axT.annotate(f"{100 * r.UL95_waste_heat_at_T:.2g}%", (T, r.UL95_waste_heat_at_T), textcoords="offset points",
                         xytext=(dx, -15), ha="center", fontsize=8.5, color=INK)
        axT.axvspan(275, 620, color=GREY, alpha=0.25, lw=0)
        axT.text(450, 2.4e-3, "waste heat fills in the\n3.4 micron band:\nsensitivity falls", ha="center", va="bottom",
                 fontsize=8, color=INK2)
        axT.set_yscale("log")
        axT.set_ylim(2e-3, 0.5)
        axT.set_xlim(30, 620)
        axT.set_xlabel("waste-heat temperature [K]")
        axT.set_title(f"f >= {f:g}: limit vs waste-heat temperature")
        if j == 0:
            axT.set_ylabel("95% upper limit on fraction")
            axT.legend(loc="upper left", fontsize=8.5)
        else:
            plt.setp(axT.get_yticklabels(), visible=False)
    fig.text(0.012, 0.012, "Not shown: Griffith+15 (WISE) find 0 of ~1e5 galaxies above f = 0.85 for waste heat near 300 K; "
             "Lacki 2016 (Planck) finds no fully shrouded galaxy near the CMB temperature among ~3 million.",
             fontsize=8, color=INK2)
    fig.subplots_adjust(left=0.06, right=0.99, top=0.93, bottom=0.17)
    fig.savefig(os.path.join(FIGS, "fig3_upper_limits.png"), dpi=150)
    plt.close(fig)


def fig_cluster():
    pc = pd.read_csv(os.path.join(DATA, "pair_correlation.csv"))
    cs = pd.read_csv(os.path.join(DATA, "clustering_summary.csv"))
    cs = cs[cs["sample"] == "all sample galaxies"].reset_index(drop=True)
    fig, axes = plt.subplots(1, 3, figsize=(13.5, 4.3))
    ax = axes[0]
    x = np.sqrt(np.maximum(pc.s_lo, 0.5) * pc.s_hi)
    ax.axhline(0, color=INK2, lw=0.8)
    ax.fill_between(x, -2 * pc.corr_null_sd, 2 * pc.corr_null_sd, color=GREY, alpha=0.5, label="permutation null, +/- 2 sd")
    ax.plot(x, pc["corr"], color=BLUE, lw=2, marker="o", ms=7, mec=SURF, label="observed")
    ax.set_xscale("log")
    ax.set_xlabel("3D pair separation [Mpc]")
    ax.set_ylabel("correlation coefficient of residuals")
    ax.set_title("Pairwise residual correlation")
    ax.legend(fontsize=8.5)
    ax = axes[1]
    xx = np.arange(len(cs))
    ax.bar(xx - 0.2, cs.n_flagged, width=0.36, color=BLUE, label="observed")
    ax.bar(xx + 0.2, cs.null_mean, width=0.36, color=GREY, label="permutation mean")
    ax.errorbar(xx + 0.2, cs.null_mean, yerr=[cs.null_mean * 0, np.maximum(cs.null_p95 - cs.null_mean, 0)], fmt="none", ecolor=INK2, lw=1, capsize=3)
    for a, b in zip(xx, cs.n_flagged):
        ax.text(a - 0.2, b + 0.1, str(int(b)), ha="center", fontsize=9)
    ax.set_xticks(xx)
    ax.set_xticklabels([f"{r:g}" for r in cs.R_Mpc])
    ax.set_xlabel("region radius R [Mpc]")
    ax.set_ylabel("groups with median residual < -0.20 dex")
    ax.set_title("Coherently dim groups (cap: null 95th pct)")
    ax.grid(axis="x", visible=False)
    ax.legend(fontsize=8.5)
    ax = axes[2]
    c5 = cs[cs.R_Mpc >= 5].reset_index(drop=True)   # R = 2 Mpc is below the redshift-space smearing scale: not shown
    x5 = np.arange(len(c5))
    ax.plot(x5, c5.UL95_fraction_of_regions, color=BLUE, lw=2, marker="o", ms=8, mec=SURF)
    for a, r in zip(x5, c5.itertuples()):
        ax.annotate(f"{100 * r.UL95_fraction_of_regions:.1f}%\n{r.n_groups} regions\nefficiency {r.injection_efficiency:.2f}",
                    (a, r.UL95_fraction_of_regions), textcoords="offset points", xytext=(0, 12), ha="center", fontsize=8.5)
    ax.set_yscale("log")
    ax.set_ylim(5e-3, 0.4)
    ax.set_xlim(-0.5, len(c5) - 0.5)
    ax.set_xticks(x5)
    ax.set_xticklabels([f"{r:g}" for r in c5.R_Mpc])
    ax.grid(axis="x", visible=False)
    ax.set_xlabel("region radius R [Mpc]")
    ax.set_ylabel("95% UL on fraction of regions")
    ax.set_title("Regions (>= 5 sample galaxies) dimmed >= 50%")
    fig.tight_layout()
    fig.savefig(os.path.join(FIGS, "fig4_clustering.png"), dpi=150)
    plt.close(fig)


def fig_plane():
    u = pd.read_csv(os.path.join(DATA, "union_sample.csv"))
    fin = pd.read_csv(os.path.join(DATA, "candidates_final.csv"))
    fin = fin[fin.f == 0.75]
    fig, axes = plt.subplots(1, 2, figsize=(12, 4.9))
    ax = axes[0]
    c = u[u.auto_clean]
    ax.scatter(c.D, c.r_bar, s=3, color=GREY, alpha=0.5, lw=0, rasterized=True, label="sample (catalogue-clean)")
    ex, un = fin[fin.status != "UNRESOLVED"], fin[fin.status == "UNRESOLVED"]
    ax.scatter(ex.D, ex.r_bar, s=30, marker="x", color=BLUE, lw=1.3, label=f"candidate, explained ({len(ex)})")
    ax.scatter(un.D, un.r_bar, s=46, marker="o", facecolor=ORANGE, edgecolor=SURF, lw=1.0, label=f"candidate, unresolved ({len(un)})")
    ax.axvline(np.log10(0.25), color=INK2, lw=1, ls=":")
    ax.axhline(-0.35, color=INK2, lw=1, ls=":")
    ax.text(0.97, -0.38, "baryon-deficient below this line", fontsize=8.5, color=INK2, va="top", ha="right")
    ax.set_xlim(-2.0, 1.0)
    ax.set_ylim(-1.4, 0.8)
    ax.set_xlabel("stellar-light Tully-Fisher residual D [dex]")
    ax.set_ylabel("baryonic Tully-Fisher residual [dex]")
    ax.set_title("Starlight deficit vs baryon deficit")
    ax.legend(loc="upper left", fontsize=8.5)
    ax = axes[1]
    vt = pd.read_csv(os.path.join(DATA, "vet_table.csv"))
    k = vt[vt.role == "control"]
    ax.scatter(k.D, k.f_warm_max, s=14, color=GREY, lw=0, label="random controls")
    m = fin.merge(vt[["uid", "f_warm_max"]], on="uid", how="left", suffixes=("", "_v"))
    fw = m.f_warm_max if "f_warm_max" in m else m.f_warm_max_v
    exm, unm = m.status != "UNRESOLVED", m.status == "UNRESOLVED"
    ax.scatter(m.D[exm], fw[exm], s=30, marker="x", color=BLUE, lw=1.3, label="candidate, explained")
    ax.scatter(m.D[unm], fw[unm], s=46, marker="o", facecolor=ORANGE, edgecolor=SURF, lw=1.0, label="candidate, unresolved")
    dd = np.linspace(-2, -0.05, 100)
    ax.plot(dd, 1 - 10 ** dd, color=INK, lw=1.5, ls="--", label="f needed to explain the deficit")
    ax.set_xlim(-2.0, 1.0)
    ax.set_ylim(0, 1.02)
    ax.set_xlabel("stellar-light Tully-Fisher residual D [dex]")
    ax.set_ylabel("largest f allowed by W3/W4 if waste heat is 100-600 K")
    ax.set_title("Mid-infrared test for warm waste heat")
    ax.legend(loc="upper right", fontsize=8.5)
    fig.tight_layout()
    fig.savefig(os.path.join(FIGS, "fig5_baryons_and_midIR.png"), dpi=150)
    plt.close(fig)


def main():
    for fn in (fig_tf, fig_resid, fig_limits, fig_cluster, fig_plane):
        fn()
        print(f"  wrote {fn.__name__}")


if __name__ == "__main__":
    main()