import itertools
import math
import numpy as np
from scipy.special import eval_hermitenorm, roots_hermitenorm


def multi_indices(d, degree):
    out = []
    def rec(i, left, cur):
        if i == d - 1:
            out.append(tuple(cur + [left]))
            return
        for k in range(left + 1):
            rec(i + 1, left - k, cur + [k])
    for total in range(degree + 1):
        rec(0, total, [])
    return out


def hermite_basis(x, indices):
    x = np.asarray(x)
    vals = []
    for a in indices:
        h = np.ones(x.shape[0])
        for j, k in enumerate(a):
            if k:
                h *= eval_hermitenorm(k, x[:, j]) / math.sqrt(math.factorial(k))
        vals.append(h)
    return np.stack(vals, axis=1)


def density_ratio(x, cov, mean=None):
    x = np.asarray(x)
    d = x.shape[1]
    mean = np.zeros(d) if mean is None else np.asarray(mean)
    sign, logdet = np.linalg.slogdet(cov)
    if sign <= 0:
        raise ValueError('covariance must be positive definite')
    inv = np.linalg.inv(cov)
    delta = x - mean
    # M_cov,mean / M_I,0; clipping only prevents rare MC overflow.
    logf = -.5 * logdet - .5*np.einsum('bi,ij,bj->b', delta, inv, delta) + .5*np.sum(x*x, axis=1)
    return np.exp(np.clip(logf, -80, 80))


def project_density(samples_w, cov, degree, mean=None):
    idx = multi_indices(samples_w.shape[1], degree)
    B = hermite_basis(samples_w, idx)
    f = density_ratio(samples_w, cov, mean)
    # Orthonormal Hermites: coefficient is E_w[f phi] = E_M[phi].
    coef = np.mean(B * f[:, None], axis=0)
    return idx, coef


def reconstruct_error(samples_w, cov, idx, coef, mean=None):
    f = density_ratio(samples_w, cov, mean)
    pred = hermite_basis(samples_w, idx) @ coef
    return float(np.sqrt(np.mean((f - pred)**2))), float(np.sqrt(np.mean(f*f)))


def heat_identity_check(rng, d=2, n=12, nz=12):
    # Product Gauss-Hermite quadrature over Z gives a deterministic check of the
    # expectation in the heat-semigroup identity at independent test points v.
    u = np.array([.35, -.25])[:d]
    t = .31
    nodes, weights = roots_hermitenorm(nz)
    weights = weights / math.sqrt(2 * math.pi)
    grid = np.array(list(itertools.product(nodes, repeat=d)))
    wg = np.array([np.prod([weights[k] for k in ix])
                   for ix in itertools.product(range(nz), repeat=d)])
    v = rng.normal(size=(n, d))
    uu = u[None, :] + math.sqrt(2*t) * grid
    lhs = np.exp(uu @ v.T - .5*np.sum(uu*uu, axis=1)[:, None]).T @ wg
    s = 1 + 2*t
    rhs = s**(-d/2) * np.exp((v @ u)/s - np.sum(u*u)/(2*s) + t*np.sum(v*v, axis=1)/s)
    return float(np.max(np.abs(lhs-rhs) / (np.abs(rhs)+1e-12)))


def run(seed=7):
    rng = np.random.default_rng(seed)
    heat_rel = heat_identity_check(rng)
    d = 4
    # Shared samples make comparisons low-noise. All eigenvalues lie in q<1 regime.
    nw, nm = 160000, 120000
    w = rng.normal(size=(nw, d))
    covs = [np.diag([1.04, .97, 1.01, .99]),
            np.diag([1.20, .82, 1.08, .94]),
            np.diag([1.55, .62, 1.18, .88])]
    qs = [float(np.max(np.abs(np.linalg.eigvalsh(c)-1))) for c in covs]
    maxdeg = 6
    results = []
    # Empirical C is calibrated only on the mildest covariance, as a practical tail scale.
    calib = covs[1]
    tails_cal = []
    projections = {}
    for N in range(maxdeg + 1):
        idx, co = project_density(w, calib, N)
        err, norm = reconstruct_error(w, calib, idx, co)
        tails_cal.append(err)
        projections[(id(calib), N)] = (idx, co)
    # C_hat is deliberately a held-out-style empirical scale, not a claimed universal constant.
    C_hat = max(tails_cal[N] / (qs[1] ** ((N+1)/2)) for N in range(maxdeg+1))
    target = .08
    fixed_N = 4
    for ci, cov in enumerate(covs):
        q = qs[ci]
        rows = []
        for N in range(maxdeg + 1):
            idx, co = project_density(w, cov, N)
            err, norm = reconstruct_error(w, cov, idx, co)
            rows.append((N, len(idx), err / norm))
        chosen = next((N for N in range(maxdeg+1) if C_hat*q**((N+1)/2) <= target), maxdeg)
        aidx, aco = project_density(w, cov, chosen)
        actual, norm = reconstruct_error(w, cov, aidx, aco)
        fidx, fco = project_density(w, cov, fixed_N)
        fixed, _ = reconstruct_error(w, cov, fidx, fco)
        results.append({'q': q, 'adaptive_degree': chosen,
                        'adaptive_coefficients': len(aidx),
                        'adaptive_rel_error': actual/norm,
                        'fixed_degree': fixed_N, 'fixed_coefficients': len(fidx),
                        'fixed_rel_error': fixed/norm, 'curve': rows})
    return {'heat_relative_error': heat_rel, 'C_hat': C_hat, 'target': target,
            'dimension': d, 'results': results}

if __name__ == '__main__':
    import json
    print(json.dumps(run(), indent=2))