#include "HistoCollector1D.h"

//_________________
HistoCollector1D::HistoCollector1D() {
  mFile = NULL;
  mQinvBins = 40;
  mQinvLo = 0.;
  mQinvHi = 0.4;
  if(!hQinvNum.empty()) {
    hQinvNum.clear();
  }
  if(!hQinvDen.empty()) {
    hQinvDen.clear();
  }

  mCorrFctnColor = 1;
  mCorrFctnLineWidth = 2;
  mCorrFctnMarkerStyle = 20;
  mCorrFctnMarkerSize = 1.8;

  mPartType = 4; //3-pions, 4-kaons
  mCentBins = 9;
  mKtBins = 24;
}

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

//_________________
void HistoCollector1D::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 HistoCollector1D::LoadData() {
  OpenFile();
  ReadHistograms();
}

//_________________
void HistoCollector1D::SetHistoStyle(TH1* histo, int charge) {
  switch(charge) {
  case 0:
    mCorrFctnColor = 2; //Red
    break;
  case 1:
    mCorrFctnColor = 4; //Blue
    break;
  case 2:
    mCorrFctnColor = 1; //Black
    break;
  default:
    mCorrFctnColor = 6; //Magenta
  };
  histo->SetLineColor(mCorrFctnColor);
  histo->SetMarkerColor(mCorrFctnColor);
  histo->SetLineWidth(mCorrFctnLineWidth);
  histo->SetMarkerSize(mCorrFctnMarkerSize);
  histo->SetMarkerStyle(mCorrFctnMarkerStyle);
}

//_________________
void HistoCollector1D::ReadHistograms() {

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

  for(int iCharge=0; iCharge<2; iCharge++) {
    for(int iCent=0; iCent<mCentBins; iCent++) {
      for(int iKt=0; iKt<mKtBins; iKt++) {
	
	if(mPartType == 4) { //Kaons
	  TH1F *hNumer = (TH1F*)mFile->Get(Form("hKaonQinvMixKt_%d_%d_Num_bin_%d",
						iCharge,iCent,iKt));
	  TH1F *hDenom = (TH1F*)mFile->Get(Form("hKaonQinvMixKt_%d_%d_Den_bin_%d",
						iCharge,iCent,iKt));
	  hQinvNum.push_back(hNumer);
	  hQinvDen.push_back(hDenom);
	}
	if(mPartType == 3) { //Pions
	  TH1F *hNumer = (TH1F*)mFile->Get(Form("hPionQinvMixKt_%d_%d_Num_bin_%d",
						iCharge,iCent,iKt));
	  TH1F *hDenom = (TH1F*)mFile->Get(Form("hPionQinvMixKt_%d_%d_Den_bin_%d",
						iCharge,iCent,iKt));
	  hQinvNum.push_back(hNumer);
	  hQinvDen.push_back(hDenom);
	}
      }
    }
  }
  std::cout << "\t [DONE]" << std::endl;
  std::cout << Form("Number of histos in numerator:   %d", hQinvNum.size()) 
	    << std::endl
	    << Form("Number of histos in denominator: %d", hQinvDen.size()) 
	    << std::endl;
}

//_________________
TH1F* HistoCollector1D::BuildCentCF(int charge, int binLo, int binHi) {

  TH1F *hNumSum = new TH1F("hNumSum","",mQinvBins, mQinvLo, mQinvHi);
  TH1F *hDenSum = new TH1F("hDenSum","",mQinvBins, mQinvLo, mQinvHi);
  TH1F *hRatio;
  hNumSum->Sumw2();
  hDenSum->Sumw2();
  TString sName = "h";

  sName += "CentCF_";
  sName += charge;
  sName += "_";
  sName += binLo;
  sName += "_";
  sName += binHi;

  hRatio = new TH1F(sName.Data(),"",mQinvBins,mQinvLo,mQinvHi);
  hRatio->Sumw2();

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

  MakeCorrFctn1D(hRatio, hNumSum, hDenSum);
  delete hNumSum;
  delete hDenSum;

  SetHistoStyle(hRatio,charge);
  return hRatio;
}

//_________________
TH1F* HistoCollector1D::BuildCentKtCF(int charge, int centBinLo, int centBinHi,
				      int ktBinLo, int ktBinHi) {

  TH1F *hNumSum = new TH1F("hNumSum","",mQinvBins, mQinvLo, mQinvHi);
  TH1F *hDenSum = new TH1F("hDenSum","",mQinvBins, mQinvLo, mQinvHi);
  TH1F *hRatio;
  TString sName = "h";

  sName += "CentKtCF_";
  sName += charge;
  sName += "_";
  sName += centBinLo;
  sName += "_";
  sName += centBinHi;
  sName += "_";
  sName += ktBinLo;
  sName += "_";
  sName += ktBinHi;

  hRatio = new TH1F(sName.Data(),"",mQinvBins,mQinvLo,mQinvHi);

  if(charge<2) { //Positive or negative charge
    for(Int_t iCent=centBinLo; iCent<=centBinHi; iCent++) {
      for(Int_t iKt=ktBinLo; iKt<=ktBinHi; iKt++) {

	hNumSum->Add(hQinvNum[charge*mCentBins*mKtBins+iCent*mKtBins+iKt]);
	hDenSum->Add(hQinvDen[charge*mCentBins*mKtBins+iCent*mKtBins+iKt]);
      } 
    }
  }
  else { //Sum of charges
    for(Int_t iCharge=0; iCharge<2; iCharge++) {
      for(Int_t iCent=centBinLo; iCent<=centBinHi; iCent++) {
	for(Int_t iKt=ktBinLo; iKt<=ktBinHi; iKt++) {
	  hNumSum->Add(hQinvNum[iCharge*mCentBins*mKtBins+iCent*mKtBins+iKt]);
	  hDenSum->Add(hQinvDen[iCharge*mCentBins*mKtBins+iCent*mKtBins+iKt]);
	}
      }
    }
  }

  MakeCorrFctn1D(hRatio,hNumSum,hDenSum);
  delete hNumSum;
  delete hDenSum;

  SetHistoStyle(hRatio,charge);
  return hRatio;
}

//_________________
void HistoCollector1D::NormalizeQinv(TH1 *h) {

  Int_t mNbins = h->GetNbinsX();
  Int_t mHNbins = mNbins/2;

  Double_t scale = h->Integral(1, mNbins);
  scale -= h->Integral(1, mHNbins);
  h->Scale(1./scale);
}

//_________________
void HistoCollector1D::MakeCorrFctn1D(TH1* hRat, TH1* hNum, TH1* hDen) {
  NormalizeQinv(hNum);
  NormalizeQinv(hDen);
  hRat->Divide(hNum,hDen);
}

//_________________
TH1F* HistoCollector1D::BuildNumerCentKtCF(int charge, int centBinLo, int centBinHi,
					   int ktBinLo, int ktBinHi) {

  TString sName = "h";
  sName += "NumCentKtCF_";
  sName += charge;
  sName += "_";
  sName += centBinLo;
  sName += "_";
  sName += centBinHi;
  sName += "_";
  sName += ktBinLo;
  sName += "_";
  sName += ktBinHi;
  TH1F *hNumSum = new TH1F("hNumSum","",mQinvBins, mQinvLo, mQinvHi);
  hNumSum->SetNameTitle(sName.Data(),sName.Data());

  if(charge<2) { //Positive or negative charge
    for(Int_t iCent=centBinLo; iCent<=centBinHi; iCent++) {
      for(Int_t iKt=ktBinLo; iKt<=ktBinHi; iKt++) {
	hNumSum->Add(hQinvNum[charge*mCentBins*mKtBins+iCent*mKtBins+iKt]);
      } 
    }
  }
  else { //Sum of charges
    for(Int_t iCharge=0; iCharge<2; iCharge++) {
      for(Int_t iCent=centBinLo; iCent<=centBinHi; iCent++) {
	for(Int_t iKt=ktBinLo; iKt<=ktBinHi; iKt++) {
	  hNumSum->Add(hQinvNum[iCharge*mCentBins*mKtBins+iCent*mKtBins+iKt]);
	}
      }
    }
  }

  return hNumSum;
}

//_________________
TH1F* HistoCollector1D::BuildDenomCentKtCF(int charge, int centBinLo, int centBinHi,
					   int ktBinLo, int ktBinHi) {

  TString sName = "h";
  sName += "DenCentKtCF_";
  sName += charge;
  sName += "_";
  sName += centBinLo;
  sName += "_";
  sName += centBinHi;
  sName += "_";
  sName += ktBinLo;
  sName += "_";
  sName += ktBinHi;
  TH1F *hDenSum = new TH1F("hDenSum","",mQinvBins, mQinvLo, mQinvHi);
  hDenSum->SetNameTitle(sName.Data(),sName.Data());

  if(charge<2) { //Positive or negative charge
    for(Int_t iCent=centBinLo; iCent<=centBinHi; iCent++) {
      for(Int_t iKt=ktBinLo; iKt<=ktBinHi; iKt++) {
	hDenSum->Add(hQinvDen[charge*mCentBins*mKtBins+iCent*mKtBins+iKt]);
      } 
    }
  }
  else { //Sum of charges
    for(Int_t iCharge=0; iCharge<2; iCharge++) {
      for(Int_t iCent=centBinLo; iCent<=centBinHi; iCent++) {
	for(Int_t iKt=ktBinLo; iKt<=ktBinHi; iKt++) {
	  hDenSum->Add(hQinvDen[iCharge*mCentBins*mKtBins+iCent*mKtBins+iKt]);
	}
      }
    }
  }

  return hDenSum;
}