#include <iostream>
#include <math.h>

#include "TFile.h"
#include "TRandom.h"
#include "TH1.h"

using namespace std;

main() {

  TFile *f = new TFile("e791.root","RECREATE");
  TH1F *h_xf = new TH1F("xf","x_{F}",100,0,1);
  TH1F *h_pzcm = new TH1F("pzcm","p_{z} in CM frame (GeV/c)",100,0,50);
  TH1F *h_pz = new TH1F("pz","p_{z} in lab frame (GeV/c)",100,0,500);
  TH1F *h_len = new TH1F("len","decay length in lab frame (cm)",100,0,2);
  TH1F *h_costh = new TH1F("costh","cos theta (D0 frame)",100,-1,1);
  TH1F *h_theta = new TH1F("theta","theta (D0 frame)",180,0,180.0);
  TH1F *h_pzk = new TH1F("pzk","p_{z}(K) in lab frame (GeV/c)",100,0,500);
  TH1F *h_pzpi = new TH1F("pzpi","p_{z}(pi) in lab frame (GeV/c)",100,0,500);
  TH1F *h_thklab = new TH1F("thklab","theta(K) (lab frame)",100,0,5.0);
  TH1F *h_thpilab = new TH1F("thpilab","theta(pi) (lab frame)",100,0,5.0);
  TH1F *h_mrec = new TH1F("mrec","Reconstructed D0 mass",100,1.6,2.0);
  TH1F *h_mrecwide = new TH1F("mrecwide","Reconstructed D0 mass",100,0.0,5.0);
  TH1F *h_mwrong = new TH1F("mwrong","Wrong reconstructed D0 mass",100,0.0,5.0);
  TH1F *h_mwrongnarrow = new TH1F("mwrongnarrow","Wrong reconstructed D0 mass",
                                  100,1.6,2.0);

  const double ebeam = 500.0;        //  500 GeV pion beam.
  const double mn = 0.9383;          //  Nucleon mass in GeV.
  const double mpi = 0.1396;         //  Pion mass in GeV.
  const double md0 = 1.864;         //  D0 mass in GeV.
  const double mk = 0.495;           //  K mass in GeV.
  const double ctd0 = 123.7e-4;      //  D0 lifetime in cm;

  double ppi = sqrt(ebeam*ebeam-mpi*mpi);
  double s = (ebeam+mn)*(ebeam+mn) - ppi*ppi;
  double ecm = sqrt(s);        //  Center-of-mass energy.
  cout << "Centre-of-mass energy = " << ecm << " GeV" << endl;

  double ed0cm = (s+md0*md0)/(2*sqrt(s));
  double pzmaxcm = sqrt(ed0cm*ed0cm-md0*md0);

  cout << "Maximum D0 momentum in c.m. frame = " << pzmaxcm << " GeV" << endl;

  double ereccm = (s-md0*md0)/(2*sqrt(s));
  double g = (ebeam+mn)/(ed0cm+ereccm);
  double gb = sqrt(g*g-1);
  double b = gb/g;

  cout << "gamma = " << g << endl;
  cout << "beta = " << b << endl;
  cout << "gammabeta = " << gb << endl;
  cout << "pzmaxlab = " << pzmaxcm*g+gb*ed0cm << endl;
  cout << "pzreclab = " << ereccm*(g-gb) << endl;

  TRandom *rnd = new TRandom();

  double nin = 0;
  const int nevt = 100000;
  for ( int i=0; i<nevt; i++ ) {
    double xf;
    do {
      xf = rnd->Uniform();
    }
    while ( 1.5*rnd->Uniform() > 6*xf*(1-xf) );
    h_xf->Fill(xf,1);

    double pzcm = pzmaxcm*xf;
    h_pzcm->Fill(pzcm,1);
    double edcm = sqrt(pzcm*pzcm+md0*md0);

    double pz = pzcm*g+edcm*gb;
    h_pz->Fill(pz,1);

    double gb = pz/md0;
    double b = sqrt(gb*gb-1)/gb;
    double g = gb/b;

    double l = rnd->Exp(gb*ctd0);
    if ( l < 1 ) nin = nin + 1;
    h_len->Fill(l,1);

    double ek = (md0*md0-mpi*mpi+mk*mk)/(2*md0);
    double epi = (md0*md0+mpi*mpi-mk*mk)/(2*md0);
    double pk = sqrt(ek*ek-mk*mk);
    double ppi = sqrt(epi*epi-mpi*mpi);
    double costh = 2*rnd->Uniform()-1;
    double theta = acos(costh);

    h_costh->Fill(costh,1);
    h_theta->Fill(theta*180.0/M_PI,1);

    double pzk = gb*ek+g*pk*costh;
    double pxk = pk*sqrt(1-costh*costh);
    double pzpi = gb*epi-g*ppi*costh;
    double pxpi = -ppi*sqrt(1-costh*costh);
    double thklab = atan(pxk/pzk);
    double thpilab = atan(-pxpi/pzpi);
    h_thklab->Fill(thklab*180.0/M_PI,1);
    h_thpilab->Fill(thpilab*180.0/M_PI,1);
    h_pzk->Fill(pzk,1);
    h_pzpi->Fill(pzpi,1);

    double pklab = sqrt(pzk*pzk+pxk*pxk);       //  K momentum in lab frame
    double ppilab = sqrt(pzpi*pzpi+pxpi*pxpi);  //  pi momentum in lab frame
    double pkres = 0.0002*pklab;                //  K momentum resolution
    double ppires = 0.0002*ppilab;                //  pi momentum resolution
    double pkrec = pklab*(1+rnd->Gaus()*pkres);   //  Reconstructed k momentum  
    double ppirec = ppilab*(1+rnd->Gaus()*ppires); //  Reconstructed pi momentum
    double pzkrec = pkrec*cos(thklab);          //  Reconstructed z and x
    double pxkrec = pkrec*sin(thklab);          //  components in the lab
    double pzpirec = ppirec*cos(thpilab);
    double pxpirec = -ppirec*sin(thpilab);
 
// 
//  Calculate reconstructed energy... and invariant mass.
//
    double ekrec = sqrt(mk*mk+pzkrec*pzkrec+pxkrec*pxkrec);
    double epirec = sqrt(mpi*mpi+pzpirec*pzpirec+pxpirec*pxpirec);
    double ekwrong = sqrt(mpi*mpi+pzkrec*pzkrec+pxkrec*pxkrec);
    double epiwrong = sqrt(mk*mk+pzpirec*pzpirec+pxpirec*pxpirec);
    double erec = ekrec+epirec;
    double ewrong = ekwrong+epiwrong;
    double pxrec = pxkrec+pxpirec;
    double pzrec = pzkrec+pzpirec;
    double mrec = sqrt(erec*erec-pxrec*pxrec-pzrec*pzrec);
    double mwrong = sqrt(ewrong*ewrong-pxrec*pxrec-pzrec*pzrec);
    h_mrec->Fill(mrec,1);
    h_mrecwide->Fill(mrec,1);
    h_mwrong->Fill(mwrong,1);
    h_mwrongnarrow->Fill(mwrong,1);
  }

  cout << "P(l<1cm) = " << nin/nevt << endl;
  f->Write();
  f->Close();
}