import pycuda.autoinit
import pycuda.driver as cuda
from pycuda.compiler import SourceModule
from pycuda import gpuarray
import numpy as np
import ROOT
import os

current_directory = os.path.split(os.path.realpath(__file__))[0]
source_directory = current_directory + '/../src'

source = open(current_directory + '/test_sample_cdf.cu').read()

mod = SourceModule(source, options=['-I' + source_directory], no_extern_c=True, cache_dir=False)

test_sample_cdf = mod.get_function('test_sample_cdf')

def compare_sampling(hist, reps=10):
    nbins = hist.GetNbinsX();
    xaxis = hist.GetXaxis()
    intg = hist.GetIntegral()
    cdf_y = np.empty(nbins+1, dtype=float)
    cdf_x = np.empty_like(cdf_y)

    cdf_x[0] = xaxis.GetBinLowEdge(1)
    cdf_y[0] = 0.0
    for i in xrange(1,len(cdf_x)):
        cdf_y[i] = intg[i]
        cdf_x[i] = xaxis.GetBinUpEdge(i)

    cdf_x_gpu = gpuarray.to_gpu(cdf_x.astype(np.float32))
    cdf_y_gpu = gpuarray.to_gpu(cdf_y.astype(np.float32))
    block =(128,1,1)
    grid = (128, 1)
    out_gpu = gpuarray.empty(shape=int(block[0]*grid[0]), dtype=np.float32)

    out_h = ROOT.TH1D('out_h', '', hist.GetNbinsX(), 
                      xaxis.GetXmin(),
                      xaxis.GetXmax())
    out_h.SetLineColor(ROOT.kGreen)

    for i in xrange(reps):
        test_sample_cdf(np.int32(i),
                        np.int32(len(cdf_x_gpu)), 
                        cdf_x_gpu, cdf_y_gpu, out_gpu, block=block, grid=grid)
        out = out_gpu.get()
        for v in out:
            out_h.Fill(v)

    prob = out_h.KolmogorovTest(hist)
    return prob, out_h 

def test_sampling():
    '''Verify that the CDF-based sampler on the GPU reproduces a binned
    Gaussian distribution'''
    f = ROOT.TF1('f_gaussian', 'gaus(0)', -5, 5)
    f.SetParameters(1.0/np.sqrt(np.pi * 2), 0.0, 1.0)
    gaussian = ROOT.TH1D('gaussian', '', 100, -5, 5)
    gaussian.Add(f)

    prob, out_h = compare_sampling(gaussian, reps=50)

    assert prob > 0.01