"""
Step 6. Figures (PNG, ../figures/). Reads only the tables written by the earlier steps.
Colours: a fixed categorical order (blue, orange, aqua, yellow) checked for colour-blind separation; text is never coloured.
"""
import json
import os

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

HERE = os.path.dirname(os.path.abspath(__file__))
DATA = os.path.join(HERE, "..", "data")
FIG = os.path.join(HERE, "..", "figures")
SURF, INK, INK2, GRID = "#fcfcfb", "#0b0b0b", "#52514e", "#e4e3df"
BLUE, ORANGE, AQUA, YELLOW, GREY = "#2a78d6", "#eb6834", "#1baf7a", "#eda100", "#8a8983"
plt.rcParams.update({"figure.facecolor": SURF, "axes.facecolor": SURF, "savefig.facecolor": SURF, "font.family": "DejaVu Sans", "font.size": 10.5,
                     "text.color": INK, "axes.labelcolor": INK2, "xtick.color": INK2, "ytick.color": INK2, "axes.edgecolor": GRID,
                     "axes.spines.top": False, "axes.spines.right": False, "axes.grid": True, "grid.color": GRID, "grid.linewidth": 0.8,
                     "axes.axisbelow": True, "legend.frameon": False, "axes.titlesize": 11.5, "axes.titleweight": "bold", "axes.titlelocation": "left"})


def save(fig, name):
    os.makedirs(FIG, exist_ok=True)
    fig.savefig(os.path.join(FIG, name), dpi=160, bbox_inches="tight")
    plt.close(fig)
    print("wrote", name)


def fig1_earth():
    pm_all = pd.read_csv(os.path.join(DATA, "pubmed_method_counts.csv"))
    ua = pd.read_csv(os.path.join(DATA, "earth_unassigned_by_biome.csv"))
    es = json.load(open(os.path.join(DATA, "earth_stats.json")))
    un = pm_all[pm_all.method_class == "union"].set_index("method").pubmed_records
    pm = pm_all[pm_all.method_class != "union"].copy()
    # split stable isotope probing: records that name DNA, RNA or sequencing are not agnostic to the genetic polymer
    i = pm.index[pm.method == "stable isotope probing"][0]
    n_na = int(un["stable isotope probing that names DNA, RNA or sequencing"]); n_all = int(pm.loc[i, "pubmed_records"])
    top, bot = pm.loc[:i - 1], pm.loc[i + 1:]
    sip = pd.DataFrame([dict(pm.loc[i], method="stable isotope probing naming DNA, RNA or sequencing", sees_non_nucleic_acid_life="no", pubmed_records=n_na),
                        dict(pm.loc[i], method="stable isotope probing, other (lipid, protein, single cell)", sees_non_nucleic_acid_life="yes", pubmed_records=n_all - n_na)])
    pm = pd.concat([top, sip, bot], ignore_index=True)

    def sees(r):
        s = r.sees_non_nucleic_acid_life
        if s.startswith("designed"):
            return 3
        if s.startswith("yes") or s.startswith("partly"):
            return 2
        if s.startswith("only"):
            return 1
        return 0
    pm["k"] = pm.apply(sees, axis=1)
    lab = ["No: needs standard nucleic acids or amino acids", "Only as unmatched leftovers", "Yes, though not used as a search", "Designed as a search"]
    col = [BLUE, ORANGE, AQUA, YELLOW]
    pm = pm.iloc[::-1].reset_index(drop=True)
    fig, ax = plt.subplots(1, 2, figsize=(14.8, 7.0), gridspec_kw=dict(width_ratios=[1.25, 1], wspace=0.78))
    a = ax[0]
    y = np.arange(len(pm))
    a.barh(y, pm.pubmed_records, height=0.62, color=[col[k] for k in pm.k], edgecolor=SURF, linewidth=1.5)
    a.set_xscale("log"); a.set_xlim(1, 4e5)
    short = {"metaproteomics (peptides matched to predicted proteins)": "metaproteomics", "untargeted metabolomics of microbial or environmental samples": "untargeted metabolomics (environmental)",
             "cryo-electron tomography of environmental or uncultivated cells": "cryo-electron tomography, uncultivated cells",
             "fluorescence in situ hybridisation of environmental microbes": "fluorescence in situ hybridisation (rRNA probes)",
             "BONCAT (amino-acid analogue incorporation)": "amino-acid analogue labelling (BONCAT)",
             "stable isotope probing naming DNA, RNA or sequencing": "isotope probing naming DNA, RNA or sequencing",
             "stable isotope probing, other (lipid, protein, single cell)": "isotope probing, other (lipid, protein, cell)"}
    a.set_yticks(y); a.set_yticklabels([short.get(m, m) for m in pm.method], fontsize=9.3)
    for yi, v in zip(y, pm.pubmed_records):
        a.text(v * 1.18, yi, f"{v:,}", va="center", fontsize=9, color=INK2)
    P = es["pubmed"]
    a.set_xlabel("PubMed records matching the method (log scale; a record count is a crude proxy for effort)\n"
                 f"Counted once each: {P['sequence_based_records_union']:,} sequence-based records against {P['polymer_agnostic_records_union_excluding_nucleic_acid_sip']:,} to "
                 f"{P['polymer_agnostic_records_union']:,} polymer-agnostic,\na ratio of roughly {P['ratio_sequence_to_agnostic_low']:.0f} to {P['ratio_sequence_to_agnostic_high']:.0f} to 1", fontsize=9.6)
    a.set_title("Almost all survey effort uses methods that cannot see\nlife without standard nucleic acids")
    a.grid(axis="y", visible=False)
    h = [plt.Rectangle((0, 0), 1, 1, color=c) for c in col]
    a.legend(h, lab, title="Could the method register a second biochemistry?", loc="upper center", fontsize=8.8, title_fontsize=9, bbox_to_anchor=(0.35, -0.2), ncol=2)
    b = ax[1]
    series = [("assembled contigs: predicted proteins with no InterPro match", "Predicted proteins with no known family (assemblies)", BLUE, 0.17),
              ("raw reads: no InterPro-matched fragment", "Raw reads with no matched fragment", ORANGE, -0.17)]
    biomes = ["Environmental", "Engineered", "Mixed", "Host-associated"]
    for meas, lbl, c, off in series:
        for i, bm in enumerate(biomes):
            r = ua[(ua.measure == meas) & (ua.biome == bm)]
            if len(r) == 0:
                continue
            r = r.iloc[0]
            yy = len(biomes) - 1 - i + off
            b.plot([100 * r.p05, 100 * r.p95], [yy, yy], color=c, lw=1.6, solid_capstyle="round", alpha=0.55)
            b.plot([100 * r.p25, 100 * r.p75], [yy, yy], color=c, lw=6, solid_capstyle="round")
            b.plot(100 * r["median"], yy, "o", ms=9, color=c, mec=SURF, mew=2, label=lbl if i == 0 else None)
            b.text(101.5, yy, f"{int(r.analyses):,} analyses,\n{int(r.studies)} studies", va="center", fontsize=8.2, color=INK2)
    b.set_yticks(range(len(biomes))); b.set_yticklabels(biomes[::-1])
    b.set_xlim(0, 100); b.set_ylim(-0.6, len(biomes) - 0.35)
    A = es["assembly_unmatched"]
    b.set_xlabel("Share with no match in the protein-family databases (%)\ndot: median over analyses; thick bar: middle half; thin line: 5th to 95th percentile\n"
                 f"All assemblies: median {100 * A['median_unweighted']:.0f}% ({100 * A['median_host_associated_weighted_x4']:.0f}% with the host-associated subsample reweighted), "
                 f"{100 * A['pooled_share_protein_weighted']:.0f}% of all proteins pooled", fontsize=9.6)
    b.set_title("A quarter to three quarters of predicted proteins match no known family\n(MGnify pipeline 5.0 holdings; these analyses come from few studies)")
    b.grid(axis="y", visible=False)
    b.legend(loc="upper center", bbox_to_anchor=(0.5, -0.26), fontsize=9, ncol=1)
    save(fig, "fig1_earth_survey.png")


def fig2_completeness():
    H = pd.read_csv(os.path.join(DATA, "completeness_by_habitat.csv"))
    rows = [("Earth: a second biosphere alive today", "Earth", "alive"), ("Earth: a second origin at any time", "Earth", "c"), ("Mars", "Mars", "c"), ("Europa", "Europa", "c"),
            ("Enceladus", "Enceladus", "c"), ("Titan", "Titan", "c"), ("Venus", "Venus", "c"), ("Carbonaceous asteroids, Ceres", "Small bodies", "c")]
    eps = [("today", "Evidence in hand, October 2026", BLUE), ("funded", "After the funded missions (about 2038)", ORANGE), ("unfunded", "After two missions not funded (2050s)", AQUA)]
    DRAWS = np.load(os.path.join(DATA, "completeness_draws.npz"))
    fig, ax = plt.subplots(1, 2, figsize=(15.5, 5.6), gridspec_kw=dict(width_ratios=[2.1, 1], wspace=0.06), sharey=True)
    a = ax[0]
    for i, (lab, hab, kind) in enumerate(rows):
        y0 = len(rows) - 1 - i
        prev = None
        for j, (ep, elab, c) in enumerate(eps):
            r = H[(H.epoch == ep) & (H.habitat == hab)].iloc[0]
            lo, md, hi = (r.c_alive_today_p05, r.c_alive_today_p50, r.c_alive_today_p95) if kind == "alive" else (r.c_p05, r.c_p50, r.c_p95)
            if prev is not None and abs(md - prev) < 1e-9:
                continue
            prev = md
            yy = y0 + (0.24 - 0.24 * j)
            a.plot([lo, hi], [yy, yy], color=c, lw=2, solid_capstyle="round", alpha=0.6)
            a.plot(md, yy, "o", ms=9, color=c, mec=SURF, mew=2, label=elab if i == 2 else None)
            if j == 0:
                a.text(hi * 1.35, yy, f"{md:.2g}" if md < 0.01 else f"{md:.3f}".rstrip("0"), va="center", fontsize=8.8, color=INK2)
    a.set_xscale("log"); a.set_xlim(1e-5, 2)
    a.set_yticks(range(len(rows))); a.set_yticklabels([r[0] for r in rows][::-1])
    a.set_xlabel("Search completeness c: chance the evidence would already show a second origin, if one happened there\n(median and 5th to 95th percentile over independent draws of my itemised judgements; log scale)")
    a.set_title("Search completeness (structured judgement, not measurement):\na few percent on Mars, a tenth of a percent or less on the ocean worlds")
    a.grid(axis="y", visible=False)
    a.legend(loc="lower left", fontsize=9, bbox_to_anchor=(0.0, 0.0))
    b = ax[1]
    for i, (lab, hab, kind) in enumerate(rows):
        if kind == "alive":
            continue
        y0 = len(rows) - 1 - i
        key = hab.replace(" ", "_")
        qo = DRAWS[f"q_today_{key}"] * DRAWS[f"o_today_{key}"]          # percentiles of the product, draw by draw
        lo, md, hi = np.quantile(qo, [0.05, 0.5, 0.95])
        b.plot([lo, hi], [y0, y0], color=GREY, lw=2, solid_capstyle="round", alpha=0.6)
        b.plot(md, y0, "o", ms=9, color=GREY, mec=SURF, mew=2)
    b.set_xscale("log"); b.set_xlim(3e-7, 2)
    b.set_xlabel("Opportunity for an origin, q × o\n(1 = an Earth-like planet over 5.6 Gyr; log scale)")
    b.set_title("How much there\nwas to find")
    b.grid(axis="y", visible=False)
    save(fig, "fig2_completeness_by_habitat.png")


def fig3_likelihood():
    L = pd.read_csv(os.path.join(DATA, "null_likelihood_curves.csv"))
    T = json.load(open(os.path.join(DATA, "completeness_summary.json")))["totals"]
    m = (L["lambda"] >= 1e-3)
    x = L["lambda"][m]
    fig, a = plt.subplots(figsize=(10.6, 6.0))
    a.axvspan(1, 1e3, color=GRID, alpha=0.55, lw=0)
    a.text(7, 0.52, "the model's\n'life starts easily' camp:\n1 to 1000 origins\nper planet", fontsize=9, color=INK2, va="center")
    a.fill_between(x, L["L_p05_today"][m], L["L_p95_today"][m], color=BLUE, alpha=0.12, lw=0)
    a.fill_between(x, L["L_p05_today_origin_persist"][m], L["L_p95_today_origin_persist"][m], color=ORANGE, alpha=0.12, lw=0)
    d1, d2, d3 = T["today"]["D_fast_mean"], T["today"]["D_fast_mean_origin_persist"], T["today"]["D_fast_model_form_mean"]
    ser = [("L_mean_today", f"Searches to date, counted per habitat (misses {100 * (1 - d1):.0f} in 100 in the fast camp)", BLUE, "-"),
           ("L_mean_today_origin_persist", f"Searches to date, each origin persisting independently (misses {100 * (1 - d2):.0f} in 100)", ORANGE, "-"),
           ("L_mean_today_model_form", f"Searches to date in the model's own form, exp(−Fλ) (misses {100 * (1 - d3):.0f} in 100)", AQUA, "-"),
           ("L_mean_ceiling_ever", f"Any life that ever arose on any body is found (misses {100 * (1 - T['ceiling_ever']['D_fast_mean']):.0f} in 100)", YELLOW, "-"),
           ("L_model_thorough", f"The model's 'thorough search', sg_frac = 0.3 (misses {100 * (1 - T['model_thorough']['D_fast']):.0f} in 100)", GREY, (0, (5, 3)))]
    for k, lab, c, ls in ser:
        a.plot(x, L[k][m], color=c, lw=2.2, ls=ls, label=lab, solid_capstyle="round")
    ends = [("L_mean_today", "per habitat"), ("L_mean_today_origin_persist", "per origin\n(persistence)"), ("L_mean_today_model_form", "model's\nown form"), ("L_mean_ceiling_ever", "everything\nfound")]
    for k, lab in ends:
        a.text(1.25e3, float(L[k].iloc[-1]), lab, va="center", fontsize=9, color=INK)
    a.set_xscale("log"); a.set_xlim(1e-3, 1e3); a.set_ylim(0, 1.03)
    a.set_xlabel("λ: expected independent origins of life per 5.6-Gyr window of an Earth-like planet\n(shaded bands: 5th to 95th percentile over independent draws of the itemised judgements)")
    a.set_ylabel("Probability of having found no second origin")
    a.set_title("What the null to date is worth depends on whether many origins in one place\nmean many chances of detection; under every reading it is far from a thorough search")
    a.legend(loc="lower left", fontsize=8.8, bbox_to_anchor=(0.0, 0.02))
    save(fig, "fig3_null_likelihood.png")


def fig4_update():
    U = pd.read_csv(os.path.join(DATA, "update_scenarios.csv")); U = U[U.convention == "site"].set_index("scenario")
    R = pd.read_csv(os.path.join(DATA, "update_range_over_completeness.csv")); R = R[(R.convention == "site") & (R.quantity == "p_alone")]
    Ch = pd.read_csv(os.path.join(DATA, "cheyava_sensitivity.csv")); Ch = Ch[Ch.sB.str.startswith("factor") & (Ch.transfer_probability_T == 0)]
    ST = [("habitat", "per habitat", BLUE, "Counted per habitat (fates fully correlated)"), ("origin_persist", "per origin (persistence)", ORANGE, "Each origin persists independently"),
          ("model_form", "model's own form, exp(-F lambda)", AQUA, "Model's own form, exp(−Fλ): every origin a separate chance (upper bound)")]
    EP = {"today": "Null to date", "funded": "Null after the funded missions (about 2038)", "unfunded": "Null after two unfunded missions as well (2050s)"}
    base = U.loc["No update (model as published)", "p_alone"]
    rows = [("Evidence in hand scored as a pure null", "ep", "today"),
            ("…Cheyava Falls scored at likelihood ratio 1.5", "lr", 1.5), ("…Cheyava Falls scored at likelihood ratio 2", "lr", 2.0), ("…Cheyava Falls scored at likelihood ratio 3", "lr", 3.0),
            ("Null after the funded missions (≈2038)", "ep", "funded"), ("Null after two unfunded missions too (2050s)", "ep", "unfunded"),
            ("Everything that survives today is found: none", "ceil", None), ("Any life that ever arose is found: none", "one", "Ceiling: any life that ever arose on any body is found (c = 1 everywhere)"),
            ("Model's 'thorough search finds none'", "one", "Model's 'thorough search finds none' (sg_frac = 0.3)")]
    fig, a = plt.subplots(figsize=(11.8, 8.2))
    a.axvline(base, color=GREY, lw=1)
    a.text(base, len(rows) - 0.42, f"model as published: {base:.3f}", fontsize=9, color=INK2, ha="center", va="bottom")
    seen = set()
    for i, (lab, kind, arg) in enumerate(rows):
        y0 = len(rows) - 1 - i
        pts = []
        if kind == "ep":
            for j, (st, nm, c, leg) in enumerate(ST):
                v = U.loc[f"{EP[arg]}; structure: {nm}", "p_alone"]
                r = R[(R.structure == st) & (R.epoch == arg)]
                pts.append((v, c, leg, (r.iloc[0].p05, r.iloc[0].p95) if len(r) else None))
        elif kind == "lr":
            for st, nm, c, leg in ST[:2]:
                v = Ch[(Ch.structure == st) & (Ch.likelihood_ratio_for_past_life_at_cheyava == arg)].p_alone.iloc[0]
                pts.append((v, c, leg, None))
        elif kind == "ceil":
            for st, nm, c, leg in ST[:2]:
                pts.append((U.loc[f"Ceiling: everything that survives or is preserved today is found; structure: {nm}", "p_alone"], c, leg, None))
        else:
            pts.append((U.loc[arg, "p_alone"], GREY, "Does not depend on the counting", None))
        n = len(pts)
        for j, (v, c, leg, rng) in enumerate(pts):
            yy = y0 + (0.27 * ((n - 1) / 2 - j) if n > 1 else 0)
            if rng:
                a.plot(rng, [yy, yy], color=c, lw=5, solid_capstyle="round", alpha=0.3)
            a.plot(v, yy, "o", ms=9, color=c, mec=SURF, mew=2, label=None if leg in seen else leg)
            seen.add(leg)
            a.text(max(v, rng[1] if rng else v) + 0.0028, yy, f"{v:.3f}", fontsize=8.6, color=INK2, va="center", ha="left")
    a.set_yticks(range(len(rows))); a.set_yticklabels([r[0] for r in rows][::-1])
    a.set_ylim(-0.6, len(rows) - 0.1)
    a.set_xlim(0.728, 0.892)
    a.grid(axis="y", visible=False)
    a.set_xlabel("P(no other civilisation has ever arisen in the Galaxy), Earth-as-given column\n"
                 "shaded bars: 5th to 95th percentile over independent factor draws only; the spread between colours (structure) and between rows (scoring of Mars) is larger")
    a.set_title("The evidence in hand moves the answer by a few hundredths at most, and its sign\ndepends on how the Cheyava Falls potential biosignature is scored")
    a.legend(loc="upper center", fontsize=8.8, bbox_to_anchor=(0.4, -0.13), ncol=2, title="How detection chances are counted when one habitat has many origins", title_fontsize=9)
    save(fig, "fig4_update_by_scenario.png")


def fig5_independence_claims():
    I = pd.read_csv(os.path.join(DATA, "independence.csv"), keep_default_na=False)
    Cl = pd.read_csv(os.path.join(DATA, "claim_value.csv"))
    F = pd.read_csv(os.path.join(DATA, "detection_forecast.csv"))
    fig, ax = plt.subplots(1, 2, figsize=(14.6, 5.8), gridspec_kw=dict(wspace=0.3))
    a = ax[0]
    M = I[(I.body == "Mars") & (I.T_earth_to_target > 0) & ~I.T_target_to_earth_label.str.contains("per-rock")]
    base_lab, fast_lab = "model's base posterior (after the null to date)", "fast camp only (life starts easily)"
    pn = {k: M[M.prior == k].p_native_origin.iloc[0] for k in (base_lab, fast_lab)}
    ser = [(base_lab, "same as T_em", AQUA, "-", f"Either route, model's base prior (native Martian origin {pn[base_lab]:.2f})"),
           (base_lab, "0", AQUA, (0, (4, 2.5)), "Earth-to-Mars seeding only, base prior (first pass)"),
           (fast_lab, "same as T_em", YELLOW, "-", f"Either route, if life starts easily (native Martian origin {pn[fast_lab]:.2f})"),
           (fast_lab, "0", YELLOW, (0, (4, 2.5)), "Earth-to-Mars seeding only, if life starts easily")]
    for prior, tl, c, ls, lab in ser:
        d = M[(M.prior == prior) & (M.T_target_to_earth_label == tl)].sort_values("T_earth_to_target")
        a.plot(d.T_earth_to_target, d.p_independent_given_life_found, color=c, lw=2.2, ls=ls, marker="o", ms=7 if ls == "-" else 0, mec=SURF, mew=1.8, label=lab)
    a.set_xscale("log"); a.set_xlim(0.008, 1.0); a.set_ylim(0, 1.05)
    a.set_xlabel("T: chance of shared ancestry through rock transfer. 'Either route' sets the chance\nthat Earth life seeded Mars and the chance that Earth's tree descends from Martian\n"
                 "life both equal to T (Worth et al. 2013: 360,000 rocks Earth to Mars and\n21 million Mars to Earth in 3.5 Gyr)", fontsize=9.4)
    a.set_ylabel("P(independent origin | life found on Mars)")
    a.set_title("Would Martian life be a second origin? That depends on\ntransfer in both directions and on the prior for a native origin")
    a.legend(loc="lower left", fontsize=8.6)
    b = ax[1]
    pd_h = F[(F.structure == "habitat") & F.searches.str.startswith("funded missions")].p_true_detection.iloc[0]
    pd_o = F[(F.structure == "origin_persist") & F.searches.str.startswith("funded missions")].p_true_detection.iloc[0]
    for stc, lab, c in (("habitat", f"Funded missions, counted per habitat (chance of a real find {100 * pd_h:.1f}%)", BLUE),
                        ("origin_persist", f"Funded missions, each origin persisting independently ({100 * pd_o:.1f}%)", ORANGE),
                        ("model", "The model's 'thorough search' (19%)", GREY)):
        d = Cl[(Cl.structure == stc) & (Cl.false_alarm_rate > 0)]
        b.plot(d.false_alarm_rate, d.p_fast_after_claim, "-o", color=c, lw=2.2, ms=8, mec=SURF, mew=2, label=lab)
    b.axvline(0.15, color=GREY, lw=1)
    b.text(0.142, 0.03, "the model's assumed\nfalse-alarm rate (0.15)", fontsize=8.8, color=INK2, ha="right", va="bottom")
    b.set_xscale("log"); b.set_xlim(0.0008, 0.4); b.set_ylim(0, 1.05)
    b.set_xlabel("Chance of an accepted but false claim of an independent origin\n(all curves under the model's base prior, P(life starts easily) = 0.19)")
    b.set_ylabel("P(life starts easily | an accepted claim)")
    b.set_title("A claim is worth only as much as its false-alarm rate\nallows (base-rate argument of Foote et al. 2023)")
    b.legend(loc="upper center", fontsize=8.8, bbox_to_anchor=(0.5, -0.22), ncol=1)
    save(fig, "fig5_independence_and_claims.png")


if __name__ == "__main__":
    fig1_earth(); fig2_completeness(); fig3_likelihood(); fig4_update(); fig5_independence_claims()