#include <math.h>
#include <iostream.h>
#include "Utils.h"
#include "FourVector.h"
#include "random.h"
#include "GenEvent.h"
#include "SimEvent.h"
#include "toymc.h"

SimEvent simulate_det (int* seedPtr, GenEvent genEvt){

  // smear jet energies

  FourVector P1_det, P2_det, P3_det, P4_det;
  P1_det = SmearJetMom(seedPtr, genEvt.P1());
  P2_det = SmearJetMom(seedPtr, genEvt.P2());
  P3_det = SmearJetMom(seedPtr, genEvt.P3());
  P4_det = SmearJetMom(seedPtr, genEvt.P4());

  //  Generate Btag values (p-value for all tracks to come from
  //  primary vertex.

  float btag1 = bTag (seedPtr, genEvt.flav1());
  float btag2 = bTag (seedPtr, genEvt.flav2());
  float btag3 = bTag (seedPtr, genEvt.flav3());
  float btag4 = bTag (seedPtr, genEvt.flav4());

  SimEvent evt (P1_det, btag1, P2_det, btag2, P3_det, btag3, P4_det, btag4);
  return evt;

}

float bTag (int* seedPtr, int flav){

  // Quick kludge for btag values.  This is the p-value
  // that all tracks in the jet originate from the primary vertex.

  float pVal;
  float r = random (*seedPtr);
  int f = abs(flav);
  if ( f <=3){                         // u, d, s
    pVal = r;
  }
  else if ( f == 4 ){                   // c
    pVal = -0.2 * log(r);  
      }
  else if ( f == 5 ){ 
    pVal = -0.05 * log(r);
  }
  if ( pVal > 1 ) { pVal = 1.; }

  return pVal;

}

FourVector SmearJetMom (int* seedPtr, FourVector P){
  
  float const Eres = 0.6;           // typical HCAL
  float E = P.t();
  float r = GaussRan(seedPtr);
  float E_det = E * (1. +  r * Eres / sqrt (E));

  // scale momenta to keep jet mass same

  float px = P.x();
  float py = P.y();
  float pz = P.z();
  float mass = P.m();;
  float pmag = sqrt (pow(px,2) + pow(py,2) + pow(pz,2));
  float f = sqrt (pow(E_det,2) - pow(mass,2)) / pmag;
  FourVector P_out (f*px, f*py, f*pz, E_det);
  return P_out;

}