#include "HistoCollector3D.h"

//_________________
HistoCollector3D::HistoCollector3D() {
  mFile = NULL;
  mQinvBins = 40;
  mQinvLo = 0.;
  mQinvHi = 0.4;
  if(!h3DNum.empty()) {
    h3DNum.clear();
  }
  if(!h3DDen.empty()) {
    h3DDen.clear();
  }
  if(!h3DQinv.empty()) {
    h3DQinv.clear();
  }

  mPartType = 4; //3-pions, 4-kaons
  mKtBins = 11;
}

//_________________
HistoCollector3D::~HistoCollector3D() {
  mFile->Close();
  delete mFile;
}

//_________________
void HistoCollector3D::OpenFile() {
  std::cout << "Opening file: " << mFileName;
  mFile = TFile::Open(mFileName);
  if(mFile) {
    std::cout << "\t [DONE]" << std::endl;
  }
  else {
    std::cout << "\t [ERROR] file does not exist" << std::endl;
  }
}

//_________________
void HistoCollector3D::LoadData() {
  OpenFile();
  ReadHistograms();
}

//_________________
void HistoCollector3D::ReadHistograms() {

  std::cout << "Retrieving histograms from file...";

  TH3F *hNumer = NULL;
  TH3F *hDenom = NULL;
  TH3F *hQinv = NULL;
  for(int iCharge=0; iCharge<2; iCharge++) {
    for(int iKt=0; iKt<mKtBins; iKt++) {
	
      if(mPartType == 4) { //Kaons

	//Numerator
	hNumer = (TH3F*)mFile->Get(Form("hKaonBPLCMSFrameKtCF_%d_Num_bin_%d",
					iCharge,iKt));
	hNumer->SetName(Form("hKaonKtCFNum_%d_%d,",iCharge,iKt));

	//Denominator
	hDenom = (TH3F*)mFile->Get(Form("hKaonBPLCMSFrameKtCF_%d_Den_bin_%d",
					iCharge,iKt));
	hDenom->SetName(Form("hKaonKtCFDen_%d_%d,",iCharge,iKt));
	
	//Qinv weight
	hQinv = (TH3F*)mFile->Get(Form("hKaonBPLCMSFrameKtCF_%d_Qinv_bin_%d",
				       iCharge,iKt));
	hQinv->SetName(Form("hKaonKtCFQinv_%d_%d,",iCharge,iKt));

	//Add histograms to the vectors
	h3DNum.push_back(hNumer);
	h3DDen.push_back(hDenom);
	h3DQinv.push_back(hQinv);
      }
      if(mPartType == 3) { //Pions

	//Numerator
	hNumer = (TH3F*)mFile->Get(Form("hPionBPLCMSFrameKtCF_%d_Num_bin_%d",
					iCharge,iKt));
	hNumer->SetName(Form("hPionKtCFNum_%d_%d,",iCharge,iKt));

	//Denominator
	hDenom = (TH3F*)mFile->Get(Form("hPionBPLCMSFrameKtCF_%d_Den_bin_%d",
					iCharge,iKt));
	hDenom->SetName(Form("hPionKtCFDen_%d_%d,",iCharge,iKt));

	//Qinv weight
	hQinv = (TH3F*)mFile->Get(Form("hPionBPLCMSFrameKtCF_%d_Qinv_bin_%d",
				       iCharge,iKt));
	hQinv->SetName(Form("hPionKtCFQinv_%d_%d,",iCharge,iKt));

	//Add histograms to the vectors
	h3DNum.push_back(hNumer);
	h3DDen.push_back(hDenom);
	h3DQinv.push_back(hQinv);
      }
    }
  }
  std::cout << "\t [DONE]" << std::endl;
  std::cout << Form("Number of histos in numerator:   %d", h3DNum.size()) 
	    << std::endl
	    << Form("Number of histos in denominator: %d", h3DDen.size()) 
	    << std::endl
	    << Form("Number of histos in qinv:        %d", h3DQinv.size()) 
	    << std::endl;
}

//_________________
TH3F* HistoCollector3D::BuildKt3DNum(int charge, int binLo, int binHi) {

  TString sName = "h";
  sName += "NumerKtCF_";
  sName += charge;
  sName += "_";
  sName += binLo;
  sName += "_";
  sName += binHi;

  TH3F *hNumSum = new TH3F(sName.Data(),sName.Data(),
			   mQinvBins,mQinvLo,mQinvHi,
			   mQinvBins,mQinvLo,mQinvHi,
			   mQinvBins,mQinvLo,mQinvHi);
  hNumSum->Sumw2();

  if(charge<2) { //Positive or negative particles
    for(int iKt=binLo; iKt<=binHi; iKt++) {
      hNumSum->Add(h3DNum[charge*mKtBins + iKt]);
    }
  }
  else { //Sum of charges
    for(int iCharge=0; iCharge<2; iCharge++) {
      for(int iKt=binLo; iKt<=binHi; iKt++) {
	hNumSum->Add(h3DNum[iCharge*mKtBins + iKt]);
      }
    }
  }

  return hNumSum;
}

//_________________
TH3F* HistoCollector3D::BuildKt3DDen(int charge, int binLo, int binHi) {

  TString sName = "h";
  sName += "DenomKtCF_";
  sName += charge;
  sName += "_";
  sName += binLo;
  sName += "_";
  sName += binHi;

  TH3F *hDenSum = new TH3F(sName.Data(),sName.Data(),
			   mQinvBins,mQinvLo,mQinvHi,
			   mQinvBins,mQinvLo,mQinvHi,
			   mQinvBins,mQinvLo,mQinvHi);
  hDenSum->Sumw2();

  if(charge<2) { //Positive or negative particles
    for(int iKt=binLo; iKt<=binHi; iKt++) {
      hDenSum->Add(h3DDen[charge*mKtBins + iKt]);
    }
  }
  else { //Sum of charges
    for(int iCharge=0; iCharge<2; iCharge++) {
      for(int iKt=binLo; iKt<=binHi; iKt++) {
	hDenSum->Add(h3DDen[iCharge*mKtBins + iKt]);
      }
    }
  }

  return hDenSum;
}

//_________________
TH3F* HistoCollector3D::BuildKt3DQinv(int charge, int binLo, int binHi) {

  TString sName = "h";
  sName += "QinvKtCF_";
  sName += charge;
  sName += "_";
  sName += binLo;
  sName += "_";
  sName += binHi;

  TH3F *hQinvSum = new TH3F(sName.Data(),sName.Data(),
			    mQinvBins,mQinvLo,mQinvHi,
			    mQinvBins,mQinvLo,mQinvHi,
			    mQinvBins,mQinvLo,mQinvHi);
  hQinvSum->Sumw2();

  if(charge<2) { //Positive or negative particles
    for(int iKt=binLo; iKt<=binHi; iKt++) {
      hQinvSum->Add(h3DQinv[charge*mKtBins + iKt]);
    }
  }
  else { //Sum of charges
    for(int iCharge=0; iCharge<2; iCharge++) {
      for(int iKt=binLo; iKt<=binHi; iKt++) {
	hQinvSum->Add(h3DQinv[iCharge*mKtBins + iKt]);
      }
    }
  }

  return hQinvSum;
}
