import random
import matplotlib.pyplot as plt
import numpy as np

# Random numerical factors
C_F = 12./5.
alpha_s = .12

def sample_stSpace(num):
    """Samples points in the triangle    0 < s < 1,  0 < t < 1-s"""
    st = []
    numTotal = 0
    numInside = 0

    while(len(st) < num):
        s,t = random.random(),random.random()
        srange = 0 < s and s < 1
        trange = 0 < t and t < 1 - s
        if(srange and trange):
            st.append([s,t])
            numInside+=1.
        numTotal+=1.
    st = np.array(st)
    # colors = eventWeight(st[:,0],st[:,1]); plt.scatter(st[:,0],st[:,1],c=colors,s = min(10,1000/num)); plt.show()

    return st

def eventWeight(s,t):
    """returns dsigma/(dsdt) in terms of s and t -- the `pdf' (not normalizable) in s and t space"""
    return (C_F * alpha_s/(2.*np.pi) ) * ( (s**2. - 2.*s + t**2. - 2.*t + 2)/(s * t) )


def getThrusts(numbins, numSamplePoints):
    """Generates events in st space and returns an np histogram of the variable that we're interested in: looking for the pdf of the function of s and t given by min(s,t,1-s-t)
    """
    st = sample_stSpace(numSamplePoints)

    s = st[:,0];  t = st[:,1]
    thrusts         = np.min(np.array([s,t,1-s-t]), axis=0)
    weights         = np.array( eventWeight(s,t)     * (3.* numbins/2.)     / numSamplePoints )
    # total weight of a bin = dsigma/dT = <dsigma/(ds dt)>_bin / bin_width  --> sum_bins dsigma/dT * bin_width = sigma_total
    # THIS IS WHERE I INSERT A FACTOR OF 2 BY HAND: the random variable goes from 0 to 1/3, so the bin width is total_width/numbins = (3 * numbins). I need a factor of 2 to make things right :'(

    squared_weights = np.array( eventWeight(s,t)**2. * (3.* numbins/2.)**2. / numSamplePoints )

    return thrusts, weights, squared_weights


def MC_thrust_Distribution_Plot(numbins, numSamplePoints):
    """Plots the analytic and MC `pdf' of the variable min(s,t,1-s-t)"""
    thrusts, weights, squared_weights = getThrusts(numbins, numSamplePoints)

    # Data for Differential Cross Section:
    diff_xsec, bins   = np.histogram(thrusts, bins=numbins, range=(0,1./3.),weights=weights)
    squared_counts, _ = np.histogram(thrusts, bins=numbins, range=(0,1./3.),weights=squared_weights)
    num_in_bin, _     = np.histogram(thrusts, bins=numbins, range=(0,1./3.))

    diff_xsec_error = np.sqrt( squared_counts - np.square(diff_xsec) ) / np.sqrt(num_in_bin)

    # Analytic Expressions:
    diff_xsec_τ_Analytic = lambda τ: C_F*alpha_s*( 3*(1+τ)*(3*τ-1)/τ + (4 + 6*τ*(τ-1))*np.log((1-2*τ)/(τ)) / ((1 - τ)*τ) ) / (2. * np.pi)
    analytic_diff_xsec = np.array( diff_xsec_τ_Analytic(np.linspace(.0001,1./3.,numbins)) )

    #MC Plot
    plt.errorbar(x=(bins[:-1]+bins[1:])/2., y=diff_xsec,
                  xerr=(bins[:-1]-bins[1:])/2., yerr=diff_xsec_error, color='orange', label='MC EMD diff_xsec')
    #Analytic plot
    plt.plot(np.linspace(.0001,1./3.,numbins), analytic_diff_xsec, 'b--', label='Analytic Result')

    plt.show()


if __name__ == "__main__":
    MC_thrust_Distribution_Plot(100, 100000)
    plt.show()