"""
Latent-heterogeneity test for protocol v4.4, section 3.

Model:

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

    tau^2  total latent heterogeneity
    lam    fraction of it tracking the phylogeny
    C      phylogenetic correlation matrix (Brownian, from branch lengths)
    sigma2 per-species spectrum-estimation variance (KNOWN, not estimated)

Null:        tau^2 = 0
Alternative: tau^2 > 0

Two things make the standard tests invalid:
  - tau^2 = 0 sits on the BOUNDARY of the parameter space
  - lam is UNIDENTIFIED under the null (it multiplies zero)
so significance comes from parametric bootstrap under H0, never from Wald or LRT.
"""

import numpy as np
from scipy.optimize import minimize_scalar, minimize


def brownian_correlation(tree_paths):
    """Phylogenetic correlation matrix from shared root-to-node path lengths.

    tree_paths[i][j] = shared path length from root for taxa i and j.
    Correlation = shared / sqrt(total_i * total_j).
    """
    V = np.asarray(tree_paths, float)
    d = np.sqrt(np.diag(V))
    return V / np.outer(d, d)


def _neg_loglik(params, R, C, sigma2, I):
    """mu is profiled out analytically (GLS), so only (tau2, lam) are searched."""
    tau2, lam = params
    if tau2 < 0 or not (0.0 <= lam <= 1.0):
        return 1e12
    S = tau2 * ((1 - lam) * I + lam * C) + np.diag(sigma2)
    try:
        L = np.linalg.cholesky(S)
    except np.linalg.LinAlgError:
        return 1e12
    one = np.ones(len(R))
    Si1 = np.linalg.solve(S, one)
    mu = float(one @ np.linalg.solve(S, R) / (one @ Si1))
    r = R - mu
    logdet = 2.0 * np.sum(np.log(np.diag(L)))
    quad = float(r @ np.linalg.solve(S, r))
    return 0.5 * (logdet + quad)


def fit(R, C, sigma2):
    """ML fit of (tau2, lam, mu). Returns dict."""
    R = np.asarray(R, float)
    sigma2 = np.asarray(sigma2, float)
    n = len(R)
    I = np.eye(n)
    scale = max(np.var(R) - np.mean(sigma2), 1e-12)
    best = None
    for lam0, t0 in ((0.9, scale), (0.1, scale)):
        res = minimize(_neg_loglik, x0=[max(t0, 1e-12), lam0],
                       args=(R, C, sigma2, I), method="Nelder-Mead",
                       options={"xatol": 1e-12, "fatol": 1e-9, "maxiter": 300})
        if best is None or res.fun < best.fun:
            best = res
    tau2, lam = float(max(best.x[0], 0.0)), float(min(max(best.x[1], 0.0), 1.0))
    S = tau2 * ((1 - lam) * I + lam * C) + np.diag(sigma2)
    one = np.ones(n)
    mu = float(one @ np.linalg.solve(S, R) / (one @ np.linalg.solve(S, one)))
    return {"tau2": tau2, "lambda": lam, "mu": mu, "nll": float(best.fun)}


def fit_null(R, sigma2):
    """Fit under tau^2 = 0: only mu, by inverse-variance weighting."""
    R = np.asarray(R, float)
    sigma2 = np.asarray(sigma2, float)
    w = 1.0 / sigma2
    return {"mu": float(np.sum(w * R) / np.sum(w)), "tau2": 0.0}


def test_tau2(R, C, sigma2, n_boot=1000, seed=0):
    """Parametric bootstrap under H0: tau^2 = 0.

    1. fit mu with tau^2 = 0
    2. simulate R* ~ N(mu, diag(sigma2))   -- each species' own uncertainty
    3. refit tau^2 on each replicate
    4. compare observed tau^2 to the null distribution
    """
    R = np.asarray(R, float)
    sigma2 = np.asarray(sigma2, float)
    obs = fit(R, C, sigma2)
    null = fit_null(R, sigma2)

    rng = np.random.default_rng(seed)
    boot = np.empty(n_boot)
    for b in range(n_boot):
        Rb = null["mu"] + rng.normal(0.0, np.sqrt(sigma2))
        boot[b] = fit(Rb, C, sigma2)["tau2"]

    p = (1.0 + np.sum(boot >= obs["tau2"])) / (n_boot + 1.0)
    return {"tau2_obs": obs["tau2"], "lambda_obs": obs["lambda"],
            "mu_obs": obs["mu"], "p_value": float(p),
            "null_tau2_95": float(np.quantile(boot, 0.95)),
            "null_tau2_median": float(np.median(boot))}
