import json
import numpy as np
from pathlib import Path


def stein_kernel(x, y, gamma):
    # Pairwise matrix for x[b,n,d], y[b,m,d], target N(0,I), RBF kernel.
    diff = x[:, :, None, :] - y[:, None, :, :]
    r2 = np.sum(diff * diff, axis=-1)
    dot = np.einsum('bid,bjd->bij', x, y)
    d = x.shape[-1]
    return np.exp(-gamma * r2) * (dot + 2 * gamma * d - (2 * gamma + 4 * gamma**2) * r2)


def stein_kernel_paired(x, y, gamma):
    # Paired h(x_i,y_i), avoiding a large pairwise matrix.
    diff = x - y
    r2 = np.sum(diff * diff, axis=-1)
    dot = np.sum(x * y, axis=-1)
    d = x.shape[-1]
    return np.exp(-gamma * r2) * (dot + 2 * gamma * d - (2 * gamma + 4 * gamma**2) * r2)


def trace_c(d, gamma):
    return (1 + 2 * gamma) * d


def trace_c2(d, gamma):
    a = 1 + 8 * gamma
    return (a ** (-d / 2.0)) * d / (a * a) * ((64 * gamma**4 + 32 * gamma**3 + 4 * gamma**2) * d + 128 * gamma**4 + 128 * gamma**3 + 80 * gamma**2 + 16 * gamma + 1)


def estimates(x, gamma):
    h = stein_kernel(x, x, gamma)
    n = x.shape[1]
    v = np.sqrt(np.maximum(h.mean(axis=(1, 2)), 0))
    u2 = (h.sum(axis=(1, 2)) - np.trace(h, axis1=1, axis2=2)) / (n * (n - 1))
    u = np.sqrt(np.maximum(u2, 0))
    return v, u, u2


def replication(d, n, gamma, reps, rng, chunk=100):
    vals_v, vals_u, vals_u2 = [], [], []
    for start in range(0, reps, chunk):
        b = min(chunk, reps - start)
        x = rng.standard_normal((b, n, d))
        v, u, u2 = estimates(x, gamma)
        vals_v.append(v); vals_u.append(u); vals_u2.append(u2)
    v = np.concatenate(vals_v); u = np.concatenate(vals_u); u2 = np.concatenate(vals_u2)
    return {
        'd': d, 'n': n, 'reps': reps,
        'V_mean': float(v.mean()), 'V_sd': float(v.std(ddof=1)),
        'U_mean': float(u.mean()), 'U_sd': float(u.std(ddof=1)),
        'U2_mean': float(u2.mean()), 'U2_sd': float(u2.std(ddof=1)),
        'U2_negative_frac': float(np.mean(u2 < 0)),
        'trace_scale': float(np.sqrt(trace_c(d, gamma) / n)),
        'hs_scale': float(np.sqrt(np.sqrt(trace_c2(d, gamma)) / n)),
        'effective_rank': float(trace_c(d, gamma)**2 / trace_c2(d, gamma)),
    }


def main():
    rng = np.random.default_rng(395)
    gamma = 0.5
    check = {}
    for d in (2, 5, 10):
        x = rng.standard_normal((200000, d)); y = rng.standard_normal((200000, d))
        hxy = stein_kernel_paired(x, y, gamma)
        hxx = np.sum(x*x, axis=1) + 2*gamma*d
        exact2 = trace_c2(d, gamma)
        emp2 = float(np.mean(hxy*hxy))
        check[str(d)] = {
            'diagonal_empirical': float(hxx.mean()), 'diagonal_exact': trace_c(d, gamma),
            'second_moment_empirical': emp2, 'second_moment_exact': exact2,
            'paired_kernel_mean': float(hxy.mean()),
            'relative_diag_error': float(abs(hxx.mean()-trace_c(d,gamma))/trace_c(d,gamma)),
            'relative_second_moment_error': float(abs(emp2-exact2)/exact2)
        }
    results = []
    for d in (2, 5, 10):
        for n in (32, 64, 128, 512):
            reps = 1000 if n <= 128 else 300
            results.append(replication(d, n, gamma, reps, rng))
    d5 = [r for r in results if r['d'] == 5]
    ns = np.array([r['n'] for r in d5], float)
    slopes = {
        'V_mean_vs_n': float(np.polyfit(np.log(ns), np.log([r['V_mean'] for r in d5]), 1)[0]),
        'U_mean_vs_n': float(np.polyfit(np.log(ns), np.log([r['U_mean'] for r in d5]), 1)[0])
    }
    out = {'gamma': gamma, 'math_check': check, 'results': results, 'loglog_slopes_d5': slopes}
    Path('results.json').write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == '__main__':
    main()