#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- -> HZ with H -> bbbar
// and Z -> qqbar.

GenEvent generate_HZ (int* seedPtr, float Ecm, float mh, float fmax_u, 
                     float fmax_d)
{

  // generate flavour for Z decay

  int flav = get_flav (seedPtr);
  float fmax;
  if ( flav%2 == 1 ) {
    fmax = fmax_u; 
  }
  else {
    fmax = fmax_d; 
  }

  // generate virtual mass of Z

  const float mz = 91.2;
  const float gammaz = 2.5;
  const float s2tw = 0.23;
  const float pi = 3.14159265;
  bool findRan = true;
  float CRan;
  while ( findRan ){
    CRan = CauchyRan(seedPtr);
    findRan = CRan > 10. || CRan < -10. ;
  }
  float mzstar = mz  + CRan*gammaz/2.;

  // generate scattering angle and decay angles for Z -> qqbar

  float phi = 2. * pi * random(*seedPtr);
  float ct, cts, phi_star;
  bool keep_trying = true;
  while ( keep_trying ){

    ct  = 2. * random(*seedPtr) - 1;
    cts = 2. * random(*seedPtr) - 1;
    phi_star = 2. * pi * random(*seedPtr);
    float diffxs = diffxs_HZ (Ecm, mh, mzstar, s2tw, 
			    flav, ct, cts, phi_star);
    float u = random(*seedPtr) * fmax;
    if ( u < diffxs ) { keep_trying = false; }

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

  //  Assign components of momentum vector for Z.

  FourVector P4_Z = twoBodyDecay(Ecm, mzstar, mh, ct, phi);
  ThreeVector betaGamma_Z = P4_Z.betaGamma();

  // Assign momentum components of fermion pairs from decay 
  // Z -> ffbar in rest frame of Z, then boost back to lab frame.  
  // Assume massless fermions.

  FourVector P4_f_star = twoBodyDecay(mzstar, 0., 0., cts, phi_star);
  FourVector P4_fbar_star = twoBodyDecay(mzstar, 0., 0., -cts, phi_star+pi);
  ThreeVector bg_to_lab (-betaGamma_Z.x(), -betaGamma_Z.y(),
                         -betaGamma_Z.z()); 
  FourVector P4_f = boostVector(P4_f_star, bg_to_lab);
  FourVector P4_fbar = boostVector(P4_fbar_star, bg_to_lab);

  // Repeat for Higgs decay (isotropic).

  FourVector P4_H = twoBodyDecay(Ecm, mh, mzstar, -ct, phi+pi);
  ThreeVector betaGamma_H = P4_H.betaGamma();
  phi_star = 2. * pi * random(*seedPtr);
  cts = 2. * random(*seedPtr) - 1;
  FourVector P4_b_star = twoBodyDecay(mh, 0., 0., cts, phi_star);
  FourVector P4_bbar_star = twoBodyDecay(mh, 0., 0., -cts, phi_star+pi);
  bg_to_lab.set (-betaGamma_H.x(), -betaGamma_H.y(),
                         -betaGamma_H.z()); 
  FourVector P4_b = boostVector(P4_b_star, bg_to_lab);
  FourVector P4_bbar = boostVector(P4_bbar_star, bg_to_lab);
  
  // Assign values to output object

  GenEvent evt(P4_f, flav, P4_fbar, -flav, P4_b, 5, P4_bbar, -5);
  evt.shuffle_jets();

  return evt;

}

float get_fmax_HZ (int* seedPtr, int num_shots, int flav, 
                   float Ecm, float mh){

  const float mz = 91.2;
  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;
    cts = 2. * random(*seedPtr) - 1;
    phi_star = 2. * pi * random(*seedPtr);
    float diffxs = diffxs_HZ (Ecm, mh, mz, s2tw, 
			    flav, ct, cts, phi_star);
    if ( diffxs > fmax ) { fmax = diffxs; }
  }

  return fmax;

}

int get_flav (int* seedPtr){

  const float s2tw = 0.23;
  const float au = -1.;
  const float vu = -1. + (8./3.)*s2tw;
  const float ad = -1.;
  const float vd =  -1. + (4./3.)*s2tw;
  const float gamma_up = pow(vu,2) + pow(au,2);
  const float gamma_down = pow(vd,2) + pow(ad,2);
  const float tot = 2.*gamma_up + 3.*gamma_down;
  const float up = gamma_up / tot;
  const float down = gamma_down / tot;

  int flav;
  float r = random(*seedPtr);
  if ( r < up ){ 
    flav = 1; }
  else if ( r < up + down ){ 
    flav = 2; }
  else if ( r < 2.*up + down ){ 
    flav = 3; }
  else if ( r < 2.*up + 2.*down ){ 
    flav = 4; }
  else {
    flav = 5; }

  return flav;

}


// Differential cross section for cos_theta, cos_theta_star and phi_star
// from V. Barger et al., Phys Rev D 49 (1994) 79, equation (13).

float diffxs_HZ (float Ecm, float mh, float mz, float s2tw, 
                 int flav, float ct, float cts, float phi_star){

  float s = pow(Ecm,2);
  float beta2 = ( 1. - ( pow(mh+mz, 2) / s ) ) * 
                ( 1. - ( pow(mh-mz, 2) / s ) );
  float gamma = 1. / sqrt (1. - beta2);

  float ct2 = pow(ct,2);
  float cts2 = pow(cts,2);
  float st2 = 1. - ct2;
  float st = sqrt(st2);
  float sts2 = 1. - cts2;
  float sts = sqrt(sts2);
  float s2t = 2. * st * ct;
  float s2ts = 2. * sts * cts;
  float cfs = cos (phi_star);
  float c2fs = cos (2.*phi_star);

  float ve = -1. + 4.*s2tw;
  float ae = -1.; 
  float vf;
  if ( flav%2 == 1 ) {                    // up type
    vf = -1. + (8./3.)*s2tw; }
  else {                                  // down type
    vf =  -1. + (4./3.)*s2tw;
  }
  float af = -1.;
  float couplings = ( (2.*ve*ae) / (pow(ve,2) + pow(ae,2)) ) * 
                    ( (2.*vf*af) / (pow(vf,2) + pow(af,2)) );

  float diffxs = st2 * sts2 
    - s2t * s2ts * cfs / (2. * gamma)
    + ((1 + ct2)*(1. + cts2)  +  st2 * sts2 * c2fs) / (2.* pow(gamma,2))
    - couplings  * (2./gamma) * (st * sts * cfs  -  ct * cts / gamma);

  return diffxs;

}

FourVector twoBodyDecay(float M, float m1, float m2, float ct, float phi){

  //  Returns four momentum of daughter 1 in rest frame of mother
  //  particle for decay M -> m1 + m2, where particle 1 has cosine of 
  //  polar angle ct and azimuthal angle phi.

  float E_1  = (pow(M,2)  +  pow(m1,2) - pow(m2,2) ) / (2.*M);
  float p_1  =   sqrt ( pow(E_1,2) - pow(m1,2) );
  float pz_1 =   p_1 * ct;
  float pt_1 =   sqrt (pow(p_1,2) - pow(pz_1,2));
  float px_1 =   pt_1 * cos(phi);
  float py_1 =   pt_1 * sin(phi);
  FourVector P4_1 (px_1, py_1, pz_1, E_1);
  return P4_1;

}