//script to add DOCAZ parameterisation to ntuple
//
//Kevin Maguire

//-----------------------------------------------------------

#include <TH2.h>
#include <TStyle.h>
#include <TTree.h>
#include <TFile.h>
#include <TCanvas.h>
#include <iostream>
#include <string>
#include <TGraph.h>
#include <TFitResult.h>
#include <TFitResultPtr.h>
#include <TVector3.h>

using namespace std;

string add(string s1, string s2, string s3 = "")
{
  string add = s1 + s2 + s3;
  cout << add << endl;
  return add;
}

void LoopForDOCAZ(string filepath, string treename, string treefolder)
{

  //options
  string D0Prefix = "D0";
  string DstPrefix = "Dstar";
  string dauplusPrefix = "dauplus";
  string dauminusPrefix = "dauminus";
  string slpiPrefix = "slowPi";
  string DOCAZ_0 = "_docazparfn_p0";
  string DOCAZ_1 = "_docazparfn_p1";
  string nPV_name = "nPV";
  string PVX_name = "PVX";
  string PVY_name = "PVY";

  //Get old file, old tree and set top branch address
  TFile *f = TFile::Open(filepath.c_str(), "UPDATE");
  if (!f) cout<<"ERROR: Ntuple file " << filepath << "not found" << endl;
  TDirectory *folder = nullptr;
  if (treefolder != "")
    {
      folder = f->GetDirectory(treefolder.c_str());
      folder->cd();
    }
  TString treepath = treefolder+treename;
  TTree *tree = (TTree*)f->Get( treepath );
  if (!tree) cout<<"ERROR: Ntuple tree path " << treepath << " not found" << endl;

  Long64_t nentries = tree->GetEntriesFast();
  cout << "Processing tree with " << nentries << " entries" << endl;

  //declare variable
  Double_t D0_PX;
  Double_t D0_PY;
  Double_t D0_PZ;
  Double_t D0_ENDVERTEX_X;
  Double_t D0_ENDVERTEX_Y;
  Double_t D0_ENDVERTEX_Z;
  Double_t D0_OWNPV_X;
  Double_t D0_OWNPV_Y;
  Double_t D0_OWNPV_Z;
  Double_t D0_MM;
  Double_t D0_P;
  Double_t dauplus_PX;
  Double_t dauplus_PY;
  Double_t dauplus_PZ;
  Double_t dauminus_PX;
  Double_t dauminus_PY;
  Double_t dauminus_PZ;
  Int_t nPV;
  Float_t PVX[200];
  Float_t PVY[200];
  
  //Set branch addresses
  tree->SetBranchStatus("*",1);
  tree->SetBranchAddress(add(D0Prefix, "_PX").c_str(),&D0_PX);
  tree->SetBranchAddress(add(D0Prefix, "_PY").c_str(),&D0_PY);
  tree->SetBranchAddress(add(D0Prefix, "_PZ").c_str(),&D0_PZ);
  tree->SetBranchAddress(add(D0Prefix, "_ENDVERTEX_X").c_str(),&D0_ENDVERTEX_X);
  tree->SetBranchAddress(add(D0Prefix, "_ENDVERTEX_Y").c_str(),&D0_ENDVERTEX_Y);
  tree->SetBranchAddress(add(D0Prefix, "_ENDVERTEX_Z").c_str(),&D0_ENDVERTEX_Z);
  tree->SetBranchAddress(add(D0Prefix, "_OWNPV_X").c_str(),&D0_OWNPV_X);
  tree->SetBranchAddress(add(D0Prefix, "_OWNPV_Y").c_str(),&D0_OWNPV_Y);
  tree->SetBranchAddress(add(D0Prefix, "_OWNPV_Z").c_str(),&D0_OWNPV_Z);
  tree->SetBranchAddress(add(D0Prefix, "_MM").c_str(),&D0_MM);
  tree->SetBranchAddress(add(D0Prefix, "_P").c_str(),&D0_P);
  tree->SetBranchAddress(add(dauplusPrefix, "_PX").c_str(),&dauplus_PX);
  tree->SetBranchAddress(add(dauplusPrefix, "_PY").c_str(),&dauplus_PY);
  tree->SetBranchAddress(add(dauplusPrefix, "_PZ").c_str(),&dauplus_PZ);
  tree->SetBranchAddress(add(dauminusPrefix, "_PX").c_str(),&dauminus_PX);
  tree->SetBranchAddress(add(dauminusPrefix, "_PY").c_str(),&dauminus_PY);
  tree->SetBranchAddress(add(dauminusPrefix, "_PZ").c_str(),&dauminus_PZ);
  tree->SetBranchAddress(nPV_name.c_str(),&nPV);
  tree->SetBranchAddress(PVX_name.c_str(),&PVX);
  tree->SetBranchAddress(PVY_name.c_str(),&PVY);

  //make new branches
  double dauplus_docazparfn_p0;
  double dauplus_docazparfn_p1;
  double dauminus_docazparfn_p0;
  double dauminus_docazparfn_p1;
  TBranch *dauplus_p0_branch = tree->Branch(add(dauplusPrefix, DOCAZ_0).c_str(), &dauplus_docazparfn_p0, add(dauplusPrefix, DOCAZ_0, "/D").c_str());
  TBranch *dauplus_p1_branch = tree->Branch(add(dauplusPrefix, DOCAZ_1).c_str(), &dauplus_docazparfn_p1, add(dauplusPrefix, DOCAZ_1, "/D").c_str());
  TBranch *dauminus_p0_branch = tree->Branch(add(dauminusPrefix, DOCAZ_0).c_str(), &dauminus_docazparfn_p0, add(dauminusPrefix, DOCAZ_0, "/D").c_str());
  TBranch *dauminus_p1_branch = tree->Branch(add(dauminusPrefix, DOCAZ_1).c_str(), &dauminus_docazparfn_p1, add(dauminusPrefix, DOCAZ_1, "/D").c_str());
  
  //decalre all
  Long64_t ientry;
  Float_t xbeam = 0; Float_t ybeam = 0;  
  Float_t npv;
  TVector3 beamlinevec;
  TVector3 childvec_dauplus;
  TVector3 childvec_dauminus;  
  TVector3 perpchild_dauplus;
  TVector3 perpchild_dauminus;
  TVector3 BP;
  int point;
  TVector3 origVec;
  double thistau;
  double docaz_dauplus;
  double docaz_dauminus;
  TFitResultPtr rplus;
  TFitResultPtr rminus;       

  for (Long64_t jentry=0; jentry < nentries;jentry++) {
  
    tree->GetEntry(jentry);

    //===========================================
    //CALCULATION WITH BRANCHES STARTS HERE
    //===========================================
    TGraph* docazvstime_dauplus_graph = docazvstime_dauplus_graph = new TGraph(12);
    TGraph* docazvstime_dauminus_graph = docazvstime_dauminus_graph = new TGraph(12);    
    
    xbeam = 0; ybeam = 0;  
    npv = 0.;
    npv = (Float_t)nPV;
    for (int pv = 0;pv < nPV;++pv) {   
      xbeam += PVX[pv]; ybeam += PVY[pv];
    }   
    xbeam = xbeam/npv; 
    ybeam = ybeam/npv;
    
    // Get the DOCAZ
    beamlinevec.Clear();
    childvec_dauplus.Clear();
    childvec_dauminus.Clear();  
    perpchild_dauplus.Clear(); 
    perpchild_dauminus.Clear(); 
    beamlinevec = TVector3(0,0,1);
    childvec_dauplus = TVector3(dauplus_PX,dauplus_PY,dauplus_PZ);
    childvec_dauminus = TVector3(dauminus_PX,dauminus_PY,dauminus_PZ);  
    perpchild_dauplus = childvec_dauplus.Cross(beamlinevec);
    perpchild_dauminus = childvec_dauminus.Cross(beamlinevec);
    
    //OK now for this one event swim
    BP.Clear();
    BP = TVector3(D0_PX,D0_PY,D0_PZ).Unit();
    point = 0;
    for (double dist = 600.;dist > -600; dist -= 1.) {

      origVec.Clear();
      origVec = TVector3(D0_ENDVERTEX_X-xbeam+dist*BP.X(),
			 D0_ENDVERTEX_Y-ybeam+dist*BP.Y(),
			 D0_ENDVERTEX_Z+dist*BP.Z());
      // Stop once behind PV in Z
      if (D0_ENDVERTEX_Z+dist*BP.Z()<D0_OWNPV_Z)
	{
	  break;
	}
      
      thistau = (10.0*D0_MM/(3.0*D0_P))*sqrt(pow(D0_ENDVERTEX_X+dist*BP.X()-D0_OWNPV_X,2)+
					     pow(D0_ENDVERTEX_Y+dist*BP.Y()-D0_OWNPV_Y,2)+   
					     pow(D0_ENDVERTEX_Z+dist*BP.Z()-D0_OWNPV_Z,2));  
      docaz_dauplus  = perpchild_dauplus.Dot(origVec)/perpchild_dauplus.Mag();
      docaz_dauminus = perpchild_dauminus.Dot(origVec)/perpchild_dauminus.Mag();
      
      if (thistau > 10.0)
	{
	  continue;
	}      
  
      docazvstime_dauplus_graph->SetPoint(point,thistau,docaz_dauplus);
      docazvstime_dauminus_graph->SetPoint(point,thistau,docaz_dauminus);
      ++point;
    }
    

    // Fit
    rplus   = docazvstime_dauplus_graph->Fit("pol1","QS");
    rminus  = docazvstime_dauminus_graph->Fit("pol1","QS"); 

    // Get the fit parameters
    dauplus_docazparfn_p0  = rplus->Value(0);
    cout << "dauplus_docazparfn_p0: " << dauplus_docazparfn_p0 << endl;
    dauplus_docazparfn_p1  = rplus->Value(1);
    dauminus_docazparfn_p0 = rminus->Value(0);
    dauminus_docazparfn_p1 = rminus->Value(1);
    
    docazvstime_dauplus_graph->Clear();
    docazvstime_dauminus_graph->Clear();

    delete docazvstime_dauplus_graph; delete docazvstime_dauminus_graph;
    //===========================================
    //CALCULATION WITH BRANCHES ENDS HERE    
    //===========================================
    
    dauplus_p0_branch->Fill(); 
    dauplus_p1_branch->Fill(); 
    dauminus_p0_branch->Fill(); 
    dauminus_p1_branch->Fill(); 

  }

  f->cd();
  if (folder) folder->cd();
  tree->Write();
  f->Close();
  delete f;

}


void AddDOCAZ(string filepath, string treename, string treefolder)
{
  LoopForDOCAZ(filepath, treename, treefolder);    
}

int main( int argc, char *argv[] )
{
  cout << argv[1] << " " << argv[2] << " " << argv[3] << endl;
  AddDOCAZ( argv[1], argv[2], argv[3] );  
  return 0;
}