#!/usr/bin/env python

"""
Calculation of Drell-Yan production of lepton pair
 p p(bar) -->  l+ l-  X

author: kkumer@phy.hr 2011-07-12 
"""


from __future__ import division # don't truncate int division like in C
import sys
from math import pi, sqrt

import numpy as np
from scipy.integrate import quad as integral_numerical


# Initialize PDFs
try:
    import lhapdf  # Les Houches Accord PDF library
    lhapdf.setVerbosity(0)
    lhapdf.initPDFSetByName('cteq66.LHgrid')
    #lhapdf.initPDFSetByName('CT10.LHgrid')
    lhapdf.initPDF(int(0));
    have_lhapdf = True
except ImportError:
    sys.stderr.write('Warning: LHPDF library missing. Using super-simple PDFs!!\n')
    have_lhapdf = False


# some constants
Nc = 3;
sw2 = 0.234;  # at MZ
cw2 = 1 - sw2;
cw = sqrt(cw2);
sw = sqrt(sw2);
alpha = 1/128.93  # at MZ
ee2 = 4*pi*alpha
MZ = 91.19;
Zwidth = 2.5;
MW = 80.403;
g2 = sqrt(ee2)/sqrt(sw2);
sLHC = (7000.+7000.)**2;
#sTEV = (980.+980.)**2;
sTEV = 1800**2
GeV2nb = 389379.;  # GeV -> nb conversion factor
GeV2fb = GeV2nb*1e6;
GeV2pb = GeV2nb*1e3;



class fermion():
    
    def __init__(self, ind, Q=None, I3=None):
        """ 
        ind - corresponding index in LHAPDF array, -1 for leptons
        """
        self.ind = ind
        self.Q = Q
        
        self.vg = -Q
        self.ag = 0
        self.vZ = (I3 - 2*sw2*Q)/(2*sw*cw)
        self.aZ = I3/(2*sw*cw)
        
        # lambda_f^ij coefficients
        self.lamp = {'gg' : self.vg**2+self.ag**2, 
                     'gZ' : self.vg*self.vZ+self.ag*self.aZ,
                     'ZZ' : self.vZ**2+self.aZ**2, 
                     'Zg' : self.vg*self.vZ+self.ag*self.aZ}        
        
    def pdf(self, x, Qfac):
        """Return PDF f(x). Not x*f(x)!"""
        if self.ind>=0:
            # "quark"
            if have_lhapdf:
                return lhapdf.xfx(float(x), float(Qfac))[self.ind]/x
            else:
                # just simple u-quark PDF, rest are zero
                if self.ind == 8:
                    return float(x)**(-0.962)*(1.-float(x))**2.36
                else:
                    return 0
        else:
            # "lepton", has no PDF
            return None


# creating particles instances
sbar = fermion(3, Q=1/3, I3=1/2)      
ubar = fermion(4, Q=-2/3, I3=-1/2) 
dbar = fermion(5, Q=1/3, I3=1/2)      
d = fermion(7, Q=-1/3, I3=-1/2)
u = fermion(8, Q=2/3, I3=1/2)
s = fermion(9, Q=-1/3, I3=-1/2)
el = fermion(-1, Q=-1, I3=1/2)  # electron

#sys.stderr.write('u.pdf(0.1, 10) = %f\n' % u.pdf(0.1, 10))

def flavourint(f, fbar, tau=10**2/1800., S=1800., antip=True):
    """Flavour's f=u,d,... contribution  to integral \int_tau^1 dx1/x1 F_f(x1, x2=tau/x1).
    
       antip -- do we have proton-antiproton scattering?
                (Tevatron: True, LHC: False)

    """
    QF = sqrt(tau*S)  # factorizaton scale
    
    if antip:
        # p-pbar collider
        return integral_numerical(lambda x1: (f.pdf(x1, QF)*f.pdf(tau/x1, QF) +
                 fbar.pdf(x1, QF)*fbar.pdf(tau/x1, QF))/x1, tau, 1, limit=100, epsrel=0.001)[0]
    else:
        # p-p collider
        assert(have_lhapdf)  # don't do this without lhapdf library!
        return integral_numerical(lambda x1: (f.pdf(x1, QF)*fbar.pdf(tau/x1, QF) +
                 fbar.pdf(x1, QF)*f.pdf(tau/x1, QF))/x1, tau, 1, limit=100, epsrel=0.001)[0]


def ampsq(lepton, f, fbar, b1='g', b2='g', M=10, S=sTEV, antip=True):
    """Squared amplitude without common prefactor.

       M -- dilepton invariant mass M_ll
       S -- Mandelstam S of proton-(anti)proton scattering
    
    """
    prop = {'g': 1/M**2, 'Z': 1/(M**2 - MZ**2 + 1j*MZ*Zwidth)}
    #sys.stderr.write('lepton.lamp[gg] = %g\n' % lepton.lamp[b1+b2])
    #sys.stderr.write('f.lamp[gg] = %g\n' % f.lamp[b1+b2])
    return ((prop[b1].conjugate() * prop[b2]).real * lepton.lamp[b1+b2] * 
                 flavourint(f, fbar, tau=M**2/S, S=S, antip=antip) * f.lamp[b1+b2])


print "Typical contribution of various flavours:"
ua = ampsq(el, u, ubar, 'g', 'g', M=10., S=1800., antip=True)
da = ampsq(el, d, dbar, 'g', 'g', M=10., S=1800., antip=True)
sa = ampsq(el, s, sbar, 'g', 'g', M=10., S=1800., antip=True)
print "u: %.2g %%" % (100*ua/(ua+da+sa))
print "d: %.2g %%" % (100*da/(ua+da+sa))
print "s: %.2g %%" % (100*sa/(ua+da+sa))


def sigma(M, S=sTEV, antip=True):
    """Cross section dsigma/dM for pp->N+N- in femtobarns."""
   
    sigma0 = 8*pi*alpha**2*M**3/(9*S)
    bosons = ['g', 'Z']
    quarks = [(u, ubar), (d, dbar), (s, sbar)]
    tot = 0
    for b1 in bosons:
        for b2 in bosons:
            for f, fbar in quarks:
                tot += ampsq(el, f, fbar, b1, b2, M, S, antip)
    return sigma0*tot*GeV2pb

# testing for two characteristic values
print '-------------------'
for M in [10., 300.]:
    print "M = %f   ds/dM = %g" % (M, sigma(M))

# doing for the whole plot
Ms = list(np.linspace(10, 80, 40)) + list(np.linspace(80, 110, 40)) +\
        list(np.linspace(110, 300, 30))
f = open('DY.dat', 'w')
for M in Ms:
    f.write('%f  %g\n' % (M, sigma(M)))
f.close()