import json, math
from pathlib import Path
import numpy as np
from scipy.stats import qmc, norm
from scipy.integrate import quad

SEED = 2108
EPS = 1e-10

# A one-step Euler probability-flow surrogate.  The monotone transport is
# T_lambda(z)=z+lambda*tanh(z), which has an exactly known change-of-variables
# Jacobian and is smooth enough for scrambled Sobol integration.
def transport(z, lam):
    return z + lam * np.tanh(z)

def det_jacobian(z, lam):
    return np.prod(1.0 + lam / np.cosh(z)**2, axis=1)

def log_phi(z):
    return -0.5*np.sum(z*z, axis=1) - z.shape[1]*0.5*math.log(2*math.pi)

def log_pi(x):
    # Known target pi=N(0,I), so the correction is an exact importance ratio.
    return log_phi(x)

def metric(x):
    return x[:, 0]**2

def estimate(z, lam, weighted):
    x = transport(z, lam)
    vals = metric(x)
    if weighted:
        # pi(T(z))*|det JT(z)| / phi(z)
        logw = log_pi(x) + np.log(det_jacobian(z, lam)) - log_phi(z)
        vals = vals * np.exp(np.clip(logw, -40, 40))
    return float(np.mean(vals))

def sobol_z(n, d, seed):
    u = qmc.Sobol(d=d, scramble=True, seed=seed).random_base2(int(round(math.log2(n))))
    return norm.ppf(np.clip(u, EPS, 1-EPS))

def iid_z(n, d, rng):
    return rng.standard_normal((n, d))

def exact_unweighted_expectation(lam):
    # E[(Z+lambda*tanh(Z))^2] = 1 + 2 lambda A + lambda^2 B.
    phi = lambda z: math.exp(-z*z/2)/math.sqrt(2*math.pi)
    A = quad(lambda z: z*math.tanh(z)*phi(z), -np.inf, np.inf, epsabs=1e-12)[0]
    B = quad(lambda z: math.tanh(z)**2*phi(z), -np.inf, np.inf, epsabs=1e-12)[0]
    return 1.0 + 2.0*lam*A + lam*lam*B, A, B

def run():
    rng = np.random.default_rng(SEED)
    ns = [32, 64, 128, 256, 512, 1024]
    reps = 48
    z0, lam0, h = np.array([[0.37, -1.1]]), 0.8, 1e-5
    num = np.zeros((2, 2))
    for j in range(2):
        zp, zm = z0.copy(), z0.copy(); zp[0,j] += h; zm[0,j] -= h
        num[:,j] = ((transport(zp,lam0)-transport(zm,lam0))/(2*h))[0]
    jac_err = abs(np.linalg.det(num) - det_jacobian(z0,lam0)[0])

    rows = []
    for lam in [0.0, 0.25, 0.5, 1.0]:
        exact_u, _, _ = exact_unweighted_expectation(lam)
        for n in ns:
            sob_u, sob_w, iid_u, iid_w = [], [], [], []
            for r in range(reps):
                z = sobol_z(n, 2, SEED + 10000*int(lam*100) + r)
                sob_u.append(estimate(z,lam,False)); sob_w.append(estimate(z,lam,True))
                zi = iid_z(n,2,rng)
                iid_u.append(estimate(zi,lam,False)); iid_w.append(estimate(zi,lam,True))
            rows.append(dict(lam=lam,n=n, exact_unweighted=exact_u,
                sobol_unweighted_mean=float(np.mean(sob_u)), sobol_unweighted_sd=float(np.std(sob_u,ddof=1)),
                sobol_weighted_mean=float(np.mean(sob_w)), sobol_weighted_sd=float(np.std(sob_w,ddof=1)),
                iid_unweighted_sd=float(np.std(iid_u,ddof=1)), iid_weighted_sd=float(np.std(iid_w,ddof=1)),
                sobol_unweighted_bias_to_target=float(np.mean(sob_u)-1),
                sobol_bias_to_exact_unweighted=float(np.mean(sob_u)-exact_u),
                sobol_weighted_bias=float(np.mean(sob_w)-1)))

    def slope(lam, field):
        rr = [x for x in rows if x['lam']==lam]
        return float(np.polyfit(np.log2([x['n'] for x in rr]), np.log2([x[field] for x in rr]), 1)[0])

    exact05, A, B = exact_unweighted_expectation(0.5)
    last = [x for x in rows if x['n']==1024]
    bias_fit = np.polyfit([x['lam'] for x in last], [x['sobol_unweighted_bias_to_target'] for x in last], 2)
    summary = {
      'jacobian_check_abs_error': float(jac_err),
      'target_expectation': 1.0,
      'analytic_unweighted_formula': '1 + 2*lambda*A + lambda^2*B',
      'analytic_A_E[z_tanh(z)]': A, 'analytic_B_E[tanh(z)^2]': B,
      'analytic_expectation_lambda_0.5': exact05,
      'slopes_sd_log2N_at_lambda_0.5': {
        'scrambled_sobol_unweighted': slope(0.5,'sobol_unweighted_sd'),
        'scrambled_sobol_weighted': slope(0.5,'sobol_weighted_sd'),
        'iid_unweighted': slope(0.5,'iid_unweighted_sd'),
        'iid_weighted': slope(0.5,'iid_weighted_sd')},
      'bias_quadratic_fit_coefficients_descending_lambda2_lambda_const': [float(v) for v in bias_fit],
      'bias_vs_lambda_n1024': [dict(lam=x['lam'], predicted_bias=x['exact_unweighted']-1,
          observed_bias=x['sobol_unweighted_bias_to_target'], weighted_bias=x['sobol_weighted_bias']) for x in last],
      'rows': rows
    }
    Path('results.json').write_text(json.dumps(summary, indent=2))
    print(json.dumps({k:v for k,v in summary.items() if k!='rows'}, indent=2))
    print('TABLE lambda=0.5: n sobol_sd iid_sd weighted_sd target_bias weighted_bias')
    for x in rows:
        if x['lam']==0.5:
            print(x['n'], '%.5g %.5g %.5g %.5g %.5g' % (x['sobol_unweighted_sd'],x['iid_unweighted_sd'],x['sobol_weighted_sd'],x['sobol_unweighted_bias_to_target'],x['sobol_weighted_bias']))

if __name__ == '__main__':
    run()