#include <set>
#include <string>
#include <iostream>
#include <iomanip>
#include <stdio.h>

#include "TStyle.h"
#include "TROOT.h"
#include "TFile.h"
#include "TLine.h"
#include "TF1.h"
#include "TGraph.h"
#include "TH1D.h"
#include "TH2D.h"
#include "TPaletteAxis.h"
#include "TBox.h"
#include "TLatex.h"
#include "TCanvas.h"
#include "TLegend.h"

#include "JLang/JToken.hh"

#include "Jeep/JPrint.hh"
#include "Jeep/JParser.hh"
#include "Jeep/JMessage.hh"
#include "Jeep/JProperties.hh"

#include "JDetector/JPMTParameters.hh"
#include "JDetector/JPMTAnalogueSignalProcessor.hh"

#include "JGizmo/JGizmoToolkit.hh"

#include "JROOT/JManager.hh"


/**
 * \file
 *
 * Program to create plots for PMT-modeling documentation.
 * \author bjjung
 */
int main(int argc, char **argv)
{
  using namespace std;
  using namespace JPP;

  typedef JToken<';'>                                  JToken_t;
  typedef JRange<Double_t>                             JRange_t;

  string                  inputFile;
  string                  outputDataDir;
  string         	  outputFigureDir;

  Double_t                Nth    = 20;
  Double_t                Nthb   = 16;

  JRange_t                thRange (0.17, 0.37);
  JRange_t                thbRange(0.0,  0.32);

  set<double>             Pselection;
  int                     Nrows = -1;
  int                     Ncols = -1;
  
  int            	  NPE;
  int                     debug;

  try {

    JProperties properties;

    properties.insert(gmake_property(Nth));
    properties.insert(gmake_property(thRange));

    properties.insert(gmake_property(Nthb));
    properties.insert(gmake_property(thbRange));

    properties.insert(gmake_property(Pselection));
    properties.insert(gmake_property(Nrows));
    properties.insert(gmake_property(Ncols));
    
    JParser<> zap("Example program to histogram time-over-threshold distributions.");

    zap['f'] = make_field(inputFile)       = JPARSER::initialised();
    zap['@'] = make_field(properties)      = JPARSER::initialised();
    zap['D'] = make_field(outputDataDir)   = "/tmp";
    zap['F'] = make_field(outputFigureDir) = "/tmp";
    zap['N'] = make_field(NPE)             = 1;
    zap['d'] = make_field(debug)           = 1;

    zap(argc, argv);
  }
  catch(const exception &error) {
    FATAL(error.what() << endl);
  }

  const int plotwidth  = 800;
  const int plotheight = 600;

  JPMTParameters parameters;
  const JPMTAnalogueSignalProcessor cpu(parameters);


  gStyle -> SetOptStat(0);

  string outputFile = outputDataDir + "/PMTmodeling.root";
  TFile out(outputFile.c_str(), "RECREATE");


  // Analogue pulse shape

  const double dT    =   0.1;
  const double Tmin  = -20.0;
  const double Tmax  =  50.0;
  const double nT    = (int) ((Tmax - Tmin) / dT);

  const double ampl  = 0.6;                 // V
  const double sigma = 3.0;                 // ns
  const double tau   = 5.0;                 // ns
  const double t_m   = sigma * sigma / tau; // ns

  const double C     = exp(-0.5 * sigma * sigma / tau / tau);

  TGraph gr_pulseGaus(nT);
  gr_pulseGaus.SetLineColor(9);
  gr_pulseGaus.SetLineWidth(3);
  gr_pulseGaus.SetLineStyle(kDashed);
  TGraph gr_pulseExp(nT);
  gr_pulseExp.SetLineColor(7);
  gr_pulseExp.SetLineWidth(3);  
  gr_pulseExp.SetLineStyle(kDashed);
  TGraph gr_pulseTot(nT);
  gr_pulseTot.SetLineColor(4);
  gr_pulseTot.SetLineWidth(4);  

  for (int i=0; i < nT; ++i) {

    const double T    = Tmin + i * dT;
    const double Gaus = ampl * exp(-0.5 * T * T / sigma / sigma);
    const double Exp  = ampl * exp(-T / tau) / C;

    gr_pulseGaus.SetPoint(i, T, Gaus);
    gr_pulseExp. SetPoint(i, T, Exp);
    gr_pulseTot. SetPoint(i, T, (T > t_m ? Exp : Gaus));
  }

  TCanvas c_pulse("c_pulse", "", plotwidth, plotheight);

  c_pulse.SetLeftMargin  (0.15);
  c_pulse.SetBottomMargin(0.15);  

  gr_pulseGaus.SetTitle(NULL);  
  gr_pulseGaus.Draw("AL");
  gr_pulseExp.Draw("L Same");
  gr_pulseTot.Draw("L Same");

  gr_pulseGaus.GetXaxis()->SetLabelFont(63);
  gr_pulseGaus.GetYaxis()->SetLabelFont(63);
  
  gr_pulseGaus.GetXaxis()->SetLabelSize(24);
  gr_pulseGaus.GetYaxis()->SetLabelSize(24);

  gr_pulseGaus.GetXaxis()->SetTitleSize(0.05);
  gr_pulseGaus.GetYaxis()->SetTitleSize(0.05);

  gr_pulseGaus.GetXaxis()->SetTitle("time [ns]");
  gr_pulseGaus.GetYaxis()->SetTitle("voltage [V]");

  const double VminView =   0.0;
  const double VmaxView =   0.7;
  const double TminView = -15.0;
  const double TmaxView =  30.0;
  
  gr_pulseGaus.GetXaxis()->SetRangeUser(TminView, TmaxView);
  gr_pulseGaus.GetYaxis()->SetRangeUser(VminView, VmaxView);

  TLine  l_tMatch(t_m,      VminView, t_m,      VmaxView);
  TLine  l_vMatch(TminView, ampl * C, TmaxView, ampl * C);
  
  l_tMatch.SetLineStyle(6);
  l_tMatch.SetLineWidth(3);
  l_tMatch.SetLineColor(13);
  
  l_vMatch.SetLineStyle(6);
  l_vMatch.SetLineWidth(3);  
  l_vMatch.SetLineColor(13);
  
  l_tMatch.Draw("Same");
  l_vMatch.Draw("Same");
  
  TLegend L_pulse(0.65,0.65,0.9,0.9);
  
  L_pulse.AddEntry(&gr_pulseGaus, "Gaus");
  L_pulse.AddEntry(&gr_pulseExp,  "Exp");
  L_pulse.AddEntry(&gr_pulseTot,  "Total");
  
  L_pulse.Draw("Same");
  
  c_pulse.SaveAs(MAKE_CSTRING(outputFigureDir << "/pulseShape.png"));
  

  // Time-over-threshold versus charge

  const double dQ   =   0.001;
  const double Qmin =   0.0     - 0.5 * dQ;
  const double Qmax = 100.0*NPE + 0.5 * dQ;
  const int    nq   = (int) ((Qmax - Qmin) / dQ);

  TH1D h_Qtot = TH1D("h_Qtot", NULL, nq, Qmin, Qmax);
  for (int i=1; i <= h_Qtot.GetNbinsX(); ++i) {

    double npe = h_Qtot.GetBinCenter(i);

    try {

      if (cpu.getThresholdDomain(npe) == cpu.ABOVE_THRESHOLD) {
	h_Qtot.SetBinContent(i,cpu.getTimeOverThreshold(npe));
      }

    } catch (const JValueOutOfRange& exception) {

      continue;
    }
  }

  TCanvas c_Qtot("c_Qtot","", plotwidth, plotheight);

  c_Qtot.SetLeftMargin  (0.15);
  c_Qtot.SetBottomMargin(0.15);  
  
  h_Qtot.Draw("L");
  h_Qtot.SetTitle(NULL);
  h_Qtot.SetLineWidth(3);
  h_Qtot.GetXaxis()->SetTitleSize(0.05);
  h_Qtot.GetYaxis()->SetTitleSize(0.05);
  h_Qtot.GetXaxis()->SetTitle("charge [npe]");
  h_Qtot.GetYaxis()->SetTitle("time-over-threshold [ns]");

  c_Qtot.SaveAs(MAKE_CSTRING(outputFigureDir << "/ToTvsQ.png"));

  TCanvas c_QtotZoom("c_QtotZoom","", plotwidth, plotheight);

  c_QtotZoom.SetLeftMargin  (0.15);
  c_QtotZoom.SetBottomMargin(0.15);  
  
  h_Qtot.Draw("L");
  h_Qtot.GetXaxis()->SetRangeUser(0.0, 5.0);

  h_Qtot.SetTitle(NULL);
  h_Qtot.GetXaxis()->SetTitleSize(0.05);  
  h_Qtot.GetYaxis()->SetTitleSize(0.05);  
  h_Qtot.GetXaxis()->SetTitle("charge [npe]");
  h_Qtot.GetYaxis()->SetTitle("Time-over-Threshold [ns]");
  

  c_QtotZoom.SaveAs(MAKE_CSTRING(outputFigureDir << "/ToTvsQzoom.png"));

  h_Qtot.Write();


  // Single photo-electron charge distribution
  
  const double thresholdMin = parameters.threshold - parameters.thresholdBand;
  const double f_UA  = parameters.gainSpread * parameters.gainSpread;

  TH1D h_SPE   = TH1D("h_SPE", NULL,   nq, Qmin, Qmax);
  TH1D h_SPEfA = TH1D("h_SPEfA", NULL, nq, Qmin, Qmax);
  TH1D h_SPEuA = TH1D("h_SPEuA", NULL, nq, Qmin, Qmax);
  
  for (int i=1; i <= h_SPE.GetNbinsX(); ++i) {

    double npe = h_SPE.GetBinCenter(i);

    double norm = 0.0;
    std::vector<double> probs(NPE+1);
    std::vector<double> norms(NPE+1);
    for (int k=0; k<=NPE; ++k) {

      const double mu    = ((NPE-k) * f_UA * parameters.gain) + (k * parameters.gain);
      const double sigma = sqrt((NPE-k) * f_UA + k) * cpu.getGainSpread(1);
      const double W     = TMath::Binomial(NPE,k) * pow(parameters.PunderAmplified,NPE-k) * pow(1-parameters.PunderAmplified,k);

      norms[k] = W * 0.5 * TMath::Erfc( (thresholdMin - mu) / sqrt(2.0) / sigma );
      probs[k] = W * TMath::Gaus(npe, mu, sigma, kTRUE);
      norm    += norms[k];
    }

    h_SPE.SetBinContent(i,cpu.getChargeProbability(npe,NPE));
    h_SPEfA.SetBinContent(i,probs[NPE]/norm);
    h_SPEuA.SetBinContent(i,probs[0]/norm);
  }

  TCanvas c_SPE("c_SPE","", plotwidth, plotheight);
  TLegend l_SPE(0.58,0.7,0.9,0.9);

  c_SPE.SetLeftMargin  (0.15);
  c_SPE.SetBottomMargin(0.15);

  const double QmaxView = 2.3 * NPE;
  const double pMaxView = 1.2 * h_SPE.GetMaximum();
  const double xtxt     = 0.7 * QmaxView;
  const double ytxt     = 0.5 * pMaxView;

  TLatex pUAtext(xtxt,ytxt,("p = " + std::to_string(parameters.PunderAmplified).substr(0,4)).c_str());

  h_SPE.Draw("L");
  h_SPE.SetLineColor(kRed);
  h_SPE.SetLineWidth(3);
  h_SPE.GetXaxis()->SetTitleSize(0.06);
  h_SPE.GetYaxis()->SetTitleSize(0.06);  
  h_SPE.GetXaxis()->SetTitle("charge [npe]");
  h_SPE.GetYaxis()->SetTitle("probability [npe^{-1}]");
  h_SPE.GetXaxis()->SetRangeUser(0.,QmaxView);
  h_SPE.GetYaxis()->SetRangeUser(0.,pMaxView);
  l_SPE.AddEntry(&h_SPE,"SPE distribution");
  h_SPEfA.Draw("Same");
  h_SPEfA.SetLineColor(1);
  h_SPEfA.SetLineWidth(2);
  h_SPEfA.GetXaxis()->SetRangeUser(0.,QmaxView);
  h_SPEfA.GetYaxis()->SetRangeUser(0.,pMaxView);
  l_SPE.AddEntry(&h_SPEfA,"Nominal Amplification");
  h_SPEuA.Draw("Same");
  h_SPEuA.SetLineColor(16);
  h_SPEuA.SetLineWidth(2);
  h_SPEuA.GetXaxis()->SetRangeUser(0.,QmaxView);
  h_SPEuA.GetYaxis()->SetRangeUser(0.,pMaxView);
  l_SPE.AddEntry(&h_SPEuA,"Underamplified");

  TBox b_thresholdBand = TBox(thresholdMin, 0.0,
  			      parameters.threshold, pMaxView);
  b_thresholdBand.SetFillColor(kOrange);
  b_thresholdBand.SetFillStyle(3001);
  b_thresholdBand.Draw("Same");
  l_SPE.AddEntry(&b_thresholdBand, "Threshold band", "f");
  
  l_SPE.Draw("Same");
  l_SPE.SetFillStyle(0);
  l_SPE.SetTextSize(0.037);
  
  //l_SPE.SetTextSize(0.025);
  //pUAtext.Draw("Same");
  
  c_SPE.SaveAs(MAKE_CSTRING(outputFigureDir << "/SPEdistr.png"));

  h_SPE.Write();
  h_SPEfA.Write();
  h_SPEuA.Write();



  // Survival probabilities for different underamplification probabilities

  const int maxNPE = 10;
  
  const int NevalPua 	= 5;
  const double maxPua 	= 0.5;
  const double minPua 	= 0.00;
  const double dPua 	= (maxPua - minPua) / ((double) NevalPua );

  TCanvas c_survival("c_survival","", plotwidth, plotheight);
  TLegend l_survival = TLegend(0.7,0.2,0.9,0.4);

  TGraph *grs_survival = new TGraph[NevalPua];
  
  for (int i=1; i<=NevalPua; ++i) {

    double p_UA = minPua + i * dPua;

    parameters.PunderAmplified		  = p_UA;
    const JPMTAnalogueSignalProcessor CPU = JPMTAnalogueSignalProcessor(parameters);

    std::string name = "PunderAmplified=" + std::to_string(p_UA);

    grs_survival[i-1] = TGraph(maxNPE);
    grs_survival[i-1].SetName(name.c_str());
    
    for (int j=1; j<=maxNPE; ++j) {

      double pSurv = CPU.getSurvivalProbability(j);
      grs_survival[i-1].SetPoint(j,j,pSurv);

    }

    if (i-1 == 0) {
      
      grs_survival[i-1].Draw("apl");
      grs_survival[i-1].GetXaxis() -> SetTitle("number of photons");
      grs_survival[i-1].GetYaxis() -> SetTitle("survival probability");
      
    } else {
      
      grs_survival[i-1].Draw("pl Same");
    }
    
    grs_survival[i-1].SetTitle("");

    int color = kCyan + i-1;
    grs_survival[i-1].SetLineColor(color);
    grs_survival[i-1].SetMarkerColor(color);
    grs_survival[i-1].SetMarkerStyle(kFullCircle);

    std::string title = "p = " + std::to_string(p_UA);
    l_survival.AddEntry(&grs_survival[i-1], title.c_str());

    grs_survival[i-1].Write();
  }

  l_survival.Draw("Same");
  c_survival.SaveAs(MAKE_CSTRING(outputFigureDir << "/pSurvival.png"));



  // Median logarithmic reduced chi-square heat-maps

  NOTICE("Extracting median logarithmic reduced chi-square values..." << endl);

  ifstream in(inputFile.c_str());

  JManager<double, TH2D> chi2scapes( new TH2D("chi2scapeP%", "",
  					      Nth,  thRange.getLowerLimit(),  thRange.getUpperLimit(),
  					      Nthb, thbRange.getLowerLimit(), thbRange.getUpperLimit()) );

  DEBUG(RIGHT(20) << "PunderAmplified" << RIGHT(20) << "threshold" << RIGHT(20) << "thresholdBand" << RIGHT(20) << "median log10(chi2)" << endl);

  for (string buffer; getline(in, buffer); ) {

    istringstream iss(buffer);

    string fileName;
    
    JToken_t p0;
    JToken_t p1;
    JToken_t p2;

    if (iss >> fileName >> p0 >> p1 >> p2) {

      TFile file(fileName.c_str());
      TH1D* h1 = (TH1D*) file.Get("chi2");

      const double P   = getValue(p0);
      const double th  = getValue(p1);
      const double thb = getValue(p2);
      
      if (h1 != NULL && (Pselection.empty() || Pselection.find(P) != Pselection.end())) {
	
	double median;
	const double cumP = 0.5;
	h1->GetQuantiles(1, &median, &cumP);

	TH2D* h2 = chi2scapes[P];
      
	h2 -> SetBinContent(h2->GetXaxis()->FindBin(th),
			    h2->GetYaxis()->FindBin(thb),
			    median);

	DEBUG(FIXED(20,3) << P << FIXED(20,3) << th << FIXED(20,3) << thb << FIXED(20,3) << median << endl);

      } else if (h1 == NULL) {

	FATAL("No chi2-histogram in " << fileName << "; skip." << endl);
      }
      
      file.Close();
      
    } else {

      FATAL("Invalid input: " << buffer << endl);
    }
  }


  gStyle->SetPadLeftMargin  (0.10);
  gStyle->SetPadRightMargin (0.20);
  gStyle->SetPadBottomMargin(0.15);
  gStyle->SetPalette(kBlueRedYellow);

  const int   Nscapes = chi2scapes.size();

  if (Ncols < 0) {
    Ncols = 3;
  }

  if (Nrows < 0) {
    Nrows = (int)((Nscapes & 1 ? Nscapes+1 : Nscapes) / Ncols);
  }

  const float cmin = 0.5;
  const float cmax = 1.5;

  const float colorBarRelWidth = 0.06;
  const float relWidth         = 1.0 + colorBarRelWidth;

  TCanvas c_chi2scapes("c_chi2scapes", "",
  		       (Int_t)(relWidth * Ncols * plotwidth), (Int_t)(Nrows * plotheight));

  c_chi2scapes.Divide(Ncols, Nrows, 0.01, 0.01);

  for (JManager<double, TH2D>::const_iterator it = chi2scapes.cbegin(); it != chi2scapes.cend(); ++it) {

    const int padNr0 = (int) distance(chi2scapes.cbegin(), it) + 1;

    c_chi2scapes.cd(padNr0);
    
    it->second->Draw("COLZ");

    it->second->GetXaxis()->CenterTitle();
    it->second->GetYaxis()->CenterTitle();
    it->second->GetZaxis()->CenterTitle();

    it->second->GetXaxis()->SetLabelFont(63);
    it->second->GetYaxis()->SetLabelFont(63);
    it->second->GetZaxis()->SetLabelFont(63);

    it->second->GetXaxis()->SetLabelSize(24);
    it->second->GetYaxis()->SetLabelSize(24);
    it->second->GetZaxis()->SetLabelSize(24);

    it->second->GetXaxis()->SetTitleFont(63);
    it->second->GetYaxis()->SetTitleFont(63);
    it->second->GetZaxis()->SetTitleFont(63);

    it->second->GetXaxis()->SetTitleSize(28);
    it->second->GetYaxis()->SetTitleSize(28);
    it->second->GetZaxis()->SetTitleSize(28);

    it->second->GetXaxis()->SetTitleOffset(1.95);
    it->second->GetYaxis()->SetTitleOffset(1.95);
    it->second->GetZaxis()->SetTitleOffset(1.95);

    it->second->GetZaxis()->SetRangeUser(cmin, cmax);

    it->second->GetXaxis()->SetTitle("Q_{0} [npe]");
    it->second->GetYaxis()->SetTitle("Q_{b} [npe]");
    it->second->GetZaxis()->SetTitle("median log_{10}#left(#chi^{2} / NDF#right)");

    c_chi2scapes.Modified();
    c_chi2scapes.Update();
  }

  c_chi2scapes.SaveAs(MAKE_CSTRING(outputFigureDir << "/chi2scapes.png"));

  out << chi2scapes;

  out.Close();
  out.Write();

  delete[] grs_survival;

  return 0;
}