"""The two phylogenetic models the manuscript compares, fitted from the deposit.

Section 6 reports a coefficient on log kappa and on the CpG shift under a
Pagel-lambda model, with the delta BIC against an intercept-only fit. Those
numbers were produced here rather than read from a file: this script takes the
deposited per-species values and the calibrated tree, fits both models by
maximum likelihood with the mean profiled out, and prints and deposits every
quantity the manuscript states.

The likelihood is the same Brownian pruning used by code/prune.py. It is a
contrast likelihood: the root contributes no term of its own, and the mean is
profiled rather than estimated. Both models here share the same random-effect
structure and differ only in fixed effects, and the fixed effects are estimated
by generalized least squares inside each likelihood evaluation, so the two
likelihoods are on the same scale and their BIC difference is interpretable.
Refitting with the root term included moves lambda and tau^2 in the fifth
decimal place and leaves the delta BIC unchanged at the reported precision;
that check is recorded in the protocol.

    python3 code/fit_primary_models.py

Writes data/primary_models.json.
"""
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")
sys.path.insert(0, HERE)


def main():
    import dendropy
    from prune import neg_ll, prepare
    from scipy.optimize import minimize

    al = pd.read_csv(os.path.join(DATA, "confirmatory_aligned.csv"))
    sc = pd.read_csv(os.path.join(DATA, "spectrum_components.csv"))
    d = al.merge(sc, on="species")
    print(f"  {len(d)} species with R, kappa and the CpG shift")

    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()}
    d = d[d["binom"].isin(labels)].reset_index(drop=True)
    tree.retain_taxa_with_labels(list(d["binom"]))
    order = [l.taxon.label.strip() for l in tree.leaf_node_iter()]
    d = d.set_index("binom").loc[order].reset_index()
    print(f"  {len(d)} resolve on the calibrated tree")

    nodes, _scale = prepare(tree, order)
    y = d["R"].to_numpy(float)
    s2 = d["sigma2"].to_numpy(float)
    z = lambda v: (v - v.mean()) / v.std()
    X = np.column_stack([z(d["kappa_log"].to_numpy(float)),
                         z(d["cpg_shift_log"].to_numpy(float))])

    def contrasts(vec, tau2, lam):
        """Independent contrasts of a vector on the same tree and variances the
        likelihood uses. Regressing contrasts of y on contrasts of X is the
        GLS solution under this covariance, which is what the fixed effects
        need: ordinary least squares would ignore the shared ancestry the
        model is there to account for."""
        tl = tau2 * lam
        extra = (tau2 * (1.0 - lam) + s2) / tl
        m, v, out = {}, {}, []
        for kind, key, b, kids, ti in nodes:
            if kind == "L":
                m[key], v[key] = vec[ti], b + extra[ti]
                continue
            cm = cv = None
            for c in kids:
                mi, vi = m[c], v[c]
                if cm is None:
                    cm, cv = mi, vi
                else:
                    sv = cv + vi
                    out.append((cm - mi) / np.sqrt(sv))
                    cm = (vi * cm + cv * mi) / sv
                    cv = cv * vi / sv
            m[key], v[key] = cm, cv + b
        return np.array(out)

    def gls(design, tau2, lam):
        """Fixed effects by regression on contrasts, and the residual the
        likelihood then evaluates."""
        cy = contrasts(y, tau2, lam)
        cx = np.column_stack([contrasts(design[:, k], tau2, lam)
                              for k in range(design.shape[1])])
        b = np.linalg.lstsq(cx, cy, rcond=None)[0]
        return b, y - design @ b

    def fit(design):
        """Profile the fixed effects out at each (tau2, lambda), so that the
        two models are compared on the same likelihood."""
        def nll(p):
            tau2, lam = np.exp(p[0]), 1 / (1 + np.exp(-p[1]))
            if not (1e-6 <= lam <= 1.0) or tau2 <= 0:
                return 1e15
            if design is None:
                r = y
            else:
                try:
                    _, r = gls(design, tau2, lam)
                except Exception:
                    return 1e15
            return neg_ll((tau2, lam), nodes, r, s2)

        best = None
        for t0 in (-6.0, -5.5, -5.0):
            for l0 in (2.0, 4.0):
                r = minimize(nll, [t0, l0], method="Nelder-Mead",
                             options={"xatol": 1e-8, "fatol": 1e-8,
                                      "maxiter": 4000})
                if best is None or r.fun < best.fun:
                    best = r
        tau2, lam = np.exp(best.x[0]), 1 / (1 + np.exp(-best.x[1]))
        ll = -best.fun
        if design is None:
            beta = []
        else:
            beta, _ = gls(design, tau2, lam)
        k = 2 + (0 if design is None else design.shape[1])
        bic = k * np.log(len(y)) - 2 * ll
        return dict(tau2=float(tau2), lam=float(lam), loglik=float(ll),
                    k=k, bic=float(bic),
                    beta=[float(b) for b in np.atleast_1d(beta)])

    print("\n  fitting the intercept-only model")
    m0 = fit(None)
    print(f"    tau2 {m0['tau2']:.8f}  lambda {m0['lam']:.6f}  "
          f"logL {m0['loglik']:.3f}")

    print("  fitting the model with log kappa and the CpG shift")
    m1 = fit(X)
    print(f"    tau2 {m1['tau2']:.8f}  lambda {m1['lam']:.6f}  "
          f"logL {m1['loglik']:.3f}")
    print(f"    beta(log kappa) {m1['beta'][0]:+.6f}   "
          f"beta(CpG shift) {m1['beta'][1]:+.6f}")

    dbic = m1["bic"] - m0["bic"]
    red = 100 * (1 - m1["tau2"] / m0["tau2"])
    print(f"\n  delta BIC {dbic:+.1f}   tau^2 reduced by {red:.1f}%")

    out = dict(n=len(d), null=m0, spectrum=m1,
               delta_bic=round(dbic, 2), tau2_reduction_pct=round(red, 2))
    json.dump(out, open(os.path.join(DATA, "primary_models.json"), "w"),
              indent=1)
    print("\n  data/primary_models.json written\n")
    return 0


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