#include <string.h>
#include <iostream>
#include <vector>
#include <fstream>

//void mc1(int iShow){
void mc1(int iShow){
  gROOT->Reset();
  gROOT->SetStyle("Plain");
  gStyle->SetOptStat(0);

  ofstream myfile_data;
  myfile_data.open ("data.csv");
  

  TRandom3 * randGen =  new TRandom3();


  TH2D * hBetaVsP = new TH2D("hBetaVsP","",200,0.0,5100,200,0.0,1.0);
  TH2D * hPVsGB = new TH2D("hPVsGB","",200,0.0,5100,200,0.0,10);
  TH1D * hMass = new TH1D("hMass","",140,0.0,1.400);
  TH1D * hMassPi = new TH1D("hMassPi","",140,0.0,1.400);
  TH1D * hMassK = new TH1D("hMassK","",140,0.0,1.400);
  TH1D * hMassPro = new TH1D("hMassPro","",140,0.0,1.400);
  TH1D * hMassSq = new TH1D("hMassSq","",140,-1.0,2.0);

  int nToGen = 1000000;

  double mPro = 938.280;
  double mK = 493.67;
  double mPi = 139.569;
  double tSmear = 0.05;
  double pSmear = 0.02;
  for (int i = 0; i<nToGen; i++) {
    //Proton
    double pPro = randGen->Uniform(300,5000);
    double ePro = sqrt(pPro*pPro + mPro*mPro);
    double betaPro = pPro/ePro;
    double length = 1.5;
    double dt = length/(betaPro*3e8);
    double pProS = pPro*(1 + randGen->Gaus(0,pSmear));
    double dtS = dt + randGen->Gaus(0,tSmear)*1e-9;
    double lengthS = length + randGen->Gaus(0,0.005);
    double betaProS = (lengthS/dtS)/3e8;

    bool keepEvent = true;

    double massProSqS = pProS*pProS*(1 - betaProS*betaProS)/betaProS;
    hMassSq->Fill(massProSqS/1000000);
    if (betaProS < 1){ 
      hBetaVsP->Fill(pProS,betaProS);

      double massProS = sqrt(massProSqS);
      hMass->Fill(massProS/1000);
      hMassPro->Fill(massProS/1000);
      double gammaProS = 1/sqrt(1 - betaProS*betaProS);
      hPVsGB->Fill(pProS,betaProS*gammaProS);

      myfile_data<<pProS;
      myfile_data<<",";
      myfile_data<<betaProS;
      myfile_data<<",";
      myfile_data<<2212;
      myfile_data<<endl;
    }else{
      keepEvent = false;
    }

    //kaon
    double pK = randGen->Uniform(200,5000);
    double eK = sqrt(pK*pK + mK*mK);
    double betaK = pK/eK;
    length = 1.5;
    dt = length/(betaK*3e8);
    double pKS = pK*(1 + randGen->Gaus(0,pSmear));
    dtS = dt + randGen->Gaus(0,tSmear)*1e-9;
    lengthS = length + randGen->Gaus(0,0.005);
    double betaKS = (lengthS/dtS)/3e8;

    double massKSqS = pKS*pKS*(1 - betaKS*betaKS)/betaKS;
    hMassSq->Fill(massKSqS/1000000);
    if (betaKS < 1){ 
      hBetaVsP->Fill(pKS,betaKS);

      double massKS = sqrt(massKSqS);
      hMass->Fill(massKS/1000);
      hMassK->Fill(massKS/1000);
      double gammaKS = 1/sqrt(1 - betaKS*betaKS);
      hPVsGB->Fill(pKS,betaKS*gammaKS);


      myfile_data<<pKS;
      myfile_data<<",";
      myfile_data<<betaKS;
      myfile_data<<",";
      myfile_data<<321;
      myfile_data<<endl;

    }else{
      keepEvent = false;
    }
    //pi
    double pPi = randGen->Uniform(200,5000);
    double ePi = sqrt(pPi*pPi + mPi*mPi);
    double betaPi = pPi/ePi;
    length = 1.5;
    dt = length/(betaPi*3e8);
    double pPiS = pPi*(1 + randGen->Gaus(0,pSmear));
    dtS = dt + randGen->Gaus(0,tSmear)*1e-9;
    lengthS = length + randGen->Gaus(0,0.005);
    double betaPiS = (lengthS/dtS)/3e8;

    double massPiSqS = pPiS*pPiS*(1 - betaPiS*betaPiS)/betaPiS;
    hMassSq->Fill(massPiSqS/1000000);
    if (betaPiS < 1){ 
      hBetaVsP->Fill(pPiS,betaPiS);

      double massPiS = sqrt(massPiSqS);
      hMass->Fill(massPiS/1000);
      hMassPi->Fill(massPiS/1000);
      double gammaPiS = 1/sqrt(1 - betaPiS*betaPiS);
      hPVsGB->Fill(pPiS,betaPiS*gammaPiS);


      myfile_data<<pPiS;
      myfile_data<<",";
      myfile_data<<betaPiS;
      myfile_data<<",";
      myfile_data<<211;
      myfile_data<<endl;

    }else{
      keepEvent = false;
    }

  }

  myfile_data.close();

  if (iShow == 1) {
    hBetaVsP->Draw("colz");
  }
  if (iShow == 2) {
    hMassPro->SetLineColor(4);
    hMassK->SetLineColor(3);
    hMassPi->SetLineColor(2);

    hMass->Draw();
    hMassPro->Draw("same");
    hMassPi->Draw("same");
    hMassK->Draw("same");
  }
  if (iShow == 3) {
    hMassSq->Draw();
  }
  if (iShow == 4) {
    hPVsGB->Draw("colz");
  }
}