"""Stability diagnostics for the phylogenetic model.

Three things a reviewer is entitled to ask about a fit that returns
lambda = 0.992 on a tree with thousands of tips, and that none of the existing
scripts answers:

  1. the likelihood profile in lambda and in tau^2, so that the reported
     maximum is shown to be a maximum and its interval is read off the curve
     rather than from an asymptotic approximation;
  2. simulation recovery — data generated on the same tree with known
     (tau^2, lambda) and refitted, to show the estimator returns what it was
     given at this sample size and with this measurement-error structure;
  3. invariance to the order in which polytomies are resolved, since the
     covariance is built by sequential pairwise merging and the claim that
     this is exact under the Brownian model deserves a test rather than an
     assertion.

Writes data/phylo_diagnostics.json and figures/FigS1.*.
"""
import json
import os
import sys

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
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")
sys.path.insert(0, HERE)

import dendropy  # noqa: E402

from prune import moment_tau2, neg_ll, prepare  # noqa: E402

RED = "#a4342a"
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 confirmatory set on the calibrated tree, in tip order.

    The tree is keyed by binomial, while the spectra carry strain and
    subspecies suffixes: Agaricus_bisporus_var._bisporus_H97 is a tip called
    Agaricus_bisporus. Joining on the wrong column drops 254 species without
    error, most of them fungi, so the join uses `binom` and asserts the count.
    """
    per = pd.read_csv(os.path.join(DATA, "cohortB_R_per_species.csv"))
    conf = pd.read_csv(os.path.join(DATA, "confirmatory_aligned.csv"))
    d = conf[["species", "binom"]].merge(per, on="species")
    tree = dendropy.Tree.get(path=os.path.join(DATA,
                                               "timetree_5053_species.nwk"),
                             schema="newick", preserve_underscores=True)
    labels = {l.taxon.label.strip() for l in tree.leaf_node_iter()}
    missing = sorted(set(d["binom"]) - labels)
    assert not missing, f"{len(missing)} binomials absent from the tree"
    tree.retain_taxa_with_labels(d["binom"].tolist())
    order = [l.taxon.label.strip() for l in tree.leaf_node_iter()]
    d = d.set_index("binom").loc[order].reset_index()
    assert len(d) == 4972, f"expected 4,972 species, got {len(d)}"
    return d, tree, order


def fit_pruned(x, s2, nodes):
    """ML fit of (tau2, lambda) by Brownian pruning, mu profiled out."""
    from scipy.optimize import minimize
    scale = max(float(np.var(x) - np.mean(s2)), 1e-12)
    best = None
    for lam0 in (0.99, 0.5):
        r = minimize(neg_ll, x0=[scale, lam0], args=(nodes, x, s2),
                     method="Nelder-Mead",
                     options={"xatol": 1e-10, "fatol": 1e-7, "maxiter": 400})
        if best is None or r.fun < best.fun:
            best = r
    return float(best.x[0]), float(min(max(best.x[1], 0.0), 1.0)), float(best.fun)


def _crossings(xs, ys, level):
    """Where the profile crosses `level`, by linear interpolation between the
    bracketing grid points; the grid then sets the resolution, not the answer."""
    order = np.argsort(xs)
    xs, ys = np.asarray(xs)[order], np.asarray(ys)[order]
    k = int(np.argmin(ys))
    def side(rng):
        prev = k
        for i in rng:
            if ys[i] >= level:
                x0, x1, y0, y1 = xs[prev], xs[i], ys[prev], ys[i]
                if y1 == y0:
                    return float(x1)
                return float(x0 + (level - y0) * (x1 - x0) / (y1 - y0))
            prev = i
        return float(xs[rng[-1]] if len(rng) else xs[k])
    return side(range(k - 1, -1, -1)), side(range(k + 1, len(xs)))


def simulate(nodes, n, tau2, lam, s2, rng):
    """A Brownian draw on the same tree: each node's value is its parent's plus
    a normal step of variance tau2*lam*branch, with tau2*(1-lam) and the
    per-species measurement variance added at the tips."""
    kids, blen, leaf = {}, {}, {}
    for kind, nid, b, ch, row in nodes:
        blen[nid] = b
        if kind == "L":
            leaf[nid] = row
        else:
            kids[nid] = ch
    root = nodes[-1][1]
    val = {root: 0.0}
    stack = [root]
    y = np.zeros(n)
    while stack:
        nd = stack.pop()
        for c in kids.get(nd, []):
            val[c] = val[nd] + rng.normal(0.0,
                                          np.sqrt(max(tau2 * lam * blen[c],
                                                      0.0)))
            stack.append(c)
        if nd in leaf:
            y[leaf[nd]] = val[nd]
    y = y + rng.normal(0.0, np.sqrt(tau2 * (1 - lam) + s2))
    return y


def main():
    from scipy.optimize import minimize_scalar
    out = {}
    d, tree, order = load()
    nodes, _ = prepare(tree, order)
    n = len(d)
    x = d["R"].to_numpy()
    s2 = (d["sigma_R"].to_numpy()) ** 2
    print(f"\n1. LIKELIHOOD PROFILE  ({n} species on the calibrated tree)\n")

    t_hat, lam_hat, nll0 = fit_pruned(x, s2, nodes)
    print(f"   maximum at lambda = {lam_hat:.4f}, tau2 = {t_hat:.4e}")
    out["lambda_hat"] = round(lam_hat, 4)
    out["tau2_hat"] = float(f"{t_hat:.4g}")

    lams = np.concatenate([np.linspace(0.980, 1.0, 41),
                           np.array([0.0, 0.5, 0.9, 0.95, 0.97])])
    lams = np.unique(lams)
    prof = []
    for lam in lams:
        r = minimize_scalar(lambda t: neg_ll([t, max(lam, 1e-6)], nodes, x, s2),
                            bounds=(1e-10, max(np.var(x) * 50, 1e-6)),
                            method="bounded")
        prof.append(float(r.fun))
    prof = np.array(prof)
    dl = 2 * (prof - prof.min())
    lo_l, hi_l = _crossings(lams, dl, 3.84)
    out["lambda_ci"] = [float(f"{lo_l:.5g}"), float(f"{hi_l:.5g}")]
    print(f"   lambda 95% profile interval  {out['lambda_ci'][0]:.4f} "
          f"to {out['lambda_ci'][1]:.4f}")

    # lambda is optimized continuously at each fixed tau2 rather than taken
    # from a grid, so the profile is the profile and not a lower envelope of
    # five slices through it.
    taus = np.geomspace(t_hat / 1.6, t_hat * 1.6, 41)
    proft = np.array([
        minimize_scalar(lambda L: neg_ll([t, min(max(L, 1e-6), 1.0)],
                                         nodes, x, s2),
                        bounds=(1e-6, 1.0), method="bounded",
                        options={"xatol": 1e-6}).fun
        for t in taus])
    dt = 2 * (proft - proft.min())
    # the crossings are located by interpolation on the profile rather than
    # read off the grid, so the interval does not depend on grid spacing
    lo_t, hi_t = _crossings(taus, dt, 3.84)
    out["tau2_ci"] = [float(f"{lo_t:.4g}"), float(f"{hi_t:.4g}")]
    print(f"   tau2   95% profile interval  {out['tau2_ci'][0]:.3e} "
          f"to {out['tau2_ci'][1]:.3e}")

    print("\n2. SIMULATION RECOVERY\n")
    # Data are generated on this same tree at known (tau2, lambda) and refitted.
    # Generation uses the tree's own covariance through its Cholesky factor,
    # built once from the pruning structure.
    rng = np.random.default_rng(7)
    rec = []
    for lam_true in (0.5, 0.9, 0.99):
        gl, gt = [], []
        for _ in range(6):
            y = simulate(nodes, n, t_hat, lam_true, s2, rng)
            t_i, l_i, _ = fit_pruned(y, s2, nodes)
            gl.append(l_i); gt.append(t_i)
        rec.append((lam_true, float(np.median(gl)),
                    float(np.min(gl)), float(np.max(gl)),
                    float(np.median(gt) / t_hat)))
        print(f"   lambda in {lam_true:.2f}  ->  recovered "
              f"{np.median(gl):.3f}  [{min(gl):.3f}, {max(gl):.3f}]"
              f"   tau2 ratio {np.median(gt)/t_hat:.2f}")
    out["recovery"] = [{"lambda_true": a, "lambda_median": round(b, 3),
                        "lambda_min": round(c, 3), "lambda_max": round(e, 3),
                        "tau2_ratio": round(f, 3)} for a, b, c, e, f in rec]

    print("\n3. POLYTOMY RESOLUTION ORDER\n")
    # The tree carries polytomies; resolving them arbitrarily but at zero
    # branch length should leave the likelihood unchanged under Brownian
    # motion. Six independent random resolutions test that.
    fits = []
    for seed in range(6):
        t2 = dendropy.Tree.get(path=os.path.join(DATA,
                                                 "timetree_5053_species.nwk"),
                               schema="newick", preserve_underscores=True)
        t2.retain_taxa_with_labels(d["binom"].tolist())
        t2.resolve_polytomies(rng=np.random.default_rng(seed))
        o2 = [l.taxon.label.strip() for l in t2.leaf_node_iter()]
        nd2, _ = prepare(t2, o2)
        x2 = d.set_index("binom").loc[o2, "R"].to_numpy()
        s22 = (d.set_index("binom").loc[o2, "sigma_R"].to_numpy()) ** 2
        ti, li, _ = fit_pruned(x2, s22, nd2)
        fits.append((li, ti))
    la = np.array([f[0] for f in fits]); ta = np.array([f[1] for f in fits])
    out["order_lambda_range"] = [round(float(la.min()), 5),
                                 round(float(la.max()), 5)]
    out["order_tau2_cv"] = float(f"{ta.std() / ta.mean():.3g}")
    print(f"   lambda across 6 resolution orders  {la.min():.5f} "
          f"to {la.max():.5f}")
    print(f"   tau2 coefficient of variation      {ta.std()/ta.mean():.3g}")

    fig, ax = plt.subplots(1, 2, figsize=(6.85, 2.4))
    m = lams >= 0.975
    ax[0].plot(lams[m], dl[m], "o-", color=DARK, lw=1.2, ms=2.5)
    ax[0].axhline(3.84, color=RED, lw=1, ls="--")
    ax[0].axvline(lam_hat, color=RED, lw=1)
    ax[0].set_xlabel("$\\lambda$")
    ax[0].set_ylabel("$2\\,\\Delta$ log-likelihood")
    ax[0].set_ylim(0, min(40, max(dl[m].max() * 1.05, 5)))
    ax[1].plot(taus, dt, "o-", color=DARK, lw=1.2, ms=2.5)
    ax[1].axhline(3.84, color=RED, lw=1, ls="--")
    ax[1].axvline(t_hat, color=RED, lw=1)
    ax[1].set_xscale("log")
    ax[1].set_xlabel("$\\tau^2$")
    ax[1].set_ylabel("$2\\,\\Delta$ log-likelihood")
    ax[1].set_ylim(0, min(40, max(dt.max() * 1.05, 5)))
    fig.tight_layout()
    for ext in ("eps", "pdf", "png"):
        fig.savefig(os.path.join(FIGS, f"FigS1.{ext}"), bbox_inches="tight")
    plt.close(fig)

    out["n_species"] = n
    json.dump(out, open(os.path.join(DATA, "phylo_diagnostics.json"), "w"),
              indent=1)
    print("\n   figures/FigS1.* and data/phylo_diagnostics.json written\n")
    return 0


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