//////////////////////////////////////////////////////////
//   This file has been automatically generated 
//     (Sun Apr 20 00:32:22 2008 by ROOT version5.17/01)
//   from TTree singmuon/single muon
//   found on file: pythia-pp-200-pt00-Maxpt00-cc-1m-207-2-ana.root
//  
//   05/02/2008 calculatge single muon A_N in open charm production
//   05/13/2008  use a function to calculate A_N(pT, xF)
//
//   11/08/2017  modefied to do heavy ion v2 analysis
//
//   04/06/2018 HF DCA analysis 
//              QFlag = 1: B->X
//                      2: D->X
//                      3: B->D->X  
//////////////////////////////////////////////////////////

#include <iostream>
#include <stdlib.h>
#include <time.h>

#include <TROOT.h>
#include <TStyle.h>

#include <TObjArray.h>
#include <TArray.h>
#include <TCanvas.h>
#include <TH1.h>
#include <TH2.h>
#include <TF1.h>
#include <TFile.h>
#include <TObject.h>
#include <TNtuple.h>
#include <TRandom.h>
#include <TRandom3.h>
#include <TLorentzVector.h>
#include <TTree.h>
#include <TStopwatch.h>
#include <TLatex.h>
#include <TGraph.h>
#include <TH1F.h>
#include <TH2F.h>
#include <math.h>


//define user function
Double_t A_N_xF(Float_t par_pT, Float_t par_xF);
Double_t A_N_y(Float_t par_pT, Float_t par_y);



void singmuon_ana_DCA(const char * InFile = "input.root" ){

//Reset ROOT and connect tree file
   gROOT->Reset();
   //   gStyle->SetOptStat("Plain");
   gStyle->SetOptStat(1111);
   gStyle->SetOptFit();
   //  gStyle->SetPadLeftMargin(0.15);
   // gStyle->SetPadBottomMargin(0.2);
   //  gStyle->SetLineWidth(1.5);



   TFile *f = (TFile*)gROOT->GetListOfFiles()->FindObject(InFile);
   if (!f) {
      f = new TFile(InFile);
   }
   TTree *singmuon = (TTree*)gDirectory->Get("singmuon");


   //----- define historgrams ------

   TCanvas*   c1 = new TCanvas("c1","cc: Single muon properties", 10,20,600,600);
   TCanvas*   c2 = new TCanvas("c2","cc: D  hadron properties", 100,200,600,600);
  
   TCanvas*   c3 = new TCanvas("c3","cc: D and muon pT correlations", 200,250,400,800);
   TCanvas*   c4 = new TCanvas("c4","cc: D and muon xF correlations", 300,300,400,400);

   TCanvas*   c5 = new TCanvas("c5","cc: D and muon V_2 correlations", 400,350,400,400);

   TCanvas*   c6 = new TCanvas("c6","cc: D and muon V_2 correlations", 450,400,400,400);

   TCanvas*   c7 = new TCanvas("c7","BB: muon DCA_XY", 500,500,400,400);
 

   TH1F*   h10 = new TH1F("h10","muon_pT",100,0.0,10.0); // 
   TH1F*   h11 = new TH1F("h11","muon_ID",40,-20.0,20.0); // 
   TH1F*   h12 = new TH1F("h12","muon_phi",100,-5.0,5.0); // 
   TH1F*   h13 = new TH1F("h13","muon_eta",100,-3.0,3.0); // 
   TH1F*   h14 = new TH1F("h14","muon_xF",100, 0,0.5); // 
   TH1F*   h15 = new TH1F("h15","muon_PID",50,500.0,550.0); // 
   TH1F*   h16 = new TH1F("h16","muon_y",100,-3.0,3.0); // 
   TH1F*   h17 = new TH1F("h17","muon_R_xy",100, 0, 2.0); // 
   TH1F*   h18 = new TH1F("h18","muon_DCA_xy",100,0.0,2.0); // 
   TH1F*   h18a = new TH1F("h18a","muon_DCA_xy > 500um",100,0.0,2.0); // 
   TH1F*   h18b = new TH1F("h18b","muon_DCA_xy > 1000um",100,0.0,2.0); // 

   TH1F*   h19 = new TH1F("h19","muon_vtx_p_sinXY",100,0.0,1.0); // 


   TH1F*   h20 = new TH1F("h20","D_pT",100,0.0,10.0); // 
   TH1F*   h21 = new TH1F("h21","D_ID",50,400.0,450.0); // 
   TH1F*   h22 = new TH1F("h22","D_phi",100,-5.0,5.0); // 
   TH1F*   h23 = new TH1F("h23","D_xF",100, 0,0.50); // 
   TH2F*   h24 = new TH2F("h24","D_xF vs mu_xF",100,-0.3,0.3, 100, -0.3,0.3); // 
   TH2F*   h25 = new TH2F("h25","D_pT vs mu_pT",100,0,10,100,0,10); // 

   TH2F*   h26 = new TH2F("h26","D_phi vs mu_phi",100,-5,5,100,-5,5); // 
   TH1F*   h27 = new TH1F("h27","D_y",100,-5.0,5.0); // 

   //mu-D pT corellation
   TH1F*   h101 = new TH1F("h101","muon_pT",100,0.0,10.0); // 
   TH1F*   h102 = new TH1F("h102","muon_pT",100,0.0,10.0); // 
   TH1F*   h103 = new TH1F("h103","muon_pT",100,0.0,10.0); // 
   TH1F*   h104 = new TH1F("h104","muon_pT",100,0.0,10.0); // 
   TH1F*   h105 = new TH1F("h105","muon_pT",100,0.0,10.0); // 
   TH1F*   h106 = new TH1F("h106","muon_pT",100,0.0,10.0); // 

   TH1F*   h201 = new TH1F("h201","D_pT",100,0.0,10.0); // 
   TH1F*   h202 = new TH1F("h202","D_pT",100,0.0,10.0); // 
   TH1F*   h203 = new TH1F("h203","D_pT",100,0.0,10.0); // 
   TH1F*   h204 = new TH1F("h204","D_pT",100,0.0,10.0); // 
   TH1F*   h205 = new TH1F("h205","D_pT",100,0.0,10.0); // 
   TH1F*   h206 = new TH1F("h206","D_pT",100,0.0,10.0); // 

   //mu-D xF corellation

   TH1F*   h141 = new TH1F("h141","muon_xF",100, 0,0.5); // 
   TH1F*   h142 = new TH1F("h142","muon_xF",100, 0,0.5); // 
   TH1F*   h151 = new TH1F("h151","muon_PID",50,400.0,450.0); // 

   TH1F*   h231 = new TH1F("h231","D_xF",100, 0,0.50); // 
   TH1F*   h232 = new TH1F("h232","D_xF",100, 0,0.50); // 
   TH1F*   h152 = new TH1F("h152","muon_PID",50,400.0,450.0); // 

   //mu-D A_N asuymmetry correlations
   TH1F*   h411p = new TH1F("h411p","D+(+411)  D_phi",100,-5.0,5.0); // 
   TH1F*   h421p = new TH1F("h421p","D0(+421)  D_phi",100,-5.0,5.0); // 
   TH1F*   h411m = new TH1F("h411m","D-(-411)  D_phi",100,-5.0,5.0); // 
   TH1F*   h421m = new TH1F("h421m","D0_bar(-421)  D_phi",100,-5.0,5.0); // 

   TH1F*   h411p_mu = new TH1F("h411p_mu","D+(+411)->mu muon_phi",100,-5.0,5.0); // 
   TH1F*   h421p_mu = new TH1F("h421p_mu","D0(+421)->mu muon_phi",100,-5.0,5.0); // 
   TH1F*   h411m_mu = new TH1F("h411m_mu","D-(-411)->mu muon_phi",100,-5.0,5.0); // 
   TH1F*   h421m_mu = new TH1F("h421m_mu","D0_bar(-421)->mu muon_phi",100,-5.0,5.0); // 

   TH1F*   h12_mup = new TH1F("h12_mup","mu+: muon_phi",100,-5.0,5.0); // 
   TH1F*   h12_mum = new TH1F("h12_mum","mu-: muon_phi",100,-5.0,5.0); // 

   TH1F*   h12_mum1 = new TH1F("h12_mum1","mu-: muon_phi: |xF|<0.05",100,-5.0,5.0); // 
   TH1F*   h12_mum2 = new TH1F("h12_mum2","mu-: muon_phi: |xF|>0.05",100,-5.0,5.0); // 


   //fitting function
   //   TF1* fit_AN = new TF1("fit_AN", "[0]*(1+[1]*sin(x))", -3.1415926, 3.1415926);
      TF1* fit_AN = new TF1("fit_AN", "[0]*(1+[1]*cos(2*x))", -3.1415926, 3.1415926);

   
   fit_AN->SetParName(0,"N");
   fit_AN->SetParName(1,"V2");

   fit_AN->SetParameter(1, 0.1);
   fit_AN->SetParLimits(1,-1,1);
   fit_AN->SetLineColor(2);
   fit_AN->SetLineWidth(2);



//Declaration of leaves types
   Float_t         IF1;
   Float_t         IF2;
   Float_t         OF1;
   Float_t         OF2;
   Float_t         X1;
   Float_t         X2;
   Float_t         Q2;
   Float_t         M;
   Float_t         px;
   Float_t         py;
   Float_t         pz;
   Float_t         eta;
   Float_t         QFlag;
   Float_t         Flavor;
   Float_t         mc_x;
   Float_t         mc_y;
   Float_t         mc_z;
   Float_t         pxP;
   Float_t         pyP;
   Float_t         pzP;
   Float_t         FlavorP;
   Float_t         ISUB;
   Float_t         CosTh;

   // Set branch addresses.
   singmuon->SetBranchAddress("IF1",&IF1);
   singmuon->SetBranchAddress("IF2",&IF2);
   singmuon->SetBranchAddress("OF1",&OF1);
   singmuon->SetBranchAddress("OF2",&OF2);
   singmuon->SetBranchAddress("X1",&X1);
   singmuon->SetBranchAddress("X2",&X2);
   singmuon->SetBranchAddress("Q2",&Q2);
   singmuon->SetBranchAddress("M",&M);
   singmuon->SetBranchAddress("px",&px);
   singmuon->SetBranchAddress("py",&py);
   singmuon->SetBranchAddress("pz",&pz);
   singmuon->SetBranchAddress("eta",&eta);
   singmuon->SetBranchAddress("QFlag",&QFlag);
   singmuon->SetBranchAddress("Flavor",&Flavor);
   singmuon->SetBranchAddress("mc_x",&mc_x);
   singmuon->SetBranchAddress("mc_y",&mc_y);
   singmuon->SetBranchAddress("mc_z",&mc_z);
   singmuon->SetBranchAddress("pxP",&pxP);
   singmuon->SetBranchAddress("pyP",&pyP);
   singmuon->SetBranchAddress("pzP",&pzP);
   singmuon->SetBranchAddress("FlavorP",&FlavorP);
   singmuon->SetBranchAddress("ISUB",&ISUB);
   singmuon->SetBranchAddress("CosTh",&CosTh);

//     This is the loop skeleton
//       To read only selected branches, Insert statements like:
// singmuon->SetBranchStatus("*",0);  // disable all branches
// TTreePlayer->SetBranchStatus("branchname",1);  // activate branchname

   Double_t E =0, EP =0;
   Double_t pt =0, ptP=0, ptot=0;
   Double_t xF =0, xFP=0;
   Double_t y  =0, yP=0;
   Double_t M_c = 4.2;

   Double_t R_xy=0, DCA_xy=0;  // vertex R_xy
   Double_t sinXY =0, cosXY2=0; 

   Double_t wgt =1; 
   Double_t phi_mu =0, phi_D=0;


   Long64_t nentries = singmuon->GetEntries();
   cout << "total entries = " << nentries << endl;
   Long64_t nbytes = 0;
   //
   for (Long64_t i=0; i<nentries;i++) {
     //for (Long64_t i=0; i<50000;i++) {
      nbytes += singmuon->GetEntry(i);

      if (i%100000==0) cout << "processed " << 100.0*i/nentries << "%" << endl;

      //     
      //--- apply PHENIX Central Arm acceptance cuts
      //
      if (fabs(eta) > 1 ) continue;
      if (fabs(Flavor) > 20 ) continue;

      //D->muon only
      //      if ( fabs(FlavorP)> 500 || fabs(FlavorP)<400 ) continue;

      //B->muon only
      if ( fabs(FlavorP)< 500 ) continue;

      pt  = sqrt(px*px + py*py);
      ptP = sqrt(pxP*pxP + pyP*pyP);

      ptot  = sqrt(px*px + py*py + pz*pz);
      R_xy = sqrt(mc_x*mc_x + mc_y*mc_y);


      if (R_xy < 0.01) continue;  // R_xy > 100um, beam size 
 
      cosXY2 = (px*mc_x + py*mc_y)*(px*mc_x + py*mc_y)/(pt*pt)/(R_xy*R_xy);
      sinXY =  sqrt(fabs(1 - cosXY2));

      if (cosXY2>1){
	cout << "!!! cosXY2 = " << cosXY2 << endl;
	cout << "sinXY = " << sinXY << endl;
      }



      DCA_xy = R_xy*sinXY;

      xF = pz/100.0;
      xFP = pzP/100.0;

      E = sqrt(px*px + py*py + pz*pz + M*M);
      y = 0.5*log((E + pz)/(E - pz));  // there are log(x)= ln(x) and log10(x)

      EP = sqrt(pxP*pxP + pyP*pyP + pzP*pzP + M_c*M_c);
      yP = 0.5*log((EP + pzP)/(EP - pzP));

      phi_mu = atan2(py,px);
      phi_D  = atan2(pyP,pxP);

      //muons

      h10->Fill(pt);
      h11->Fill(Flavor);
      h12->Fill(phi_mu);
      h13->Fill(eta);
      h14->Fill(fabs(xF));
      h15->Fill(fabs(FlavorP));
      h16->Fill(y);

      h17->Fill(R_xy);
      h18->Fill(DCA_xy);
      if (DCA_xy > 0.05) {  // trigger DCA cut > 0.5mm  
	h18a->Fill(DCA_xy);
      }

     if (DCA_xy > 0.1) {  // trigger DCA cut > 1mm  
	h18b->Fill(DCA_xy);
      }

      h19->Fill(sinXY);


      //D mesons
      h20->Fill(ptP);
      h21->Fill(fabs(FlavorP));
      h22->Fill(phi_D);
      h23->Fill(fabs(xFP));
      h24->Fill(xF, xFP);
      h25->Fill(pt, ptP);

      h26->Fill(phi_mu, phi_D);
      h27->Fill(yP);

      //mu-D correlations

      if (pt>1.0 && pt < 1.5){
	h101->Fill(pt);
	h201->Fill(ptP);
      }
      else if (pt>1.5 && pt < 2.0){
	h102->Fill(pt);
	h202->Fill(ptP);
      }
      else if(pt>2.0 && pt < 2.5){
	h103->Fill(pt);
	h203->Fill(ptP);
      }
      else if (pt>2.5 && pt < 3.0){
	h104->Fill(pt);
	h204->Fill(ptP);
      }
      else if (pt>3.0 && pt < 3.85){
	h105->Fill(pt);
	h205->Fill(ptP);
      }
      else if (pt>3.85 && pt < 5.0){
	h106->Fill(pt);
	h206->Fill(ptP);
      }
      else {};

      //xF correlation

      if (fabs(xF)>0 && fabs(xF) < 0.05) {
	h141->Fill(fabs(xF));
	h231->Fill(fabs(xFP));
	h151->Fill(fabs(FlavorP));
      }
      else if (fabs(xF)>0.05 && fabs(xF) <0.2){
	h142->Fill(fabs(xF));
	h232->Fill(fabs(xFP));
	h152->Fill(fabs(FlavorP));
      }
      else {};

      //
      //calcualte single muon A_N from D A_N
      //

      Float_t AN_Dm = A_N_y(ptP,yP),  AN_D0bar = A_N_y(ptP, yP); // at xF = 0.15
      Float_t AN_Dp = -A_N_y(ptP, yP), AN_D0 = 0;


      if (FlavorP == -411){ // D-
	//	wgt = 1 + AN_Dm*sin(phi_D);
		wgt = 1 + AN_Dm*cos(2*phi_D);

	h411m->Fill(phi_D,wgt);
	h411m_mu->Fill(phi_mu,wgt);

      }
      else if (FlavorP == -421){ // D0bar
	// 	wgt = 1 + AN_D0bar*sin(phi_D);
	 	wgt = 1 + AN_D0bar*cos(2*phi_D);

	h421m->Fill(phi_D,wgt);
	h421m_mu->Fill(phi_mu,wgt);
      }
      else if (fabs(FlavorP)==431){ 
	wgt = 1.0;

	h411p->Fill(phi_D,wgt);
	h411p_mu->Fill(phi_mu,wgt);

	h421p->Fill(phi_D,wgt);
	h421p_mu->Fill(phi_mu,wgt);
      }
      else {
	wgt = 1.0;

	h411p->Fill(phi_D,wgt);
	h411p_mu->Fill(phi_mu,wgt);

	h421p->Fill(phi_D,wgt);
	h421p_mu->Fill(phi_mu,wgt);
      }

      //mu+ and mu-
      if (Flavor > 0){ // 11 = e-, 13 = mu- 
	h12_mum->Fill(phi_mu,wgt);

	if (fabs(xF)<0.05){
	  h12_mum1->Fill(phi_mu,wgt);
	}
	else {
	  h12_mum2->Fill(phi_mu,wgt);
	}

      }
      else {
	h12_mup->Fill(phi_mu,wgt);
      }


   } // end of event loop


   //make plots

   // single muons
   c1->cd();
   c1->Divide(3,3);

   c1->cd(1)->SetLogy();
   h10->Draw();

   c1->cd(2);
   h11->Draw();

   c1->cd(3);
   h12->Draw();

   c1->cd(4);
   h13->Draw();

   c1->cd(5)->SetLogy();
   h14->Draw();

   c1->cd(6);
   h15->Draw();

   c1->cd(7)->SetLogy();
   h17->Draw();

   c1->cd(8)->SetLogy();
   h18->Draw();

   c1->cd(9)->SetLogy();
   //   h18a->Draw();
   h19->Draw();




   // single D mesons
   c2->cd();
   c2->Divide(2,4);

   c2->cd(1)->SetLogy();
   h20->Draw();

   c2->cd(2);
   h21->Draw();

   c2->cd(3);
   h22->Draw();

   c2->cd(4)->SetLogy();
   h23->Draw();

   c2->cd(5);
   h24->Draw();

   c2->cd(6);
   h25->Draw();

   c2->cd(7);
   //   h26->Draw();
   h16->Draw();
   c2->cd(8);
   h27->Draw();


   //muon - D pT correlations
   c3->cd();
   c3->Divide(2,6);

   c3->cd(1);
   h101->Draw();
   c3->cd(2);
   h201->Draw();

   c3->cd(3);
   h102->Draw();
   c3->cd(4);
   h202->Draw();

   c3->cd(5);
   h103->Draw();
   c3->cd(6);
   h203->Draw();

   c3->cd(7);
   h104->Draw();
   c3->cd(8);
   h204->Draw();

   c3->cd(9);
   h105->Draw();
   c3->cd(10);
   h205->Draw();

   c3->cd(11);
   h106->Draw();
   c3->cd(12);
   h206->Draw();

   //muon - D xF correlations
   c4->cd();
   c4->Divide(3,2);

   c4->cd(1);
   h141->Draw();
   c4->cd(2);
   h231->Draw();
   c4->cd(3);
   h151->Draw();

   c4->cd(4);
   h142->Draw();
   c4->cd(5);
   h232->Draw();
   c4->cd(6);
   h152->Draw();

   //muon-D A_N corellation
   c5->cd();
   c5->Divide(4,3);

   //D-
   c5->cd(1);
   h411m->Fit("fit_AN","R");
   h411m->Draw();

   c5->cd(2);
   h421m->Fit("fit_AN","R");
   h421m->Draw();

   c5->cd(3);
   h411p->Fit("fit_AN","R");
   h411p->Draw();

   c5->cd(4);
   h421p->Fit("fit_AN","R");
   h421p->Draw();

   //mu-/+
   c5->cd(5);
   h411m_mu->Fit("fit_AN","R");
   h411m_mu->Draw();

   c5->cd(6);
   h421m_mu->Fit("fit_AN","R");
   h421m_mu->Draw();

   c5->cd(7);
   h411p_mu->Fit("fit_AN","R");
   h411p_mu->Draw();

   c5->cd(8);
   h421p_mu->Fit("fit_AN","R");
   h421p_mu->Draw();

   //mu-, mu+
   c5->cd(9);
   h12_mum->Fit("fit_AN","R");
   h12_mum->Draw();

   c5->cd(10);
   h12_mup->Fit("fit_AN","R");
   h12_mup->Draw();

   // |xF|<0.05
   c5->cd(11);
   h12_mum1->Fit("fit_AN","R");
   h12_mum1->Draw();

   c5->cd(12);
   h12_mum2->Fit("fit_AN","R");
   h12_mum2->Draw();

   //for output
   c6->cd();
   c6->Divide(2,2);

   c6->cd(1);
   h411m->Draw();

   c6->cd(2);
   h411p->Draw();

   c6->cd(3);
   h411m_mu->Draw();

   c6->cd(4);
   h411p_mu->Draw();


   //--- c7 trigger DCA study 
   c7->cd();
   c7->Divide(2,2);

   c7->cd(1)->SetLogy();
   h10->Draw();

   c7->cd(2)->SetLogy();
   h18->Draw();

   c7->cd(3)->SetLogy();
   h18a->Draw();

   c7->cd(4)->SetLogy();
   h18b->Draw();


   //--- print output files ----

   c1->Print("c1.pdf");
   c2->Print("c2.pdf");
   c3->Print("c3.pdf");
   c4->Print("c4.pdf");
   c5->Print("c5.pdf");
   
   c6->Print("c6.gif");
   c6->Print("c6.pdf");

   c7->Print("c7.gif");
   c7->Print("c7.pdf");

   
}

Double_t A_N_xF(Float_t par_pT, Float_t par_xF){

  //  Double_t ANX = 0.1*par_pT + 0.1*par_xF;
    Double_t ANX = 0.2;

  return ANX;
 
}

Double_t A_N_y(Float_t x, Float_t y){

  Double_t ANX =0.2;

  /*
  if (y < 0.5) {
    ANX = 0.0762 + 0.0411*x + 0.0416*x*x - 0.0197*x*x*x + 0.0021*x*x*x*x; // 
  }
  else if ( y > 0.5 && y < 1.5 ){
    ANX = 0.0762 + 0.0411*x + 0.0416*x*x - 0.0197*x*x*x + 0.0021*x*x*x*x; // y=1
  }
  else if ( y > 1.5 && y < 2.5 ){
    ANX = 0.0762 + 0.0411*x + 0.0416*x*x - 0.0197*x*x*x + 0.0021*x*x*x*x; //y=2
  }
  else if ( y > 2.5 && y < 3.0 ){
    ANX = 0.0762 + 0.0411*x + 0.0416*x*x - 0.0197*x*x*x + 0.0021*x*x*x*x; // y =3
  }
  else if ( y > 3.0 && y < 4.0 ){
    ANX = 0.0762 + 0.0411*x + 0.0416*x*x - 0.0197*x*x*x + 0.0021*x*x*x*x; // y = 4
  }
  else {
    ANX = 0;
  }
  */
  
  return ANX;
 
}