"""Figures for the timing study.  Every function reads only objects produced by run_all.py."""
from __future__ import annotations

import os

import numpy as np
import pandas as pd

import config as cfg

cfg.style()
import matplotlib.pyplot as plt  # noqa: E402
from matplotlib.lines import Line2D  # noqa: E402

from events_table import EVENTS  # noqa: E402
from timing_lib import bmax_hardsteps, weighted_quantile  # noqa: E402

C = cfg.C
GREY = "#8c8b86"
LN10 = np.log(10.0)


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


# ---------------------------------------------------------------------------------
def fig1_timeline():
    ev = {e["key"]: e for e in EVENTS}
    fig, ax = plt.subplots(figsize=(11.5, 6.2))
    rows = {"window": 6.0, "life": 4.7, "transition": 3.4, "environment": 2.15, "order": 1.0, "order2": 0.3}
    x = lambda age: -age            # plot time increasing to the right; 0 = today

    # habitable window bar
    ax.barh(rows["window"], 4.404 + 0.9, left=x(4.404), height=0.28, color="#d9e6f5", edgecolor="none")
    ax.barh(rows["window"], 1.84 - 0.9, left=0.9, height=0.28, color="#edf2f8", edgecolor="none")
    ax.plot([x(4.51), x(4.404)], [rows["window"]] * 2, color=C["blue"], lw=2)
    ax.text(x(4.46), rows["window"] + 0.28, "start: 4.51 Ga (Moon impact)\nto 4.404 Ga (zircon water)",
            ha="center", va="bottom", fontsize=8, color=cfg.INK2)
    for key, lab, dy in [("end_ck92", "0.9", 0.26), ("end_or21", "1.08", -0.56), ("end_gha24", "1.73", 0.26),
                         ("end_hmw26", "1.84", -0.56)]:
        ax.plot(-ev[key]["age"], rows["window"], "o", color=C["blue"], ms=6, mec=cfg.SURFACE, mew=1.2)
        ax.text(-ev[key]["age"], rows["window"] + dy, lab, ha="center", va="bottom", fontsize=8, color=cfg.INK2)
    ax.text(1.37, rows["window"] + 0.6, "end of habitable window: four published\nestimates (Gyr from now)",
            ha="center", va="bottom", fontsize=8, color=cfg.INK2)

    def pt(key, row, color):
        e = ev[key]
        ax.plot([x(e["hi"]), x(e["lo"])], [row, row], color=color, lw=2, solid_capstyle="round")
        ax.plot(x(e["age"]), row, "o", ms=7, mec=color, mew=1.5, mfc=color)

    def lab(xx, yy, txt, ha="center", va="bottom"):
        ax.text(xx, yy, txt, ha=ha, va=va, fontsize=8, color=cfg.INK)

    pt("luca", rows["life"], C["orange"])
    lab(x(4.2), rows["life"] + 0.26, "LUCA 4.2\n(molecular clock, 4.33-4.09)")
    pt("dresser", rows["life"], C["orange"])
    lab(x(3.40), rows["life"], "Dresser stromatolites 3.48 (accepted fossil evidence)", ha="left", va="center")
    pt("oxyphoto", rows["transition"], C["aqua"])
    lab(x(2.99), rows["transition"] + 0.26, "O$_2$ photosynthesis 2.99")
    pt("leca", rows["transition"], C["aqua"])
    lab(x(1.735), rows["transition"] + 0.26, "eukaryotes (LECA) 1.74")
    pt("metazoa", rows["transition"], C["aqua"])
    lab(x(0.70), rows["transition"] + 0.26, "animals 0.60")
    pt("sapiens", rows["transition"], C["aqua"])
    lab(0.07, rows["transition"], "humans", ha="left", va="center")
    for key in ("goe", "noe", "poe"):
        e = ev[key]
        ax.barh(rows["environment"], e["hi"] - e["lo"], left=x(e["hi"]), height=0.3, color=C["violet"])
    lab(x(2.34), rows["environment"] + 0.24, "Great Oxidation\n2.46-2.22")
    lab(x(0.82), rows["environment"] + 0.24, "Neoproterozoic\n0.80-0.54")
    lab(x(0.41), rows["environment"] - 0.24, "Paleozoic\n0.42-0.40", va="top")

    # expected order statistics for n = 4 hard steps
    for row, F, txt in [(rows["order"], 0.9, "if the window ends in 0.9 Gyr"), (rows["order2"], 1.73, "if the window ends in 1.73 Gyr")]:
        T = 4.404 + F
        exp = np.arange(1, 5) / 5 * T - 4.404
        ax.plot([x(4.404), F], [row, row], color=GREY, lw=1)
        ax.plot(exp, [row] * 4, "|", color=cfg.INK, ms=14, mew=2)
        ax.text(F + 0.06, row, txt, ha="left", va="center", fontsize=8, color=cfg.INK2)
    ax.plot([x(4.2), x(1.735), x(0.60), 0], [rows["order"] + 0.36] * 4, "v", color=C["aqua"], ms=6)
    ax.text(x(4.404) - 0.05, rows["order"] + 0.36, "observed", ha="right", va="center", fontsize=8, color=cfg.INK2)

    ax.axvline(0, color=cfg.INK2, lw=0.8, ls=(0, (3, 3)))
    ax.set_yticks([rows["window"], rows["life"], rows["transition"], rows["environment"], 0.65])
    ax.set_yticklabels(["habitable\nwindow", "earliest\nlife", "biological\ntransitions", "oxygen\ngates",
                        "expected times,\n4 hard steps"])
    ax.set_ylim(-0.2, 7.5)
    ax.set_xlim(-4.85, 3.3)
    ticks = np.arange(-4.5, 2.01, 0.5)
    ax.set_xticks(ticks)
    ax.set_xticklabels(["today" if abs(t) < 1e-9 else (f"{-t:g}" if t < 0 else f"+{t:g}") for t in ticks])
    ax.set_xlabel("age (Ga before present)  |  future (Gyr from now)")
    ax.grid(axis="y", visible=False)
    ax.set_title("Earth's timeline as used in this study: dated events, oxygen gates and the habitable window")
    return _save(fig, "fig1_timeline.png")


# ---------------------------------------------------------------------------------
def fig2_abiogenesis(lam, curves, moody_tab, moody_summ=None):
    fig, axs = plt.subplots(1, 3, figsize=(14.5, 4.9), gridspec_kw={"wspace": 0.3})
    ax = axs[0]
    order = ["1 later hard step(s)", "2 later hard step(s)", "4 later hard step(s)",
             "1 later hard step + 1.5 Gyr delay", "pacing, deadline = today",
             "1 later step at 1.4 per Gyr (not hard)"]
    labels = ["1 later hard step (Kipping's limit)", "2 later hard steps", "4 later hard steps",
              "1 hard step + 1.5 Gyr fixed delay", "life must start before today (no hard steps)",
              "1 later step that is not hard (1.4 per Gyr)"]
    cols = cfg.SERIES[:5] + [C["violet"]]
    for k, (key, lab) in enumerate(zip(order, labels)):
        ax.plot(lam, curves[key], color=cols[k], label=lab, ls=(0, (5, 2)) if k == 5 else "-")
        dy = {"1 later hard step + 1.5 Gyr delay": 0.45, "2 later hard step(s)": -0.45}.get(key, 0.0)
        ax.text(lam[-1] * 1.25, curves[key][-1] + dy, f"{curves[key][-1]:.1f}", va="center", fontsize=8.5)
    ax.set_xscale("log")
    ax.set_xlim(lam[0], lam[-1] * 6)
    ax.set_ylim(0, 33)
    ax.set_xlabel("abiogenesis rate (events per Gyr)")
    ax.set_ylabel("likelihood relative to rate -> 0")
    ax.set_title("a  The likelihood has two plateaus")
    ax.legend(fontsize=7.6, loc="upper left")

    ax = axs[1]
    m = np.arange(1, 7)
    scen = [("LUCA 4.2 Ga, start 4.404 Ga, end +0.9 Gyr", 0.204, 4.404, 5.304),
            ("LUCA 4.2 Ga, start 4.404 Ga, end +1.73 Gyr", 0.204, 4.404, 6.134),
            ("LUCA 4.2 Ga, start 4.51 Ga, end +0.9 Gyr", 0.31, 4.51, 5.41),
            ("fossils 3.48 Ga, start 4.404 Ga, end +0.9 Gyr", 0.924, 4.404, 5.304)]
    for k, (lab, tl, tobs, T) in enumerate(scen):
        b = bmax_hardsteps(m, tl, tobs, T)
        ax.plot(m, b, "-o", color=cfg.SERIES[k], ms=5, label=lab)
    ax.axhline(10, color=cfg.INK2, lw=0.9, ls=(0, (4, 3)))
    ax.text(6.05, 10.3, "10:1 ('strong')", ha="right", va="bottom", fontsize=8, color=cfg.INK2)
    ax.set_xlabel("number of hard steps after abiogenesis")
    ax.set_ylabel("plateau ratio, fast : slow abiogenesis,\nif every later step is hard")
    ax.set_ylim(0, 19.5)
    ax.set_title("b  All-hard-steps value B(m)")
    ax.text(6.05, 18.9, "not a ceiling: 19.5 if the one later step has rate 1.4 per Gyr,\n"
            "21.6 if life only had to start before today (panel a)", ha="right", va="top", fontsize=7.6,
            color=cfg.INK2)
    ax.legend(fontsize=7.6, loc="upper right", bbox_to_anchor=(1.0, 0.86))

    ax = axs[2]
    post = moody_tab[moody_tab.stage == "posterior"]
    rng = np.random.default_rng(1)
    xmax = 100.0
    for k, (col, lab) in enumerate([("BF_kipping_t0=4.404", "habitable from 4.404 Ga"),
                                    ("BF_kipping_t0=4.51", "habitable from 4.51 Ga")]):
        v = np.sort(post[col].values)
        y = 1 - k + rng.uniform(-0.22, 0.22, len(v))
        inside = v <= xmax
        ax.plot(v[inside], y[inside], "o", color=cfg.SERIES[k], ms=6, mec=cfg.SURFACE, mew=1.0, label=lab)
        if (~inside).any():                      # point values that diverge: draw at the edge
            ax.plot(np.full((~inside).sum(), xmax * 0.97), y[~inside], ">", color=cfg.SERIES[k], ms=7,
                    mec=cfg.SURFACE, mew=1.0)
            ax.text(xmax * 0.9, 1 - k + 0.34, f"{(~inside).sum()} off scale, up to {v.max():.0f}:\n"
                    "node copies within a few Myr\nof the assumed start (artefact\nof dividing by a tiny gap)",
                    ha="right", va="bottom", fontsize=7.4, color=cfg.INK2)
        ax.plot([np.median(v)] * 2, [1 - k - 0.34, 1 - k + 0.34], color=cfg.INK, lw=1.6)
        ax.text(np.median(v) * 1.06, 1 - k - 0.36, f"median {np.median(v):.1f}", ha="left", va="top", fontsize=8)
    ax.axvline(10, color=cfg.INK2, lw=0.9, ls=(0, (4, 3)))
    ax.text(10.6, -0.66, "10:1", fontsize=8, color=cfg.INK2, va="bottom")
    if moody_summ is not None:
        mm = moody_summ["marginalised"]
        a = mm["paper main (ILN+GBM) | t0 ~ U(4.404, 4.50) | future mix"]["B_marg_m1"]
        b = mm["all 38 posterior estimates | t0 ~ U(4.404, 4.50) | future mix"]["B_marg_m1"]
        ax.text(0.03, 0.02, f"with date uncertainty integrated out:\n{a:.1f} (paper's two main runs), {b:.1f} (all {len(post)} variants)",
                transform=ax.transAxes, ha="left", va="bottom", fontsize=8, color=cfg.INK,
                bbox=dict(facecolor=cfg.SURFACE, edgecolor="none", pad=1.5))
    ax.set_xscale("log")
    ax.set_xlim(3.2, xmax * 1.08)
    ax.set_ylim(-1.15, 2.45)
    ax.set_yticks([])
    ax.grid(axis="y", visible=False)
    ax.set_xlabel("T / (2 t$_L$') evaluated at each variant's mean LUCA age")
    ax.set_title(f"c  {len(post)} LUCA ages in the repository: point values")
    ax.legend(fontsize=8, loc="upper left", bbox_to_anchor=(0.0, 1.0))
    return _save(fig, "fig2_abiogenesis_bayes_factor.png")


# ---------------------------------------------------------------------------------
def fig3_hardsteps(nlike, gof):
    fig, axs = plt.subplots(1, 2, figsize=(13.5, 4.8), gridspec_kw={"wspace": 0.28, "width_ratios": [1, 1.25]})
    ax = axs[0]
    futs = [c for c in nlike.case.unique() if c.startswith("future")]
    for k, c in enumerate(futs):
        d = nlike[nlike.case == c]
        d = d[d.n <= 12]
        y = d.like_last / d.like_last.max()
        ax.plot(d.n, y, "-o", ms=4.5, color=cfg.SERIES[k], label="window ends in " + c.replace("future ", "").split(" (")[0]
                + " Gyr (" + c.split("(")[1].rstrip(")").split(";")[0] + ")")
    ax.set_xlabel("number of hard steps n")
    ax.set_ylabel("likelihood of n (relative to best n)")
    ax.set_ylim(0, 1.08)
    ax.set_xticks(range(1, 13))
    ax.set_title("a  Time remaining barely constrains n")
    ax.legend(fontsize=8, loc="lower center")

    ax = axs[1]
    sets = list(dict.fromkeys(gof.candidate_set))
    short = ["2 steps: life, humans", "3 steps: O$_2$ photosynthesis, eukaryotes, humans",
             "4 steps, 2021 dates: life, eukaryotes, sex, humans", "4 steps, 2024-26 dates: life, eukaryotes, animals, humans",
             "5 steps: life, O$_2$ photosynthesis, eukaryotes, animals, humans"]
    for i, (s, lab) in enumerate(zip(sets, short)):
        y = len(sets) - 1 - i
        d = gof[(gof.candidate_set == s) & gof.future.str.startswith("0.9")].iloc[0]
        u = np.array([float(v) for v in d.u.split()])
        n = len(u)
        exp = np.arange(1, n + 1) / (n + 1)
        ax.plot([0, 1], [y, y], color=cfg.GRID, lw=5, solid_capstyle="butt")
        ax.plot(exp, [y] * n, "|", color=cfg.INK, ms=17, mew=1.8)
        ax.plot(u, [y] * n, "o", color=C["blue"], ms=8, mec=cfg.SURFACE, mew=1.3)
        d2 = gof[(gof.candidate_set == s) & gof.future.str.startswith("1.73")].iloc[0]
        u2 = np.array([float(v) for v in d2.u.split()])
        ax.plot(u2, [y - 0.24] * n, "o", color=C["orange"], ms=6, mec=cfg.SURFACE, mew=1.0)
        ax.text(0.0, y + 0.3, lab, fontsize=8.3, va="bottom", ha="left")
    ax.set_yticks([])
    ax.set_ylim(-0.6, len(sets) - 0.25)
    ax.set_xlim(-0.01, 1.01)
    ax.set_xlabel("fraction of the habitable window elapsed when the step occurred")
    ax.grid(axis="y", visible=False)
    ax.set_title("b  Observed step times against the hard-steps expectation")
    h = [Line2D([], [], marker="o", ls="", color=C["blue"], ms=8, mec=cfg.SURFACE, label="observed (window ends +0.9 Gyr)"),
         Line2D([], [], marker="o", ls="", color=C["orange"], ms=6, mec=cfg.SURFACE, label="observed (window ends +1.73 Gyr)"),
         Line2D([], [], marker="|", ls="", color=cfg.INK, ms=13, mew=1.8, label="expected: k/(n+1)")]
    ax.legend(handles=h, fontsize=8, loc="lower right", ncol=3, bbox_to_anchor=(1.0, -0.24))
    return _save(fig, "fig3_hard_steps.png")


# ---------------------------------------------------------------------------------
def fig4_snyder_beattie(t2):
    fig, ax = plt.subplots(figsize=(6.6, 6.0))
    pub = t2[["pub_1", "pub_2", "pub_3"]].values.ravel()
    e = t2[["ours_epow_1", "ours_epow_2", "ours_epow_3"]].values.ravel()
    t = t2[["ours_10pow_1", "ours_10pow_2", "ours_10pow_3"]].values.ravel()
    lim = (1.5e-4, 1.6)
    ax.plot(lim, lim, color=cfg.INK2, lw=0.9, ls=(0, (4, 3)))
    ax.plot(pub, t, "s", color=C["orange"], ms=7, mec=cfg.SURFACE, mew=1.2,
            label="thresholds as printed: $10^{-2}$, $10^{-12}$, $10^{-24}$")
    ax.plot(pub, e, "o", color=C["blue"], ms=7, mec=cfg.SURFACE, mew=1.2,
            label="thresholds $e^{-2}$, $e^{-12}$, $e^{-24}$")
    ax.set_xscale("log")
    ax.set_yscale("log")
    ax.set_xlim(*lim)
    ax.set_ylim(*lim)
    ax.set_xlabel("published posterior weight (Snyder-Beattie et al. 2021, Table 2)")
    ax.set_ylabel("our recomputation with the stated model and prior")
    ax.text(0.25, 0.14, "1:1", fontsize=8.5, color=cfg.INK2, rotation=45, ha="center")
    ax.set_title("Table 2 is reproduced only with natural-log thresholds")
    ax.legend(fontsize=8.5, loc="lower right")
    return _save(fig, "fig4_snyder_beattie_table2.png")


# ---------------------------------------------------------------------------------
def fig5_fbio(samples, curves):
    fig, axs = plt.subplots(1, 2, figsize=(14, 5.2), gridspec_kw={"wspace": 0.2, "width_ratios": [1.45, 1]})
    ax = axs[0]
    fams = [("A", "A   pacing,\nEarth typical"), ("A'", "A'  pacing,\ncentre unknown"),
            ("B2", "B   2 hard steps,\nmoderate prior"), ("B3", "B   3 hard steps,\nmoderate prior"),
            ("C4", "C   4 hard steps,\nwide prior"), ("C5", "C   5 hard steps,\nwide prior"),
            ("D", "D   very wide\nabiogenesis prior")]
    lo_clip = -32.0
    qs = [0.05, 0.25, 0.5, 0.75, 0.95]
    for i, (key, lab) in enumerate(fams):
        y = len(fams) - 1 - i
        lp, lq, w, (ls, ws) = samples[key]
        qp = np.quantile(lp / LN10, qs)
        qq = weighted_quantile(lq / LN10, w, qs)
        qsia = weighted_quantile(ls / LN10, ws, qs)
        for q, yy, col in [(qp, y + 0.26, GREY), (qq, y, C["blue"]), (qsia, y - 0.26, C["orange"])]:
            c = np.maximum(q, lo_clip)
            ax.plot([c[0], c[4]], [yy, yy], color=col, lw=1.6, solid_capstyle="butt")
            ax.plot([c[1], c[3]], [yy, yy], color=col, lw=6, solid_capstyle="butt")
            ax.plot(c[2], yy, "o", color=cfg.SURFACE, ms=4.5, mec=col, mew=1.6, zorder=5)
            if q[0] < lo_clip:
                ax.annotate("", xy=(lo_clip - 1.0, yy), xytext=(lo_clip, yy),
                            arrowprops=dict(arrowstyle="->", color=col, lw=1.3))
                ax.text(lo_clip + 0.2, yy + 0.13, f"25% at {q[1]:.0f}, 5% at {q[0]:.0f}", fontsize=7.5,
                        color=cfg.INK2, va="bottom")
    ax.set_yticks([len(fams) - 1 - i for i in range(len(fams))])
    ax.set_yticklabels([lab for _, lab in fams], fontsize=8.6, ha="left")
    ax.tick_params(axis="y", pad=92, length=0)
    ax.set_ylim(-0.55, len(fams) - 0.45)
    ax.set_xlim(lo_clip - 1.3, 0.8)
    ax.set_xlabel("log$_{10}$ f$_{bio}$  (probability of technological intelligence within the habitable window)")
    ax.grid(axis="y", visible=False)
    ax.set_title("a  f$_{bio}$ under four prior families and two anthropic conventions")
    h = [Line2D([], [], color=GREY, lw=6, label="prior"),
         Line2D([], [], color=C["blue"], lw=6, label="posterior, likelihood conditioned on success (main analysis)"),
         Line2D([], [], color=C["orange"], lw=6, label="posterior, not conditioned (self-indication)"),
         Line2D([], [], color=cfg.INK2, lw=0, label="dot: median; thick: 25-75%; thin: 5-95%")]
    ax.legend(handles=h, fontsize=8.2, loc="upper center", ncol=2, bbox_to_anchor=(0.42, -0.13), handlelength=1.6)

    ax = axs[1]
    for k, (F, lab) in enumerate([(0.9, "window ends in 0.9 Gyr"), (1.73, "window ends in 1.73 Gyr")]):
        f, rel = curves[F]
        ax.plot(f, rel, color=cfg.SERIES[k], label=lab)
    ax.set_xlim(0, 1.0)
    ax.set_ylim(0, 1.08)
    ax.set_xlabel("f$_{bio}$ (linear scale)")
    ax.set_ylabel("profile likelihood (relative to maximum)")
    ax.text(0.03, 0.66, "flat to within 4% for every f$_{bio}$ < 0.1,\nall the way down to f$_{bio}$ = 0:\n"
            "the timeline cannot exclude\nany small value", fontsize=8.8, va="top")
    ax.text(0.47, 0.31, "only values near 1 are\ndisfavoured (we arrived late);\nthis holds for constant-rate\nsteps, not for pacing", fontsize=8.8, ha="left", va="top")
    ax.set_title("b  Conditioned likelihood alone (two-step model)")
    ax.legend(fontsize=8.5, loc="lower left")
    return _save(fig, "fig5_fbio.png")