from mynumpy import *
chi2 = scipy.stats.chi2
import lhsmdu
import halton
import sobol_lib

def iid(nsamps,M,ndims):
    return randn(nsamps,M,ndims)

def antithetic(nsamps,M,ndims):
    # just flip half the samples
    z = randn(nsamps,M,ndims)
    if M > 1:
        z[:,-(M//2):,:] = -z[:,:M//2,:]
    return z

def stratified01(nsamps,M):
    v = zeros((nsamps,M))
    for m in range(M):
        v[:,m] = rand(nsamps)/M + m/M
    return v

def strat_radius(nsamps,M,ndims):
    # generate samples from the sphere
    u = randn(nsamps,M,ndims)
    u = u / expand_dims(np.sqrt(np.sum(u**2,axis=2)),axis=2)
    # generate radii via inverse CDF of chi distribution
    v = stratified01(nsamps,M)
    r = sqrt(chi2.ppf(v,ndims))
    z = expand_dims(r,axis=2) * u
    return z

def latin(ndims,nsamps):
    k = array(lhsmdu.sample(ndims, nsamps))
    k = k.T[k.T[:,1].argsort()].T
    return k

def latin_cart(nsamps,M,ndims):
    # first get 0-1 samples
    rez = zeros((nsamps,M,ndims))
    for i in range(nsamps):
        zi = array(lhsmdu.sample(ndims,M)).T
        rez[i,:,:] = zi
    return scipy.stats.norm.ppf(rez)

def qmc_cart(nsamps,M,ndims):
    #sobol_seq.i4_sobol_generate(3, 5)
    rez0 = sobol_lib.i4_sobol_generate(ndims,M,1)
    #rez0 = halton.halton_sequence(1,M,ndims)
    rez0 = expand_dims(rez0.T,axis=0)
    #rez0 = rand(nsamps,M,ndims)
    shifts = rand(nsamps,1,ndims)
    rez = np.mod(rez0 + shifts,1.0)
    return scipy.stats.norm.ppf(rez)

def qmc_cartp1(nsamps,M,ndims):
    #sobol_seq.i4_sobol_generate(3, 5)
    rez0 = sobol_lib.i4_sobol_generate(ndims+1,M,1)
    #rez0 = halton.halton_sequence(1,M,ndims)
    rez0 = expand_dims(rez0.T,axis=0)
    #rez0 = rand(nsamps,M,ndims)
    shifts = rand(nsamps,1,ndims+1)
    rez = np.mod(rez0 + shifts,1.0)
    
    # radii
    rez_radius = rez[:,:,0 ]
    r = sqrt(chi2.ppf(rez_radius,ndims))
    
    # samples on the sphere
    rez_rest   = rez[:,:,1:]
    u = scipy.stats.norm.ppf(rez_rest)
    u = u / expand_dims(np.sqrt(np.sum(u**2,axis=2)),axis=2)

    z = expand_dims(r,axis=2) * u
    return z

def anti_qmc_cartp1(nsamps,M,ndims):
    if M==1:
        z = qmc_cartp1(nsamps,M,ndims)
    else:
        assert(np.allclose(M//2,M/2))
        z = qmc_cartp1(nsamps,M//2,ndims)
        z = np.concatenate([z,-z],axis=1)
    return z

def qmc_p1(nsamps,M,ndims):
    # generate nsamps shifted radiuses [0,1]
    radius0 = np.arange(0,1,1/M)
    radius  = radius0 + np.random.rand(nsamps,M)/M

    # convert to real radius
    r = sqrt(chi2.ppf(radius,ndims))
    
    # generate spherical samples
    u = randn(nsamps,M,ndims)
    u = u / expand_dims(np.sqrt(np.sum(u**2,axis=2)),axis=2)

    z = expand_dims(r,axis=2) * u
    return z
