#! /usr/bin/python2
########################################################################
#
# Einfuehrung in die Theoretische Teilchenphysik
# Karlsruher Institut fuer Technologie, WS 2017/18
#
# Sheet 12, Exercise 1, Higgs-Produktion via W Fusion
#
# Author: Michael Rauch <michael.rauch@kit.edu>
# Last change: January 16 2018
#

import sys
import random
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages

# input values
## number of integration points
numpoints=int(1e8)
## ILC as example
sqrts=500.
## Higgs mass
mh=125.
## Fermi constant for the weak coupling
GF=1.16638e-5
## W mass
mw=80.385

# plots
xbinrange = { 
  'x3': (0,1),
  'x4': (0,1),
  'pt_nu': (0,500), 
  'pt_nubar': (0,250), 
  'pt_H': (0,500), 
  'y_nu': (-5,5),
  'y_nubar': (-5,5),
  'y_H': (-2.5,2.5),
  'm_nunubar': (0,1000),
  'deltay_nunubar': (-10,10),
  'deltaphi_nunubar': (-np.pi,np.pi)
}
nbins = 100
bins = { }
step = { }
distributions = { }
val = { }
for key in xbinrange.keys():
  low,up = xbinrange[key]
  bins[key],step[key] = np.linspace(low,up,nbins+1,retstep=True)
  distributions[key] = np.zeros(nbins+2)

# derived quantities
## vacuum expectation value of the Higgs field
v=1./np.sqrt(np.sqrt(2.)*GF)
## weak coupling constant
g=np.sqrt(4*np.sqrt(2)*GF*mw**2)
## kinematic lower bound on x_3
onemsqs=1-(mh/sqrts)**2

# global rotation by the azimuthal angle phi_3 around the z axis
## The matrix element is independent of this angle, so the integral has been
## taken and the corresponding factor 2 pi put into the prefactor of the
## phase-space integral. This only randomizes the momenta, so we do not have a
## preferred direction in the x-y-plane.
def azrotate(mom,phi):
  mom[1],mom[2] = np.cos(phi)*mom[1]-np.sin(phi)*mom[2], np.cos(phi)*mom[2]+np.sin(phi)*mom[1]

# dot product for four vectors
def dotprod(mom1,mom2):
  return mom1[0]*mom2[0]-mom1[1]*mom2[1]-mom1[2]*mom2[2]-mom1[3]*mom2[3]

# momentum array
## p_1 and p_2 is fixed, p_3, p_4 and p_5 will be filled per phase-space point
mom = [
  np.array([sqrts/2.,0,0, sqrts/2.],float),
  np.array([sqrts/2.,0,0,-sqrts/2.],float),
  np.array([0,0,0,0],float),
  np.array([0,0,0,0],float),
  np.array([0,0,0,0],float)
]

crosssection = 0
error = 0

# loop over phase-space points
for pt in range(numpoints):
  if (pt+1)%10000==0:
    sys.stderr.write("\rEvent {} / {}".format(pt+1,numpoints))
  # global factor from phase-space integral
  weight = sqrts**2/(512*np.pi**4)
  # convert random numbers to phase space integration variables
  # keep track of the jacobians from transforming the integrals to integrals on
  # the unit hypercube
  x3 = onemsqs*random.random()
  weight *= onemsqs
  x4 = onemsqs-x3 + random.random()*(1-(mh/sqrts)**2/(1.-x3) - (onemsqs-x3))
  weight *= 1-(mh/sqrts)**2/(1.-x3) - (onemsqs-x3)
  costheta3 = 2.*random.random()-1.
  weight *= 2.
  phi4 = 2.*np.pi*random.random()
  weight *= 2.*np.pi
  # derived variables
  sintheta3 = np.sqrt(1.-costheta3**2)
  costheta4 = 1+(2.*(onemsqs-x3-x4))/(x3*x4)
  sintheta4 = np.sqrt(1.-costheta4**2)

  # generate momenta
  mom[2][0] = sqrts/2.*x3
  mom[2][1] = mom[2][0]*sintheta3
  mom[2][2] = 0.
  mom[2][3] = mom[2][0]*costheta3
  mom[3][0] = sqrts/2.*x4
  mom[3][1] = mom[3][0]*(sintheta3*costheta4+costheta3*sintheta4*np.cos(phi4))
  mom[3][2] = mom[3][0]*sintheta4*np.sin(phi4)
  mom[3][3] = mom[3][0]*(costheta3*costheta4-sintheta3*sintheta4*np.cos(phi4))
  # global azimuthal-angle rotation
  phi3 = 2.*np.pi*random.random()
  azrotate(mom[2],phi3)
  azrotate(mom[3],phi3)

  # calculate p_5 via global 4-momentum conservation
  mom[4] = mom[0]+mom[1]-mom[2]-mom[3]

  # momenta of the exchanged W bosons
  q1 = mom[0]-mom[2]
  q2 = mom[1]-mom[3]

  # matrix element as calculated in the exercise
  me = g**8*v**2*(
    dotprod(mom[0],mom[3])*dotprod(mom[1],mom[2])
    /(dotprod(q1,q1)-mw**2)**2
    /(dotprod(q2,q2)-mw**2)**2
  )/4.
  weight *= me
  # flux factor
  weight /= (4.*dotprod(mom[0],mom[1]))
  weight /= numpoints
  # 1/(hbar c)^2 - convert from GeV^2 (natural units) to femtobarn
  weight *= 3.89379e11

  crosssection += weight
  error += weight**2

  # store values for histograms
  val['x3'] = x3
  val['x4'] = x4
  val['pt_nu'] = np.sqrt(mom[2][1]**2+mom[2][2]**2)
  val['pt_nubar'] = np.sqrt(mom[3][1]**2+mom[3][2]**2)
  val['pt_H']  = np.sqrt(mom[4][1]**2+mom[4][2]**2)
  val['y_nu'] = np.log((mom[2][0]+mom[2][3])/(mom[2][0]-mom[2][3]))/2.
  val['y_nubar'] = np.log((mom[3][0]+mom[3][3])/(mom[3][0]-mom[3][3]))/2.
  val['y_H']  = np.log((mom[4][0]+mom[4][3])/(mom[4][0]-mom[4][3]))/2.
  val['m_nunubar'] = np.sqrt(2.*dotprod(mom[2],mom[3]))
  val['deltay_nunubar'] = val['y_nu']-val['y_nubar']
  val['deltaphi_nunubar'] = np.arctan2(mom[2][2],mom[2][1])-np.arctan2(mom[3][2],mom[3][1])
  if val['deltaphi_nunubar'] > np.pi:
    val['deltaphi_nunubar'] = val['deltaphi_nunubar'] - 2*np.pi
  elif val['deltaphi_nunubar'] < -np.pi:
    val['deltaphi_nunubar'] = val['deltaphi_nunubar'] + 2*np.pi

  for key in distributions.keys():
    distributions[key][np.digitize(val[key],bins[key])] += weight

# error = sqrt( (<x^2>-<x>^2)/N )
error = np.sqrt(error-crosssection**2/numpoints)

sys.stderr.write("\n")
print("Cross section: {} +/- {} fb".format(crosssection,error))

pdfout = PdfPages('WFusionHiggs.pdf')

for key in distributions.keys():
  plt.hist(bins[key][0:nbins]+0.5*step[key], bins=nbins, range=xbinrange[key], weights=distributions[key][1:nbins+1])
  plt.xlabel(key)
# save to pdf
  pdfout.savefig()
# show on screen
#  plt.show()
  plt.close()

pdfout.close()