#include "SiPixelLorentzAngleTree_.C"
#include <map>
#include <fstream>
#include <Riostream.h>
#include "tdrstyle.C"

void Draw(TH1F *h, int iLayer){

  if (h == 0)
    return;
  
  TCanvas *cc= new TCanvas("cc", "cc", 600, 600);
    
  h->UseCurrentStyle();
  h->GetXaxis()->SetTitle("production depth (#mum)");
  h->GetYaxis()->SetTitle("average drift of electrons (#mum)");
  h->SetAxisRange(0., 285);
  h->SetMarkerStyle(20);
  h->SetLineColor(1);  
  h->Draw("e0");

  TLatex *detLabel = new TLatex();  
  detLabel->SetNDC();
  detLabel->SetTextSize(0.055);	
  detLabel->SetTextAlign(10);
  string detLabelTxt="Barrel Pixel Detector";

  if (iLayer>=1)
    detLabelTxt=string(Form("Layer %d",iLayer));
  
  detLabel->DrawLatex(0.2, 0.85, detLabelTxt.c_str());

  TLatex *fitResult = new TLatex();
  fitResult->SetNDC();
  fitResult->SetTextSize(0.045);	
  fitResult->SetTextAlign(10);  
  double LA = h->GetFunction("ff")->GetParameter(1);
  double LA_err = h->GetFunction("ff")->GetParError(1);  
  string fitResultTxt="#bf{tan(#it{#theta_{LA}}) = %.4f #pm %.4f}";
  fitResult->DrawLatex(0.20,0.78,(TString)Form(fitResultTxt.c_str(),LA, LA_err));
  
  TLegend* leg = new TLegend(0.2, 0.64, .5, .76);
  leg->SetBorderSize(0);
  leg->SetTextFont(42);
  leg->SetTextSize(0.045);
  TF1 * ff= h->GetFunction("ff");
  leg->AddEntry(h, "Collision 2017", "lep");
  leg->AddEntry(ff, "Fit result", "l");
  leg->Draw();
  
  TLatex *cmsPrel = new TLatex();
  cmsPrel->SetTextSize(0.055);	
  cmsPrel->SetLineWidth(2);
  cmsPrel->SetNDC();    
  cmsPrel->SetTextAlign(11);
  cmsPrel->DrawLatex(gPad->GetLeftMargin(),1-0.85*gPad->GetTopMargin(),
		     "#scale[1.0]{CMS}#scale[0.9]{#it{#bf{ Work in progress}}}");	  
  cc->SaveAs(Form("LA_layer%d.pdf",iLayer));  
  delete detLabel;
  delete fitResult;
  delete cmsPrel;
  delete leg;
  delete cc;
}

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

float calc_residual(float trackhit_x, float rechit_x, float trackhit_y, float rechit_y){
  return TMath::Sqrt(pow(trackhit_x - rechit_x, 2)+
		     pow(trackhit_y - rechit_y, 2));      
}

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

int rejectEvent_BpixCollision(SiPixelLorentzAngleTree_ *an){
  
  bool large_pix = false;

  for (int j = 0; j < an->npix; j++){
    int colpos = static_cast<int>(an->colpix[j]);
    if (an->rowpix[j] == 0  || an->rowpix[j] == 79 ||
        an->rowpix[j] == 80 || an->rowpix[j] == 159 ||
        colpos % 52 == 0    || colpos % 52 == 51 ){
      large_pix = true;
    }
  }

  float residual = calc_residual(an->trackhit_x, an->rechit_x,
                                 an->trackhit_y, an->rechit_y);
  if (large_pix) return 1;
  if (an->pt < 3) return 1;
  if (an->clust_size_y < 4) return 1;
  if (an->clust_charge >= 1000) return 1;
  if (residual >= 0.01) return 1;
  if ((an->chi2/an->ndof) >= 2) return 1;
  if ((an->chi2/an->ndof) <= 0.2) return 1;

  return 0;
}

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

int makeFileList(string infile, int isList, vector<TString> &filenames){
  
  if (isList != 0){
    ifstream in(infile.c_str());
    
    if (!in.good()) {
      cout <<  "[Util::makeFileList] Error reading file: " << infile << endl;
      return 1;
    }
    
    string buffer;    
    while(!in.eof()){
      getline(in, buffer);        
      if (buffer.size () == 0) continue;    
      TString filename = buffer;   
      filenames.push_back(filename);
    }
  }
  else{
    filenames.push_back(infile);
  }
  
  return 0;
}
//----------------------------------------------------------------------------------------------------------
int GetMean(TH2F* h_adc, TH2F* h_adc2, TH2F* h_noadc, TH1F* &h_mean){

  int hist_drift_=h_adc->GetNbinsX();
  int hist_depth_=h_adc->GetNbinsY();
  float min_drift_=h_adc->GetXaxis()->GetBinLowEdge(1);
  float max_drift_=h_adc->GetXaxis()->GetBinUpEdge(hist_drift_);
  float min_depth_=h_adc->GetYaxis()->GetBinLowEdge(1);
  float max_depth_=h_adc->GetYaxis()->GetBinUpEdge(hist_depth_);

  static int count = 0;
  static TH1F * h_drift_depth_adc_slice_=0;
  
  if (h_drift_depth_adc_slice_ != 0)
    delete h_drift_depth_adc_slice_;
  
  h_drift_depth_adc_slice_= new
    TH1F("h_slice","h_slice", hist_drift_, min_drift_, max_drift_);
  
  for( int i = 1; i <= hist_depth_; i++){
    double npix = 0;
    
    h_drift_depth_adc_slice_->Reset("ICE");
    
    // determine sigma and sigma^2 of the adc counts and average adc counts
    //loop over bins in drift width

    for( int j = 1; j<= hist_drift_; j++){
      if(h_noadc->GetBinContent(j, i) >= 1){
	double adc_error2 = (h_adc2->GetBinContent(j,i) - h_adc->GetBinContent(j,i)*h_adc->GetBinContent(j, i) / h_noadc->GetBinContent(j, i)) /  h_noadc->GetBinContent(j, i);
	
	//cout << "no adc counts" << endl;
	h_drift_depth_adc_slice_->SetBinContent(j, h_adc->GetBinContent(j,i));
	h_drift_depth_adc_slice_->SetBinError(j, sqrt(adc_error2));
	npix += h_noadc->GetBinContent(j,i);	
      }
      else{
	h_drift_depth_adc_slice_->SetBinContent(j, h_adc->GetBinContent(j,i));
	h_drift_depth_adc_slice_->SetBinError(j, 0);
      }
    } // end loop over bins in drift width
    
    double mean = h_drift_depth_adc_slice_->GetMean(1); 
    double error = 0;
    if(npix != 0){
      error = h_drift_depth_adc_slice_->GetRMS(1) / sqrt(npix);
    }
    
    h_mean->SetBinContent(i, mean);
    h_mean->SetBinError(i, error);
    
  }// end loop over bins in depth 
  
  return 0;
}

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

int FitMeanHisto(TH1F* &h_mean){
  TF1 *ff = new TF1("ff","[0] + [1]*x", 50., 235.);    
  ff->SetParName(0,"offset");
  ff->SetParName(1,"tan#theta_{LA}");
  ff->SetParameter(0,0);
  ff->SetParameter(1,0.4);
  ff->SetLineColor(2);
  ff->SetLineWidth(3);
  int retVal = h_mean->Fit(ff,"ERQ");
  delete ff;
  return retVal;
}

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

int 
LALayer(string infile, int isList, string outfile, long int nMaxEvents=-1){

  setTDRStyle();  
  
  const int nLayers=4;
  const int hist_drift_ = 400;
  const int hist_depth_ = 100;
  const double min_drift_ = -1000;
  const double max_drift_ = 1000;
  const double min_depth_ = -100;
  const double max_depth_ = 400;
  const double width_ = 0.0285;
  
  vector<TString> filenames;  
  
  if (makeFileList(infile, isList, filenames) == 1)
    return 1;
  
  TFile *f = new TFile(outfile.c_str(),"RECREATE");  
  
  TH2F *h_drift_depth_adc[nLayers+1];
  TH2F *h_drift_depth_adc2[nLayers+1];
  TH2F *h_drift_depth_noadc[nLayers+1];
  TH1F *h_layers_mean[nLayers+1];

  for (int iLayer = 0; iLayer<=nLayers; iLayer++){
    h_drift_depth_adc[iLayer] = 
      new TH2F(Form("h_drift_depth_adc_L%d",iLayer), 
	       Form("h_drift_depth_adc_L%d",iLayer), 
	       hist_drift_ , min_drift_, max_drift_, 
	       hist_depth_, min_depth_, max_depth_);
    h_drift_depth_adc2[iLayer] = 
      new TH2F(Form("h_drift_depth_adc2_L%d",iLayer),
	       "h_drift_depth_adc2", 
	       hist_drift_ , min_drift_, max_drift_, 
	       hist_depth_, min_depth_, max_depth_);
    h_drift_depth_noadc[iLayer] = 
      new TH2F(Form("h_drift_depth_noadc_L%d",iLayer), 
	       ";drift [#mum];depth [#mum]",
	       hist_drift_ , min_drift_, max_drift_, 
	       hist_depth_, min_depth_, max_depth_);

    h_layers_mean[iLayer] =
      new TH1F(Form("meanL%d", iLayer),Form("meanL%d", iLayer),	       
	       hist_depth_, min_depth_, max_depth_);
}

  TChain *ch1 = new TChain("SiPixelLorentzAngleTree_");
  for (int i=0;i<filenames.size();++i){
    ch1->Add(filenames[i]);
  }  
  
  SiPixelLorentzAngleTree_ *an = new SiPixelLorentzAngleTree_(ch1);
  long int nentries = an->fChain->GetEntries();  
  cout << "nentries: " << nentries << endl;  
  
  int nentriesLoop = nentries; 
  
  if (nMaxEvents != -1)
    nentriesLoop = nMaxEvents;
  int goodHits = 0;
  
  for(long int ientrie = 0 ; ientrie < nentriesLoop; ientrie++){        
    int nbytes = an->LoadTree(ientrie);
    int retVal = an->GetEntry(ientrie);
    int rn = an->run;
    
    if (ientrie%10000==0)  
      cout << "Entry: " << ientrie <<  "  " 
      	   << an->run << " " << an->event << " " <<  endl;
    
    if (rejectEvent_BpixCollision(an))
      continue;

    for (int j = 0; j <  an->npix; j++){
      float dx = (an->xpix[j]  - (an->trackhit_x - width_/2. / TMath::Tan(an->trackhit_alpha))) * 10000.;
      float dy = (an->ypix[j]  - (an->trackhit_y - width_/2. / TMath::Tan(an->trackhit_beta))) * 10000.;
      float depth = dy * tan(an->trackhit_beta);
      float drift = dx - dy * tan(an->trackhit_gamma_);
      h_drift_depth_adc[0]->Fill(drift, depth, an->adc[j]);
      h_drift_depth_adc2[0]->Fill(drift, depth, an->adc[j]*an->adc[j]);
      h_drift_depth_noadc[0]->Fill(drift, depth);
      h_drift_depth_adc[an->layer]->Fill(drift, depth, an->adc[j]);
      h_drift_depth_adc2[an->layer]->Fill(drift, depth, an->adc[j]*an->adc[j]);
      h_drift_depth_noadc[an->layer]->Fill(drift, depth);      
    }    
  }
  
  for (int iLayer = 0; iLayer<=nLayers; iLayer++){
    GetMean(h_drift_depth_adc[iLayer], h_drift_depth_adc2[iLayer],
	    h_drift_depth_noadc[iLayer], h_layers_mean[iLayer]);
    FitMeanHisto(h_layers_mean[iLayer]);    
    Draw(h_layers_mean[iLayer], iLayer);
  }
  
  for (int iLayer = 0; iLayer<=nLayers; iLayer++){
    h_drift_depth_adc[iLayer]->Write();
    h_drift_depth_adc2[iLayer]->Write();
    h_layers_mean[iLayer]->Write();

  }
  
  f->Close();
  return 0;
  
}