"""The three statistics diverge: R rises with transition bias while the
standardized separation and the rank both fall.

Reads the deposited per-species values and writes figures/Fig2.* and
data/divergence.json. No external download is needed.
"""
import json
import os

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

# Springer wants embedded outline fonts in vector figures, and a column width
# of 84 or 174 mm so that nothing is rescaled at typesetting.
plt.rcParams.update({"ps.fonttype": 42, "pdf.fonttype": 42})
_MM = 1 / 25.4
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")
FIGS = os.path.join(ROOT, "figures")
N_ENS = 10000

RED = "#a4342a"
GREY = "#7a7a7a"
DARK = "#2b2b2b"
plt.rcParams.update({"font.size": 8, "axes.labelsize": 8,
                     "xtick.labelsize": 7, "ytick.labelsize": 7,
                     "axes.spines.top": False, "axes.spines.right": False,
                     "figure.dpi": 150})


def load():
    """The 4,972 species that resolved on the calibrated tree.

    cohortB_R_per_species.csv holds all 5,053 focal species; the confirmatory
    set is the subset that resolved against TimeTree, and every figure in the
    paper is computed on that subset.
    """
    per = pd.read_csv(os.path.join(DATA, "cohortB_R_per_species.csv"))
    conf = pd.read_csv(os.path.join(DATA, "confirmatory_aligned.csv"))
    comp = pd.read_csv(os.path.join(DATA, "spectrum_components.csv"))
    d = conf[["species"]].merge(per, on="species").merge(comp, on="species")
    assert len(d) == 4972, f"expected 4,972 species, got {len(d)}"
    d["CV"] = d["sd_null"] / d["Phi_null"]
    d["beat"] = (d["p_emp"] * (N_ENS + 1) - 1).round().astype(int)
    return d


def main():
    d = load()
    out = {"n_species": int(len(d))}

    # --- the rank is not a property of the code -------------------------
    out["beat_min"] = int(d["beat"].min())
    out["beat_max"] = int(d["beat"].max())
    out["beat_median"] = float(d["beat"].median())
    out["pct_none_better"] = round(float(100 * (d["beat"] == 0).mean()), 1)
    print("\n1. THE RANK IS NOT A PROPERTY OF THE CODE\n")
    print(f"   randomized codes that beat the standard code   "
          f"{out['beat_min']} to {out['beat_max']} of {N_ENS}")
    print(f"   species for which none is better               "
          f"{int((d['beat'] == 0).sum())} ({out['pct_none_better']}%)")
    print(f"   median                                         "
          f"{out['beat_median']:.0f}")

    # --- the three statistics diverge ----------------------------------
    print("\n2. THE THREE STATISTICS DIVERGE\n")
    dec = pd.qcut(d["kappa_log"], 10, labels=False)
    lo, hi = dec == 0, dec == 9
    d["none"] = (d["beat"] == 0).astype(float) * 100
    for lab, col in (("R", "R"), ("z", "z"), ("CV", "CV"), ("beat", "beat"),
                     ("none", "none")):
        a, b = float(d[col][lo].mean()), float(d[col][hi].mean())
        out[f"{lab}_lo_decile"] = round(a, 4)
        out[f"{lab}_hi_decile"] = round(b, 4)
        out[f"{lab}_pct_change"] = round(100 * (b / a - 1), 1)
        print(f"   {lab:<5} {a:8.4f} -> {b:8.4f}   {100*(b/a-1):+6.1f}%")
    for lab, col in (("R", "R"), ("z", "z"), ("CV", "CV"), ("beat", "beat")):
        r = float(np.corrcoef(d["kappa_log"], d[col])[0, 1])
        out[f"corr_kappa_{lab}"] = round(r, 3)
    out["corr_R_z"] = round(float(np.corrcoef(d["R"], d["z"])[0, 1]), 3)
    print(f"\n   corr with log kappa:  R {out['corr_kappa_R']:+.2f}   "
          f"z {out['corr_kappa_z']:+.2f}   CV {out['corr_kappa_CV']:+.2f}   "
          f"beat {out['corr_kappa_beat']:+.2f}")
    print(f"   corr(R, z) {out['corr_R_z']:+.3f}")

    # --- the figure ----------------------------------------------------
    q = pd.qcut(d["kappa_log"], 10, labels=False)
    g = d.groupby(q).agg(k=("kappa_log", "mean"), R=("R", "mean"),
                         z=("z", "mean"), CV=("CV", "mean"),
                         beat=("beat", "mean"),
                         none=("beat", lambda x: 100 * (x == 0).mean()))
    fig, ax = plt.subplots(1, 3, figsize=(174 * _MM, 57 * _MM))

    a = ax[0]
    a.plot(g["k"], 100 * g["R"], "o-", color=RED, ms=3.5, lw=1.3)
    a.set_xlabel("log $\\kappa$, decile mean")
    a.set_ylabel("advantage $R$ (%)")
    a.set_title("more protection", fontsize=8, color=DARK)

    a = ax[1]
    a.plot(g["k"], g["z"], "o-", color=DARK, ms=3.5, lw=1.3)
    a.set_xlabel("log $\\kappa$, decile mean")
    a.set_ylabel("standardized separation $z$")
    a.set_title("less exceptional", fontsize=8, color=DARK)

    a = ax[2]
    a.plot(g["k"], g["none"], "o-", color=GREY, ms=3.5, lw=1.3)
    a.set_xlabel("log $\\kappa$, decile mean")
    a.set_ylabel("species with no better code (%)")
    a.set_title("less rare", fontsize=8, color=DARK)

    fig.tight_layout()
    for ext in ("eps", "pdf", "png"):
        fig.savefig(os.path.join(FIGS, f"Fig2.{ext}"), bbox_inches="tight")
    plt.close(fig)

    json.dump(out, open(os.path.join(DATA, "divergence.json"), "w"), indent=1)
    print(f"\n   figures/Fig2.* and data/divergence.json written\n")
    return 0


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