"""Figures.  Reads the tables written by part1.py and part2.py (plus the input definitions in classic.py).

Colours: the first three slots of a categorical palette validated for colour-vision deficiency on all pairs
(blue, orange, aqua); every series is also labelled directly, and text is always in ink colours.
"""
import json
import os

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

import classic
from config import DATA, FIGS, SEED

BLUE, ORANGE, AQUA = "#2a78d6", "#eb6834", "#1baf7a"
INK, INK2, GRID, SURF = "#0b0b0b", "#52514e", "#dddcd6", "#fcfcfb"

plt.rcParams.update({
    "figure.facecolor": SURF, "axes.facecolor": SURF, "savefig.facecolor": SURF,
    "font.size": 9.5, "axes.edgecolor": GRID, "axes.labelcolor": INK2, "text.color": INK,
    "xtick.color": INK2, "ytick.color": INK2, "axes.spines.top": False, "axes.spines.right": False,
    "axes.grid": True, "grid.color": GRID, "grid.linewidth": 0.6, "axes.titlesize": 10.5,
    "axes.titleweight": "bold", "axes.titlelocation": "left", "legend.frameon": False,
})


def _save(fig, name):
    path = os.path.join(FIGS, name)
    fig.savefig(path, dpi=170, bbox_inches="tight")
    plt.close(fig)
    return path


def fig_timeline():
    """Transition dates (each published estimate) against the conventional oxygenation-event gates and the
    gates that follow from physiological thresholds under low and high mid-Proterozoic oxygen."""
    lags = pd.read_csv(os.path.join(DATA, "part1_lags.csv"))
    rows = [("eukaryotes", "Eukaryotes (LECA)"), ("animals", "Animals (crown Metazoa)"),
            ("humans", "Humans")]
    fig, ax = plt.subplots(figsize=(9.6, 4.6))
    y = 0.0
    yt, yl = [], []
    prop = {"eukaryotes": (2.22, 2.46), "animals": classic.ENV["noe_window"][:2], "humans": classic.ENV["poe"][:2]}
    conv_name = {"eukaryotes": "Great Oxidation", "animals": "Neoproterozoic oxygenation window",
                 "humans": "Palaeozoic oxygenation"}
    for key, name in rows:
        est = classic.TRANSITIONS[key]["estimates"]
        y0 = y
        for (lab, kind, mode, lo, hi, _, _) in est:
            ax.plot([lo, hi], [y, y], color=BLUE, lw=5, solid_capstyle="butt")
            if key != "humans":
                ax.text(lo - 0.03, y, lab, va="center", ha="left", fontsize=7.5, color=INK2)
            y += 0.55
        # conventional oxygenation-event gate
        ax.plot(prop[key], [y, y], color=ORANGE, lw=5, solid_capstyle="butt")
        ax.text(max(prop[key]) + 0.03, y, "conventional gate: " + conv_name[key], va="center", ha="right",
                fontsize=7.5, color=INK2)
        y += 0.55
        for lab, sel in (("threshold gate, low mid-Proterozoic O2", '"midprot": "low"'),
                         ("threshold gate, high mid-Proterozoic O2", '"midprot": "high"')):
            r = lags[(lags.gate_version == "physio") & lags.physio.str.contains(sel, regex=False)
                     & lags.physio.str.contains("animal_sponge") & (lags.transition == key)]
            g = float(r["gate_age_median_Ga"].iloc[0])
            ax.plot([g], [y], marker="D", ms=7, color=AQUA, mec=SURF, mew=1.2)
            ax.text(g + 0.05, y, lab, va="center", ha="right", fontsize=7.5, color=INK2)
            y += 0.55
        yt.append(0.5 * (y0 + y - 0.55)); yl.append(name)
        y += 0.7
        ax.axhline(y - 0.6, color=GRID, lw=0.8)
    ax.set_yticks(yt); ax.set_yticklabels(yl)
    ax.set_xlim(2.75, -0.05); ax.set_ylim(y - 0.6, -0.5)
    ax.grid(axis="y", visible=False)
    ax.set_xlabel("Age (billion years ago)")
    ax.set_title("Dated transitions against candidate oxygen gates")
    h = [plt.Line2D([], [], color=BLUE, lw=5, label="transition date (one bar per published estimate)"),
         plt.Line2D([], [], color=ORANGE, lw=5, label="conventional oxygenation-event gate (the animal window is "
                    "this project's construction, not named by Mills et al. 2025)"),
         plt.Line2D([], [], color=AQUA, marker="D", ms=7, lw=0, label="gate from a physiological O2 threshold "
                    "(median), sponge-grade animals")]
    ax.legend(handles=h, loc="lower center", bbox_to_anchor=(0.5, -0.36), ncol=1, fontsize=8)
    return _save(fig, "fig1_transitions_vs_gates.png")


def fig_bayes():
    bf = pd.read_csv(os.path.join(DATA, "part1_bayes_factors.csv"))
    nsteps = {"C3": "3 steps", "C4": "4 steps", "C5": "5 steps"}
    lab = bf["config"] + " [" + bf["set"].map(nsteps) + "]"
    for old, new in (("physiological gates, ", "threshold gates: "),
                     ("animal_mills25", "animals need 0.1-1% PAL"), ("animal_sponge", "animals need 0.5-4% PAL"),
                     ("animal_catling_circ", "animals need 4.7-47% PAL"), ("animal_catling", "animals need 4.7% PAL"),
                     ("logU_1_1000", "log-uniform 1-1000 Myr"), ("logU_10_2000", "log-uniform 10-2000 Myr"),
                     ("U_0_500", "uniform 0-500 Myr"), ("future (0.9,)", "0.9 Gyr of habitability left"),
                     ("future (1.73,)", "1.73 Gyr left"), ("future (0.9, 1.08, 1.73)", "0.9, 1.08 or 1.73 Gyr left"),
                     ("mid-Proterozoic all", "mid-Proterozoic O2 0.1-10% PAL"),
                     ("mid-Proterozoic low", "mid-Proterozoic O2 0.1-1% PAL"),
                     ("mid-Proterozoic high", "mid-Proterozoic O2 1-10% PAL"),
                     (" (after last ocean-vaporising impact)", ""), ("all, sponge, ", ""),
                     (" (both models conditioned on its date)", ""),
                     ("conventional gates, three steps counted", "conventional gates, steps counted"),
                     ("C4 strict (photosynthesis prompt after life)", "photosynthesis prompt after life if easy"),
                     ("C5 strict", "photosynthesis prompt after life if easy"), (", C5, ", ", "),
                     (" (Mills et al. 2025: animals excluded only before ~2.2 Ga)", " (Mills et al. 2025 reading)"),
                     (" (Dahl et al. 2010)", ""),
                     ("animal gate uncertain: equal mix of Neoproterozoic window, Tonian onset and Great Oxidation",
                      "animal gate uncertain: equal mix of three readings"),
                     ("timing-study inputs (Kay LECA; animals uniform 0.574-0.80 Ga)", "timing-study dates"),
                     ("timing-study inputs and its lag prior (log-uniform 10-3000 Myr)",
                      "timing-study dates and lag prior 10-3000 Myr")):
        lab = lab.str.replace(old, new, regex=False)
    # rows from top to bottom in table order, with a gap and a header before each group
    ypos, yhead, y = [], [], 0.0
    prev = None
    for gname in bf["group"]:
        if gname != prev:
            y += 1.3 if prev is not None else 0.6
            yhead.append((y - 0.75, gname))
            prev = gname
        ypos.append(y)
        y += 1.0
    yy = np.array(ypos)
    fig, ax = plt.subplots(figsize=(9.8, 0.25 * y + 1.7))
    ax.axvspan(1 / 3.0, 3.0, color=GRID, alpha=0.45, lw=0)
    ax.axvline(1.0, color=INK2, lw=0.8)
    lo = 3e-3
    ax.hlines(yy, bf["BF_G_Hg"].clip(lower=lo), bf["BF_G_H"].clip(lower=lo), color=GRID, lw=1.2, zorder=1)
    ax.scatter(bf["BF_G_H"].clip(lower=lo), yy, s=40, color=BLUE, edgecolor=SURF, lw=1, zorder=3,
               label="gating : hard steps (ungated, Carter)")
    ax.scatter(bf["BF_G_Hg"].clip(lower=lo), yy, s=40, color=ORANGE, marker="s", edgecolor=SURF, lw=1, zorder=3,
               label="gating : hard steps that also wait for the gate")
    for yv, (a_, b_) in zip(yy, zip(bf["BF_G_H"], bf["BF_G_Hg"])):
        if a_ < lo:
            ax.text(lo * 1.25, yv, "exactly 0 (every animal date precedes this gate)", fontsize=7, va="center",
                    color=INK2)
    for yh, gname in yhead:
        ax.text(lo, yh, gname.upper(), fontsize=7.2, color=INK2, fontweight="bold", va="center", ha="left")
    ax.set_xscale("log")
    ax.set_xlim(lo * 0.8, 60)
    ax.set_yticks(yy); ax.set_yticklabels(lab, fontsize=7.2)
    ax.set_ylim(y - 0.3, -0.9)
    ax.grid(axis="y", visible=False)
    ax.set_xlabel("Bayes factor (log scale); > 1 favours gating.  Grey band: 1/3 to 3, not worth a mention")
    ax.set_title("Bayes factor for gating over hard steps: no single value, the answer follows the analysis choice")
    ax.legend(loc="lower center", bbox_to_anchor=(0.3, 1.02), ncol=2, fontsize=8, columnspacing=1.2)
    return _save(fig, "fig2_bayes_factor_sensitivity.png")


def fig_posterior():
    p = pd.read_csv(os.path.join(DATA, "part3_posterior_hard_steps.csv"))
    p = p.iloc[::-1].reset_index(drop=True)
    lab = (p["config"].str.replace(", hard gated, flat", "", regex=False)
           .str.replace(", hard steps gated, photosynthesis flat", "", regex=False)
           .str.replace(", hard steps ungated, photosynthesis flat", ", hard steps ungated", regex=False)
           .str.replace(", hard steps gated, photosynthesis strict", ", photosynthesis prompt if easy", regex=False))
    fig, axs = plt.subplots(1, 3, figsize=(11.6, 0.36 * len(p) + 1.7), sharey=True,
                            gridspec_kw=dict(wspace=0.08, width_ratios=[1.15, 1, 1]))
    yy = np.arange(len(p))
    ax = axs[0]
    ax.axvspan(1 / 3.0, 3.0, color=GRID, alpha=0.45, lw=0)
    ax.axvline(1.0, color=INK2, lw=0.8)
    ax.errorbar(p["BF_k0_vs_kge1"], yy, xerr=[p["BF_k0_vs_kge1"] * (1 - np.exp(-2 * p["BF_k0_vs_kge1_lnse"])),
                                               p["BF_k0_vs_kge1"] * (np.exp(2 * p["BF_k0_vs_kge1_lnse"]) - 1)],
                fmt="o", color=INK, ms=6, mec=SURF, ecolor=INK, elinewidth=1.2)
    ax.set_xscale("log"); ax.set_xlim(0.02, 5)
    ax.set_yticks(yy); ax.set_yticklabels(lab, fontsize=7.6)
    ax.set_xlabel("Bayes factor, no hard steps : at least one\n(grey band 1/3 to 3)")
    ax.set_title("What the dates say")
    ax = axs[1]
    for col, mk, tag, name in ((BLUE, "o", "", "candidate-step prior"), (ORANGE, "s", "imp_", "model's implied prior")):
        ax.axvline(float(p[tag + "prior_k0"].iloc[0]), color=col, lw=1, ls="--")
        ax.errorbar(p[tag + "post_k0"], yy + (0.13 if tag else -0.13), xerr=2 * p[tag + "post_k0_mcse"], fmt=mk,
                    color=col, ms=5.5, mec=SURF, ecolor=col, elinewidth=1.1, label=name)
    ax.set_xlim(0, 0.52)
    ax.set_xlabel("Posterior probability of no hard steps\n(dashed: priors 0.30 and 0.35)")
    ax.set_title("P(no hard steps)")
    ax.legend(fontsize=7.6, loc="lower center", bbox_to_anchor=(0.5, 1.07), ncol=1, borderaxespad=0.0)
    ax = axs[2]
    for col, mk, tag in ((BLUE, "o", ""), (ORANGE, "s", "imp_")):
        ax.axvline(float(p[tag + "prior_mean_k"].iloc[0]), color=col, lw=1, ls="--")
        ax.plot(p[tag + "post_mean_k"], yy + (0.13 if tag else -0.13), mk, color=col, ms=5.5, mec=SURF)
    ax.set_xlim(0.6, 3.0)
    ax.set_xlabel("Posterior mean number of hard steps\n(of 4 candidates; dashed: prior means)")
    ax.set_title("Mean number of hard steps")
    for a_ in axs:
        a_.grid(axis="y", visible=False)
        a_.set_ylim(-0.6, len(p) - 0.4)
    fig.suptitle("What the dated record does to the model's prior on the number of hard steps", x=0.01, ha="left",
                 fontsize=11, fontweight="bold", y=1.04)
    return _save(fig, "fig3_posterior_hard_steps.png")


def fig_phanerozoic():
    c = pd.read_csv(os.path.join(DATA, "curves_1myr.csv"))
    inn = pd.read_csv(os.path.join(DATA, "part2_per_innovation.csv"))
    fig, axs = plt.subplots(2, 1, figsize=(10, 6.6), sharex=True, gridspec_kw=dict(height_ratios=[1.25, 1], hspace=0.06))
    ax = axs[0]
    a = c["age_Ma"]
    ax.fill_between(a, c["gcso_p16"], c["gcso_p84"], color=BLUE, alpha=0.16, lw=0)
    ax.plot(a, c["gcso_median"], color=BLUE, lw=2)
    m = a <= 800
    ax.plot(a[m], c["copse_standard"][m], color=ORANGE, lw=2)
    ax.fill_between(a, c["charcoal_lo"], c["charcoal_hi"], color=AQUA, alpha=0.16, lw=0)
    ax.plot(a, c["charcoal_mid"], color=AQUA, lw=2)
    ax.axhline(20.95, color=INK2, lw=0.8, ls=":")
    ax.text(-4, 19.3, "present\nlevel", fontsize=7.5, color=INK2, ha="left", va="top")
    ax.text(520, 26.5, "isotope mass balance\n(GEOCARBSULFOR)", fontsize=8, color=INK, ha="center")
    ax.plot([510, 495], [25.8, 15.0], color=BLUE, lw=0.8)
    ax.text(585, 15.5, "forward model\n(COPSE)", fontsize=8, color=INK, ha="center")
    ax.plot([585, 577], [14.9, 12.4], color=ORANGE, lw=0.8)
    ax.text(92, 34.0, "charcoal\nenvelope", fontsize=8, color=INK, ha="center")
    ax.plot([92, 100], [33.4, 29.0], color=AQUA, lw=0.8)
    ax.set_ylabel("Atmospheric O2 (per cent)")
    ax.set_ylim(0, 38)
    ax.set_title("Three reconstructions of Phanerozoic oxygen and the first appearances of 31 innovations")
    ax = axs[1]
    inn = inn.sort_values("origin50", ascending=False).reset_index(drop=True)
    lanes = 8
    for i, r in inn.iterrows():
        yv = lanes - (i % lanes)
        col = INK if r["o2_class"] == "high" else INK2
        ax.plot([r["fad_old"] + r["ext95"], r["fad_young"]], [yv, yv], color=col, lw=2.2, solid_capstyle="butt")
        ax.plot([r["origin50"]], [yv], marker="o" if r["o2_class"] == "high" else "s", ms=4.5, color=col, mec=SURF,
                mew=0.8)
        ax.text(r["fad_young"] - 4, yv, r["id"].replace("_", " "), fontsize=6.8, va="center", ha="left",
                color=INK2)
    ax.set_ylim(0.3, lanes + 0.7)
    ax.set_yticks([])
    ax.grid(axis="y", visible=False)
    ax.set_xlim(610, -45)
    ax.set_xlabel("Age (million years ago)\nbars: first-appearance bracket plus 95% range extension;  "
                  "circles: high oxygen demand;  squares: plant or other")
    return _save(fig, "fig4_phanerozoic_oxygen_and_origins.png")


def fig_nulls():
    nd = pd.read_csv(os.path.join(DATA, "part2_null_primary.csv"))
    pw = pd.read_csv(os.path.join(DATA, "part2_power.csv"))
    pn = pd.read_csv(os.path.join(DATA, "part2_power_vs_n.csv"))
    fig = plt.figure(figsize=(10.2, 6.6))
    gs = fig.add_gridspec(2, 6, hspace=0.62, wspace=1.0)
    cols = {"O2 GEOCARBSULFOR median": BLUE, "O2 COPSE standard": ORANGE, "O2 charcoal midpoint": AQUA}
    short = {"O2 GEOCARBSULFOR median": "GEOCARBSULFOR", "O2 COPSE standard": "COPSE", "O2 charcoal midpoint": "Charcoal"}
    for j, (cn, col) in enumerate(cols.items()):
        ax = fig.add_subplot(gs[0, 2 * j:2 * j + 2])
        d = nd[nd.curve == cn]
        obs = float(d[d.shift_Myr == 0]["mean_rise"].iloc[0])
        null = d[d.shift_Myr > 0]["mean_rise"].to_numpy()
        ax.hist(null, bins=28, color=col, alpha=0.55, edgecolor=SURF, lw=0.6)
        ax.axvline(obs, color=INK, lw=2)
        p = (np.sum(null >= obs) + 1) / (len(null) + 1)
        ax.set_title("%s\nobserved %.2f (line), null %.2f, p = %.2f" % (short[cn], obs, null.mean(), p),
                     fontsize=8.4)
        ax.set_xlabel("Mean O2 change over the 20 Myr\nbefore origin (percentage points)", fontsize=8)
        if j == 0:
            ax.set_ylabel("Circular shifts")
        ax.grid(axis="x", visible=False)
    for j, tau in enumerate((10.0, 30.0)):
        ax = fig.add_subplot(gs[1, 3 * j:3 * j + 3])
        for true, test, col, ls, lab in (
                ("GEOCARBSULFOR median", "GEOCARBSULFOR median", BLUE, "-", "right curve (GEOCARBSULFOR)"),
                ("COPSE standard", "COPSE standard", ORANGE, "-", "right curve (COPSE)"),
                ("GEOCARBSULFOR median", "COPSE standard", BLUE, "--", "truth GEOCARBSULFOR, tested on COPSE"),
                ("COPSE standard", "GEOCARBSULFOR median", ORANGE, "--", "truth COPSE, tested on GEOCARBSULFOR")):
            d = pw[(pw.true_curve == true) & (pw.test_curve == test) & ((pw.mean_lag_Myr == tau) | (pw.gated_fraction == 0))]
            d = d.sort_values("gated_fraction")
            ax.plot(d["gated_fraction"], d["power_rise_20Myr"], color=col, ls=ls, lw=2, marker="o", ms=4.5,
                    mec=SURF, label=lab)
        # the same test with fewer, independent origins (the real ones are clustered): open markers at 0.5
        for cname, col, dx in (("GEOCARBSULFOR median", BLUE, -0.012), ("COPSE standard", ORANGE, 0.012)):
            d = pn[(pn.curve == cname) & (pn.mean_lag_Myr == tau) & (pn.n_independent_origins < 27)]
            ax.plot(0.5 + dx + 0 * d["power_rise_20Myr"], d["power_rise_20Myr"], "o", ms=5, mfc=SURF, mec=col,
                    mew=1.4, label="right curve, only 14 or 9 independent origins" if cname.startswith("GEO") else None)
        ax.axhline(0.05, color=INK2, lw=0.8, ls=":")
        ax.set_ylim(0, 1.03); ax.set_xlim(-0.03, 1.03)
        ax.set_xlabel("Fraction of the innovations that are truly oxygen-gated")
        ax.set_title("Power (an upper bound), mean gate-to-origin lag %d Myr" % tau, fontsize=9.5)
        if j == 0:
            ax.set_ylabel("Probability of detection (alpha 0.05)")
        else:
            ax.legend(fontsize=7.6, loc="upper center", bbox_to_anchor=(-0.12, -0.24), ncol=3, handlelength=2.6,
                      columnspacing=1.4)
    fig.suptitle("Do the innovations follow rises in oxygen?  Observed statistic against the circular-shift null, "
                 "and what the test could have seen", x=0.01, ha="left", fontsize=10.5, fontweight="bold", y=0.975)
    return _save(fig, "fig5_phanerozoic_null_and_power.png")


def run_all():
    os.makedirs(FIGS, exist_ok=True)
    return [fig_timeline(), fig_bayes(), fig_posterior(), fig_phanerozoic(), fig_nulls()]


if __name__ == "__main__":
    for p in run_all():
        print(p)