#!/usr/bin/env python3
"""Step 5.  Figures (PNG) from the derived tables in data/."""
import os

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

from common import DATA, FIGS, load_json

# palette (categorical slots in fixed order; neutral inks)
BLUE, ORANGE, AQUA, YELLOW, MAGENTA, GREEN, VIOLET, RED = (
    "#2a78d6", "#eb6834", "#1baf7a", "#eda100", "#e87ba4", "#008300", "#6250d6", "#e34948")
SEQ = ["#86b6ef", "#3987e5", "#1c5cab", "#0d366b"]      # one-hue ramp for ordered quantities
INK, INK2, MUTED, GRID, SURF = "#0b0b0b", "#52514e", "#898781", "#e1e0d9", "#fcfcfb"
KMS = 1.0227121650537077

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


def delay_label(td):
    return "no delay" if td == 0 else f"delay 10$^{{{int(round(np.log10(td * 1e6)))}}}$ yr"


def fig1():
    s1 = load_json("s1_density_kinematics.json")
    dens = pd.read_csv(os.path.join(DATA, "density_vs_radius.csv"))
    lf = pd.read_csv(os.path.join(DATA, "luminosity_function_kinematics.csv"))
    fig, ax = plt.subplots(1, 2, figsize=(11, 4.2))
    a = ax[0]
    V = 4 / 3 * np.pi * dens.R_pc**3
    a.plot(dens.R_pc, dens.dens_sources, "-o", color=MUTED, ms=4, label="GCNS sources (individual Gaia entries)")
    a.plot(dens.R_pc, dens.n_systems_table3 / V, "-o", color=ORANGE, ms=4, label="systems, GCNS binary table only")
    a.errorbar(dens.R_pc, dens.dens_systems, yerr=dens.dens_systems_err, fmt="-o", color=BLUE, ms=5,
               label="systems, + co-moving pair search (used)")
    a.errorbar([10 * 0.93], [s1["tenpc"]["dens_systems_stellar"]], yerr=[s1["tenpc"]["dens_systems_stellar_err"]],
               fmt="s", color=INK, ms=6, label="10 pc sample, stellar systems (Reylé+ 2021)")
    a.axhspan(s1["n_sys_adopted"] - s1["n_sys_adopted_err"], s1["n_sys_adopted"] + s1["n_sys_adopted_err"],
              color=BLUE, alpha=0.12, lw=0, label="adopted: 50 pc value × 1.045 ± 0.045 for missing systems")
    a.set_xscale("log"); a.set_xticks([10, 20, 30, 50, 100]); a.set_xticklabels(["10", "20", "30", "50", "100"])
    a.set_ylim(0.058, 0.092)
    a.set_xlabel("radius of Sun-centred sphere [pc]"); a.set_ylabel("number density [pc$^{-3}$]")
    a.set_title("Local density of star systems")
    a.legend(loc="upper left")
    b = ax[1]
    k = (lf.MG_lo >= 2) & (lf.MG_lo <= 16)
    x = lf.MG_lo[k] + 0.5
    b.errorbar(x, lf.mean_v[k], yerr=lf.mean_v_err[k], fmt="-o", color=BLUE, ms=5,
               label=r"all systems: $(4/\pi)\langle v_T\rangle$, no RV needed")
    kk = k & (lf.N_rv >= 100)
    b.plot(lf.MG_lo[kk] + 0.5, lf.v3d_rv[kk], "-s", color=ORANGE, ms=4, label=r"systems with clean Gaia RV: $\langle |v| \rangle$ (3D)")
    kk = k & ((lf.N - lf.N_rv) >= 80) & (lf.MG_lo <= 12)
    b.plot(lf.MG_lo[kk] + 0.5, lf.v_from_vT_norv[kk], "--^", color=AQUA, ms=4, lw=1.5,
           label=r"systems without clean Gaia RV: $(4/\pi)\langle v_T\rangle$")
    b.axhline(s1["mean_speed_adopted_kms"], color=INK2, lw=1, ls=":")
    b.text(2.1, s1["mean_speed_adopted_kms"] + 0.6, f"all systems <50 pc: {s1['mean_speed_adopted_kms']:.1f} km/s",
           color=INK2, fontsize=8)
    b.set_xlabel("absolute magnitude M$_G$"); b.set_ylabel("mean speed relative to the Sun [km/s]")
    b.set_title("Speeds: stars lacking RVs are kinematically colder")
    b.set_ylim(25, 65); b.legend(loc="lower right")
    fig.tight_layout()
    fig.savefig(os.path.join(FIGS, "fig1_density_speed.png"))
    plt.close(fig)


def fig2():
    s1 = load_json("s1_density_kinematics.json")
    s2 = load_json("s2_census.json")
    s3 = load_json("s3_exposure.json")
    cen = pd.read_csv(os.path.join(DATA, "census_obs_vs_pred.csv"))
    enc = pd.read_csv(os.path.join(DATA, "census_encounters_100pc.csv"))
    m = cen[(cen.T_Myr == 1.0) & (cen.ruwe_cut == 0) & np.isinf(cen.vmax)]
    fig, ax = plt.subplots(1, 2, figsize=(11, 4.6))
    a = ax[0]
    d = np.sort(enc[(enc.t_ph.abs() < 1.0) & (enc.d_ph < 10)].d_ph.values)
    a.step(d, np.arange(1, len(d) + 1), where="post", color=BLUE, label="observed (clean Gaia DR3 RVs)")
    a.plot(m.d_pc, m.N_pred, "o-", color=ORANGE, ms=5, lw=1.5,
           label="expected for the same stars (consistency check)")
    a.plot(m.d_pc, m.N_pred_all_systems, "s--", color=MUTED, ms=4, lw=1.2,
           label="expected for all systems (with or without RV)")
    xx = np.array([0.4, 10]); a.plot(xx, m.N_pred.iloc[-1] * (xx / 10) ** 2, ":", color=INK2, lw=1,
                                     label="exact d² scaling")
    a.set_xscale("log"); a.set_yscale("log"); a.set_xlim(0.4, 11); a.set_ylim(0.8, 3000)
    a.set_xlabel("perihelion distance d [pc]"); a.set_ylabel("systems passing within d  (|t| < 1 Myr, now < 100 pc)")
    a.set_xticks([0.5, 1, 2, 5, 10]); a.set_xticklabels(["0.5", "1", "2", "5", "10"])
    a.set_title("Direct census versus expectation")
    a.legend(loc="upper left")
    b = ax[1]
    mb = s3["martinez_barbosa_2017_gamma1pc"][1]
    items = [("García-Sánchez+ 2001\n(Hipparcos, corrected)", 11.7, 1.3, MUTED, None),
             ("Martínez-Barbosa+ 2017\n(N-body, no migration)", mb, 0.0, MUTED, None),
             ("Bailer-Jones+ 2018\n(Gaia DR2, corrected)", 19.7, 2.2, MUTED, None),
             ("Fernandez-Puig+ 2026\n(25 pc sample)", 10.6, 4.5, MUTED, None),
             ("This work: raw clean-RV counts\n(lower bound, no correction)", s2["lower_bound_gamma1pc"],
              s2["lower_bound_gamma1pc_err"], ORANGE, "≥ "),
             ("This work: star systems\n(Gaia DR3/GCNS, RV-free)", s1["gamma1pc_adopted"], s1["gamma1pc_adopted_err"],
              BLUE, None)]
    for i, (lab, v, e, c, pre) in enumerate(items):
        y = len(items) - 1 - i
        b.errorbar([v], [y], xerr=[e] if e else None, fmt="o", color=c, ms=7, capsize=3, lw=2)
        txt = f"{v:.1f} ± {e:.1f}" if e else f"{v:.1f}"
        b.text(v, y + 0.24, (pre or "") + txt, ha="center", color=INK, fontsize=8.5)
    b.set_yticks(range(len(items))); b.set_yticklabels([it[0] for it in items][::-1])
    b.set_ylim(-0.6, len(items) - 0.2); b.set_xlim(0, 25)
    b.set_xlabel("encounters within 1 pc of the Sun per Myr"); b.grid(axis="y", visible=False)
    b.set_title("Encounter rate within 1 pc")
    fig.tight_layout()
    fig.savefig(os.path.join(FIGS, "fig2_census_rate.png"))
    plt.close(fig)


def fig3():
    exp = pd.read_csv(os.path.join(DATA, "exposure_table.csv"))
    con = pd.read_csv(os.path.join(DATA, "fermi_constraint.csv"))
    fig, ax = plt.subplots(1, 2, figsize=(11, 4.6))
    a = ax[0]
    for name, c in (("O2 atmosphere (since GOE)", BLUE), ("High O2/O3 (since NOE)", ORANGE),
                    ("Land vegetation (red edge)", AQUA)):
        e = exp[exp.era == name]
        a.fill_between(e.d_pc, e.N_low, e.N_high, color=c, alpha=0.15, lw=0)
        a.plot(e.d_pc, e.N_central, "-o", color=c, ms=4, label=f"{name}, {e.T_Myr.iloc[0] / 1000:.2f} Gyr")
    e = exp[(exp.era.str.startswith("Radio")) & (exp.d_pc >= 2)]
    a.plot(e.d_pc, e.N_central, "s", color=INK, ms=5, label="Radio era (100 yr): systems inside d today")
    a.set_xscale("log"); a.set_yscale("log")
    a.set_xticks([0.1, 0.5, 1, 2, 5, 10]); a.set_xticklabels(["0.1", "0.5", "1", "2", "5", "10"])
    a.set_xlabel("closest-approach distance d [pc]"); a.set_ylabel("distinct star systems passing within d")
    a.set_title("Systems that passed while Earth showed a biosphere")
    a.legend(loc="upper left")
    b = ax[1]
    Ls = [1.0, 10.0, 100.0, 470.0, 2330.0]
    cols = ["#b7d3f6", "#86b6ef", "#3987e5", "#1c5cab", "#0d366b"]
    for L, c in zip(Ls, cols):
        s = con[(con.persistence_Myr == L) & (con.passers == "any speed")]
        lab = f"{L:.0f} Myr" if L < 1000 else f"{L / 1000:.2f} Gyr"
        b.plot(s.d_pc, s.x95, "-o", color=c, ms=4, label=f"presence lasts {lab}")
    for L, ls, lab in ((2330.0, "--", "2.33 Gyr"), (470.0, ":", "470 Myr")):
        s = con[(con.persistence_Myr == L) & (con.passers == "v<30 km/s")]
        b.plot(s.d_pc, s.x95, ls, color=ORANGE, lw=1.6, label=f"{lab}, only passers slower than 30 km/s")
    b.set_xscale("log"); b.set_yscale("log"); b.set_ylim(5e-7, 2)
    b.set_xticks([0.1, 0.5, 1, 2, 5, 10]); b.set_xticklabels(["0.1", "0.5", "1", "2", "5", "10"])
    b.set_xlabel("hop distance d the civilization can cross [pc]")
    b.set_ylabel("95% upper limit on fraction x of passing systems")
    b.set_title("Local Fermi constraint (no evident presence here)")
    b.legend(loc="lower left", fontsize=8)
    fig.tight_layout()
    fig.savefig(os.path.join(FIGS, "fig3_exposure_constraint.png"))
    plt.close(fig)


def fig4():
    st = pd.read_csv(os.path.join(DATA, "static_percolation.csv"))
    sim = pd.read_csv(os.path.join(DATA, "settlement_runs.csv"))
    cur = pd.read_csv(os.path.join(DATA, "settlement_curves.csv"))
    s4 = load_json("s4_settlement.json")
    fig, ax = plt.subplots(2, 2, figsize=(11, 8.6))
    a = ax[0, 0]
    s = st[st.R_pc == 40]
    a.fill_between(s.d_hop, s.poisson_largest_frac - 2 * s.poisson_sd, s.poisson_largest_frac + 2 * s.poisson_sd,
                   color=MUTED, alpha=0.25, lw=0)
    a.plot(s.d_hop, s.poisson_largest_frac, "--", color=MUTED, lw=1.5, label="random (Poisson) field, same density")
    a.plot(s.d_hop, s.largest_frac, "-", color=BLUE, label="real Gaia systems within 40 pc")
    a.axvline(s4["static"]["d_c_infinite_poisson_adopted_density"], color=INK2, lw=1, ls=":")
    a.text(s4["static"]["d_c_infinite_poisson_adopted_density"] - 0.05, 0.9, "infinite-field\nthreshold",
           ha="right", color=INK2, fontsize=8)
    a.set_xlabel("hop length d$_{hop}$ [pc]"); a.set_ylabel("fraction of systems in largest connected cluster")
    a.set_title("Static star field: connected only for hops > 2.1–2.3 pc")
    a.set_xlim(0.6, 4); a.legend(loc="lower right")

    b = ax[0, 1]
    lim = sim[sim.dv_max_kms == 30]
    for td, c in zip(sorted(lim.t_delay_Myr.unique()), SEQ):
        s = lim[lim.t_delay_Myr == td]
        b.plot(s.d_hop_pc, s.T2_kinetic_Myr * 1e3, "-", color=c, lw=1.4)
        b.plot(s.d_hop_pc, s.T2_sim_Myr * 1e3, "o", color=c, ms=6, label=delay_label(td))
    anyv = sim[np.isinf(sim.dv_max_kms) & (sim.t_delay_Myr == 0)]
    b.plot(anyv.d_hop_pc, anyv.T2_kinetic_Myr * 1e3, "-", color=MUTED, lw=1.2)
    b.plot(anyv.d_hop_pc, anyv.T2_sim_Myr * 1e3, "s", color=MUTED, ms=5,
           label="no delay, ships match any speed\n(lower bound on the doubling time)")
    b.set_xscale("log"); b.set_yscale("log")
    b.set_xticks([0.5, 1, 2, 3, 5]); b.set_xticklabels(["0.5", "1", "2", "3", "5"])
    b.xaxis.set_minor_formatter(matplotlib.ticker.NullFormatter())
    b.set_xlabel("hop length d$_{hop}$ [pc]"); b.set_ylabel("doubling time, 10 → 1000 systems [kyr]")
    b.plot([], [], "-", color=INK2, lw=1.4, label="lines: kinetic theory; points: runs")
    b.set_ylim(25, 4e4)
    b.set_title("Doubling time with a 30 km/s rendezvous limit")
    b.legend(loc="upper right", ncol=2, fontsize=8)

    c_ = ax[1, 0]
    cc = cur[(cur.kind == "cube") & (cur.d_hop == 2.0) & (cur.dv == 30)]
    for td, col in zip(sorted(cc.t_delay.unique()), SEQ):
        s = cc[cc.t_delay == td]
        first = True
        for _, g in s.groupby("run"):
            c_.plot(g.t, np.maximum(g.n, 1), "-", color=col, lw=1.5, label=delay_label(td) if first else None)
            first = False
    c_.set_yscale("log"); c_.set_xscale("log"); c_.set_xlim(0.05, 10)
    c_.set_xlabel("time since seed [Myr]"); c_.set_ylabel("settled systems")
    c_.set_title("Settled population, d$_{hop}$ = 2 pc, 30 km/s limit, single seeds")
    c_.legend(loc="upper left")

    d_ = ax[1, 1]
    series = [("slab", 3.0, 30.0, BLUE, "-", "along U (radial)"),
              ("slab_axis1", 3.0, 30.0, ORANGE, "-", "along V (rotation)"),
              ("slab_axis2", 3.0, 30.0, AQUA, "-", "along W (vertical)")]
    for kind, dh, dv, col, ls, lab in series:
        s = cur[(cur.kind == kind) & (cur.d_hop == dh) & (cur.dv == dv) & (cur.t_delay == 0.0)]
        if len(s):
            d_.plot(s.t, 0.5 * (s.f50p + s.f50m), ls, color=col, label=f"30 km/s limit, d$_{{hop}}$ = 3 pc, {lab}")
    s = cur[(cur.kind == "slab") & (cur.d_hop == 1.0) & np.isinf(cur.dv) & (cur.t_delay == 0.0)]
    if len(s):
        d_.plot(s.t, 0.5 * (s.f50p + s.f50m), "--", color=MUTED, lw=1.5,
                label="ships match any speed, d$_{hop}$ = 1 pc, along U")
    tt = np.array([0.0, 25.0])
    d_.plot(tt, 30 * KMS * tt, ":", color=INK2, lw=1)
    d_.text(24.6, 30 * KMS * 24.6 - 30, "30 km/s (ship speed)", color=INK2, fontsize=8, ha="right", va="top")
    d_.set_xlim(0, 25); d_.set_ylim(0, 1100)
    d_.set_xlabel("time since seed [Myr]"); d_.set_ylabel("position of half-settled front [pc]")
    d_.set_title("Spreading in a 3 kpc slab (no delay)")
    d_.legend(loc="upper center", bbox_to_anchor=(0.5, -0.16), ncol=2, fontsize=8)
    fig.tight_layout()
    fig.savefig(os.path.join(FIGS, "fig4_settlement.png"))
    plt.close(fig)


def fig5():
    """How little capability is needed: growth time to 1e11 systems versus the
    rendezvous limit and versus the settleable fraction."""
    kd = pd.read_csv(os.path.join(DATA, "kinetic_vs_dvcap.csv"))
    fk = pd.read_csv(os.path.join(DATA, "f_scaling_kinetic.csv"))
    fr = pd.read_csv(os.path.join(DATA, "f_scaling_runs.csv"))
    sim = pd.read_csv(os.path.join(DATA, "settlement_runs.csv"))
    vr = pd.read_csv(os.path.join(DATA, "voyager_class_runs.csv"))
    ND = np.log2(1e11)
    hops = [0.5, 1.0, 2.0, 5.0]
    fig, ax = plt.subplots(1, 2, figsize=(11, 4.6))
    a = ax[0]
    fin = kd[np.isfinite(kd.dv_max_kms)]
    for d, c in zip(hops, SEQ):
        a.plot(fin.dv_max_kms, fin[f"T_grow_1e11_Myr_d{d:g}"], "-", color=c, label=f"d$_{{hop}}$ = {d:g} pc (kinetic theory)")
        s = sim[(sim.dv_max_kms == 30) & (sim.d_hop_pc == d) & (sim.t_delay_Myr == 0)]
        a.plot([30.0], ND * s.T2_sim_Myr * 1e3 / 1e3, "o", color=c, ms=6, mec=SURF, mew=1.2)
        s = vr[vr.d_hop_pc == d] if len(vr) else vr
        if len(s):
            a.plot([10.0], ND * s.T2_sim_Myr, "o", color=c, ms=6, mec=SURF, mew=1.2)
    a.plot([], [], "o", color=INK2, ms=6, label="simulation")
    for y, lab in ((1e3, "1 Gyr"), (1e4, "10 Gyr (age of the disc)")):
        a.axhline(y, color=INK2, lw=1, ls=":")
        a.text(102, y * 1.12, lab, color=INK2, fontsize=8, ha="right")
    a.axvline(30, color=INK2, lw=1, ls=":")
    a.text(29, 2.4e4, "30 km/s (lead rule)", color=INK2, fontsize=8, ha="right", va="center")
    a.set_xscale("log"); a.set_yscale("log"); a.set_xlim(4.5, 105); a.set_ylim(0.08, 5e4)
    a.set_xticks([5, 10, 20, 30, 50, 100]); a.set_xticklabels(["5", "10", "20", "30", "50", "100"])
    a.xaxis.set_minor_formatter(matplotlib.ticker.NullFormatter())
    a.set_xlabel("largest velocity change a ship can make to stop at a passing star [km/s]")
    a.set_ylabel("growth time to 10$^{11}$ systems at the local density [Myr]")
    a.set_title("Growth versus rendezvous limit (all systems settleable)")
    a.legend(loc="lower left", fontsize=8)
    b = ax[1]
    for d, c in zip(hops, SEQ):
        s = fk[(fk.dv_max_kms == 30) & (fk.d_hop_pc == d) & (fk.t_delay_Myr == 0)].sort_values("f")
        b.plot(s.f, s.T_grow_1e11_Myr, "-", color=c, label=f"d$_{{hop}}$ = {d:g} pc (kinetic theory)")
        r = fr[(fr.kind == "cube") & (fr.d_hop_pc == d)]
        if len(r):
            b.plot(r.f, ND * r.T2_sim_Myr, "o", color=c, ms=6, mec=SURF, mew=1.2)
        s1_ = sim[(sim.dv_max_kms == 30) & (sim.d_hop_pc == d) & (sim.t_delay_Myr == 0)]
        b.plot([1.0], ND * s1_.T2_sim_Myr, "o", color=c, ms=6, mec=SURF, mew=1.2)
    b.plot([], [], "o", color=INK2, ms=6, label="simulation")
    for y, lab in ((1e3, "1 Gyr"), (1e4, "10 Gyr (age of the disc)")):
        b.axhline(y, color=INK2, lw=1, ls=":")
        b.text(1.05, y * 1.12, lab, color=INK2, fontsize=8, ha="right")
    b.set_xscale("log"); b.set_yscale("log"); b.set_xlim(0.0095, 1.1); b.set_ylim(0.08, 5e4)
    b.set_xticks([0.01, 0.03, 0.1, 0.3, 1]); b.set_xticklabels(["0.01", "0.03", "0.1", "0.3", "1"])
    b.xaxis.set_minor_formatter(matplotlib.ticker.NullFormatter())
    b.set_xlabel("settleable fraction of star systems f")
    b.set_ylabel("growth time to 10$^{11}$ systems at the local density [Myr]")
    b.set_title("Growth versus settleable fraction (30 km/s limit, no delay)")
    b.legend(loc="lower left", fontsize=8)
    fig.tight_layout()
    fig.savefig(os.path.join(FIGS, "fig5_capability.png"))
    plt.close(fig)


def main():
    fig1(); fig2(); fig3()
    if os.path.exists(os.path.join(DATA, "settlement_runs.csv")):
        fig4(); fig5()
    print("figures written to", FIGS)


if __name__ == "__main__":
    main()