"""Recompute the external replication from the independently inferred spectra.

Ramos-Almodovar et al. (2026) deposit fitted posterior summaries per species.
Their trinucleotide layer collapses onto the same 96 categories used here, so
the pipeline runs on them unchanged: same reference exposure, same frozen
ensemble, same estimand.

Writes data/external_replication_per_species.csv and, from it,
data/external_replication.json and data/replication_species.csv. The aggregate
figures reported in the manuscript are therefore recomputable from a deposited
per-species table rather than asserted.

    python3 code/replication.py /path/to/posterior_summaries_directory

Without the directory the script rebuilds the aggregates from the deposited
per-species CSV, which is what the audit does.
"""
import glob
import gzip
import json
import os
import sys

import numpy as np
import pandas as pd

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
DATA = os.path.join(ROOT, "data")
PER = os.path.join(DATA, "external_replication_per_species.csv")

ALT = {"A": ["C", "G", "T"], "C": ["A", "G", "T"]}
COMP = {"A": "U", "C": "G", "G": "C", "U": "A", "T": "A"}


def _rc(s):
    return "".join(COMP[b] for b in reversed(s))


def read_spectrum(path):
    """The trinucleotide layer, collapsed to the 96 pyrimidine-focal
    categories used throughout."""
    with gzip.open(path, "rt") as fh:
        d = json.load(fh)
    root = list(d)[0]
    out = {}
    for ctx, val in d[root]["2"]["p_vec"].items():
        cu = ctx.replace("T", "U")
        for alt, pv in zip(ALT[ctx[1]], val["mean"]):
            au = alt.replace("T", "U")
            key = (cu, au) if cu[1] == "C" else (_rc(cu), COMP[au])
            out[key] = out.get(key, 0.0) + pv
    return out


def per_species(folder):
    """R, z, CV, rank and kappa for each species, against the frozen
    ensemble under the same reference exposure."""
    sys.path.insert(0, HERE)
    from codecore import (AA20, POLAR_REQUIREMENT, STANDARD_CODE,
                          scalar_metric_matrix)

    keys = [tuple(q) for q in json.load(open(os.path.join(DATA,
                                                         "edge_keys.json")))]
    cats = sorted(set(keys))
    col = np.array([cats.index(k) for k in keys])
    fr, to = np.load(os.path.join(DATA, "edge_meta.npy"))
    pq = np.load(os.path.join(DATA, "edge_pq.npy"))
    codons = sorted(STANDARD_CODE)
    ci = {c: i for i, c in enumerate(codons)}
    ai = {a: i for i, a in enumerate(AA20 + ["*"])}
    dm = scalar_metric_matrix(POLAR_REQUIREMENT, True)
    D = np.zeros((21, 21))
    for a in AA20:
        for b in AA20:
            D[ai[a], ai[b]] = dm[(a, b)]
    ens = json.load(open(os.path.join(DATA, "confirmatory_ensemble.json")))
    perm = [ci[c] for c in ens["codon_order"]]
    codes = np.zeros((ens["n"], 64), dtype=np.int8)
    for i, sc in enumerate(ens["codes"]):
        for k, aa in enumerate(sc):
            codes[i, perm[k]] = ai[aa]
    sgc = np.array([ai[STANDARD_CODE[c]] for c in codons])

    files = sorted(glob.glob(os.path.join(
        folder, "*ALL*final_posterior_summaries*json.gz")))
    rows = []
    for f in files:
        name = os.path.basename(f).split("_non_coding")[0]
        spec = read_spectrum(f)
        p = np.array([spec.get(k, 0.0) for k in cats])
        p = p / p.sum()
        w = p[col] * pq
        w = w / w.sum()
        phis = float((w * D[sgc[fr], sgc[to]]).sum())
        phie = (D[codes[:, fr], codes[:, to]] * w[None, :]).sum(axis=1)
        mu, sd = float(phie.mean()), float(phie.std())
        f_ = lambda t: p[[i for i, c in enumerate(cats)
                          if c[0][1] == t[0] and c[1] == t[1]]].sum()
        kap = float((f_(("C", "U")) + f_(("U", "C")))
                    / (f_(("C", "A")) + f_(("C", "G"))
                       + f_(("U", "A")) + f_(("U", "G"))))
        rows.append(dict(species=name, kappa=kap, R=1 - phis / mu,
                         z=(mu - phis) / sd, CV=sd / mu,
                         rank=int((phie < phis).sum()),
                         p_emp=(1 + int((phie <= phis).sum())) / (ens["n"] + 1)))
    return pd.DataFrame(rows).sort_values("species").reset_index(drop=True)


def aggregate(d):
    """The figures the manuscript reports, from the per-species table."""
    lo = d["kappa"] <= d["kappa"].quantile(1 / 3)
    hi = d["kappa"] >= d["kappa"].quantile(2 / 3)
    out = {"n_species": int(len(d)),
           "beat_min": int(d["rank"].min()),
           "beat_max": int(d["rank"].max()),
           "beat_median": float(d["rank"].median()),
           "pct_none_better": round(100 * float((d["rank"] == 0).mean()), 1)}
    lk = np.log(d["kappa"])
    for lab, col in (("R", "R"), ("z", "z"), ("CV", "CV"), ("beat", "rank")):
        a, b = float(d[col][lo].mean()), float(d[col][hi].mean())
        out[f"{lab}_lo_tercile"] = round(a, 4)
        out[f"{lab}_hi_tercile"] = round(b, 4)
        out[f"{lab}_pct_change"] = round(100 * (b / a - 1), 1)
        out[f"corr_kappa_{lab}"] = round(
            float(np.corrcoef(lk, d[col])[0, 1]), 3)
    out["corr_R_z"] = round(float(np.corrcoef(d["R"], d["z"])[0, 1]), 3)
    # Rank is a right-skewed count and one extreme species is enough to flip a
    # Pearson correlation over this many points, so it is summarized ordinally.
    from scipy import stats as _st
    out["rank_spearman_kappa"] = round(
        float(_st.spearmanr(lk, d["rank"])[0]), 3)
    out["rank_median_lo"] = float(d["rank"][lo].median())
    out["rank_median_hi"] = float(d["rank"][hi].median())
    out["pct_none_lo"] = round(100 * float((d["rank"][lo] == 0).mean()), 1)
    out["pct_none_hi"] = round(100 * float((d["rank"][hi] == 0).mean()), 1)
    return out


def main():
    if len(sys.argv) > 1 and os.path.isdir(sys.argv[1]):
        d = per_species(sys.argv[1])
        d.to_csv(PER, index=False)
        print(f"  {len(d)} species written to "
              f"{os.path.relpath(PER, ROOT)}")
    elif os.path.isfile(PER):
        d = pd.read_csv(PER)
        print(f"  {len(d)} species read from the deposited table")
    else:
        print("  no per-species table and no source directory given")
        return 1
    json.dump(aggregate(d),
              open(os.path.join(DATA, "external_replication.json"), "w"),
              indent=1)
    d[["species"]].to_csv(os.path.join(DATA, "replication_species.csv"),
                          index=False)
    print("  data/external_replication.json and "
          "data/replication_species.csv rewritten")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
