import json, math
import numpy as np
from scipy.special import roots_hermitenorm, eval_hermitenorm
from hermite_bottleneck import multi_indices, heat_identity_check


def one_dim_coeffs(var, maxdeg, nodes=80):
    # E_{N(0,var)}[He_k(X)/sqrt(k!)] by Gaussian quadrature.
    x, wt = roots_hermitenorm(nodes)
    wt = wt / math.sqrt(2 * math.pi)
    y = math.sqrt(var) * x
    return np.array([np.sum(wt * eval_hermitenorm(k, y)) / math.sqrt(math.factorial(k))
                     for k in range(maxdeg + 1)])


def exact_report(vars_, maxdeg=12):
    d = len(vars_)
    one = [one_dim_coeffs(v, maxdeg) for v in vars_]
    allidx = multi_indices(d, maxdeg)
    coeff = np.array([np.prod([one[j][a[j]] for j in range(d)]) for a in allidx])
    # Parseval norm of the density ratio, computed independently in closed form.
    S = np.diag(vars_)
    norm2 = np.linalg.det(S) ** -1 * np.linalg.det(2*np.linalg.inv(S)-np.eye(d)) ** -.5
    rows = []
    for N in range(maxdeg + 1):
        keep = np.array([sum(a) <= N for a in allidx])
        err = math.sqrt(max(0., norm2 - float(np.sum(coeff[keep]**2)))) / math.sqrt(norm2)
        rows.append((N, int(keep.sum()), err))
    q = max(abs(v - 1) for v in vars_)
    return q, norm2, rows


def run():
    rng = np.random.default_rng(11)
    # The semigroup formula is checked independently by Monte Carlo.
    heat = heat_identity_check(rng, n=12, nz=12)
    cases = [[1.04, .97, 1.01, .99], [1.20, .82, 1.08, .94],
             [1.55, .62, 1.18, .88]]
    reports = []
    for v in cases:
        q, norm2, rows = exact_report(v)
        reports.append({'variances': v, 'q': q, 'norm2': norm2,
                        'curve': [{'degree': n, 'coefficients': k, 'relative_tail': e}
                                  for n,k,e in rows]})
    # Fit the observed geometric envelope on degrees 2..8; report ratios rather than
    # asserting a universal constant (the statement allows C(Sigma,u,d)).
    ratios = []
    for r in reports:
        q = r['q']
        for row in r['curve']:
            if row['relative_tail'] > 1e-12 and q > 0:
                ratios.append(row['relative_tail'] / q**((row['degree']+1)/2))
    # Practical degree choice using a conservative empirical C from all cases.
    C = max(ratios)
    target = .08
    for r in reports:
        q = r['q']
        chosen = next((x['degree'] for x in r['curve'] if C*q**((x['degree']+1)/2) <= target), 12)
        r['chosen_degree'] = chosen
        r['chosen_coefficients'] = next(x['coefficients'] for x in r['curve'] if x['degree']==chosen)
        r['chosen_tail'] = next(x['relative_tail'] for x in r['curve'] if x['degree']==chosen)
        fixed = next(x for x in r['curve'] if x['degree']==4)
        r['fixed_degree_4_tail'] = fixed['relative_tail']
        r['fixed_degree_4_coefficients'] = fixed['coefficients']
    return {'heat_relative_error': heat, 'empirical_envelope_C': C,
            'target': target, 'reports': reports}

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