//  
//  landauFit.C
//  Obtaining the MPV from a fit of the ADC counts distribution
//
//  Created by Christian Elsasser on 31.05.13.
//  University of Zurich, elsasser@cern.ch
//  Copyright only for commercial use
//
//  Compile with:
//  g++ -o landauFit landauFit.C `root-config --glibs --cflags` -lRooFit -lRooFitCore -lEG
//

//  General include
#include <iostream>
#include <fstream>
#include <vector>
#include <map>
#include <stdlib.h>
#include <limits>
#include <string>
#include <time.h>


//  ROOT include
#include "../../include/incBasicROOT.h"
#include "../../include/incDrawROOT.h"
#include "../../include/incHistROOT.h"
#include "../../include/incIOROOT.h"
#include "../../include/incRooFit.h"


#include "../../include/basicROOTTools.h"
#include "../../include/basicRooFitTools.h"


#include "../../include/lhcbstyle.C"


#include "deplTool.h"


// -d detector (TT/IT)
// -l layer (TTaX,TTaU,TTbV,TTbX, T[1-3]{X1,U,V,X2})
// -s read-out sector number
// -f fill number
// -c Calibration step (Odin step) [0-65]
// -r consider only hit in a distance closer than r mm from the beam axis
// -v number of summed up strips
// -ns Do not save the value in the root file
// -ss Perform a temporary save


int landauFit(int argc, char* argv[]);


int main(int argc, char* argv[]){
  

  Printf("Starting...");

  // if i want to use this, set $LHCBSTYLE to either 0 or 1
  TStyle* style = lhcbStyle();
  
  int argc_copy = argc;
  
  TApplication* theApp = new TApplication("Analysis", &argc, argv);
  argc = theApp->Argc();
  argv = theApp->Argv();
  
  bool runBatch = false;
  
  for (int i = 1; i<argc; i++) {
    if (strcmp(argv[i],"-b")==0) {
      runBatch = true;
      gROOT->SetBatch();
      Printf("Set BATCH mode...");
    }
  }
  

  
  int exit = landauFit(argc,argv);
  
  
  
  Printf("End");
  if (!runBatch) {
    theApp->Run(!runBatch);
  }
  
  
  return exit;
}


//  Executable method
int landauFit(int argc, char* argv[]){
  
  TVectorD v_res(5);
  
  gStyle->SetPadLeftMargin(0.20);
  gStyle->SetTitleOffset(1.5,"Y");
  gStyle->SetTitleOffset(1.0,"X");

  // To point it to the right directory:
  // export DISK=/disk/data1/hep/elena
  const char* diskVar = std::getenv ("DISK");
  Printf("%s",diskVar);
  
  int fill      = 2083;
  TString det("TT");
  TString lay("TTaU");
  int sector    = 2634;
  int caliStep  = 0;
  int val       = 3;
  bool notSave  = false;
  bool temSave  = false;
  double radius = -10.0;
  
  for (int i=1; i<argc; i++) {
    if (strcmp(argv[i],"-d")==0 && i+1<argc) {
      det = argv[++i];
    }else if (strcmp(argv[i],"-l")==0 && i+1<argc) {
      lay = argv[++i];
    }else if (strcmp(argv[i],"-s")==0 && i+1<argc) {
      sector = std::atoi(argv[++i]);
    }else if (strcmp(argv[i],"-f")==0 && i+1<argc) {
      fill = std::atoi(argv[++i]);
    }else if (strcmp(argv[i],"-v")==0 && i+1<argc) {
      val = std::atoi(argv[++i]);
    }else if (strcmp(argv[i],"-c")==0 && i+1<argc) {
      caliStep = std::atoi(argv[++i]);
    }else if (strcmp(argv[i],"-r")==0 && i+1<argc) {
      radius = std::atof(argv[++i]);
    }else if (strcmp(argv[i],"-ns")==0) {
      notSave = true;
    }else if (strcmp(argv[i],"-ss")==0) {
      temSave = true;
    }
  }
  
  
  
  int voltStep  = caliStep/6;
  int timeStep  = caliStep%6;
  
  double time   = STTool::GetTime(caliStep, fill);
  double volt   = 0.0;
  
  
  if (det.EqualTo("TT")) {
    volt = STTool::GetTTVoltage(caliStep);
  }else{
    volt = STTool::GetITVoltage(caliStep);
  }
  
  // Name of the vector where it is stored
  TString vn_result(Form("v_%d_val%d",caliStep,val));
  
  // Directory name and file names for the plots
  TString dn_graph(Form("%s/data/ST/graphics/%s/%s/%d/%d",
                        diskVar,det.Data(),lay.Data(),sector,fill));
  TString fn_graph_s(Form("%s/signal_%d_val%d",
                          dn_graph.Data(),caliStep,val));
  TString fn_graph_n(Form("%s/noise_%d_val%d",
                          dn_graph.Data(),caliStep,val));
  TString fn_graph_a(Form("%s/alt_%d_val%d",
                          dn_graph.Data(),caliStep,val));
  
  // Director name and file name for the data
  TString dn_output(Form("%s/data/ST/Aging",diskVar));
  TString fn_output(Form("%s/CCEScan.root",dn_output.Data()));
  Printf("%s", fn_output.Data());
  TString fn_dummy(Form("%s/check",dn_output.Data()));
  
  // Name if it is not directly written into the main file
  if (temSave) {
    fn_output= TString(Form("%s/temp_%s_%s_%d_%d_r%d_step%d.root",
                            dn_output.Data(),det.Data(),lay.Data(),
                            sector,fill,(int)radius, (int)caliStep));
  }

  
  // Set the radius to an irrelevant value
  if (radius<0.0) {
    radius = 2000.0;
  }else{
    vn_result  +=Form("_r%d",(int)radius);
    fn_graph_s +=Form("_r%d",(int)radius);
  }
  
  fn_graph_s += ".pdf";
  fn_graph_n += ".pdf";
  
  
  Printf("=========================================");
  Info("landauFit","Analyse: ");
  Printf("               Sector:    %d",sector);
  Printf("               Radius:    %4.2f mm",radius);
  Printf("               Fill:      %d",fill);
  Printf("               Detector:  %s",det.Data());
  Printf("               Layer:     %s",lay.Data());
  Printf("               Voltage:   %4.2f V",volt);
  Printf("               Time:      %4.2f ns",time);
  
  
  // Obtain directory, file and tree names for the ADC counts
  TString dn_data(Form("%s/data/ST/Aging_Tuples/%s/%s"
                        ,diskVar,det.Data(),lay.Data()));
  TString fn_data(Form("%s/%d.root",
                       dn_data.Data(),fill));
  TString tn_sig(Form("STADCTrackMonitor/HitInfo/%s",
                      lay.Data()));
  TString tn_bkg(Form("STADCTrackMonitor/NoiseInfo/%s",
                      lay.Data()));
  
  
  

  TFile* f_data       = TFile::Open(fn_data.Data());
  
  
  if (!f_data) {
    Error("landauFit","Data file does not exist!");
    Printf("=========================================\n");
    return EXIT_FAILURE;
  }
  
  TTree* t_sig = (TTree*)f_data->Get(tn_sig.Data());
  if (!t_sig) {
    Error("landauFit","Signal tree not found!");
    Printf("=========================================\n");
    return EXIT_FAILURE;
  }
  
  TTree* t_bkg = (TTree*)f_data->Get(tn_bkg.Data());
  if (!t_bkg) {
    Error("landauFit","Background tree not found!");
    Printf("=========================================\n");
    return EXIT_FAILURE;
  }

  
  
  TH1F* h_sig = new TH1F("h_sig","Signal histo",
                         208,-32.0,176.0);
  TH1F* h_bkg = new TH1F("h_bkg","Background histo",
                         208,-32.0,176.0);

  // Select calibration and read-out sector
  TCut cu_bkg(Form("(odinStep==%d) && (sector==%d)",
                   caliStep,sector));
  
  // Select only first run in fill if you consider calibration step 0
  if (caliStep==0) {
    cu_bkg = cu_bkg && TCut(Form("(runNumber==%d)",STTool::zeroRunMap[fill]));
  }
  
  // Select tracks based on their chi^2 and ghost prob
  //TCut cu_sel = STTool::cu_track_TT;
  //if (strcmp(det.Data(),"TT")==0) {
  //  cu_sel = STTool::cu_track_IT;
  //}
  TCut cu_sel = STTool::get_track_selection(det, fill);
  
  // Get full selection for signal tracks/hits including radius
  TCut cu_sig = cu_bkg && cu_sel && TCut(Form("(TMath::Sqrt(xHit*xHit+yHit*yHit)<%f)",radius));

  
  // Fill Histograms

    t_sig->Draw(Form("val%d>>h_sig",val),cu_sig,"goff");
    t_bkg->Draw(Form("val%d>>h_bkg",val),cu_bkg,"goff");


  
  delete t_sig;
  delete t_bkg;
  
  // Skip histogram if there are too few entries in the histograms
  if (h_sig->GetEntries()<10 || h_bkg->GetEntries()<10) {
    Error("landauFit","Too few entries in histograms. Skip this step!");
    return EXIT_FAILURE;
  }
  
  // ADC value
  RooRealVar* adc = new RooRealVar("adc","Signal height",
                                   h_sig->GetXaxis()->GetXmin(),
                                   h_sig->GetXaxis()->GetXmax(),
                                   "ADC value");
  
  

  // RooFit histograms
  RooDataHist* hr_bkg = new  RooDataHist("hr_bkg","Histogram of raw noise",
                                         *adc,RooFit::Import(*h_bkg));
  RooDataHist* hr_sig = new  RooDataHist("hr_sig","Histogram of raw signal",
                                         *adc,RooFit::Import(*h_sig));
  
  // PDF for noise
  RooRealVar* nMu   = new RooRealVar("nMu","Mean of raw noise distribution",0.0,-10.0,10.0,"");
  RooRealVar* nSig1 = new RooRealVar("nSig1","Sigma 1 of raw noise distribution",3.0,0.5,6.0,"");
  RooRealVar* nSig2 = new RooRealVar("nSig2","Sigma 2 of raw noise distribution",5.0,0.5,25.0,"");
  RooRealVar* nFrac = new RooRealVar("nFrac","Fractition of Gaussian 1 to total noise distribution",0.5,0.0,1.0);
  
  RooGaussian* nG1  = new RooGaussian("nG1","Gaussian 1 of raw noise distribution",*adc,*nMu,*nSig1);
  RooGaussian* nG2  = new RooGaussian("nG2","Gaussian 2 of raw noise distribution",*adc,*nMu,*nSig2);
  
  RooAddPdf* nPDF   = new RooAddPdf("nPDF","Total distribution of raw noise",
                                    RooArgList(*nG1,*nG2),
                                    RooArgList(*nFrac));

  
  
  // Fit and plot noise sample
  RooFitResult* nRes = nPDF->fitTo(*hr_bkg,RooFit::Save());
  RooPlot* nFrame    = adc->frame(RooFit::Title(Form("Noise distribution for %4.0f V and %4.1f ns",volt,time)));
  hr_bkg->plotOn(nFrame,RooFit::DataError(RooAbsData::SumW2),RooFit::DataError(RooAbsData::SumW2));
  nPDF->plotOn(nFrame,RooFit::LineColor(kBlue));
  
  
  // average noise value
  v_res(2)     = nSig1->getVal()*nFrac->getVal()+nSig2->getVal()*(1.0-nFrac->getVal());
  // uncertainty on average noise
  v_res(3)     = TMath::Sqrt(TMath::Power(nSig1->getError()*nFrac->getVal(),2.0)+
                             TMath::Power(nSig2->getError()*(1.0-nFrac->getVal()),2.0)+
                             TMath::Power((nSig1->getVal()-nSig2->getVal())*nFrac->getError(),2.0));
  
  
  // PDF for signal
  
  // Constrain parameters fixed in the noise fit
  //nMu->setConstant(kTRUE);
  //nSig1->setConstant(kTRUE);
  //nSig2->setConstant(kTRUE);
  nFrac->setConstant(kTRUE);
  
  RooGaussian* sG1 = new RooGaussian("sG1","Gaussian 1",*adc,*nMu,*nSig1);
  RooGaussian* sG2 = new RooGaussian("sG2","Gaussian 2",*adc,*nMu,*nSig2);
  
  RooAddPdf* sG    = new RooAddPdf("sG","Gaussian smearing",
                                   RooArgList(*sG1,*sG2),
                                   RooArgList(*nFrac));
  
  RooAddPdf* nG    = new RooAddPdf("nG","Distribution of raw noise",
                                   RooArgList(*nG1,*nG2),
                                   RooArgList(*nFrac));
  
  
  RooRealVar* mpv  = new RooRealVar("mpv","Most probable ADC value",
                                    TMath::Max(10.0,h_sig->GetBinCenter(h_sig->GetMaximumBin()))
                                    ,1.0,60.0,"");
  RooRealVar* spread  = new RooRealVar("spread","Sigma of Landau distribution"
                                       ,3.0,0.1,10.0,"");
  
  RooRealVar* ampv = new RooRealVar("ampv","Most probable ADC value",
                                    TMath::Max(10.0,h_sig->GetBinCenter(h_sig->GetMaximumBin()))
                                    ,1.0,60.0,"");
  RooRealVar* aspread = new RooRealVar("aspread","Sigma of Landau distribution"
                                       ,3.0,0.1,10.0,"");
  
  // Variables for the Landau coming from photo conversion
  RooFormulaVar* scaledMpv    = new RooFormulaVar("scaledMpv","2*mpv",RooArgList(*mpv));
  RooFormulaVar* scaledSpread = new RooFormulaVar("scaledSpread","2*spread",RooArgList(*spread));
  
  
  RooLandau* landauDist       = new RooLandau("landauDist","Landau distribution of ADC values",*adc,*mpv,*spread);
  RooLandau* landauDistScaled = new RooLandau("landauDistScaled","Landau distribution of conversion",*adc,*scaledMpv,*scaledSpread);
  
  RooLandau* altlandauDist    = new RooLandau("altlandauDist","AlternativeLandau distribution of ADC values",*adc,*ampv,*aspread);
  
  RooRealVar* cFrac           = new RooRealVar("cFrac","Fraction of photo conversion",0.025);

  // Set fraction from photon conversion to zero
  if (det.EqualTo("IT")) {
    cFrac->setVal(0.0);
  }
  
  RooRealVar* sFrac           = new RooRealVar("sFrac","Fraction of signal hits",0.8,0.0,1.0);
  RooRealVar* aFrac           = new RooRealVar("aFrac","Fraction of signal hits",0.8,0.0,1.0);
  
  adc->setBins(10000,"cache") ;
  
  RooFFTConvPdf* landauConv         = new RooFFTConvPdf("landauConv","Landau (x) gauss",*adc,*landauDist,*sG);
  RooFFTConvPdf* landauConvPhoto    = new RooFFTConvPdf("landauConvPhoto","Landau Conv (x) gauss",*adc,*landauDistScaled,*sG);

  RooAddPdf* sTot             = new RooAddPdf("sTot","Total signal distribution",*landauConvPhoto,*landauConv,*cFrac);
  RooAddPdf* sPDF             = new RooAddPdf("sPDF","Total distribution",*sTot,*nG,*sFrac);
  
  // Alternative distribution only done from a single Landau
  RooFFTConvPdf* altlandauConv      = new RooFFTConvPdf("altlandauConv","Landau Conv (x) gauss",*adc,*altlandauDist,*sG);
  RooAddPdf* aPDF             = new RooAddPdf("aPDF","Alternative total distribution",*altlandauConv,*nG,*aFrac);
  
  
  // Fit the distribution
  RooFitResult* sRes = sPDF->fitTo(*hr_sig,RooFit::Save());
  
  
  
  Info("landauFit","Baseline result:");
  Printf("               mpv        = %8.4f+/-%8.4f",mpv->getVal(),mpv->getError());
  Printf("               spread     = %8.4f+/-%8.4f",spread->getVal(),spread->getError());
  Printf("               conv Frac  = %8.4f+/-%8.4f",cFrac->getVal(),cFrac->getError());
  Printf("               noise Frac = %8.4f+/-%8.4f",1.0 - sFrac->getVal(),sFrac->getError());
  
  // mpv value: 1st place in result vector
  v_res(0)     = mpv->getVal();
  // mpv uncertainty: 2nd place in result vector
  v_res(1)     = mpv->getError();
  
  
  
  
  RooPlot* sFrame = adc->frame(RooFit::Title(Form("Signal distribution for %4.0f V and %4.1f ns",volt,time)));
  hr_sig->plotOn(sFrame,RooFit::DataError(RooAbsData::SumW2),RooFit::DataError(RooAbsData::SumW2));
  sPDF->plotOn(sFrame,RooFit::LineColor(kBlue));
  sPDF->plotOn(sFrame,RooFit::Components(*nG),RooFit::LineColor(kBlue),RooFit::LineStyle(kDashed));
  sPDF->plotOn(sFrame,RooFit::Components(*landauConv),RooFit::LineColor(kRed),RooFit::LineStyle(9));
  sPDF->plotOn(sFrame,RooFit::Components(*landauConvPhoto),RooFit::LineColor(kRed),RooFit::LineStyle(kDashed));

  
  Info("landauFit","Results: ");
  Printf("   MPV:              %5.2f+/-%5.2f ADC counts",v_res(0),v_res(1));
  Printf("   Noise:            %5.2f+/-%5.2f ADC counts",v_res(2),v_res(3));

  
  if (nRes->status()==0) {
    Info("landauFit","Signal fit status: %d",nRes->status());
  }else{
    Warning("landauFit","Signal fit status: %d",nRes->status());
  }
  
  if (sRes->status()==0) {
    Info("landauFit","Noise fit status:  %d",sRes->status());
  }else{
    Warning("landauFit","Noise fit status:  %d",sRes->status());
  }
  

  // Plot the results
  
  TCanvas* c_sig = new TCanvas("c_sig","Signal Canvas",800,600);
  c_sig->SetLeftMargin(0.20);
  SwapParenthesis(sFrame);
  ReplaceSubstring(sFrame,"Events","Hits");
  ReplaceSubstring(sFrame,"1 ","");
  sFrame->GetYaxis()->SetTitleOffset(1.5);
  sFrame->Draw();
  
  TCanvas* c_bkg = new TCanvas("c_bkg","Background Canvas",800,600);
  c_bkg->SetLeftMargin(0.20);
  nFrame->GetXaxis()->SetLimits(-32.0,32.0);
  SwapParenthesis(nFrame);
  ReplaceSubstring(nFrame,"Events","Hits");
  ReplaceSubstring(nFrame,"1 ","");
  nFrame->GetYaxis()->SetTitleOffset(1.5);
  nFrame->Draw();
  

  // Save data and plots
  
  TFile* f_output;
  
  while (!gSystem->Exec(Form("ls %s &> /dev/null",fn_dummy.Data()))) {
    Info("landauFit","Waiting until file is closed");
    gSystem->Sleep(1000);
  }
  gSystem->Exec(Form("touch %s &> /dev/null",fn_dummy.Data()));
  
  f_output = new TFile(fn_output.Data(),"UPDATE");
  
  
  TString dn_result = TString::Format("%s/%s/%d/%d",
                                      det.Data(),
                                      lay.Data(),
                                      sector,
                                      fill);
  
  // Get right directory
  bool isThere = f_output->GetListOfKeys()->Contains(det.Data());
  if (!isThere) {
    f_output->mkdir(det.Data());
  }
  
  isThere = f_output->GetListOfKeys()->Contains(Form("%s/%s",
                                                     det.Data(),
                                                     lay.Data()));
  if (!isThere) {
    f_output->mkdir(Form("%s/%s",
                         det.Data(),
                         lay.Data()));
  }
  
  isThere = f_output->GetListOfKeys()->Contains(Form("%s/%s/%d",
                                                     det.Data(),
                                                     lay.Data(),
                                                     sector));
  if (!isThere) {
    f_output->mkdir(Form("%s/%s/%d",
                         det.Data(),
                         lay.Data(),
                         sector));
  }
  
  isThere = f_output->GetListOfKeys()->Contains(Form("%s/%s/%d/%d",
                                                    det.Data(),
                                                    lay.Data(),
                                                    sector,
                                                    fill));
  if (!isThere) {
    f_output->mkdir(Form("%s/%s/%d/%d",
                         det.Data(),
                         lay.Data(),
                         sector,
                         fill));
  }
  
  f_output->cd(Form("%s/%s/%d/%d",
                    det.Data(),
                    lay.Data(),
                    sector,
                    fill));
  if(!notSave){
    gDirectory->Delete(Form("%s;*",
                            vn_result.Data()));
    
    v_res.Write(Form("%s",
                     vn_result.Data()));
    
    Info("landauFit","Result written to %s",fn_output.Data());
  }
  
  f_output->Close();
  gSystem->Exec(Form("rm -f %s &> /dev/null",fn_dummy.Data()));
  
  if (!notSave) {
    // Save picture
    gSystem->Exec(Form("mkdir -p %s",dn_graph.Data()));
    gSystem->Exec(Form("rm %s &> /dev/null",fn_graph_s.Data()));
    gSystem->Exec(Form("rm %s &> /dev/null",fn_graph_n.Data()));
    c_sig->SaveAs(fn_graph_s.Data());
    c_bkg->SaveAs(fn_graph_n.Data());
  }
  
  
  Printf("=========================================\n");

  return EXIT_SUCCESS;
}