import json, math, random
import numpy as np
from scipy.special import ndtri
from scipy.stats import qmc

SEED = 2437
np.random.seed(SEED); random.seed(SEED)


def sobol_base(n, d):
    return qmc.Sobol(d=d, scramble=False, seed=SEED).random_base2(int(round(math.log2(n))))


def qmc_estimates(n, d, reps, fn):
    x = sobol_base(n, d); out = []
    rng = np.random.default_rng(SEED + 17*d + n)
    for _ in range(reps):
        out.append(float(np.mean(fn((x + rng.random(d)) % 1.0))))
    return np.asarray(out)


def iid_estimates(n, d, reps, fn):
    rng = np.random.default_rng(SEED + 991*d + n)
    return np.asarray([float(np.mean(fn(rng.random((n,d))))) for _ in range(reps)])


def smooth(u):
    return np.prod(1.0 + 0.7*np.cos(2*np.pi*u), axis=1)


def nonsmooth(u):
    return (u[:,0] < 0.37).astype(float)


def toy_checks():
    unbiased = qmc_estimates(16, 2, 20000, smooth)
    ns = [4, 8, 16, 32, 64, 128, 256]
    smooth_rows=[]; rough_rows=[]
    for n in ns:
        q=qmc_estimates(n,2,512,smooth); i=iid_estimates(n,2,512,smooth)
        smooth_rows.append((n,float(np.sqrt(np.mean((q-1)**2))),float(np.sqrt(np.mean((i-1)**2)))))
        q=qmc_estimates(n,2,512,nonsmooth); i=iid_estimates(n,2,512,nonsmooth)
        rough_rows.append((n,float(np.sqrt(np.mean((q-.37)**2))),float(np.sqrt(np.mean((i-.37)**2)))))
    dim_rows=[]
    for d in [1,2,4,8,16]:
        q=qmc_estimates(64,d,512,smooth); i=iid_estimates(64,d,512,smooth)
        dim_rows.append((d,float(np.sqrt(np.mean((q-1)**2))),float(np.sqrt(np.mean((i-1)**2)))))
    def slope(rows,col):
        return float(np.polyfit(np.log([r[0] for r in rows]),np.log([max(r[col],1e-15) for r in rows]),1)[0])
    return {'unbiasedness':{'target':1.0,'mean':float(unbiased.mean()),'absolute_error':float(abs(unbiased.mean()-1)),'shift_sd':float(unbiased.std())},'smooth_scaling':{'rows_n_qmc_rmse_iid_rmse':smooth_rows,'qmc_loglog_slope':slope(smooth_rows,1),'iid_loglog_slope':slope(smooth_rows,2)},'nonsmooth_control':{'rows_n_qmc_rmse_iid_rmse':rough_rows,'qmc_loglog_slope':slope(rough_rows,1),'iid_loglog_slope':slope(rough_rows,2)},'dimension':{'rows_d_qmc_rmse_iid_rmse':dim_rows}}


def neural_gradient_check():
    import torch
    torch.manual_seed(SEED); dtype=torch.float64
    x0=torch.linspace(-1.,1.,48,dtype=dtype).reshape(-1,1); y=torch.sin(2.4*x0)+.15*x0
    net=torch.nn.Sequential(torch.nn.Linear(1,12),torch.nn.Tanh(),torch.nn.Linear(12,1)).double()
    state={k:v.detach().clone() for k,v in net.state_dict().items()}
    def grad_for(u):
        net.load_state_dict(state); net.zero_grad()
        z=torch.tensor(ndtri(np.clip(u,1e-6,1-1e-6)),dtype=dtype).reshape(-1,1)
        pred=net(x0)+.25*z.mean(); loss=((pred-y)**2).mean(); loss.backward()
        return np.concatenate([p.grad.detach().numpy().ravel() for p in net.parameters()]),float(loss.detach())
    n=16; reps=192; base=sobol_base(n,1); rng=np.random.default_rng(SEED+777)
    gq=[];gi=[];lq=[];li=[]
    for _ in range(reps):
        g,l=grad_for((base+rng.random(1))%1);gq.append(g);lq.append(l)
        g,l=grad_for(rng.random((n,1)));gi.append(g);li.append(l)
    gq=np.asarray(gq);gi=np.asarray(gi);qvar=float(np.mean(np.var(gq,axis=0,ddof=1)));ivar=float(np.mean(np.var(gi,axis=0,ddof=1)))
    return {'n':n,'reps':reps,'gradient_component_variance_qmc':qvar,'gradient_component_variance_iid':ivar,'variance_reduction':1-qvar/ivar,'loss_mean_qmc':float(np.mean(lq)),'loss_mean_iid':float(np.mean(li)),'loss_mean_difference':float(np.mean(lq)-np.mean(li))}


def training_check():
    import torch
    torch.manual_seed(SEED); dtype=torch.float64
    x=torch.linspace(-1.,1.,64,dtype=dtype).reshape(-1,1)
    target=lambda xx,zz: torch.sin(2.2*xx)+.25*zz
    base_model=torch.nn.Sequential(torch.nn.Linear(2,16),torch.nn.Tanh(),torch.nn.Linear(16,1)).double()
    init={k:v.detach().clone() for k,v in base_model.state_dict().items()}
    n,steps=16,120; sob=sobol_base(n,2); rng=np.random.default_rng(SEED+314)
    curves=[]
    for mode in ['qmc','iid']:
        model=torch.nn.Sequential(torch.nn.Linear(2,16),torch.nn.Tanh(),torch.nn.Linear(16,1)).double(); model.load_state_dict(init)
        opt=torch.optim.Adam(model.parameters(),lr=.025); vals=[]
        for step in range(steps):
            if mode=='qmc': u=(sob+rng.random(2))%1
            else: u=rng.random((n,2))
            z=torch.tensor(ndtri(np.clip(u[:,1],1e-6,1-1e-6)),dtype=dtype).reshape(-1,1)
            # Use the first coordinate to select fixed data points, while the second is latent noise.
            idx=np.floor(u[:,0]*len(x)).astype(int); xx=x[idx]; yy=target(xx,z)
            pred=model(torch.cat([xx,z],1)); loss=((pred-yy)**2).mean()
            opt.zero_grad();loss.backward();opt.step();vals.append(float(loss.detach()))
        curves.append(vals)
    return {'steps':steps,'final_loss_qmc':curves[0][-1],'final_loss_iid':curves[1][-1],'best_loss_qmc':min(curves[0]),'best_loss_iid':min(curves[1]),'trajectory_first_last':{'qmc':[curves[0][0],curves[0][-1]],'iid':[curves[1][0],curves[1][-1]]}}


def main():
    result={'seed':SEED,'toy':toy_checks(),'neural_gradient':neural_gradient_check(),'training':training_check()}
    with open('results.json','w') as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()