import json
import math
import numpy as np


def sqrt_quantile_grid(u, a, n, delta=0.0):
    u = np.asarray(u, float); a = np.maximum(np.asarray(a, float), 0) + delta
    w = np.sqrt(a)
    c = np.zeros_like(u)
    c[1:] = np.cumsum((w[1:] + w[:-1]) * np.diff(u) / 2)
    targets = np.linspace(0, c[-1], n + 1)
    return np.interp(targets, c, u)


def piecewise_cost(grid, u, a):
    # Dense Gauss-Legendre quadrature per interval avoids coarse-grid bias.
    nodes, weights = np.polynomial.legendre.leggauss(12)
    lo, hi = grid[:-1], grid[1:]
    x = (lo[:, None] + hi[:, None]) / 2 + (hi[:, None] - lo[:, None]) * nodes[None, :] / 2
    integ = np.sum(np.interp(x, u, a) * weights[None, :], axis=1) * (hi-lo) / 2
    return .5 * np.sum(integ * (hi-lo))


def empirical_density(u, temporal_amp=1.0, spatial_amp=0.0, seed=0, batch=512):
    """Toy analogue of E[||d_u s||^2 + lambda ||J_x s||_F^2].
    The score is s(x,u)=q(u)x + p(u) tanh(x), with x~N(0,1).
    Finite differences and exact Jacobian are used as a pilot estimate.
    """
    rng = np.random.default_rng(seed)
    x = rng.normal(size=(batch, 1))
    q = 0.35 + 0.7 * u + 1.5 * np.exp(-((u-.70)/.09)**2)
    dq = 0.7 + 1.5 * np.exp(-((u-.70)/.09)**2) * (-2*(u-.70)/(.09**2))
    p = 0.25 * np.sin(5*u)
    dp = 1.25 * np.cos(5*u)
    temporal = np.mean((dq*x + dp*np.tanh(x))**2, axis=0)
    spatial = np.mean((q + p*(1-np.tanh(x)**2))**2, axis=0)
    return temporal_amp * temporal + spatial_amp * spatial


def euler_variable_ode(grid, k=18.0):
    """y'=-r(u)y, with rapid rate around .7; exact y(1) is known."""
    r = lambda z: 0.35 + k*np.exp(-((z-.70)/.075)**2)
    y = 1.0
    for lo, hi in zip(grid[:-1], grid[1:]):
        y *= 1.0 - r((lo+hi)/2)*(hi-lo)
    exact = np.exp(-(.35 + k*.075*np.sqrt(np.pi)/2 * (
        math.erf((1-.70)/.075) - math.erf((0-.70)/.075))))
    return abs(y-exact), y, exact


def main():
    u = np.linspace(0, 1, 20001)
    # Density with one narrow high-error region is the mechanism test.
    a = 1.0 + 80.0*np.exp(-((u-.70)/.075)**2)
    rows = []
    for n in [8, 16, 32, 64, 128]:
        uni = np.linspace(0, 1, n+1)
        ad = sqrt_quantile_grid(u, a, n, delta=1e-10)
        cuni = piecewise_cost(uni, u, a)
        cad = piecewise_cost(ad, u, a)
        lower = .5*np.trapz(np.sqrt(a), u)**2/n
        rows.append({'N': n, 'uniform_cost': cuni, 'adaptive_cost': cad,
                     'predicted_lower_bound': lower,
                     'adaptive_over_bound': cad/lower,
                     'uniform_over_adaptive': cuni/cad,
                     'max_equalized_cost_ratio': (np.max(np.interp((ad[:-1]+ad[1:])/2,u,a)*np.diff(ad)**2)/
                                                   np.min(np.interp((ad[:-1]+ad[1:])/2,u,a)*np.diff(ad)**2))})

    # Prediction 1: adaptive cost tends to lower bound and interval costs equalize.
    # Prediction 2: local widths follow inverse square-root density (log correlation).
    n = 64; ad = sqrt_quantile_grid(u, a, n); mids=(ad[:-1]+ad[1:])/2
    widths=np.diff(ad); amid=np.interp(mids,u,a)
    width_pred=1/np.sqrt(amid)
    corr=float(np.corrcoef(np.log(widths), np.log(width_pred))[0,1])
    # Prediction 3: multiplying density by c changes cost by c but not grid.
    scale_rows=[]
    ref=sqrt_quantile_grid(u,a,n)
    for c in [.25, 1., 4., 16.]:
        g=sqrt_quantile_grid(u,c*a,n)
        scale_rows.append({'scale':c, 'grid_max_difference':float(np.max(abs(g-ref))),
                           'cost_ratio_to_reference':piecewise_cost(g,u,c*a)/piecewise_cost(ref,u,a),
                           'predicted_cost_ratio':c})

    # Empirical-pilot robustness: noisy density estimates should preserve the allocation.
    # Each estimate is positive and smoothed by a small floor, as in the proposal.
    robustness = []
    true_grid = sqrt_quantile_grid(u, a, 32)
    for rel_noise in [0.0, 0.05, 0.10, 0.20, 0.40]:
        rng = np.random.default_rng(1000 + int(rel_noise*100))
        pilot = np.maximum(a[::200] * (1 + rel_noise*rng.normal(size=len(a[::200]))), 1e-8)
        pu = u[::200]
        est_grid = sqrt_quantile_grid(pu, pilot, 32, delta=1e-8*np.mean(pilot))
        robustness.append({'relative_noise': rel_noise,
                           'grid_linf_error': float(np.max(abs(est_grid-true_grid))),
                           'cost_ratio_using_true_a': float(piecewise_cost(est_grid,u,a)/piecewise_cost(true_grid,u,a))})

    # Mini scheduler experiment: same Euler method, fixed NFE.
    ode=[]
    for n in [8, 16, 32, 64]:
        ug=np.linspace(0,1,n+1)
        # use normalized empirical density; only its shape matters
        eg=sqrt_quantile_grid(u,a,n)
        eu=euler_variable_ode(ug); ea=euler_variable_ode(eg)
        ode.append({'N':n, 'uniform_abs_error':eu[0], 'adaptive_abs_error':ea[0],
                    'error_ratio_adaptive_over_uniform':ea[0]/eu[0]})

    report={'predictions':[
        {'claim':'adaptive cost approaches C*=L^2/(2N)', 'observed':rows},
        {'claim':'width proportional to a^{-1/2}', 'observed_log_correlation':corr,
         'tolerance':'correlation > 0.98 on this smooth density'},
        {'claim':'a -> c*a leaves grid invariant and scales cost by c', 'observed':scale_rows}
    ], 'pilot_robustness': robustness, 'ode_mini_experiment':ode,
    'density_summary':{'min':float(a.min()),'max':float(a.max()),'peak_u':float(u[a.argmax()])}}
    print(json.dumps(report, indent=2))

    # Mechanism checks (strict, reproducible).
    assert rows[-1]['adaptive_over_bound'] < 1.01
    assert corr > .98
    assert max(x['grid_max_difference'] for x in scale_rows) < 1e-9
    assert max(abs(x['cost_ratio_to_reference']-x['predicted_cost_ratio']) for x in scale_rows) < 1e-10

if __name__ == '__main__':
    main()