import random
import numpy as np


def testHistFn2D():
    numSample   = 10000
    xs = np.array( [random.random() for i in range(numSample)] )
    ys = np.array( [random.random() for i in range(numSample)] )
    numbins       = 10
    test_fn_1    = xs + ys

    # MC, the histogram becomes a pdf by weighting the points by 1/numSample * (1/size of a bin)
    h0,b0 = np.histogram(test_fn_1, numbins, weights=numbins/(2.*numSample)*np.ones(len(xs)))
    plt.errorbar(x=(b0[:-1]+b0[1:])/2., y=h0, xerr=(b0[:-1]-b0[1:])/2.);

    # Actual pdf for x + y
    pdf = lambda xi: (xi<1)*xi + (xi>1)*(2.-xi)
    plt.plot(np.linspace(0.,2.,40), pdf(np.linspace(0.,2.,40)))

    # Summing MC histogram, i.e. sum_bins pdf(bin)*(size of bin), should return one
    print( np.sum( h0 )*2./numbins )

    #compare the two
    plt.show()

def testHistFn3D():
    numSample   = 10000
    xs = np.array( [random.random() for i in range(numSample)] )
    ys = np.array( [random.random() for i in range(numSample)] )
    zs = np.array( [random.random() for i in range(numSample)] )
    numbins       = 10
    test_fn_1    = xs + ys + zs

    # MC
    h0,b0 = np.histogram(test_fn_1, numbins, weights=numbins/(3.*numSample)*np.ones(len(xs)))
    plt.errorbar(x=(b0[:-1]+b0[1:])/2., y=h0, xerr=(b0[:-1]-b0[1:])/2.);

    # Analytic
    pdf = lambda xi: (4./9.)*( (xi<3./2.)*xi + (xi>3./2.)*(3.-xi))
    plt.plot(np.linspace(0.,3.,40), pdf(np.linspace(0.,3.,40)))

    # Should return 1
    print( np.sum( h0 )*3./numbins )

    # Compare
    plt.show()


if __name__ == "__main__":
   testHistFn2D()
   testHistFn3D()