"""
Analyses reported in the manuscript whose code was not previously deposited.

Four blocks, each printing the values that appear in the text so they can be
checked against it directly:

  1. Exact enumeration of all C(64,3) = 41,664 stop-codon placements
     (Results section 6.1; Supplementary Table S6)
  2. Sensitivity to the CUG-clade reassignment
     (Results section 6.2; Supplementary Tables S7 and S10)
  3. Which stop codon a nonsense mutation creates, by quintile of kappa
     (Results section 6.1; Supplementary Table S13)
  4. Curvature in the out-of-clade intercept bias
     (Results section 5.5; Supplementary Table S12)

All four blocks need the third-party CORAL supplementary table. Block 1 will
run without it, but on the reference exposure alone rather than the
spectrum-averaged weights the manuscript reports, and the numbers it prints will
not match Supplementary Table S6. It says so when that happens. Point
CORAL_XLSX at the table to reproduce the published values.

    python3 code/deposited_analyses.py
    CORAL_XLSX=/path/to/media-1.xlsx python3 code/deposited_analyses.py
"""

import itertools
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")
CORAL = os.environ.get("CORAL_XLSX", os.path.join(DATA, "media-1.xlsx"))

sys.path.insert(0, HERE)
from codecore import SENSE, STANDARD_CODE, BASES              # noqa: E402
from context3 import EXPOSURE3, canonical3, context_exposure3  # noqa: E402

PUR = {"A", "G"}
CUG_GENERA = ["Candida", "Debaryomyces", "Meyerozyma", "Scheffersomyces",
              "Lodderomyces", "Clavispora", "Spathaspora", "Babjeviella",
              "Cephaloascus", "Metschnikowia", "Suhomyces", "Priceomyces",
              "Wickerhamia", "Hyphopichia", "Yamadazyma", "Teunomyces",
              "Danielozyma"]
YEAST_CLASSES = ["Saccharomycetes", "Pichiomycetes", "Schizosaccharomycetes",
                 "Taphrinomycetes"]


def spectra():
    """Per-species substitution rates, aggregated as described in Methods §5."""
    import coral
    counts = pd.read_pickle(os.path.join(DATA, "coral_counts.pkl"))
    abund = pd.read_pickle(os.path.join(DATA, "coral_abund.pkl"))
    return coral.spectra_from_sums(counts, abund)


def confirmatory():
    conf = pd.read_csv(os.path.join(DATA, "confirmatory_aligned.csv"))
    comp = pd.read_csv(os.path.join(DATA, "spectrum_components.csv"))
    return conf.merge(comp, on="species")


def with_taxonomy(df):
    raw = pd.read_excel(CORAL, sheet_name="CORALMutationsDB")
    for col in ("Clade", "phylum", "class"):
        df[col] = df["species"].map(raw.groupby("taxa1")[col].first())
    return df


# --------------------------------------------------------------- block 1
def stop_placements(spec=None, species=None):
    """Exact enumeration of every three-codon stop set.

    Edge weights are exposure times substitution rate, averaged over the
    confirmatory species when a spectrum is supplied. Without one the reference
    exposure alone is used, which gives a different ranking — the manuscript
    reports the spectrum-averaged version.
    """
    codons = sorted(STANDARD_CODE)
    idx = {c: i for i, c in enumerate(codons)}
    pair = np.zeros((64, 64))
    for c in codons:
        for j in (0, 1, 2):
            exp = EXPOSURE3.get((c, j)) or context_exposure3(c, j)
            for ctx, pq in exp.items():
                for b in BASES:
                    if b == c[j]:
                        continue
                    if spec is None:
                        rate = 1.0
                    else:
                        k = canonical3(ctx, b)
                        rate = float(np.mean([spec[s].get(k, 0.0)
                                              for s in species]))
                    pair[idx[c], idx[c[:j] + b + c[j + 1:]]] += pq * rate
    out_total = pair.sum(axis=1)

    def burden(stops):
        sense = [i for i in range(64) if i not in stops]
        return pair[np.ix_(sense, list(stops))].sum() / out_total[sense].sum()

    def spread(stops):
        return sum(sum(x != y for x, y in zip(codons[a], codons[b]))
                   for a, b in itertools.combinations(stops, 2))

    std = tuple(sorted(idx[c] for c in ("UAA", "UAG", "UGA")))
    b_std, s_std = burden(std), spread(std)
    vals, spreads = [], []
    for combo in itertools.combinations(range(64), 3):
        vals.append(burden(combo))
        spreads.append(spread(combo))
    vals, spreads = np.array(vals), np.array(spreads)

    print("\n1. STOP-CODON PLACEMENT, EXACT ENUMERATION")
    if spec is None:
        print("   *** WARNING: no spectrum supplied. Edge weights are the")
        print("   *** reference exposure alone, not the spectrum-averaged")
        print("   *** weights used in the manuscript. The values below will")
        print("   *** NOT match Supplementary Table S6. Set CORAL_XLSX.")
    print(f"   placements evaluated            {len(vals):,}")
    print(f"   standard set burden             {b_std:.5f}")
    print(f"   distribution  min/median/max    {vals.min():.5f} / "
          f"{np.median(vals):.5f} / {vals.max():.5f}")
    print(f"   percentile of the standard set  {100 * (vals < b_std).mean():.2f}")
    print(f"   placements with lower burden    {int((vals < b_std).sum()):,}")
    same = spreads == s_std
    print(f"   placements of equal spread ({s_std})  {int(same.sum()):,}")
    print(f"   percentile within that group    "
          f"{100 * (vals[same] < b_std).mean():.1f}")
    print(f"   corr(spread, burden)            "
          f"{np.corrcoef(spreads, vals)[0, 1]:+.3f}")


# --------------------------------------------------------------- block 2
def cug_sensitivity(spec, al):
    """Cost of the standard code against the CUG-clade code, per species."""
    from codecore import POLAR_REQUIREMENT, scalar_metric_matrix
    d = scalar_metric_matrix(POLAR_REQUIREMENT, True)
    cug = dict(STANDARD_CODE)
    cug["CUG"] = "S"
    keys, pq, cs, cc = [], [], [], []
    for c in SENSE:
        for j in (0, 1, 2):
            for ctx, w in EXPOSURE3[(c, j)].items():
                for b in BASES:
                    if b == c[j]:
                        continue
                    c2 = c[:j] + b + c[j + 1:]
                    if STANDARD_CODE[c2] == "*":
                        continue
                    keys.append(canonical3(ctx, b))
                    pq.append(w)
                    cs.append(d[(STANDARD_CODE[c], STANDARD_CODE[c2])])
                    cc.append(d[(cug[c], cug[c2])])
    pq, cs, cc = np.array(pq), np.array(cs), np.array(cc)
    delta = []
    for sp in al["species"]:
        s = spec[sp]
        w = np.array([s.get(k, 0.0) for k in keys]) * pq
        w /= w.sum()
        delta.append(float(w @ cc - w @ cs))
    al = al.assign(delta=delta)
    al["cug"] = al["species"].str.split("_").str[0].isin(CUG_GENERA)
    yz = al[al["class"].isin(YEAST_CLASSES)]

    print("\n2. THE CUG REASSIGNMENT UNDER YEAST SPECTRA")
    print(f"   {'group':<26}{'n':>6}{'mean d':>10}{'median':>10}{'% d<0':>9}")
    for lab, g in [("CUG clade", yz[yz["cug"]]),
                   ("other yeasts", yz[~yz["cug"]]),
                   ("non-yeast fungi",
                    al[(al["Clade"] == "Fungi") &
                       ~al["class"].isin(YEAST_CLASSES)])]:
        print(f"   {lab:<26}{len(g):>6}{g['delta'].mean():>10.4f}"
              f"{g['delta'].median():>10.4f}{100 * (g['delta'] < 0).mean():>8.1f}%")
    try:
        from scipy import stats
        p = stats.mannwhitneyu(yz[yz["cug"]]["delta"],
                               yz[~yz["cug"]]["delta"]).pvalue
        print(f"   Mann-Whitney, CUG clade vs other yeasts: p = {p:.3f}")
    except ImportError:
        pass
    print("\n   sensitivity with the CUG clade removed:")
    m = ~al["species"].str.split("_").str[0].isin(CUG_GENERA)
    for lab, g in [("full", al), ("excluding CUG clade", al[m])]:
        b = np.polyfit(g["kappa_log"], g["R"], 1)[0]
        print(f"     {lab:<22}n={len(g):<6}mean R {g['R'].mean():.4f}  "
              f"sd {g['R'].std(ddof=1):.4f}  beta_k {b:+.5f}")


# --------------------------------------------------------------- block 3
def stop_identity(spec, al):
    """Share of nonsense burden creating each stop codon, by quintile of kappa."""
    by = {"UAA": [[], []], "UAG": [[], []], "UGA": [[], []]}
    ts = tv = 0.0
    per_stop = {k: [0.0, 0.0] for k in by}
    for c in SENSE:
        for j in (0, 1, 2):
            for ctx, pq in EXPOSURE3[(c, j)].items():
                for b in BASES:
                    if b == c[j]:
                        continue
                    c2 = c[:j] + b + c[j + 1:]
                    if STANDARD_CODE[c2] != "*":
                        continue
                    by[c2][0].append(canonical3(ctx, b))
                    by[c2][1].append(pq)
                    is_ts = (c[j] in PUR) == (b in PUR)
                    per_stop[c2][0 if is_ts else 1] += pq
                    if is_ts:
                        ts += pq
                    else:
                        tv += pq

    print("\n3. WHICH STOP CODON A NONSENSE MUTATION CREATES")
    print(f"   sense->stop exposure: transitions {ts:.1f}, "
          f"transversions {tv:.1f}, tv/ts {tv / ts:.2f}")
    for st in ("UAA", "UAG", "UGA"):
        a, b = per_stop[st]
        print(f"     {st}: routes {len(by[st][0])}, ts {a:.1f}, tv {b:.1f}, "
              f"tv/ts {b / a:.2f}")

    frac = {}
    for st, (ks, ws) in by.items():
        ws = np.array(ws)
        frac[st] = np.array([float((np.array([spec[s].get(k, 0.0)
                                              for k in ks]) * ws).sum())
                             for s in al["species"]])
    tot = sum(frac.values())
    for st in frac:
        al[st] = frac[st] / tot
    al["tv_share"] = np.nan  # filled below for the transversion decomposition

    q = pd.qcut(al["kappa_log"], 5, labels=False)
    print(f"\n   {'quintile':<10}{'mean kappa':>12}{'UAA':>9}{'UAG':>9}{'UGA':>9}")
    for i, g in al.groupby(q, observed=True):
        print(f"   {i + 1:<10}{np.exp(g['kappa_log']).mean():>12.2f}"
              f"{g['UAA'].mean():>9.3f}{g['UAG'].mean():>9.3f}"
              f"{g['UGA'].mean():>9.3f}")
    for st in ("UAA", "UAG", "UGA"):
        print(f"   corr(log kappa, {st} share) = "
              f"{np.corrcoef(al['kappa_log'], al[st])[0, 1]:+.3f}")


# --------------------------------------------------------------- block 4
def curvature(al):
    """Leave-one-clade-out with higher-order terms in log kappa."""
    BIG = ["Fish", "Plants", "Birds", "Mammals", "Fungi", "Insects", "Reptiles"]
    al = al.assign(fold=np.where(al["Clade"].isin(BIG), al["Clade"], "Other"))
    R = al["R"].values
    w = 1.0 / al["sigma2"].values
    k = al["kappa_log"].values
    c = al["cpg_shift_log"].values

    def loco(cols):
        Z = np.column_stack(cols)
        Z = (Z - Z.mean(0)) / Z.std(0)
        rows = []
        for f in sorted(al["fold"].unique()):
            tr = (al["fold"] != f).values
            te = ~tr
            A = np.column_stack([np.ones(tr.sum()), Z[tr]])
            W = np.diag(w[tr])
            b = np.linalg.solve(A.T @ W @ A, A.T @ W @ R[tr])
            e = (b[0] + Z[te] @ b[1:]) - R[te]
            rows.append((f, te.sum(), np.sqrt((e ** 2).mean()), e.mean()))
        d = pd.DataFrame(rows, columns=["fold", "n", "rmse", "bias"])
        wm = lambda col: np.sqrt(np.average(d[col] ** 2, weights=d["n"]))
        return d, wm("rmse"), wm("bias")

    print("\n4. CURVATURE IN THE OUT-OF-CLADE INTERCEPT BIAS")
    rows = []
    for f in sorted(al["fold"].unique()):
        tr = (al["fold"] != f).values
        te = ~tr
        b0 = np.average(R[tr], weights=w[tr])
        rows.append((te.sum(), np.sqrt(((b0 - R[te]) ** 2).mean())))
    n0, r0 = (np.array(x) for x in zip(*rows))
    base = np.sqrt(np.average(r0 ** 2, weights=n0))
    print(f"   {'model':<24}{'RMSE':>9}{'|bias|':>9}{'gain':>9}"
          f"{'bias share':>12}")
    print(f"   {'intercept only':<24}{base:>9.5f}{'-':>9}{'-':>9}{'-':>12}")
    for lab, cols in [("linear (confirmatory)", [c, k]),
                      ("plus kappa^2", [c, k, k ** 2]),
                      ("plus kappa^2 and ^3", [c, k, k ** 2, k ** 3])]:
        d, rm, bi = loco(cols)
        print(f"   {lab:<24}{rm:>9.5f}{bi:>9.5f}"
              f"{100 * (1 - rm / base):>8.1f}%{100 * bi ** 2 / rm ** 2:>11.0f}%")
    dlin, _, _ = loco([c, k])
    dq, _, _ = loco([c, k, k ** 2])
    km = al.groupby("fold")["kappa_log"].mean()
    print(f"\n   {'clade':<10}{'mean log k':>12}{'bias linear':>14}"
          f"{'bias +k^2':>12}")
    for f in dlin["fold"]:
        print(f"   {f:<10}{km[f]:>12.3f}"
              f"{dlin[dlin['fold'] == f]['bias'].iloc[0]:>+14.5f}"
              f"{dq[dq['fold'] == f]['bias'].iloc[0]:>+12.5f}")
    print(f"   corr(clade mean log kappa, bias) = "
          f"{np.corrcoef(km[dlin['fold']], dlin['bias'])[0, 1]:+.3f}")


def main():
    have = os.path.isfile(CORAL)
    print("=" * 72)
    print("DEPOSITED ANALYSES")
    print("=" * 72)
    print(f"CORAL table: {'found' if have else 'NOT FOUND — blocks 2-4 skipped'}")

    if not have:
        stop_placements()
        print("\nBlocks 2-4 need the CORAL supplementary table. Set CORAL_XLSX.")
        return 0

    al = with_taxonomy(confirmatory())
    spec = spectra()
    stop_placements(spec, list(al["species"]))
    cug_sensitivity(spec, al.copy())
    stop_identity(spec, al.copy())
    curvature(al.copy())
    return 0


if __name__ == "__main__":
    sys.exit(main())
