"""
polyatree.py -- POLYA TREES: a nonparametric prior on a distribution (from scratch).

Backs the notebooks in  "Polya Trees".

A Dirichlet process is a prior over DISCRETE distributions. A Polya tree (Lavine 1992;
Ferguson 1974) is a prior over CONTINUOUS ones -- a random density -- built by recursively
splitting the sample space and randomising how mass flows down the splits.

Recipe. Map the data to [0,1] through the CDF of a centring distribution G0, u_i = G0(x_i).
Split [0,1] dyadically: level m has 2^m intervals, each parent splitting into two children.
At every split, the fraction of the parent's mass going LEFT is a Beta random variable

    Y_eps ~ Beta(alpha_m, alpha_m),   alpha_m = c * m^2   (Lavine's canonical choice),

independent across splits. alpha_m = c m^2 makes the random measure absolutely continuous (a
genuine density), and the concentration c controls how tightly it hugs G0: c -> infinity gives
back G0 exactly, small c allows wild departures. Because the splits are at G0's dyadic quantiles
(uniform on the u-scale), the prior is CENTRED on G0: E[density] = G0.

Conjugacy makes it trivial. With counts n_{eps0}, n_{eps1} of observations falling in the two
children of a node, the posterior is another Polya tree:

    Y_eps | data ~ Beta(alpha_m + n_{eps0}, alpha_m + n_{eps1}).

So the posterior density is available in closed form -- no MCMC, just Beta draws for the bands.

GOODNESS OF FIT. Centring on a fitted parametric G0 turns the Polya tree into a test of that
family: if G0 is right the data are uniform on the u-scale (every Y ~ 1/2) and the tree adds
nothing; if not, the branch counts pull the Y's away from 1/2. The Bayes factor against G0 is a
closed-form sum over nodes -- large values are evidence the parametric family is wrong.
"""

import numpy as np
from scipy.special import betaln


# --------------------------------------------------------------------------- #
#  fit (store dyadic counts of the CDF-transformed data)                        #
# --------------------------------------------------------------------------- #

def polyatree_fit(x, cdf, pdf, c=1.0, M=8):
    """Centre a Polya tree of depth M on the distribution with the given cdf/pdf, concentration c.
    Stores the counts of the transformed data in every dyadic interval up to level M."""
    x = np.asarray(x, float)
    u = np.clip(cdf(x), 1e-12, 1 - 1e-12)
    counts = [np.bincount(np.minimum((u * 2 ** m).astype(int), 2 ** m - 1),
                          minlength=2 ** m).astype(float) for m in range(M + 1)]
    return dict(counts=counts, c=c, M=M, cdf=cdf, pdf=pdf)


def _leaf_masses(pt, draw=False, rng=None):
    """mass in each of the 2^M leaf intervals (posterior mean, or a posterior draw)."""
    M, c = pt["M"], pt["c"]; mass = np.ones(1)
    for m in range(M):
        a = c * (m + 1) ** 2; nc = pt["counts"][m + 1]
        n0 = nc[0::2]; n1 = nc[1::2]
        y = rng.beta(a + n0, a + n1) if draw else (a + n0) / (2 * a + n0 + n1)
        nm = np.empty(2 ** (m + 1)); nm[0::2] = mass * y; nm[1::2] = mass * (1 - y)
        mass = nm
    return mass


# --------------------------------------------------------------------------- #
#  posterior density on a grid                                                  #
# --------------------------------------------------------------------------- #

def polyatree_density(pt, xgrid, ndraws=0, rng=None):
    """posterior density on xgrid: mean, and (if ndraws>0) an (ndraws, len(grid)) sample."""
    M = pt["M"]; xgrid = np.asarray(xgrid, float)
    u = np.clip(pt["cdf"](xgrid), 1e-12, 1 - 1e-12)
    leaf = np.minimum((u * 2 ** M).astype(int), 2 ** M - 1)
    g0 = pt["pdf"](xgrid)
    mean = _leaf_masses(pt)[leaf] * (2 ** M) * g0
    if ndraws > 0:
        D = np.empty((ndraws, len(xgrid)))
        for d in range(ndraws):
            D[d] = _leaf_masses(pt, draw=True, rng=rng)[leaf] * (2 ** M) * g0
        return mean, D
    return mean, None


# --------------------------------------------------------------------------- #
#  goodness of fit: Bayes factor against the centring family G0                 #
# --------------------------------------------------------------------------- #

def polyatree_logBF(pt):
    """log Bayes factor for H1 (Polya tree) vs H0 (data ~ G0 exactly), summed over all nodes.
    log BF > 0 favours a departure from G0; > ~2.3 (BF>10) is strong evidence against it."""
    M, c = pt["M"], pt["c"]; lbf = 0.0
    for m in range(M):
        a = c * (m + 1) ** 2; nc = pt["counts"][m + 1]
        n0 = nc[0::2]; n1 = nc[1::2]; n = n0 + n1
        lbf += np.sum(n * np.log(2) + betaln(a + n0, a + n1) - betaln(a, a))
    return lbf


# --------------------------------------------------------------------------- #
#  simulation helpers                                                           #
# --------------------------------------------------------------------------- #

def simulate_mixture(n, weights, means, sds, rng):
    weights = np.asarray(weights, float); weights /= weights.sum()
    k = rng.choice(len(weights), size=n, p=weights)
    return rng.normal(np.asarray(means)[k], np.asarray(sds)[k])

def mixture_pdf(grid, weights, means, sds):
    weights = np.asarray(weights, float); weights /= weights.sum()
    g = np.asarray(grid, float); out = np.zeros_like(g)
    for w, m, s in zip(weights, means, sds):
        out += w * np.exp(-0.5 * ((g - m) / s) ** 2) / (s * np.sqrt(2 * np.pi))
    return out