#include "cmath"

string str1[4] = { "MonteCarlo/ntpl_jpsi_incl.root", "Bes/umcinc.root", "Bes/udata.root", "mcsig_09jpsi_phieta.root" };
string str2[4] = { "MCBadger", "MCBes", "DataBes", "NewMCBadger" };

void SetISN(TH1D* his);
void CoreBg(TF1* sum, TF1* graph1, TF1* graph2);
void DrwHis(TH1D* his, TF1* total, TF1* core, TF1* bg);

static Double_t Mpip = 0.140;
static Double_t Mpim = 0.140;
static Double_t mass_pi0 = 0.135;
static Double_t Meta = 0.548;

static Double_t inf, sup, nor;
static Double_t eff = 0.011; // Efficiency of selection
static Double_t Br1 = 0.3941; // Br(eta -> 2 gamma)
static Double_t Br2 = 0.1524; // Br(phi -> pi0 pi+ pi-)

void script( int nu = 1 )
{

  TCanvas* c1 = new TCanvas("Dalitz","Dalitz", 1500, 10, 1000, 1000);
  TCanvas* c2 = new TCanvas("Masses","Masses", 500, 10, 1000, 1000);

  string str = "/home/pogodin/Data_Files/" + str1[nu];

  TFile* f = new TFile(str.c_str());

  TNtupleD* data;
  gDirectory->GetObject("one/a5C", data);

  double ch2;
  double Ppip;
  double Cpip;
  double phipip;
  double Ppim;
  double Cpim;
  double phipim;
  double Ppi0;
  double Cpi0;
  double phipi0;
  double P2g;
  double C2g;
  double phi2g;
  double M2pi3;
  double M2gg;
  double M2rec;
  double dec;
  double deceta;
  double decphi;
  double M2pi0;

  data->SetBranchAddress("ch2", &ch2);
  data->SetBranchAddress("Ppip", &Ppip);
  data->SetBranchAddress("Cpip", &Cpip);
  data->SetBranchAddress("phipip", &phipip);
  data->SetBranchAddress("Ppim", &Ppim);
  data->SetBranchAddress("Cpim", &Cpim);
  data->SetBranchAddress("phipim", &phipim);
  data->SetBranchAddress("Ppi0", &Ppi0);
  data->SetBranchAddress("Cpi0", &Cpi0);
  data->SetBranchAddress("phipi0", &phipi0);
  data->SetBranchAddress("P2g", &P2g);
  data->SetBranchAddress("C2g", &C2g);
  data->SetBranchAddress("phi2g", &phi2g);
  data->SetBranchAddress("M2pi3", &M2pi3);
  data->SetBranchAddress("M2gg", &M2gg);
  data->SetBranchAddress("M2rec", &M2rec);
  data->SetBranchAddress("dec", &dec);
  data->SetBranchAddress("deceta", &deceta);
  data->SetBranchAddress("decphi", &decphi);
  data->SetBranchAddress("M2pi0", &M2pi0);

  int n = data->GetEntries();

  TH2D *Dalitz = new TH2D("Dalitz","Dalitz plot;M_{#gamma#gamma}^{2}, GeV^{2}/c^{4};M_{#pi^{+}#pi^{-}#pi^{0}}^{2}, GeV^{2}/c^{4}", 1000, 0, 1, 1300, 0, 1.3);

  TH1D *Chi = new TH1D("Chi","#Chi^{2};#Chi^{2};Entries/1", 150, 0, 150);

  TH1D *MggTH = new TH1D("MggTH","M_{#eta};M_{#gamma#gamma}, GeV/c^{2};Entries/0.001GeV/c^{2}", 100, 0.5, 0.6);

  TH1D *Momega = new TH1D("Momega","M_{#omega};M_{#pi^{+}#pi^{-}#pi^{0}}, Gev/c^{2};Entries/0.001GeV/c^{2}", 80, 0.74, 0.82);

  TH1D *Mphi = new TH1D("Mphi","M_{#phi};M_{#pi^{+}#pi^{-}#pi^{0}}, GeV/c^{2};Entries/0.001GeV/c^{2}", 100, 0.97, 1.07);

  TH1D *Mpi0 = new TH1D("Mpi0","M_{#pi^{0}};M_{#pi^{0}}, GeV/c^{2};Entries/0.001GeV/c^{2}}", 100, 0.130, 0.140);

  for ( int i = 0; i < n; i++)
  {
    data->GetEntry(i);

    Chi->Fill(ch2);

    if ( M2gg < 1. && M2pi3 < 1.3 && ch2 < 100 )
      Dalitz->Fill( M2gg, M2pi3 );

    if ( M2pi3 > 0.5 && M2pi3 < 1.1 && ch2 < 100 )
      MggTH->Fill( sqrt(M2gg) );
  }

  SetISN(MggTH);

  TF1 *core = new TF1("core","gausn", inf, sup);
  TF1 *bg = new TF1("bg","pol1", inf, sup);

  TF1 *total = new TF1("tot","gausn(0)*0.001 + pol1(3)", inf, sup);//0.547 - 0.02, 0.547 + 0.02);

  //total->Print(); cout << endl;

  //total->SetParameter(0, 32.23);
  total->SetParameter(1, 0.546);
  total->SetParameter(2, 0.009);

  total->SetParName(0, "N");
  total->SetParName(1, "#mu");
  total->SetParName(2, "#sigma");
  total->SetParName(3, "Const 1");
  total->SetParName(4, "Const 2");

  MggTH->Fit("tot", "", "", 0.5, 0.6);//0.547 - 0.02, 0.547 + 0.02);

  CoreBg(total, core, bg);

  Double_t w = 0.008;// total->GetParameter(2);
  Double_t M2Eta = pow( 0.548, 2 );// pow(total->GetParameter(1), 2);
  Double_t M2Pi02 = pow( 0.135, 2 );

  TVector3* pip = new TVector3;
  TVector3* pim = new TVector3;
  TVector3* pi0 = new TVector3;
  TVector3* gg = new TVector3;

  TLorentzVector* lpip = new TLorentzVector();
  TLorentzVector* lpim = new TLorentzVector();
  TLorentzVector* lpi0 = new TLorentzVector();
  TLorentzVector* lgg = new TLorentzVector();

  for ( int i = 0; i < n; i++ )
  {
    data->GetEntry(i);

      pip->SetMagThetaPhi(Ppip, acos(Cpip), phipip);
      pim->SetMagThetaPhi(Ppim, acos(Cpim), phipim);
      pi0->SetMagThetaPhi(Ppi0, acos(Cpi0), phipi0);
      gg->SetMagThetaPhi(P2g, acos(C2g), phi2g);
      
      lpip->SetXYZM(pip->X(), pip->Y(), pip->Z(), Mpip);
      lpim->SetXYZM(pim->X(), pim->Y(), pim->Z(), Mpim);
      lpi0->SetXYZM(pi0->X(), pi0->Y(), pi0->Z(), mass_pi0);
      lgg->SetXYZM(gg->X(), gg->Y(), gg->Z(), Meta);

    if ( (*lpi0 + *lgg).M() > 0.9 && (*lpi0 + *lgg).M() < 1.45 ) continue;
    if ( (*lpip + *lgg).M() > 0.9 && (*lpip + *lgg).M() < 1.45 ) continue;
    if ( (*lpim + *lgg).M() > 0.9 && (*lpim + *lgg).M() < 1.45 ) continue;

    if ( abs( sqrt(M2gg) - sqrt(M2Eta) ) < 3*w && ch2 < 100)
      {
        if ( sqrt(M2pi3) > 0.74 && sqrt(M2pi3) < 0.84 )
          Momega->Fill( sqrt(M2pi3) );
        else
          Mphi->Fill( sqrt(M2pi3) );
      }

    if ( abs( sqrt(M2pi0) - sqrt(M2Pi02) )  < 3*w && ch2 < 100 )
      Mpi0->Fill( sqrt(M2pi0) );

  }

  SetISN(Momega);

  TF1 *core2 = new TF1("core2","[0]*TMath::BreitWigner(x,[1],[2])", inf, sup);
  TF1 *bg2 = new TF1("bg2","pol1", inf, sup);

  TF1 *total2 = new TF1("tot2","core2*0.001 + bg2", inf, sup);//0.782 - 0.02, 0.782 + 0.02);

  //total2->Print(); cout << endl;

  total2->SetParameter(1, 0.782);
  total2->SetParameter(2, 0.006);
  total2->SetParameter(3, 300);
  //total2->SetParameter(5, 0.001);

  total2->SetParName(0, "N");
  total2->SetParName(1, "M");//"#mu");
  total2->SetParName(2, "#Gamma");//"#sigma");
  total2->SetParName(3, "Const 1");
  total2->SetParName(4, "Const 2");

  total2->SetParLimits(1, 0.779, 0.786);
  total2->SetParLimits(2, 0.000, 0.020);

  Momega->Fit("tot2", "", "", 0.74, 0.82);//0.782 - 0.02, 0.782 + 0.02);

  CoreBg(total2, core2, bg2);

  SetISN(Mphi);

  //TF1 *core3 = new TF1("core3","[0]*TMath::BreitWigner(x,[1],[2])", inf, sup);
  //TF1Convolution *core3_conv = new TF1Convolution("[0]*TMath::BreitWigner(x,[1],[2])", "gausn(3)", inf, sup);
  TF1 *core3 = new TF1("core3", "NSUM(*TMath::BreitWigner(), gausn())", inf, sup);//, core3_conv->GetNpar());
  TF1 *bg3 = new TF1("bg3","pol1", inf, sup);

  TF1 *total3 = new TF1("tot3","core3*0.001 + bg3", inf, sup);

  //total3->Print(); cout << endl;

  total3->SetParameter(1, 1.020);
  total3->SetParameter(2, 0.0016);

  total3->SetParName(0, "N");
  total3->SetParName(1, "M");//"#mu");
  total3->SetParName(2, "#Gamma");//"#sigma");
  total3->SetParName(3, "Const 1");
  total3->SetParName(4, "Const 2");

  total3->SetParLimits(2, 0.009, 0.020);

  Mphi->Fit("tot3", "", "", 0.97, 1.07);//1.020 - 0.02, 1.020 + 0.02);

  CoreBg(total3, core3, bg3);

  string n1, n2;

  c1->cd();

  gStyle->SetOptStat(10);

  Dalitz->Draw();

  n1 = "/home/pogodin/Art_Files/Dalitz" + str2[nu] +  ".pdf";

  c1->SaveAs(n1.c_str());

  gStyle->SetOptFit(112);

  c2->Clear();

  c2->Divide(2,2);

  c2->cd();

  c2->cd(1);
  gPad->SetGrid();
  Mpi0->Draw();

  c2->cd(2);
  DrwHis(MggTH, total, core, bg);

  c2->cd(3);
  DrwHis(Momega, total2, core2, bg2);

  c2->cd(4);
  DrwHis(Mphi, total3, core3, bg3);

  n2 = "/home/pogodin/Art_Files/Graph" + str2[nu] + ".pdf";

  c2->SaveAs(n2.c_str());

  // Br(J/psi -> phi omega) = N_phi / ( N_J/psi * e(phi eta) * Br( eta - > 2 gamma) * Br( phi -> pi+ pi- pi0) )
  cout << "\nBr(J/psi -> phi omega) = " << core3->GetParameter(0)/0.001 / ( n * eff /* Br1 * Br2*/ ) << "\n";

}

void SetISN(TH1D* his)
{
  inf = his->GetXaxis()->GetXmin();
  sup = his->GetXaxis()->GetXmax();
  nor = (sup - inf)/his->GetNbinsX();
}

void CoreBg(TF1* sum, TF1* graph1, TF1* graph2)
{
  graph1->SetParameter(0, sum->GetParameter(0)*nor);

  for ( Int_t i = 1; i < graph1->GetNpar(); i++)
    graph1->SetParameter(i, sum->GetParameter(i));

  for ( Int_t i = 0; i < sum->GetNpar() - graph1->GetNpar(); i++ )
    graph2->SetParameter(i, sum->GetParameter(graph1->GetNpar() + i));

  graph1->SetLineColor(4);
  graph1->SetLineStyle(2);

  graph2->SetLineColor(620);
  graph2->SetLineStyle(2);
}

void DrwHis(TH1D* his, TF1* total, TF1* core, TF1* bg)
{
  gPad->SetGrid();
  his->SetMinimum(0);

  his->Draw("E");

  TLegend* legend = new TLegend(0.1,0.7,0.4,0.9);
  legend->SetHeader("Designations","C");
  legend->AddEntry(his,"Data");
  legend->AddEntry(total, "BW + Bg");
  legend->AddEntry(core, "BW");
  legend->AddEntry(bg, "Background");

  legend->Draw();
  core->Draw("same");
  bg->Draw("same");
}