import json, math, random
from pathlib import Path
import numpy as np

SEED = 3112


def cov_long_memory(n, alpha):
    i = np.arange(n)
    return (1.0 + np.abs(i[:, None] - i[None, :])) ** (-alpha)


def correlated_gaussian(n, alpha, rng):
    C = cov_long_memory(n, alpha)
    return np.linalg.cholesky(C + 1e-10 * np.eye(n)) @ rng.standard_normal(n)


def hermite_prob(x, q):
    if q == 1: return x
    if q == 2: return x*x - 1.0
    if q == 3: return x*x*x - 3.0*x
    raise ValueError('q must be 1, 2, or 3')


def hermite_unit(x, q):
    # For standard Gaussian input, E[He_q(G)^2] = q!.
    return hermite_prob(x, q) / math.sqrt(math.factorial(q))


def slope(x, y):
    return float(np.polyfit(np.log(np.asarray(x)), np.log(np.asarray(y)), 1)[0])


def scaling_check(alpha=0.4, q=1, max_n=512, trials=160, rng=None):
    # No per-path centering: centering would force sum(z)==0 and destroy
    # exactly the phenomenon being measured.
    rng = np.random.default_rng(SEED + q) if rng is None else rng
    ns = np.array([16, 32, 64, 128, 256, max_n])
    variances = []
    means = []
    for n in ns:
        sums = []
        for _ in range(trials):
            g = correlated_gaussian(n, alpha, rng)
            sums.append(np.sum(hermite_unit(g, q)))
        variances.append(np.var(sums, ddof=1))
        means.append(np.mean(sums))
    H = 1.0 - alpha*q/2.0
    observed_H = slope(ns, variances) / 2.0
    rms_scaled = [float(np.sqrt(v) / n**H) for n, v in zip(ns, variances)]
    # For Gaussian g, Cov(He_q(g_i)/sqrt(q!), He_q(g_j)/sqrt(q!))=rho(|i-j|)^q.
    exact_variances = []
    for n in ns:
        rho = (1.0 + np.arange(n)) ** (-alpha)
        exact_variances.append(float(n + 2.0 * sum((n-j) * rho[j]**q for j in range(1, n))))
    exact_H = slope(ns, exact_variances) / 2.0
    return {'q': q, 'theory_H': H, 'observed_H': observed_H, 'exact_finite_n_H': exact_H,
            'ns': ns.tolist(), 'sum_variance': [float(v) for v in variances],
            'exact_sum_variance': exact_variances,
            'scaled_sum_rms': rms_scaled,
            'rms_ratio_last_first': rms_scaled[-1] / rms_scaled[0],
            'sum_means': [float(x) for x in means]}


def residual_stability(alpha=0.4, n_trials=300, dims=32, rng=None):
    rng = np.random.default_rng(SEED + 99) if rng is None else rng
    depths = [32, 64, 128, 256]
    out = {}
    for kind in ['critical_correlated', 'iid_sqrt', 'correlated_inverse_sqrt', 'correlated_wrong_inverse']:
        vals = []
        for L in depths:
            final = []
            for _ in range(n_trials):
                if kind == 'critical_correlated':
                    z = correlated_gaussian(L, alpha, rng)
                    lam = L ** (-(1-alpha/2))
                elif kind == 'iid_sqrt':
                    z = rng.standard_normal(L)
                    lam = L ** -0.5
                elif kind == 'correlated_inverse_sqrt':
                    z = correlated_gaussian(L, alpha, rng)
                    lam = L ** -0.5
                else:
                    z = correlated_gaussian(L, alpha, rng)
                    lam = L ** -1.0
                h = rng.standard_normal(dims)
                for zl in z:
                    h = h + lam * zl * np.tanh(h)
                final.append(np.mean(h*h))
            vals.append((float(np.mean(final)), float(np.std(final, ddof=1)/math.sqrt(n_trials))))
        out[kind] = {'depths': depths, 'mean_final_variance': [x[0] for x in vals],
                     'sem': [x[1] for x in vals]}
    return out


def main():
    random.seed(SEED); np.random.seed(SEED)
    checks = [scaling_check(q=q) for q in (1, 2)]
    stability = residual_stability()
    report = {'seed': SEED, 'alpha': 0.4, 'checks': checks, 'stability': stability}
    Path('results.json').write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))

if __name__ == '__main__':
    main()