import json, math
import numpy as np
from scipy.integrate import quad
from shape_private_noise import select_shape, hockey_divergence, sample_noise, moment


def monte_carlo_moments(p, b, seed=7, n=500000):
    z = sample_noise(np.random.default_rng(seed), p, b, n)
    return float(np.mean(np.abs(z))), float(np.mean(z*z)), float(np.std(z))


def numerical_divergence_mc(p, b, d, eps, seed=9, n=1000000):
    # Importance sampling under P: D_e(P||Q)=E_P[(1-exp(eps+log q-log p))_+].
    rng = np.random.default_rng(seed)
    x = sample_noise(rng, p, b, n)
    lp = math.log(p) - math.log(2*b) - math.lgamma(1/p) - (np.abs(x)/b)**p
    lq = math.log(p) - math.log(2*b) - math.lgamma(1/p) - (np.abs(x-d)/b)**p
    val = np.maximum(0, 1-np.exp(np.minimum(700, eps+lq-lp)))
    return float(np.mean(val)), float(np.std(val)/math.sqrt(n))


def main():
    grid=[1,1.25,1.5,1.75,2,2.5,3,4,6,8]
    regimes=[(0.1,1e-5),(0.5,1e-5),(1,1e-5),(2,1e-5),(1,1e-3),(1,1e-2),(5,1e-5)]
    sweep=[]
    for eps,delta in regimes:
        best, rows=select_shape(epsilon=eps,delta=delta,grid=grid)
        lap=next(x for x in rows if x['p']==1)
        gauss=next(x for x in rows if x['p']==2)
        sweep.append({'epsilon':eps,'delta':delta,'best':best,
                      'laplace_moment':lap['moment'],'gaussian_moment':gauss['moment'],
                      'gain_vs_laplace':1-best['moment']/lap['moment'],
                      'gain_vs_gaussian':1-best['moment']/gauss['moment']})
    b,rows=select_shape(epsilon=1,delta=1e-5,grid=grid)
    moment_check=[]
    for r in rows[:3]:
        mc1,mc2,sd=monte_carlo_moments(r['p'],r['b'])
        moment_check.append({'p':r['p'],'formula_m1':moment(r['p'],r['b'],1),
          'mc_m1':mc1,'formula_m2':moment(r['p'],r['b'],2),'mc_m2':mc2,'std':sd,
          'quad_divergence':r['divergence'],'mc_divergence':numerical_divergence_mc(r['p'],r['b'],1,1)})
    result={'sweep':sweep,'checks':moment_check}
    open('sweep_results.json','w').write(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))
main()