#include <iostream.h>
#include <math.h>
#include <cfortran/cfortran.h>
#include "cfortran/hbook.h"
#include "ThreeVector.h"
#include "FourVector.h"
#include "Utils.h"
#include "random.h"
#include "toymc.h"

// Generates a four-jet event from e+e- -> ZZ 

GenEvent generate_ZZ (int* seedPtr, float Ecm, float fmax)
{

  const float mz = 91.2;
  const float gammaz = 2.5;
  const float pi = 3.14159265;

  // Get virtual Z masses

  bool findRan = true;
  float CRan;
  while ( findRan ){
    CRan = CauchyRan(seedPtr);
    findRan = CRan > 10. || CRan < -10. ;
  }
  float mzstar1 = mz  + CRan*gammaz/2.;

  findRan = true;
  while ( findRan ){
    CRan = CauchyRan(seedPtr);
    findRan = CRan > 10. || CRan < -10. ;
  }
  float mzstar2 = mz  + CRan*gammaz/2.;  

  // Get scattering angle using on-shell Zs.

  float phi = 2. * pi * random(*seedPtr);
  float ct;
  bool keep_trying = true;
  while ( keep_trying ){
    ct  = 2. * random(*seedPtr) - 1;
    float diffxs = diffxs_ZZ (Ecm, mz, ct);
    float u = random(*seedPtr) * fmax;
    if ( u < diffxs ) { keep_trying = false; }
  }

  //  Assign components of momentum vector for first Z and its
  // decay products.

  float cts, phi_star;
  cts = 2. * random(*seedPtr) - 1;
  phi_star = 2. * pi * random(*seedPtr);

  FourVector P4_1 = twoBodyDecay(Ecm, mzstar1, mzstar2, ct, phi);
  ThreeVector betaGamma_Z1 = P4_1.betaGamma();

  FourVector P4_f_star_1 = twoBodyDecay(mzstar1, 0., 0., cts, phi_star);
  FourVector P4_fbar_star_1 = twoBodyDecay(mzstar1, 0., 0., -cts, phi_star+pi);
  ThreeVector bg_to_lab (-betaGamma_Z1.x(), -betaGamma_Z1.y(),
                         -betaGamma_Z1.z()); 
  FourVector P4_f_1 = boostVector(P4_f_star_1, bg_to_lab);
  FourVector P4_fbar_1 = boostVector(P4_fbar_star_1, bg_to_lab);

  // Repeat for other Z decay (take also isotropic and uncorrelated to first).

  FourVector P4_2 = twoBodyDecay(Ecm, mzstar2, mzstar1, -ct, phi+pi);
  ThreeVector betaGamma_Z2 = P4_2.betaGamma();
  phi_star = 2. * pi * random(*seedPtr);
  cts = 2. * random(*seedPtr) - 1;
  FourVector P4_f_star_2 = twoBodyDecay(mzstar2, 0., 0., cts, phi_star);
  FourVector P4_fbar_star_2 = twoBodyDecay(mzstar2, 0., 0., -cts, phi_star+pi);
  bg_to_lab.set (-betaGamma_Z2.x(), -betaGamma_Z2.y(),
                         -betaGamma_Z2.z()); 
  FourVector P4_f_2 = boostVector(P4_f_star_2, bg_to_lab);
  FourVector P4_fbar_2 = boostVector(P4_fbar_star_2, bg_to_lab);
  
  // Generate quark flavours

  int flav1 = get_flav (seedPtr);
  int flav2 = get_flav (seedPtr);

  // Assign values to output object

  GenEvent evt(P4_f_1, flav1, P4_fbar_1, -flav1, 
               P4_f_2, flav2, P4_fbar_2, -flav2);

  // diagnostic histograms

  HF1 (101, phi, 1.);
  HF1 (102, ct, 1.);
  HF1 (103, phi_star, 1.);
  HF1 (104, cts, 1.);
  HF1 (106, mzstar1, 1.);
  HF1 (106, mzstar2, 1.);

  return evt;

}

float get_fmax_ZZ (int* seedPtr, int num_shots, float Ecm, float mz){

  const float s2tw = 0.23;
  const float pi = 3.14159265;

  float ct, cts, phi_star;
  float fmax = 0.;

  for (int i = 0; i<num_shots; i++){
    ct  = 2. * random(*seedPtr) - 1;
    float diffxs = diffxs_ZZ (Ecm, mz, ct);
    if ( diffxs > fmax ) { fmax = diffxs; }
  }
  return fmax;
}

// Differential cross section for cos_theta, from Mikaelian
// et al., PRD 19 (1979).

float diffxs_ZZ (float Ecm, float mz, float ct){

  float s = pow(Ecm,2);
  float mz2 = pow(mz,2);
  float Ee = Ecm/2.;
  float EZ = Ecm/2.;
  float pe = Ee;
  float pZ = sqrt (pow(EZ,2)  - mz2);
  float t = mz2 - 2.*(Ee*EZ - pe*pZ*ct);
  float u = 2.*mz2 - s - t;

  float diffxs = t/u  +  u/t  +  (4.*mz2*s)/(u*t)  -
                 pow(mz,4) * ( 1./pow(t,2)  +  1./pow(u,2) );
  return diffxs;

}