"""
O(n) likelihood for the latent-heterogeneity model on a large tree.

Model (protocol v5, section 3):

    R_obs ~ N( mu*1 ,  tau^2 [ (1-lam) I + lam C ]  +  diag(sigma2) )

With n ~ 5,000 a dense Cholesky is hopeless. But the model IS a tree covariance
with extended terminal branches:

    S = tau^2*lam*C  +  diag( tau^2*(1-lam) + sigma2 )

so adding ( tau^2(1-lam) + sigma2_s ) / (tau^2*lam) to each tip's terminal edge
(in units of C) recovers a pure Brownian covariance, and Felsenstein pruning
evaluates the likelihood in O(n) with mu profiled out analytically.

Polytomies are handled by sequential pairwise merging, which is exact for the
Brownian model.
"""

import numpy as np


def prepare(tree, tip_order):
    """Post-order arrays from a dendropy tree, with branch lengths scaled so
    mean root-to-tip = 1. Every node is keyed by id(node); leaves additionally
    carry the row index of their tip in `tip_order`."""
    depths = [l.distance_from_root() for l in tree.leaf_node_iter()]
    scale = float(np.mean(depths))
    idx = {name: i for i, name in enumerate(tip_order)}
    nodes = []
    for nd in tree.postorder_node_iter():
        b = (nd.edge.length or 0.0) / scale
        if nd.is_leaf():
            nodes.append(("L", id(nd), b, [], idx[nd.taxon.label.strip()]))
        else:
            nodes.append(("I", id(nd), b, [id(c) for c in nd.child_node_iter()], -1))
    return nodes, scale


def loglik(nodes, x, extra, tau2lam):
    """Brownian pruning with mu profiled out. `extra` inflates terminal edges."""
    m = {}
    v = {}
    ll = 0.0
    LOG2PI = np.log(2 * np.pi)
    for kind, key, b, kids, ti in nodes:
        if kind == "L":
            m[key] = x[ti]
            v[key] = 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:
                s = cv + vi
                ll += -0.5 * (LOG2PI + np.log(tau2lam * s)
                              + (cm - mi) ** 2 / (tau2lam * s))
                w1, w2 = 1.0 / cv, 1.0 / vi
                cm = (cm * w1 + mi * w2) / (w1 + w2)
                cv = 1.0 / (w1 + w2)
        m[key] = cm
        v[key] = cv + b
    rk = nodes[-1][1]
    return ll, m[rk], v[rk]


def neg_ll(params, nodes, x, sigma2):
    tau2, lam = params
    if tau2 <= 0 or not (1e-6 <= lam <= 1.0):
        return 1e15
    tl = tau2 * lam
    extra = (tau2 * (1.0 - lam) + sigma2) / tl
    ll, _, _ = loglik(nodes, x, extra, tl)
    return 1e15 if not np.isfinite(ll) else -ll


def moment_tau2(x, sigma2):
    """Moment estimator: total dispersion minus mean measurement variance."""
    return max(0.0, float(np.var(x, ddof=1) - np.mean(sigma2)))
