"""Part 2: how repeatable is the road to a hominin-sized brain?

(a) number of phylogenetically independent lineages within a given factor of the human value of
    relative brain size (EQ) and of cortical/pallial neuron number;
(b) the upper envelope of relative brain size through the Cenozoic, overall and by clade;
(c) waiting-time estimates for a second lineage to reach hominin-level brain 'hardware';
(d) hardware versus technology: a transparent (judgement-based) filter table.

Inputs : derived tables in data/ written by extract_data.py and build_transcribed.py
Outputs: data/part2_*.csv, data/part2_summary.json, figures/fig3..fig5
"""
import csv
import json
import os
import collections
import numpy as np
from scipy import stats

import phylo
import plotstyle

HERE = os.path.dirname(os.path.abspath(__file__))
DATA = os.path.join(HERE, "..", "data")
FIG = os.path.join(HERE, "..", "figures")
SEED = 20261011
NBOOT = 2000

# Endocast volume (cm3) -> brain mass (g).  The cetacean (Waugh & Thewissen 2021) and proboscidean
# (Benoit et al. 2019) tables already carry the authors' own clade-specific brain masses (about 0.75-0.80
# and 0.92 g per cm3).  For the two volume-only sources:
#   hominins (Du et al. 2018): brain mass = 1.147 * ECV^0.976 (Ruff, Trinkaus & Holliday 1997 Nature 387:173,
#       as quoted in a secondary source; 0.96-0.99 g per cm3 over 400-1,400 cm3);
#   Paleogene placentals (Bertrand et al. 2022): 1.0 g per cm3 (no clade-specific conversion exists; brain
#       mass scales isometrically with endocranial volume in primates, Isler et al. 2008 J Hum Evol 55:967).
# Sensitivity ("k092"): 0.92 g per cm3 for both, the value implied by the proboscidean table.
def ecv_to_g(ecv, kind, mode="primary"):
    if mode == "k092":
        return 0.92 * ecv
    if mode == "k1":
        return 1.0 * ecv
    if kind == "hominin":
        return 1.147 * ecv ** 0.976
    return 1.0 * ecv


FACTORS_EQ = [1.15, 1.25, 1.5, 2.0, 3.0]
SLICES_MA = [10.0, 20.0, 30.0]   # time slices used to define 'deep' lineages
FACTORS_N = [2.0, 3.0, 5.0, 10.0]
HUMAN_NCX = 16.3e9          # cortical neurons, isotropic fractionator (Azevedo et al. 2009)
EPOCHS = [("Paleocene", 66.0, 56.0), ("Eocene", 56.0, 33.9), ("Oligocene", 33.9, 23.03),
          ("Miocene", 23.03, 5.33), ("Plio-Pleistocene", 5.33, 0.0117), ("Recent", 0.0117, -1.0)]
AFROTHERIA = {"Proboscidea", "Hyracoidea", "Sirenia", "Afrosoricida", "Macroscelidea", "Tubulidentata"}


def eqj(brain_g, body_g):
    """Jerison (1973) encephalization quotient."""
    return brain_g / (0.12 * body_g ** 0.667)


def ncx_rule(brain_g, order):
    """Cortical neuron number predicted from brain mass by the isotropic-fractionator scaling rules
    quoted in Avelino-de-Souza et al. (2025, Table 2; from Herculano-Houzel and co-workers):
    primates e^17.125 M^0.933; afrotherians e^16.965 M^0.651; cetartiodactyls e^17.052 M^0.600.
    The cetartiodactyl rule is used here as a generic non-primate rule for all other orders."""
    if order == "Primates":
        return np.exp(17.125) * brain_g ** 0.933
    if order in AFROTHERIA:
        return np.exp(16.965) * brain_g ** 0.651
    return np.exp(17.052) * brain_g ** 0.600


def rd(path):
    return list(csv.DictReader(open(os.path.join(DATA, path), encoding="utf8")))


def wr(path, rows, cols=None):
    cols = cols or list(rows[0].keys())
    with open(os.path.join(DATA, path), "w", newline="", encoding="utf8") as f:
        w = csv.DictWriter(f, fieldnames=cols, extrasaction="ignore")
        w.writeheader()
        w.writerows(rows)


def fnum(x):
    try:
        v = float(x)
        return v
    except (TypeError, ValueError):
        return np.nan


# ------------------------------------------------------------------ data
def measured_ncx():
    """Cortical neuron numbers measured with the isotropic fractionator (species -> count), from the
    Massen et al. (2021) compilation and the values transcribed from open full texts."""
    out = {}
    for r in rd("neurons_mammals_massen2021.csv"):
        out[r["species"]] = float(r["cortex_neurons"])
    for r in rd("neurons_highend_transcribed.csv"):
        if r["kind"] == "IF_measured":
            out[r["species"]] = float(r["cortical_neurons_1e9"]) * 1e9
    return out


def load_extant(include_flagged=False, human_brain_g=None):
    rows = rd("mammals_extant_burger2019.csv")
    meas = measured_ncx()
    out = []
    for r in rows:
        if r["flag"] and not include_flagged:
            continue
        r = dict(r)
        r["brain_g"] = float(r["brain_g"])
        r["body_g"] = float(r["body_g"])
        if human_brain_g and r["species"] == "Homo_sapiens":
            r["brain_g"] = float(human_brain_g)
        r["well"] = r.get("well_sampled", "0") == "1"
        r["eq"] = eqj(r["brain_g"], r["body_g"])
        r["ncx_rule"] = float(ncx_rule(r["brain_g"], r["order"]))
        # measured isotropic-fractionator count where one exists, otherwise the rule prediction
        r["ncx_is_measured"] = r["species"] in meas
        r["ncx"] = float(meas.get(r["species"], r["ncx_rule"]))
        out.append(r)
    lb = np.log10([r["brain_g"] for r in out])
    lm = np.log10([r["body_g"] for r in out])
    sl, ic = stats.linregress(lm, lb)[:2]
    for r, a, b in zip(out, lb, lm):
        r["eq75"] = float(10 ** (a - (ic + sl * b)))
    return out, (float(sl), float(ic))


def load_fossils(ecv_mode="primary"):
    """One table of fossil (and sub-fossil) brain/body estimates with ages. Returns list of dicts with
    source, clade, species, age_Ma, brain_g, body_g, eq, eq_alt (alternative body mass, cetaceans).
    ecv_mode selects the endocast-volume-to-brain-mass conversion for the volume-only sources."""
    out = []
    for r in rd("fossil_placentals_bertrand2022.csv"):
        age = fnum(r["age_Ma"])
        if not np.isfinite(age):
            continue
        if "Archaeoceti" in r["order"]:
            continue  # the same archaeocete specimens are in the Waugh & Thewissen table
        b, m = ecv_to_g(float(r["ecv_cm3"]), "generic", ecv_mode), float(r["body_g"])
        clade = "Primates" if r["order"].startswith("Primates") else "other mammals"
        out.append(dict(source="Bertrand2022", clade=clade, species=r["species"], age_Ma=age, brain_g=b,
                        body_g=m, eq=eqj(b, m), eq_alt=eqj(b, m), order=r["order"]))
    for r in rd("fossil_cetaceans_waugh2021.csv"):
        b = float(r["brain_g"])
        mw = fnum(r["body_kg_waugh"]) * 1000
        mb = fnum(r["body_kg_boessenecker"]) * 1000
        e1 = eqj(b, mw)
        e2 = eqj(b, mb) if np.isfinite(mb) else e1
        out.append(dict(source="Waugh2021", clade="Cetacea", species=r["species"].replace(" ", "_"),
                        age_Ma=float(r["age_Ma"]), brain_g=b, body_g=mw, eq=e1, eq_alt=e2, order=r["family"]))
    for r in rd("fossil_proboscideans_benoit2019.csv"):
        if float(r["age_Ma"]) == 0:
            continue
        b, m = float(r["brain_g"]), float(r["body_g"])
        out.append(dict(source="Benoit2019", clade=r["clade"], species=r["species"], age_Ma=float(r["age_Ma"]),
                        brain_g=b, body_g=m, eq=eqj(b, m), eq_alt=eqj(b, m), order=r["clade"]))
    body = {r["taxon"]: float(r["body_kg_assumed"]) * 1000 for r in rd("hominin_body_mass_assumed.csv")}
    for r in rd("hominin_ecv_du2018.csv"):
        b = ecv_to_g(float(r["ecv_cm3"]), "hominin", ecv_mode)
        m = body[r["taxon"]]
        out.append(dict(source="Du2018", clade="Hominini", species=r["taxon"], age_Ma=float(r["age_Ma"]),
                        brain_g=b, body_g=m, eq=eqj(b, m), eq_alt=eqj(b, m), order="Primates"))
    return out


def epoch_of(age):
    for nm, a, b in EPOCHS:
        if a >= age > b:
            return nm
    return "Mesozoic" if age > 66 else "Recent"


# ------------------------------------------------------------------ (a) lineage counts
def deep_lineages(ext, root, tipset, metric, thr, t_cut):
    """Count 'deep' lineages over a threshold: cut the timetree t_cut Myr ago; every lineage crossing the
    cut defines a clade.  A clade counts if a non-human member is over the threshold ('any'), or, more
    robustly against single-record noise, if a member with >= 3 brains measured is over it or at least two
    members are ('robust').  Origins among those clades are then counted by parsimony on the backbone."""
    by = {r["species"]: r for r in ext if r["species"] in tipset}
    pr = phylo.prune(root, set(by))
    clades = phylo.slice_clades(pr, t_cut)
    res = {}
    for mode in ("any", "robust", "n3plus"):
        state, hit = {}, []
        for cl in clades:
            mem = [by[sp] for sp in cl if sp != "Homo_sapiens"]
            over = [r for r in mem if r[metric] >= thr]
            if mode == "any":
                ok = len(over) >= 1
            elif mode == "n3plus":      # strictest: a species with >= 3 brains measured must be over
                ok = any(r["well"] for r in over)
            else:
                ok = any(r["well"] for r in over) or len(over) >= 2
            state[cl[0]] = int(ok)
            if ok:
                hit.append(sorted({r["family"] for r in over}))
        bb = phylo.prune(pr, set(state))
        o = phylo.origins(bb, state) if sum(state.values()) else dict(gains_min=0, gains_max=0)
        res[mode] = dict(n_clades=len(clades), n_clades_over=len(hit), origins_min=o["gains_min"],
                         origins_max=o["gains_max"], families_per_clade=hit)
    return res


def lineage_counts(ext, trees, summary):
    human = [r for r in ext if r["species"] == "Homo_sapiens"][0]
    rows = []
    fam_lists = {}
    deep = {}
    for metric, factors in (("eq", FACTORS_EQ), ("eq75", FACTORS_EQ), ("brain_g", [1.0, 2.0, 3.0]),
                            ("ncx", FACTORS_N)):
        href = HUMAN_NCX if metric == "ncx" else human[metric]
        for f in factors:
            thr = href / f
            hit = [r for r in ext if r[metric] >= thr and r["species"] != "Homo_sapiens"]
            hitw = [r for r in hit if r["well"]]
            gen = collections.defaultdict(list)
            for r in ext:
                if r["species"] != "Homo_sapiens":
                    gen[(r["family"], r["genus"])].append(r[metric])
            gmed = {k: float(np.median(v)) for k, v in gen.items()}
            ghit = [k for k, v in gmed.items() if v >= thr]
            row = dict(metric=metric, within_factor=f, threshold=thr, n_other_species=len(hit),
                       n_other_genera=len({r["genus"] for r in hit}),
                       n_other_families=len({r["family"] for r in hit}),
                       n_other_orders=len({r["order"] for r in hit}),
                       n_species_n3plus=len(hitw), n_families_n3plus=len({r["family"] for r in hitw}),
                       n_genera_by_genus_median=len(ghit),
                       n_families_by_genus_median=len({k[0] for k in ghit}))
            if metric == "ncx":
                row["n_species_measured_count"] = len([r for r in hit if r["ncx_is_measured"]])
                row["species_measured_count"] = ";".join(sorted(r["species"] for r in hit if r["ncx_is_measured"]))
            fam_lists[f"{metric}_x{f}"] = sorted({r["family"] for r in hit})
            fam_lists[f"{metric}_x{f}_n3plus"] = sorted({r["family"] for r in hitw})
            for tname, (root, tipset) in trees.items():
                keep = {r["species"] for r in ext} & tipset
                pr = phylo.prune(root, keep)
                # species-level parsimony (kept as a robustness number only: it is inflated by record noise)
                st = {r["species"]: int(r[metric] >= thr) for r in ext if r["species"] in keep}
                o = phylo.origins(pr, st)
                st0 = dict(st)
                st0["Homo_sapiens"] = 0
                o0 = phylo.origins(pr, st0)
                row[f"{tname}_species_in_tree"] = int(sum(st.values()))
                row[f"{tname}_origins_min_incl_human"] = o["gains_min"]
                row[f"{tname}_origins_max_incl_human"] = o["gains_max"]
                row[f"{tname}_origins_noloss_incl_human"] = o["gains_noloss"]
                row[f"{tname}_other_origins_min"] = o0["gains_min"]
                row[f"{tname}_other_origins_max"] = o0["gains_max"]
            # deep lineages (primary tree only)
            root, tipset = trees["alvarez2022"]
            for t_cut in SLICES_MA:
                d = deep_lineages(ext, root, tipset, metric, thr, t_cut)
                row[f"deep{t_cut:.0f}Ma_clades_any"] = d["any"]["n_clades_over"]
                row[f"deep{t_cut:.0f}Ma_clades_robust"] = d["robust"]["n_clades_over"]
                row[f"deep{t_cut:.0f}Ma_clades_n3plus"] = d["n3plus"]["n_clades_over"]
                row[f"deep{t_cut:.0f}Ma_origins_robust"] = f"{d['robust']['origins_min']}-{d['robust']['origins_max']}"
                deep[f"{metric}_x{f}_{t_cut:.0f}Ma"] = d
            rows.append(row)
    wr("part2_lineage_counts.csv", rows, cols=sorted({k for r in rows for k in r}, key=lambda k: (k not in (
        "metric", "within_factor", "threshold"), k)))
    summary["lineage_counts"] = rows
    summary["families_over_threshold"] = fam_lists
    summary["deep_lineages"] = deep
    # joint criterion: large in both relative and absolute terms (EQ can be raised by a small body alone)
    joint = {}
    for f in (2.0, 3.0):
        hit = [r for r in ext if r["species"] != "Homo_sapiens" and r["eq"] >= human["eq"] / f
               and r["brain_g"] >= human["brain_g"] / f]
        joint[f"x{f}"] = dict(n_species=len(hit), n_species_n3plus=len([r for r in hit if r["well"]]),
                              families=sorted({r["family"] for r in hit}),
                              species=sorted(r["species"] for r in hit))
    summary["joint_EQ_and_brain_mass"] = joint
    return rows


def neuron_table(ext, summary):
    """Lineage-level table of cortical/pallial neuron numbers: measured (IF), predicted from the IF
    scaling rules, and stereological counts where they exist."""
    tr = rd("neurons_highend_transcribed.csv")
    birds = [r for r in rd("neurons_birds_reptiles.csv") if r["pallium_neurons"] and r["cls"] == "Aves"]
    seen, b2 = set(), []
    for r in sorted(birds, key=lambda r: -float(r["pallium_neurons"])):
        key = r["species"].replace("scandiaca", "scandiacus")
        if key in seen:
            continue
        seen.add(key)
        b2.append(r)
    rows = []
    for r in tr:
        rows.append(dict(lineage=r["lineage"], species=r["species"], kind=r["kind"],
                         neurons_1e9=float(r["cortical_neurons_1e9"]), brain_g=r["brain_g"],
                         ratio_to_human_IF=float(r["cortical_neurons_1e9"]) * 1e9 / HUMAN_NCX,
                         source=r["source"], verified=r["verified"]))
    for r in b2[:8]:
        rows.append(dict(lineage="Aves", species=r["species"], kind="IF_measured (pallium)",
                         neurons_1e9=float(r["pallium_neurons"]) / 1e9, brain_g=r["brain_g"],
                         ratio_to_human_IF=float(r["pallium_neurons"]) / HUMAN_NCX, source=r["source"],
                         verified="V"))
    wr("part2_neuron_table.csv", rows)
    # accuracy of the generic rule against measured IF counts (Massen compilation + transcribed)
    chk = []
    for r in rd("neurons_mammals_massen2021.csv"):
        order = "Primates" if r["order"] == "Primates" else r["order"]
        pred = ncx_rule(float(r["brain_g"]), order)
        chk.append((r["species"], order, float(r["cortex_neurons"]), float(pred)))
    for r in tr:
        if r["kind"] == "IF_measured" and r["brain_g"]:
            order = {"Hominini": "Primates", "Panini": "Primates", "Cercopithecidae": "Primates",
                     "Proboscidea": "Proboscidea"}.get(r["lineage"], "Cetartiodactyla")
            chk.append((r["species"] + "*", order, float(r["cortical_neurons_1e9"]) * 1e9,
                        float(ncx_rule(float(r["brain_g"]), order))))
    ratio = np.array([p / m for _, _, m, p in chk])
    summary["neuron_rule_check"] = dict(
        n=len(chk), median_pred_over_measured=float(np.median(ratio)),
        p10=float(np.percentile(ratio, 10)), p90=float(np.percentile(ratio, 90)),
        rms_log10=float(np.sqrt(np.mean(np.log10(ratio) ** 2))),
        primates=[dict(species=sp, measured_1e9=m / 1e9, predicted_1e9=pr_ / 1e9, ratio=pr_ / m)
                  for sp, od, m, pr_ in chk if od == "Primates"],
        note="rule applied to all orders; 22 species of the Massen et al. compilation plus 7 transcribed; "
             "the primate rule over-predicts the measured human count")
    # how many bird species and which orders exceed human/10
    nb = [(r["species"], float(r["pallium_neurons"])) for r in b2 if float(r["pallium_neurons"]) >= HUMAN_NCX / 10]
    summary["birds_over_tenth_human"] = [dict(species=s, pallium_neurons=v) for s, v in nb]
    summary["bird_max_pallium_neurons"] = dict(species=b2[0]["species"], value=float(b2[0]["pallium_neurons"]))
    return rows


# ------------------------------------------------------------------ (b) envelope
def envelope(ext, fos, summary, rng):
    human = [r for r in ext if r["species"] == "Homo_sapiens"][0]
    allrows = []
    for r in fos:
        allrows.append(dict(r, epoch=epoch_of(r["age_Ma"])))
    for r in ext:
        clade = {"Cetacea": "Cetacea", "Proboscidea": "Proboscidea", "Primates": "Primates"}.get(r["order"], "other mammals")
        if r["species"] == "Homo_sapiens":
            clade = "Hominini"
        allrows.append(dict(source="Burger2019", clade=clade, species=r["species"], age_Ma=0.0, brain_g=r["brain_g"],
                            body_g=r["body_g"], eq=r["eq"], eq_alt=r["eq"], order=r["order"], epoch="Recent",
                            well=r["well"]))
    wr("part2_brain_size_through_time.csv", allrows,
       ["source", "clade", "species", "age_Ma", "epoch", "brain_g", "body_g", "eq", "eq_alt", "order"])

    def binstats(rows, key="eq"):
        v = np.array([r[key] for r in rows])
        if len(v) == 0:
            return None
        return dict(n=int(len(v)), median=float(np.median(v)), p90=float(np.percentile(v, 90)), max=float(v.max()),
                    max_taxon=rows[int(np.argmax(v))]["species"])

    table = []
    for nm, a, b in EPOCHS:
        # hominins are summarised per taxon in summary["hominin_history"] (per-cranium EQ with a fixed
        # species body mass would overstate their spread)
        for clade in ("all non-hominin", "Cetacea", "Proboscidea", "Primates", "other mammals"):
            sel = [r for r in allrows if r["epoch"] == nm and
                   ((clade == "all non-hominin" and r["clade"] != "Hominini") or r["clade"] == clade)]
            s = binstats(sel)
            if not s:
                continue
            s2 = binstats(sel, "eq_alt")
            table.append(dict(epoch=nm, start_Ma=a, end_Ma=max(b, 0.0), clade=clade, **s, max_alt_body_mass=s2["max"]))
    wr("part2_envelope_by_epoch.csv", table)
    summary["envelope_by_epoch"] = [t for t in table if t["clade"] == "all non-hominin"]

    # slope of ln(max EQ) of non-hominins against time, epoch midpoints (Recent at 0)
    def slope_from(rows, key, stat="max"):
        xs, ys = [], []
        for nm, a, b in EPOCHS:
            v = [r[key] for r in rows if r["epoch"] == nm]
            if not v:
                continue
            xs.append(0.0 if nm == "Recent" else 0.5 * (a + max(b, 0)))
            ys.append(np.log(max(v) if stat == "max" else np.percentile(v, 90)))
        lrr = stats.linregress(-np.array(xs), np.array(ys))
        return lrr.slope, lrr.intercept, xs, ys, lrr.stderr

    nonh = [r for r in allrows if r["clade"] != "Hominini" and r["age_Ma"] <= 66]
    res = {}
    for key in ("eq", "eq_alt"):
        for stat in ("max", "p90"):
            sl, ic, xs, ys, se = slope_from(nonh, key, stat)
            boots = []
            byep = collections.defaultdict(list)
            for r in nonh:
                byep[r["epoch"]].append(r)
            for _ in range(NBOOT):
                samp = []
                for nm, lst in byep.items():
                    idx = rng.integers(0, len(lst), len(lst))
                    samp.extend(lst[i] for i in idx)
                boots.append(slope_from(samp, key, stat)[0])
            res[f"{key}_{stat}"] = dict(slope_ln_per_Myr=float(sl), slope_se_ols=float(se), intercept_lnEQ_today=float(ic),
                                        boot_ci95=[float(np.percentile(boots, 2.5)), float(np.percentile(boots, 97.5))],
                                        doubling_time_Myr=float(np.log(2) / sl), epoch_mid_Ma=xs,
                                        ln_stat=[float(y) for y in ys])
    summary["envelope_slope"] = res

    # sensitivities of the slope of the epoch maxima (referee points M1, M5, m4, m5) ----------------
    sens = {}
    byep = collections.defaultdict(list)
    for r in nonh:
        byep[r["epoch"]].append(r)
    # (i) every epoch rarefied to n = 10 specimens (the smallest bin has 10)
    rs = []
    for _ in range(NBOOT):
        samp = []
        for nm, lst in byep.items():
            idx = rng.choice(len(lst), min(10, len(lst)), replace=False)
            samp.extend(lst[i] for i in idx)
        rs.append(slope_from(samp, "eq")[0])
    sens["rarefied_n10_per_epoch"] = dict(mean=float(np.mean(rs)), p2_5=float(np.percentile(rs, 2.5)),
                                          p97_5=float(np.percentile(rs, 97.5)))
    # (ii) without the Recent bin (fossils only)
    sens["without_Recent_bin"] = float(slope_from([r for r in nonh if r["epoch"] != "Recent"], "eq")[0])
    # (iii) Recent maximum taken from species with >= 3 brains measured only
    sens["Recent_max_from_n3plus_records"] = float(slope_from(
        [r for r in nonh if r["epoch"] != "Recent" or r.get("well")], "eq")[0])
    # (iv) endocast volumes of the volume-only sources converted at 0.92 g per cm3
    alt = [dict(r, epoch=epoch_of(r["age_Ma"])) for r in load_fossils("k092") if r["clade"] != "Hominini"]
    alt += [r for r in nonh if r["epoch"] == "Recent"]
    sens["endocasts_at_0p92_g_per_cm3"] = float(slope_from(alt, "eq")[0])
    summary["envelope_slope_sensitivity"] = sens
    # which clade holds the maximum in each epoch, and how clade-targeted the bins are
    comp = []
    for nm, a, b in EPOCHS:
        lst = byep.get(nm, [])
        if not lst:
            continue
        top = max(lst, key=lambda r: r["eq"])
        cnt = collections.Counter(r["clade"] for r in lst)
        comp.append(dict(epoch=nm, n=len(lst), max_taxon=top["species"], max_clade=top["clade"], max_eq=float(top["eq"]),
                         n_by_clade=dict(cnt)))
    summary["envelope_bin_composition"] = comp
    # first appearance in this (clade-targeted) sample of a non-hominin within a factor 2 and 3 of the human EQ
    fa = {}
    for f in (3.0, 2.0, 1.5):
        v = [r for r in nonh if r["eq"] >= human["eq"] / f and r["age_Ma"] > 0]
        if v:
            o = max(v, key=lambda r: r["age_Ma"])
            fa[f"x{f}"] = dict(age_Ma=float(o["age_Ma"]), taxon=o["species"], eq=float(o["eq"]))
    summary["first_fossil_within_factor_of_human_EQ"] = fa

    # sample-size (rarefaction) check: expected maximum in random subsamples of n = 20
    rare = []
    for nm, a, b in EPOCHS:
        v = np.array([r["eq"] for r in nonh if r["epoch"] == nm])
        if len(v) < 5:
            continue
        n = min(20, len(v))
        mx = [v[rng.choice(len(v), n, replace=False)].max() for _ in range(NBOOT)]
        rare.append(dict(epoch=nm, n_available=int(len(v)), n_subsample=n, mean_max_subsample=float(np.mean(mx)),
                         raw_max=float(v.max())))
    summary["rarefied_max_n20"] = rare

    # clade histories -------------------------------------------------------------
    cet = [r for r in allrows if r["clade"] == "Cetacea"]
    cf = [r for r in cet if r["source"] == "Waugh2021"]

    def mx(rows, lo, hi, key="eq"):
        v = [(r[key], r["species"], r["age_Ma"]) for r in rows if lo < r["age_Ma"] <= hi]
        return max(v) if v else (np.nan, "", np.nan)

    e_eoc = mx(cf, 33.9, 60)
    e_mio = mx(cf, 5.33, 23.03)
    e_plio = mx(cf, 0.0117, 5.33)
    ext_cet = sorted([(r["eq"], r["species"]) for r in ext if r["order"] == "Cetacea"], reverse=True)
    ext_cet_w = sorted([(r["eq"], r["species"]) for r in ext if r["order"] == "Cetacea" and r["well"]], reverse=True)
    cet_hist = dict(
        max_archaeocete_EQ=dict(value=float(e_eoc[0]), taxon=e_eoc[1], age_Ma=float(e_eoc[2])),
        max_Miocene_EQ=dict(value=float(e_mio[0]), taxon=e_mio[1], age_Ma=float(e_mio[2])),
        max_PlioPleistocene_fossil_EQ=dict(value=float(e_plio[0]), taxon=e_plio[1], age_Ma=float(e_plio[2])),
        max_extant_EQ=dict(value=float(ext_cet[0][0]), taxon=ext_cet[0][1]),
        max_extant_EQ_n3plus=dict(value=float(ext_cet_w[0][0]), taxon=ext_cet_w[0][1]),
        rate_16Ma_to_extant_max_n3plus=float(np.log(ext_cet_w[0][0] / e_mio[0]) / e_mio[2]),
        rate_47_to_16Ma=float(np.log(e_mio[0] / mx(cf, 44, 60)[0]) / (47.0 - e_mio[2])),
        rate_16Ma_to_extant_max=float(np.log(ext_cet[0][0] / e_mio[0]) / e_mio[2]),
        rate_16Ma_to_4Ma_fossil=float(np.log(e_plio[0] / e_mio[0]) / (e_mio[2] - e_plio[2])))
    # delphinoid 'plateau' test: regression of ln EQ on age among fossil delphinoids (<= 18 Ma) and extant
    delph_f = [r for r in cf if r["order"] in ("Kentriodontidae", "Delphinidae", "Phocoenidae")]
    delph_e = [r for r in ext if r["family"] in ("Delphinidae", "Phocoenidae", "Monodontidae")]
    xa = np.array([r["age_Ma"] for r in delph_f] + [0.0] * len(delph_e))
    ya = np.log(np.array([r["eq"] for r in delph_f] + [r["eq"] for r in delph_e]))
    lr = stats.linregress(-xa, ya)
    mw = stats.mannwhitneyu([r["eq"] for r in delph_f], [r["eq"] for r in delph_e], alternative="two-sided")
    cet_hist["delphinoid_trend"] = dict(
        n_fossil=len(delph_f), n_extant=len(delph_e), slope_ln_per_Myr=float(lr.slope),
        slope_ci95=[float(lr.slope - 1.96 * lr.stderr), float(lr.slope + 1.96 * lr.stderr)], p=float(lr.pvalue),
        median_fossil_EQ=float(np.median([r["eq"] for r in delph_f])),
        median_extant_EQ=float(np.median([r["eq"] for r in delph_e])), mannwhitney_p=float(mw.pvalue),
        note="fossil body masses are estimated from skeletal proxies, extant ones are weighed; the two are not strictly comparable")
    # same test with crown-group fossils only (Delphinidae + Phocoenidae; kentriodontids are stem delphinoids)
    crown_f = [r for r in delph_f if r["order"] in ("Delphinidae", "Phocoenidae")]
    xa = np.array([r["age_Ma"] for r in crown_f] + [0.0] * len(delph_e))
    ya = np.log(np.array([r["eq"] for r in crown_f] + [r["eq"] for r in delph_e]))
    lrc = stats.linregress(-xa, ya)
    cet_hist["delphinoid_trend_crown_fossils_only"] = dict(
        n_fossil=len(crown_f), n_extant=len(delph_e), slope_ln_per_Myr=float(lrc.slope),
        slope_ci95=[float(lrc.slope - 1.96 * lrc.stderr), float(lrc.slope + 1.96 * lrc.stderr)], p=float(lrc.pvalue))
    # with the alternative (Boessenecker) fossil body masses
    ya = np.log(np.array([r["eq_alt"] for r in delph_f] + [r["eq"] for r in delph_e]))
    xa = np.array([r["age_Ma"] for r in delph_f] + [0.0] * len(delph_e))
    lra = stats.linregress(-xa, ya)
    cet_hist["delphinoid_trend_alt_body_mass"] = dict(slope_ln_per_Myr=float(lra.slope), p=float(lra.pvalue))
    cet_hist["plateau_summary"] = ("fossil cetacean maximum EQ is flat between the middle Miocene and the Pliocene "
                                   "(rate_16Ma_to_4Ma_fossil); the positive delphinoid slope compares fossils with "
                                   "skeletal mass estimates against weighed living animals and is not significant with "
                                   "crown-group fossils only: the data cannot distinguish a slow rise from a plateau "
                                   "since about 15 Ma (Marino et al. 2004 describe the second increase as complete by 15 Ma)")
    # EQ is not independent of body size inside a clade: among living delphinoids ln EQ falls with ln body mass
    lb_ = np.log([r["body_g"] for r in delph_e]); le_ = np.log([r["eq"] for r in delph_e])
    lbr_ = np.log([r["brain_g"] for r in delph_e])
    l1 = stats.linregress(lb_, le_); l2 = stats.linregress(lb_, lbr_)
    summary["eq_vs_body_within_delphinoids"] = dict(
        n=len(delph_e), slope_lnEQ_on_lnbody=float(l1.slope), p=float(l1.pvalue),
        brain_body_slope=float(l2.slope), brain_body_slope_se=float(l2.stderr),
        note="the within-clade brain-body slope is well below the 0.667 built into EQ, so small-bodied species "
             "have high EQ; a lineage can raise EQ by shrinking its body")
    summary["cetacean_history"] = cet_hist

    pro = [r for r in allrows if r["clade"] == "Proboscidea"]
    pf = sorted([(r["age_Ma"], r["species"], r["eq"], r["brain_g"]) for r in pro if r["source"] == "Benoit2019"], reverse=True)
    summary["proboscidean_history"] = dict(
        fossils=[dict(age_Ma=a, species=s, eq=float(e), brain_g=b) for a, s, e, b in pf],
        extant=[dict(species=r["species"], eq=r["eq"], brain_g=r["brain_g"]) for r in ext if r["order"] == "Proboscidea"])

    hom = [r for r in allrows if r["clade"] == "Hominini" and r["source"] == "Du2018"]
    xa = np.array([r["age_Ma"] for r in hom])
    lr = stats.linregress(-xa, np.log([r["brain_g"] for r in hom]))  # brain mass from ECV (Ruff et al. 1997)
    lre = stats.linregress(-xa, np.log([r["eq"] for r in hom]))
    bysp = collections.defaultdict(list)
    for r in hom:
        bysp[r["species"]].append(r)
    dolphin_max = ext_cet[0][0]
    tax = []
    for k, v in bysp.items():
        tax.append(dict(taxon=k, n=len(v), mean_age_Ma=float(np.mean([r["age_Ma"] for r in v])),
                        min_age_Ma=float(min(r["age_Ma"] for r in v)), max_age_Ma=float(max(r["age_Ma"] for r in v)),
                        median_brain_g=float(np.median([r["brain_g"] for r in v])), body_kg_assumed=v[0]["body_g"] / 1000,
                        eq=float(eqj(np.median([r["brain_g"] for r in v]), v[0]["body_g"]))))
    tax.sort(key=lambda d: -d["mean_age_Ma"])
    above = [d for d in tax if d["eq"] > dolphin_max]
    above_w = [d["taxon"] for d in tax if d["eq"] > ext_cet_w[0][0]]
    summary["hominin_history"] = dict(
        n_crania=len(hom), ecv_rate_ln_per_Myr=float(lr.slope), ecv_rate_se=float(lr.stderr),
        eq_rate_ln_per_Myr=float(lre.slope), eq_rate_se=float(lre.stderr), by_taxon=tax,
        fossil_taxa_with_median_EQ_above_max_extant_dolphin=[d["taxon"] for d in above],
        fossil_taxa_with_median_EQ_above_max_extant_dolphin_n3plus=above_w,
        oldest_such_taxon_age_range_Ma=[above[0]["min_age_Ma"], above[0]["max_age_Ma"]] if above else None,
        human_EQ=float(human["eq"]),
        note="hominin EQ uses assumed species body masses (flag L; +-30 % body mass = -+19 % EQ) and brain mass from "
             "endocranial volume by the Ruff et al. (1997) formula; which fossil hominin taxa exceed 'the dolphin "
             "level' depends on those assumptions and on which dolphin record is trusted")

    prim_f = [r for r in fos if r["clade"] == "Primates"]
    prim_e = sorted([(r["eq"], r["species"]) for r in ext if r["order"] == "Primates" and r["species"] != "Homo_sapiens"], reverse=True)
    pmax = max(prim_f, key=lambda r: r["eq"])
    summary["primate_history"] = dict(
        max_Paleogene_EQ=dict(value=float(pmax["eq"]), taxon=pmax["species"], age_Ma=float(pmax["age_Ma"])),
        max_extant_nonhuman_EQ=dict(value=float(prim_e[0][0]), taxon=prim_e[0][1]),
        rate=float(np.log(prim_e[0][0] / pmax["eq"]) / pmax["age_Ma"]))

    # whole-distribution shift: median EQ of Paleocene placentals vs Recent mammals
    pal = np.array([r["eq"] for r in nonh if r["epoch"] == "Paleocene"])
    eoc = np.array([r["eq"] for r in nonh if r["epoch"] == "Eocene" and r["source"] == "Bertrand2022"])
    rec = np.array([r["eq"] for r in ext if r["species"] != "Homo_sapiens"])
    summary["median_shift"] = dict(median_Paleocene=float(np.median(pal)), median_Eocene=float(np.median(eoc)),
                                   median_Recent=float(np.median(rec)),
                                   drift_ln_per_Myr=float(np.log(np.median(rec) / np.median(pal)) / 61.0),
                                   n_Paleocene=int(len(pal)), n_Eocene=int(len(eoc)), n_Recent=int(len(rec)),
                                   note="NOT a like-for-like trend: the Paleocene sample is of archaic mammals, the Miocene and "
                                        "Plio-Pleistocene bins are almost all cetaceans and proboscideans, and the Recent sample "
                                        "is taxonomically different; quoted for the hindcast starting point only")
    return allrows


# ------------------------------------------------------------------ (c) waiting time
ORDERS_RATE = ("Primates", "Cetacea", "Carnivora", "Rodentia", "Chiroptera", "Artiodactyla")


def bm_rates(ext, trees, summary):
    """Brownian-motion rates (variance of ln trait per Myr) from independent contrasts: the plain estimate,
    a robust (median-based) one, and -- added in revision, referee point M5 -- a noise-corrected one from a
    REML fit of Brownian motion plus independent tip error (single-record measurement noise)."""
    by = {r["species"]: r for r in ext}
    out = {}
    for tname, (root, tipset) in trees.items():
        keep = {r["species"] for r in ext} & tipset
        if tname == "upham2019":
            keep -= {"Lissodelphis_borealis"}  # misplaced inside Balaenopteridae in the distributed MCC tree
        pr = phylo.prune(root, keep)
        res = {}
        groups = {"Mammalia": keep}
        for od in ORDERS_RATE:
            groups[od] = {s for s in keep if by[s]["order"] == od}
        tau2 = {}
        for g, sub in groups.items():
            p2 = phylo.prune(pr, sub)
            for trait in ("eq", "brain_g"):
                x = {s: np.log(by[s][trait]) for s in sub}
                cons, s2, _ = phylo.contrasts(p2, x)
                c = np.array([o["contrast"] for o in cons])
                d = dict(n=len(sub), sigma2=float(s2),
                         sigma2_robust=float((np.median(np.abs(c)) / 0.6745) ** 2),
                         excess_kurtosis=float(stats.kurtosis(c)))
                if tname == "alvarez2022":
                    if g == "Mammalia":
                        me = phylo.bm_with_tip_error(p2, x)
                        tau2[trait] = me["tau2"]
                        d.update(sigma2_noise_corrected=me["sigma2"], sigma2_nc_lo95=me["sigma2_lo95"],
                                 sigma2_nc_hi95=me["sigma2_hi95"], tip_error_sd=me["tip_error_sd"],
                                 tip_error_sd_lo95=float(np.sqrt(me["tau2_lo95"])),
                                 tip_error_sd_hi95=float(np.sqrt(me["tau2_hi95"])),
                                 dloglik_vs_no_error=me["dloglik_vs_no_error"])
                    else:
                        # tip-error variance fixed at the all-mammal estimate (too few tips to fit it per order)
                        me = phylo.bm_rate_given_tip_error(p2, x, tau2[trait])
                        d.update(sigma2_noise_corrected=me["sigma2"], sigma2_nc_lo95=me["sigma2_lo95"],
                                 sigma2_nc_hi95=me["sigma2_hi95"])
                res[f"{g}_{trait}"] = d
        if tname == "alvarez2022":
            # independent cross-check with sister-species pairs only: E[d^2] = sigma2 * T + 2 tau2
            from scipy import optimize
            for trait in ("eq", "brain_g"):
                x = {s: np.log(by[s][trait]) for s in keep}
                T, D = phylo.sister_pairs(pr, x)

                def nll(th):
                    var = np.exp(th[0]) * T + 2 * np.exp(th[1])
                    return 0.5 * np.sum(np.log(var) + D * D / var)
                o = optimize.minimize(nll, [np.log(0.002), np.log(0.02)], method="Nelder-Mead",
                                      options=dict(xatol=1e-6, fatol=1e-8, maxiter=4000))
                res[f"sister_pairs_{trait}"] = dict(
                    n_pairs=int(len(T)), sigma2=float(np.exp(o.x[0])), tip_error_sd=float(np.sqrt(np.exp(o.x[1]))),
                    rms_diff_pairs_under_2Myr=float(np.sqrt(np.mean(D[T < 2] ** 2))) if (T < 2).sum() else None,
                    n_pairs_under_2Myr=int((T < 2).sum()),
                    rms_diff_pairs_20_to_200Myr=float(np.sqrt(np.mean(D[(T >= 20) & (T < 200)] ** 2))),
                    n_pairs_20_to_200Myr=int(((T >= 20) & (T < 200)).sum()),
                    note="sister-species pairs only; one all-mammal rate, used for every order in the "
                         "'sister-pair' waiting-time rows")
            # estimator check by simulation on the same tree: Brownian motion at the fitted rate plus tip noise
            # of the fitted size; which estimator recovers the true rate?
            rs = np.random.default_rng(SEED)
            s2_true, tau_true = res["Mammalia_eq"]["sigma2_noise_corrected"], res["Mammalia_eq"]["tip_error_sd"]
            sims = []
            for _ in range(8):
                val = {id(pr): 0.0}
                xs = {}
                for nd in phylo.iter_nodes(pr):
                    for c in nd.children:
                        val[id(c)] = val[id(nd)] + rs.normal(0, np.sqrt(s2_true * c.length))
                    if nd.is_tip():
                        xs[nd.name] = val[id(nd)] + rs.normal(0, tau_true)
                me = phylo.bm_with_tip_error(pr, xs)
                T, D = phylo.sister_pairs(pr, xs)
                o = optimize.minimize(lambda th: 0.5 * np.sum(np.log(np.exp(th[0]) * T + 2 * np.exp(th[1])) + D * D / (
                    np.exp(th[0]) * T + 2 * np.exp(th[1]))), [np.log(0.002), np.log(0.02)], method="Nelder-Mead",
                    options=dict(xatol=1e-6, fatol=1e-8, maxiter=4000))
                sims.append((me["sigma2_no_error"], me["sigma2"], float(np.exp(o.x[0]))))
            sims = np.array(sims)
            res["estimator_check_simulation"] = dict(
                true_sigma2=s2_true, true_tip_error_sd=tau_true, n_sims=len(sims),
                plain_range=[float(sims[:, 0].min()), float(sims[:, 0].max())],
                whole_tree_noise_corrected_range=[float(sims[:, 1].min()), float(sims[:, 1].max())],
                sister_pair_range=[float(sims[:, 2].min()), float(sims[:, 2].max())],
                note="simulated data with the fitted rate and tip noise reproduce the inflated plain estimate; the "
                     "whole-tree fit recovers the true rate tightly, the sister-pair estimator scatters widely")
        out[tname] = res
    summary["bm_rates"] = out
    return out


def p_any(d, sig2, t):
    """P(at least one of the independent Brownian walkers at distances d (ln units) below a threshold has
    crossed it by time t) -- reflection principle, no drift. d: array; sig2: array or scalar; t: array."""
    d = np.asarray(d, float)[:, None]
    s = np.sqrt(np.asarray(sig2, float).reshape(-1, 1) * np.asarray(t, float)[None, :])
    pi = 2.0 * stats.norm.sf(d / s)
    pi = np.where(d <= 0, 1.0, pi)
    return 1.0 - np.exp(np.sum(np.log1p(-np.clip(pi, 0, 1 - 1e-15)), axis=0))


def t_at(p, t, level=0.5):
    idx = np.where(p >= level)[0]
    return float(t[idx[0]]) if len(idx) else float("inf")


RATE_KEYS = {"standard": "sigma2", "robust": "sigma2_robust", "noise-corrected": "sigma2_noise_corrected",
             "sister-pair": None}


def waiting(ext, rates, summary):
    """Drift-free Brownian first-passage 'waiting times'.  These are model outputs for a biosphere that
    already contains large-brained mammals; the model has no ceiling, no extinction and no drift, and it
    fails a hindcast of the Cenozoic (below), so the numbers bound nothing empirically."""
    human = [r for r in ext if r["species"] == "Homo_sapiens"][0]
    others = [r for r in ext if r["species"] != "Homo_sapiens"]
    t = np.logspace(-2, 3.5, 600)
    R = rates["alvarez2022"]
    curves = {}
    res = {}

    def sig_for(r, trait, which):
        if which == "sister-pair":
            return R[f"sister_pairs_{trait}"]["sigma2"]
        key = f"{r['order']}_{trait}"
        k = key if key in R else f"Mammalia_{trait}"
        return R[k][RATE_KEYS[which]]

    def run(label, rows, d, trait, which):
        s2 = np.array([sig_for(r, trait, which) for r in rows])
        p = p_any(d, s2, t)
        curves[label] = p
        res[label] = dict(n_lineages=len(rows), T10_Myr=t_at(p, t, 0.1), T50_Myr=t_at(p, t, 0.5), T90_Myr=t_at(p, t, 0.9),
                          P_10Myr=float(np.interp(10, t, p)), P_100Myr=float(np.interp(100, t, p)),
                          P_1000Myr=float(np.interp(1000, t, p)))

    # brain mass a primate needs for 16.3e9 cortical neurons under the primate rule, and the masses a
    # non-primate would need under the other two rules
    m_primate = float(np.exp((np.log(HUMAN_NCX) - 17.125) / 0.933))
    m_cetart = float(np.exp((np.log(HUMAN_NCX) - 17.052) / 0.600))
    m_afro = float(np.exp((np.log(HUMAN_NCX) - 16.965) / 0.651))
    big = max(others, key=lambda r: r["brain_g"])
    summary["if_rule_required_brain_mass"] = dict(
        primate_rule_g=m_primate, cetartiodactyl_rule_g=m_cetart, afrotherian_rule_g=m_afro,
        largest_brain_in_dataset=dict(species=big["species"], brain_g=big["brain_g"]),
        primate_rule_prediction_for_human=dict(
            at_burger_brain_mass=float(ncx_rule(human["brain_g"], "Primates")),
            at_1460_g=float(ncx_rule(1460.0, "Primates")), measured=HUMAN_NCX),
        note="under the isotropic-fractionator rules a non-primate needs a brain several times heavier than any "
             "known brain, so that route is effectively closed and no waiting time is quoted for it")

    for which in RATE_KEYS:
        d = np.log(human["eq"] / np.array([r["eq"] for r in others]))
        run(f"EQ, all mammals ({which} rate)", others, d, "eq", which)
        wl = [r for r in others if r["well"]]
        d = np.log(human["eq"] / np.array([r["eq"] for r in wl]))
        run(f"EQ, all mammals, n>=3 records only ({which} rate)", wl, d, "eq", which)
        land = [r for r in others if r["order"] != "Cetacea"]
        d = np.log(human["eq"] / np.array([r["eq"] for r in land]))
        run(f"EQ, non-cetaceans ({which} rate)", land, d, "eq", which)
        landw = [r for r in land if r["well"]]
        d = np.log(human["eq"] / np.array([r["eq"] for r in landw]))
        run(f"EQ, non-cetaceans, n>=3 records only ({which} rate)", landw, d, "eq", which)
        # cortical neurons via the primate IF rule: a primate needs brain mass m_primate (rule target) or,
        # since the rule over-predicts the measured human count, the human brain mass itself
        prim = [r for r in others if r["order"] == "Primates"]
        bm_ = np.array([r["brain_g"] for r in prim])
        run(f"IF neurons, another primate, rule target {m_primate:.0f} g ({which} rate)", prim, np.log(m_primate / bm_), "brain_g", which)
        run(f"IF neurons, another primate, human brain mass ({which} rate)", prim, np.log(human["brain_g"] / bm_), "brain_g", which)
    summary["waiting_time_diffusion"] = res
    allw = [v["T50_Myr"] for k, v in res.items()]
    summary["waiting_time_range"] = dict(
        min_T50_Myr=float(min(allw)), max_T50_Myr=float(max(allw)),
        max_T90_Myr=float(max(v["T90_Myr"] for v in res.values())),
        note="range of median first-passage times over criteria, record sets and the four rate estimates")

    # hindcast: can drift-free diffusion reproduce the Cenozoic rise?  Start 61 Myr ago with N lineages at the
    # Paleocene median EQ and ask for the probability that any reaches (i) half the human EQ, (ii) the human EQ.
    ms = summary["median_shift"]
    hind = {}
    for which in RATE_KEYS:
        s2 = R["sister_pairs_eq"]["sigma2"] if which == "sister-pair" else R["Mammalia_eq"][RATE_KEYS[which]]
        for nlin in (100, 5000):
            for lab, target in (("half_human", human["eq"] / 2), ("human", human["eq"])):
                d = np.full(nlin, np.log(target / ms["median_Paleocene"]))
                hind[f"{which}_N{nlin}_{lab}"] = float(p_any(d, s2, np.array([61.0]))[0])
    summary["hindcast_driftless_diffusion"] = hind

    # gaps to the human value (no extrapolation of the envelope is made: the series switches clade at every
    # step and ends on an island dwarf elephant, so its slope is not a forecast)
    ch = summary["cetacean_history"]
    summary["gap_to_human_EQ"] = dict(
        ln_gap_top_single_record=float(np.log(human["eq"] / ch["max_extant_EQ"]["value"])),
        ln_gap_top_n3plus_record=float(np.log(human["eq"] / ch["max_extant_EQ_n3plus"]["value"])))

    # tail extrapolation: is the human EQ an outlier of the mammalian distribution?
    x = np.log(np.array([r["eq"] for r in others]))
    tail = {}
    for u_eq in (1.5, 2.0, 2.5, 3.0):
        u = np.log(u_eq)
        ex = x[x > u] - u
        c, _, sc = stats.genpareto.fit(ex, floc=0)
        tail[f"u_EQ_{u_eq}"] = dict(
            n_exceed=int(len(ex)), mean_excess=float(ex.mean()),
            expected_species_ge_human_exponential=float(len(ex) * np.exp(-(np.log(human["eq"]) - u) / ex.mean())),
            gpd_shape=float(c), gpd_upper_endpoint_EQ=float(np.exp(u - sc / c)) if c < 0 else None,
            expected_species_ge_human_gpd=float(len(ex) * stats.genpareto.sf(np.log(human["eq"]) - u, c, 0, sc)))
    summary["tail_extrapolation"] = tail
    return t, curves


# ------------------------------------------------------------------ (d) hardware vs technology
FILTERS = [
    # lineage, EQ ratio source species, terrestrial, manipulator, fire-compatible habitat, cumulative culture, Myr at/near level
    dict(lineage="Hominins (Homo)", terrestrial=1, grasping_manipulator=1, fire_possible=1, cumulative_culture=1,
         approx_Myr_at_level=2, note="the one success; large brains ~2 Myr, technology accelerating for <0.1 Myr"),
    dict(lineage="Other great apes", terrestrial=1, grasping_manipulator=1, fire_possible=1, cumulative_culture=0,
         approx_Myr_at_level=15, note="tool use and traditions, no ratcheting culture; brain mass ~1/3 of human"),
    dict(lineage="Elephants", terrestrial=1, grasping_manipulator=1, fire_possible=1, cumulative_culture=0,
         approx_Myr_at_level=20, note="trunk is a capable manipulator; modern brain mass since the early Miocene (Benoit et al. 2019)"),
    dict(lineage="Delphinoid cetaceans", terrestrial=0, grasping_manipulator=0, fire_possible=0, cumulative_culture=0,
         approx_Myr_at_level=15, note="highest non-human EQ; vocal learning and traditions; aquatic, no manipulators, no fire"),
    dict(lineage="Capuchins / atelids / macaques", terrestrial=1, grasping_manipulator=1, fire_possible=1, cumulative_culture=0,
         approx_Myr_at_level=10, note="EQ within x2 of humans but absolute brain mass < 1/10"),
    dict(lineage="Parrots, corvids, owls", terrestrial=1, grasping_manipulator=1, fire_possible=1, cumulative_culture=0,
         approx_Myr_at_level=20, note="pallial neurons up to ~1/9 of human in <25 g brains; tool manufacture in corvids; flight limits mass"),
]


def hardware_vs_technology(summary):
    wr("part2_hardware_vs_technology.csv", FILTERS)
    n_hw = len(FILTERS) - 1
    n_ok = sum(1 for f in FILTERS[1:] if f["terrestrial"] and f["grasping_manipulator"] and f["fire_possible"])
    expo = sum(f["approx_Myr_at_level"] for f in FILTERS[1:] if f["terrestrial"] and f["grasping_manipulator"])
    summary["hardware_vs_technology"] = dict(
        n_other_lineages_listed=n_hw, n_passing_physical_filters=n_ok,
        exposure_lineage_Myr_without_technology=expo,
        rate_upper_limit_95_per_lineage_Myr=3.0 / expo,
        note="hand-coded judgement table (not a measurement); exposure times are rounded literature values; the "
             "lineages are not hominin-level on neuron number, so even the upper limit is generous")
    # What the exposure does and does not constrain.  Let h be the hazard (per lineage-Myr) with which a
    # land-living, manipulating, large-brained lineage starts cumulative technology.  Zero events in `expo`
    # lineage-Myr gives h < 3/expo (95 %); there is NO lower limit, so P = 1 - exp(-h L T) over T = 1000 Myr
    # can be anything from ~0 to ~1.  The grid shows how the mean of P under a log-uniform prior on
    # [h_min, 3/expo] is set by the arbitrary floor h_min.
    hmax = 3.0 / expo
    T = 1000.0
    grid = {}
    for L in (1, 3, 10):
        for hmin in (1e-10, 1e-8, 1e-7, 1e-6, 1e-5):
            h = np.logspace(np.log10(hmin), np.log10(hmax), 4001)
            pr = 1.0 - np.exp(-h * L * T)
            grid[f"L{L}_hmin{hmin:g}"] = dict(mean_P=float(np.mean(pr)), P_at_hmin=float(pr[0]), P_at_hmax=float(pr[-1]),
                                               h_for_P50=float(np.log(2) / (L * T)))
    summary["technology_probability_sketch"] = dict(
        hazard_upper_limit_95=hmax, T_Myr=T, grid=grid,
        mean_P_range=[min(g["mean_P"] for g in grid.values()), max(g["mean_P"] for g in grid.values())],
        note="arithmetic on an ignorance prior, NOT an empirical estimate: the data give an upper limit on the hazard "
             "and no lower limit, so the probability is not bounded below; the mean is set by the assumed floor")


# ------------------------------------------------------------------ figures
def figures(ext, allrows, lin, ntab, t, curves, summary):
    plt = plotstyle.apply()
    P = plotstyle
    human = [r for r in ext if r["species"] == "Homo_sapiens"][0]
    os.makedirs(FIG, exist_ok=True)

    # ---- Fig 3: brain size through time
    fig, ax = plt.subplots(figsize=(9.2, 5.6))
    groups = [("other mammals", P.GREY, "o", "other mammals (Bertrand et al. 2022; extant)"),
              ("Primates", P.AQUA, "o", "non-hominin primates"),
              ("Cetacea", P.BLUE, "o", "cetaceans (Waugh & Thewissen 2021; extant)"),
              ("Proboscidea", P.ORANGE, "s", "proboscideans (Benoit et al. 2019; extant)"),
              ("Hominini", P.RED, "D", "hominins (Du et al. 2018, taxon medians; H. sapiens)")]
    rng = np.random.default_rng(1)
    for clade, col, mk, lab in groups:
        sel = [r for r in allrows if r["clade"] == clade and r["age_Ma"] <= 66]
        x = np.array([r["age_Ma"] for r in sel])
        y = np.array([r["eq"] for r in sel])
        if clade == "Hominini":  # one point per fossil taxon (median ECV, assumed body mass) plus H. sapiens
            hh = summary["hominin_history"]["by_taxon"]
            x = np.array([d["mean_age_Ma"] for d in hh] + [0.0])
            y = np.array([d["eq"] for d in hh] + [human["eq"]])
        x = np.where(x == 0, -rng.uniform(0.6, 3.4, len(x)), x)  # spread the Recent sample right of 0
        small = clade == "other mammals"
        ax.scatter(x, y, s=7 if small else 20, c=col, marker=mk, alpha=0.35 if small else 0.85,
                   edgecolors="none" if small else P.SURFACE, linewidths=0.5, label=lab, zorder=2 if small else 3)
    # envelope of non-hominins (running record) using the fitted epoch maxima
    es = summary["envelope_slope"]["eq_max"]
    xx = np.array(es["epoch_mid_Ma"])
    ax.plot(xx, np.exp(es["ln_stat"]), color=P.INK, lw=1.2, marker="_", ms=10, zorder=4)
    xf = np.linspace(66, 0, 50)
    ax.plot(xf, np.exp(es["intercept_lnEQ_today"] - es["slope_ln_per_Myr"] * xf), color=P.INK, lw=1.0, ls="--", zorder=4)
    ax.annotate(f"non-hominin epoch maxima (solid)\nand log-linear fit (dashed):\ndoubling every {es['doubling_time_Myr']:.0f} Myr on average",
                xy=(52, np.exp(es["intercept_lnEQ_today"] - es["slope_ln_per_Myr"] * 52)), xytext=(64.5, 1.35),
                fontsize=8.5, color=P.INK2, arrowprops=dict(arrowstyle="-", color=P.INK2, lw=0.7))
    ax.axhline(human["eq"], color=P.RED, lw=0.9, ls=":")
    ax.text(65.5, human["eq"] * 1.06, f"Homo sapiens EQ = {human['eq']:.1f}", color=P.INK2, fontsize=8.5, va="bottom")
    ax.axhline(human["eq"] / 2, color=P.GREY, lw=0.8, ls=":")
    ax.text(65.5, human["eq"] / 2 * 1.06, "half the human value", color=P.INK2, fontsize=8.5, va="bottom")
    ax.annotate("hominin taxa\n(medians)", xy=(1.9, 4.0), xytext=(9.5, 2.35), fontsize=8, color=P.INK2,
                arrowprops=dict(arrowstyle="-", color=P.INK2, lw=0.7))
    ch = summary["cetacean_history"]
    ax.annotate("Kentriodon\n(16 Ma)", xy=(16, ch["max_Miocene_EQ"]["value"]), xytext=(25, 4.4), fontsize=8, color=P.INK2,
                arrowprops=dict(arrowstyle="-", color=P.INK2, lw=0.7))
    pf = [f for f in summary["proboscidean_history"]["fossils"] if f["species"] == "Palaeoloxodon_falconeri"][0]
    ax.annotate("dwarf elephant\nP. falconeri", xy=(pf["age_Ma"], pf["eq"]), xytext=(11.5, 8.2), fontsize=8, color=P.INK2,
                arrowprops=dict(arrowstyle="-", color=P.INK2, lw=0.7))
    ax.set_yscale("log")
    ax.set_xlim(67, -4.5)
    ax.set_ylim(0.04, 11)
    ax.set_yticks([0.05, 0.1, 0.2, 0.5, 1, 2, 5, 10])
    ax.set_yticklabels(["0.05", "0.1", "0.2", "0.5", "1", "2", "5", "10"])
    ax.axvspan(0, -4.5, color=P.GRID, alpha=0.45, lw=0)
    ax.text(-2.2, 0.046, "living\nspecies", ha="center", va="bottom", fontsize=8, color=P.INK2)
    ax.set_xlabel("Age (million years ago)")
    ax.set_ylabel("Encephalization quotient (Jerison), log scale")
    chh = summary["cetacean_history"]
    for key, lab, dy in (("max_extant_EQ", "top single record", 1.0), ("max_extant_EQ_n3plus", "top species with 3+ brains", 1.0)):
        ax.plot([-4.3, -0.2], [chh[key]["value"]] * 2, color=P.BLUE, lw=1.3, zorder=5)
    ax.annotate(f"living dolphins: top single record {chh['max_extant_EQ']['value']:.1f},\ntop species with 3+ brains {chh['max_extant_EQ_n3plus']['value']:.1f}",
                xy=(-0.2, chh["max_extant_EQ_n3plus"]["value"]), xytext=(33, 7.9), fontsize=8, color=P.INK2,
                arrowprops=dict(arrowstyle="-", color=P.INK2, lw=0.7))
    ax.set_title("Relative brain size rose in several clades through the Cenozoic\n(the record-holder changes clade at most steps; fossil bins are small and clade-targeted)")
    ax.legend(loc="lower right", bbox_to_anchor=(0.93, 0.02), fontsize=8, handletextpad=0.3)
    fig.tight_layout()
    fig.savefig(os.path.join(FIG, "fig3_brain_size_through_time.png"))
    plt.close(fig)

    # ---- Fig 4: lineage counts and neuron numbers
    fig, (a1, a2) = plt.subplots(1, 2, figsize=(11.5, 5.2), gridspec_kw=dict(width_ratios=[1, 1.25]))
    rows = [r for r in lin if r["metric"] == "eq" and r["within_factor"] in (1.5, 2.0, 3.0)]
    xs = np.arange(len(rows))
    w = 0.26
    series = [("deep20Ma_clades_robust", P.BLUE, "lineages separate for 20+ Myr (noise-robust)"),
              ("n_families_n3plus", P.AQUA, "families, species with 3+ brains measured"),
              ("n_other_families", P.ORANGE, "families, all records")]
    for k, (key, col, lab) in enumerate(series):
        vals = [r[key] for r in rows]
        a1.bar(xs + (k - 1) * w, vals, width=w * 0.92, color=col, label=lab)
        for x, v_ in zip(xs, vals):
            a1.text(x + (k - 1) * w, v_ + 0.25, str(v_), ha="center", fontsize=9, color=P.INK)
    a1.set_xticks(xs)
    a1.set_xticklabels([f"÷{r['within_factor']:g}  (EQ ≥ {r['threshold']:.1f})\n{r['n_other_species']} species\n"
                        f"({r['n_species_n3plus']} with 3+ brains)" for r in rows], fontsize=8.5)
    a1.set_xlabel("Threshold: human EQ divided by …")
    a1.set_ylabel("Number of non-human lineages at or above threshold")
    a1.set_title("Other mammal lineages near the\nhuman relative brain size")
    a1.legend(loc="upper left")
    a1.grid(axis="x", visible=False)

    # neuron dot plot
    order = ["Homo_sapiens", "Gorilla_gorilla", "Pongo_pygmaeus", "Pan_troglodytes", "Loxodonta_africana",
             "Balaenoptera_physalus", "Orcinus_orca", "Globicephala_melas", "Balaenoptera_acutorostrata",
             "Phocoena_phocoena", "Papio_anubis", "Macaca_mulatta", "Giraffa_camelopardalis", "Ara_ararauna",
             "Strix_nebulosa", "Corvus_corax"]
    names = {"Homo_sapiens": "human", "Gorilla_gorilla": "gorilla", "Pongo_pygmaeus": "orangutan",
             "Pan_troglodytes": "chimpanzee", "Loxodonta_africana": "African elephant",
             "Balaenoptera_physalus": "fin whale", "Orcinus_orca": "killer whale",
             "Globicephala_melas": "long-finned pilot whale", "Balaenoptera_acutorostrata": "minke whale",
             "Phocoena_phocoena": "harbour porpoise", "Papio_anubis": "baboon", "Macaca_mulatta": "rhesus macaque",
             "Giraffa_camelopardalis": "giraffe", "Ara_ararauna": "blue-and-yellow macaw",
             "Strix_nebulosa": "great grey owl", "Corvus_corax": "raven"}
    birds = {r["species"]: float(r["pallium_neurons"]) / 1e9 for r in rd("neurons_birds_reptiles.csv") if r["pallium_neurons"]}
    styles = {"IF_measured": (P.BLUE, "o", "measured, isotropic fractionator"),
              "IF_rule_predicted": (P.BLUE, "o", "predicted from brain mass (IF scaling rule)"),
              "stereology": (P.ORANGE, "D", "measured, stereology")}
    done = set()
    for i, sp in enumerate(order):
        y = len(order) - 1 - i
        recs = [r for r in ntab if r["species"] == sp and r["lineage"] != "Aves"]
        if sp in birds and not recs:
            recs = [dict(kind="IF_measured", neurons_1e9=birds[sp])]
        vals = [r["neurons_1e9"] for r in recs]
        if len(vals) > 1:
            a2.plot([min(vals), max(vals)], [y, y], color=P.GRID, lw=2, zorder=1)
        for r in recs:
            kind = r["kind"].split(" ")[0]
            col, mk, lab = styles[kind]
            hollow = kind == "IF_rule_predicted"
            a2.scatter(r["neurons_1e9"], y, s=52, marker=mk, facecolors=P.SURFACE if hollow else col, edgecolors=col,
                       linewidths=1.6, zorder=3, label=lab if kind not in done else None)
            done.add(kind)
    a2.set_yticks(range(len(order)))
    a2.set_yticklabels([names[s] for s in reversed(order)])
    a2.set_xscale("log")
    a2.set_xlim(0.8, 60)
    a2.set_xticks([1, 2, 5, 10, 20, 50])
    a2.set_xticklabels(["1", "2", "5", "10", "20", "50"])
    for v, lab in ((16.3, "human"), (16.3 / 3, "÷3"), (16.3 / 10, "÷10")):
        a2.axvline(v, color=P.INK2, lw=0.8, ls=":")
        a2.text(v, len(order) - 0.35, lab, ha="center", fontsize=8.5, color=P.INK2)
    a2.set_ylim(-0.7, len(order) + 0.3)
    a2.set_xlabel("Cortical (birds: pallial) neurons, billions, log scale")
    a2.set_title("Cortical neuron numbers depend\non the counting method")
    a2.legend(loc="lower right", fontsize=8)
    a2.grid(axis="y", visible=False)
    fig.tight_layout()
    fig.savefig(os.path.join(FIG, "fig4_lineages_near_human.png"))
    plt.close(fig)

    # ---- Fig 5: waiting time
    fig, ax = plt.subplots(figsize=(9.2, 6.0))
    mp = summary["if_rule_required_brain_mass"]["primate_rule_g"]
    sel = [("EQ, all mammals ({} rate)", P.BLUE, "relative brain size (EQ), any mammal"),
           ("EQ, non-cetaceans ({} rate)", P.AQUA, "EQ, land mammals only"),
           ("IF neurons, another primate, rule target %.0f g ({} rate)" % mp, P.ORANGE,
            "cortical neurons (primate IF rule), another primate")]
    for key, col, lab in sel:
        ax.plot(t, curves[key.format("noise-corrected")], color=col, ls="-", label=lab)
        ax.plot(t, curves[key.format("standard")], color=col, ls=":", lw=1.2)
    ax.plot([], [], color=P.INK2, ls="-", label="solid: rate corrected for record noise")
    ax.plot([], [], color=P.INK2, ls=":", lw=1.2, label="dotted: uncorrected rate (inflated by noise)")
    for v, lab, yy in ((66, "Cenozoic\nso far", 0.03), (540, "since first\nanimals", 0.16), (1000, "habitable\ntime left", 0.03)):
        ax.axvline(v, color=P.INK2, lw=0.8, ls="--")
        ax.text(v * (0.93 if v < 1000 else 1.07), yy, lab, ha="right" if v < 1000 else "left", va="bottom", fontsize=8, color=P.INK2)
    hc = summary["hindcast_driftless_diffusion"]
    ax.text(0.06, 0.98, "Caveats: no ceiling, no extinction,\nno drift; starts from today's large-\n"
            "brained mammals; EQ can rise by\nbody shrinkage alone. The same\nmodel fails a hindcast of the\n"
            f"Cenozoic (P = {hc['standard_N5000_half_human']:.3f} uncorrected,\n"
            f"{hc['noise-corrected_N5000_half_human']:.0e} corrected).",
            fontsize=7.6, color=P.INK2, va="top", ha="left")
    ax.set_xscale("log")
    ax.set_xlim(0.05, 3000)
    ax.set_ylim(0, 1.02)
    ax.set_xticks([0.1, 1, 10, 100, 1000])
    ax.set_xticklabels(["0.1", "1", "10", "100", "1000"])
    ax.set_xlabel("Time from now (million years), log scale")
    ax.set_ylabel("P(some other lineage has reached the human value)")
    ax.set_title("Drift-free diffusion model: first-passage time for a second lineage to reach\nthe human value, starting from living mammals (a model output, not a forecast)")
    ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.13), ncol=2, fontsize=8.5)
    fig.tight_layout()
    fig.savefig(os.path.join(FIG, "fig5_waiting_time.png"))
    plt.close(fig)


# ------------------------------------------------------------------ main
def main():
    rng = np.random.default_rng(SEED)
    summary = {}
    ext, ols = load_extant()
    summary["n_extant_species"] = len(ext)
    summary["flagged_excluded"] = [r["species"] for r in rd("mammals_extant_burger2019.csv") if r["flag"]]
    summary["ols_brain_body"] = dict(slope=ols[0], intercept_log10=ols[1])
    human = [r for r in ext if r["species"] == "Homo_sapiens"][0]
    summary["human"] = dict(eq=human["eq"], eq75=human["eq75"], brain_g=human["brain_g"], body_g=human["body_g"])
    nonhum = sorted([r for r in ext if r["species"] != "Homo_sapiens"], key=lambda r: -r["eq"])
    top = nonhum[:12]
    summary["top_nonhuman_EQ"] = [dict(species=r["species"], family=r["family"], eq=r["eq"], ratio_to_human=r["eq"] / human["eq"],
                                       n_brain=r["n_brain"], well_sampled=r["well"]) for r in top]
    topw = [r for r in nonhum if r["well"]][:8]
    summary["top_nonhuman_EQ_n3plus"] = [dict(species=r["species"], family=r["family"], eq=r["eq"],
                                              ratio_to_human=r["eq"] / human["eq"], n_brain=r["n_brain"]) for r in topw]
    summary["n_species_n3plus"] = int(sum(r["well"] for r in ext))
    summary["n_top25_nonhuman_without_sample_size"] = int(sum(1 for r in nonhum[:25] if str(r["n_brain"]).strip() == ""))
    # sensitivity to the human record: Burger et al. give 1,250 g at 65.1 kg (n = 7); 1,350 g is the more usual mean
    e1350 = eqj(1350.0, human["body_g"])
    summary["human_brain_1350g_sensitivity"] = dict(
        human_eq=float(e1350), top_single_record_ratio=float(nonhum[0]["eq"] / e1350),
        top_n3plus_ratio=float(topw[0]["eq"] / e1350),
        n_within={f"x{f}": int(sum(1 for r in nonhum if r["eq"] >= e1350 / f)) for f in FACTORS_EQ},
        n_within_n3plus={f"x{f}": int(sum(1 for r in nonhum if r["well"] and r["eq"] >= e1350 / f)) for f in FACTORS_EQ})

    trees = {}
    for tname in ("alvarez2022", "upham2019"):
        root = phylo.parse_newick(open(os.path.join(DATA, f"mammal_tree_pruned_{tname}.nwk")).read())
        trees[tname] = (root, {x.name for x in phylo.tips(root)})
    summary["n_in_tree"] = {k: len({r["species"] for r in ext} & v[1]) for k, v in trees.items()}

    lin = lineage_counts(ext, trees, summary)
    ntab = neuron_table(ext, summary)
    fos = load_fossils()
    summary["n_fossil_rows"] = dict(collections.Counter(r["source"] for r in fos))
    allrows = envelope(ext, fos, summary, rng)
    rates = bm_rates(ext, trees, summary)
    t, curves = waiting(ext, rates, summary)
    hardware_vs_technology(summary)

    # independent cross-check of the top delphinid EQ values (specimen-level data, different compilation)
    mm = []
    for r in rd("cetaceans_muller2019.csv"):
        e_ = eqj(float(r["brain_g"]), float(r["body_kg"]) * 1000)
        mm.append(dict(species=r["species"], n_specimens=int(r["n_specimens"]), eq=float(e_),
                       ratio_to_human=float(e_ / human["eq"])))
    mm.sort(key=lambda d: -d["eq"])
    summary["crosscheck_muller2019"] = mm

    # endocast-conversion sensitivity for hominin taxa (0.92 g per cm3 instead of the Ruff et al. formula)
    body = {r["taxon"]: float(r["body_kg_assumed"]) * 1000 for r in rd("hominin_body_mass_assumed.csv")}
    ecvs = collections.defaultdict(list)
    for r in rd("hominin_ecv_du2018.csv"):
        ecvs[r["taxon"]].append(float(r["ecv_cm3"]))
    summary["hominin_eq_by_conversion"] = {
        k: dict(median_ecv_cm3=float(np.median(v)),
                eq_ruff1997=float(eqj(ecv_to_g(np.median(v), "hominin"), body[k])),
                eq_0p92=float(eqj(0.92 * np.median(v), body[k])), eq_1to1=float(eqj(np.median(v), body[k])))
        for k, v in ecvs.items()}

    # sensitivity: keep the flagged records
    ext2, _ = load_extant(include_flagged=True)
    h2 = [r for r in ext2 if r["species"] == "Homo_sapiens"][0]
    sens = {}
    for f in FACTORS_EQ:
        sens[f"x{f}"] = len([r for r in ext2 if r["eq"] >= h2["eq"] / f and r["species"] != "Homo_sapiens"])
    summary["sensitivity_keep_flagged_n_other_species"] = sens

    json.dump(summary, open(os.path.join(DATA, "part2_summary.json"), "w"), indent=1, default=float)
    figures(ext, allrows, lin, ntab, t, curves, summary)
    return summary


if __name__ == "__main__":
    s = main()
    for k in ("human", "top_nonhuman_EQ", "top_nonhuman_EQ_n3plus", "human_brain_1350g_sensitivity",
              "joint_EQ_and_brain_mass", "envelope_slope", "envelope_slope_sensitivity", "envelope_bin_composition",
              "first_fossil_within_factor_of_human_EQ", "rarefied_max_n20", "cetacean_history",
              "eq_vs_body_within_delphinoids", "hominin_history", "hominin_eq_by_conversion", "primate_history",
              "median_shift", "bm_rates", "if_rule_required_brain_mass", "waiting_time_diffusion",
              "waiting_time_range", "hindcast_driftless_diffusion", "gap_to_human_EQ", "tail_extrapolation",
              "neuron_rule_check", "hardware_vs_technology", "technology_probability_sketch"):
        print("==", k)
        print(json.dumps(s[k], indent=1, default=float)[:6000])