#include "TLorentzVector.h"
#include "TRandom3.h"
#include "TH1D.h"
#include "TF1.h"
#include "TGenPhaseSpace.h"
#include "TCanvas.h"
#include "TStyle.h"
#include "TVector3.h"

static const Double_t Me = 0.511;  
static const Double_t Q = 2.5;	     
static const Double_t Mr = 19559;  
static const Double_t MI = Mr+Q+Me;  



/**
   maximum momentum 
 */

double p3Max(double m, double m1, double m2, double m3) { 
  double a = (m1+m2-m3);
  double b = (m1+m2+m3);
  double pmax = sqrt( (m*m -a*a)*(m*m - b*b) )/2/m;
  return pmax; 
}


/** 
 this is the correct function from excact phase space calculation
*/
double pDist( double * x, double * par  ) {
  
  double m = MI;
  double m1 = Mr;
  double m2 = 0;
  double m3 = Me;
  double p3 = x[0];
  double e3 = sqrt(p3*p3 + m3*m3);
  double pmax = p3Max(m,m1,m2,m3);
  if (p3 > pmax) return 0; 

  double a = ( m*m - 2*m*e3 + m3*m3 - (m1-m2)*(m1-m2) );
  double b = ( m*m - 2*m*e3 + m3*m3 - (m1+m2)*(m1+m2) );
  double c = 2* (m*m + m3*m3 - 2*m*e3);
  double fval = p3*p3/2./e3*sqrt(a*b)/c;
  return par[0]*fval;
}

double eDist( double * x, double * par  ) {
  
  double m3 = Me;
  double e3 = x[0];
  if ( e3 < m3) return 0;
  double p3 = sqrt(e3*e3 - m3*m3);
  double xp[1]; xp[0] = p3; 
  double pfval = pDist(xp,par);
  double fval = pfval*e3/p3;
  return fval;
}


void doit(UInt_t NrOfEvents = 1000000)   //***Number of events can be changed here or when you want to run by "doit(...)"
{
  if (gRandom) delete gRandom;
  gRandom = new TRandom3(0); 
  TLorentzVector* p[3] = {0,0,0};
  TGenPhaseSpace* rg = new TGenPhaseSpace();
  TLorentzVector* p0 = new TLorentzVector(0,0,0,MI); 
  Double_t masses[3] = {0,Me,Mr};
  rg->SetDecay((*p0), 3, masses);
 
  TCanvas* betacan=new TCanvas("bcan","event generator for decay",300,10,900,800);  //***defining CANVAS
   
  TH1D* E1 = new TH1D("E1","E_1",200,0,3.5);
  TH1D* E1_2 = new TH1D("E1_2","E_1 (2)",200,2.5,3.1);
  TH1D* E2 = new TH1D("E2","E_2",200,0,3.5);
  TH1D* P1 = new TH1D("P1","P_1",200,0,3.0);
  TH1D* P1_2 = new TH1D("P1","P_1",200,2.,3.);

    
  betacan->Divide(2,2);   
  gStyle->SetPalette(1);          
    
  Double_t weight1;
  Double_t weight2;
  Double_t Eb,Pb;
     
  for (UInt_t i = 0; i<NrOfEvents ; i++)   
  {
      weight1 = rg->Generate();     
      
     for(UInt_t j = 0; j<3; j++) p[j] = rg->GetDecay(j);     
      
    Eb=p[1]->E();
    Pb=p[1]->P();

    E1->Fill(Eb,weight1);
    weight2=weight1*p[0]->E()*p[1]->E();
    E2->Fill(Eb,weight2);	
    E1_2->Fill(Eb,weight1);
    P1->Fill(Pb,weight1);
    P1_2->Fill(Pb,weight1);
    }
 
  betacan->cd(1);
  E1->Draw();
  betacan->cd(2); 
  E1_2->Draw();
  betacan->cd(3);
  P1->Draw();
  betacan->cd(4); 
  P1_2->Draw();
 

  // fit functions 
  gStyle->SetOptFit(1111);          
  TF1 * pFunc = new TF1("pDist",pDist,0.,1.2,1);    
  TF1 * eFunc = new TF1("EDist",eDist,0.,1.3,1);
  pFunc->SetParameter(0,1.0);
  eFunc->SetParameter(0,1.0);

  E1->Fit(eFunc);
  E1_2->Fit(eFunc);
  P1->Fit(pFunc);
  P1_2->Fit(pFunc);

}

void testEvent2() { 
  doit();
}