#include "km3net-dataformat/online/JDAQ.hh"
#include "JDetector/JDetector.hh"
#include "JDetector/JDetectorToolkit.hh"
#include "Jeep/JParser.hh"

#include "TF2.h"
#include "TFile.h"
#include "TH1D.h"
#include "TH2D.h"
#include "TMath.h"

#include <iostream>
#include <string>


namespace FITL1DTSLICES {

  /* Calculate the log-poisson
   *
   * This function has some assumptions: since the input data from JMonitorL1dt use SN datastream
   * most background is strongly suppressed. This means many bins on both data and model will
   * have no entries: n = 0 or nhat = 0. This will result in infinities when calculating the 
   * log-poisson. To avoid the infities the options below are used.
   * This has been crosschecked with L1 datastream which gives a significant background and the 
   * shapes of the likelihood distributions are similar.
   */
  double logPoison(double n, double nhat, double logP_min = -999999.9) {
    if (nhat < 0.0 || n < 0.0) {
      FATAL("logPoisson: n (" << n <<") or nhat (" << nhat << ") is < 0.0" << std::endl);
    }
    if (n == 0.0 && nhat == 0.0) {
      return 0.0; // ok.
    }
    if (n == 0.0 && nhat > 0.0) {
      return -1*nhat; // log( lab^0 * e^-lab / 0! ) = log( e^-lab ) = -lab, ok.
    }
    if (n >= 0.0 && nhat == 0.0) {
      return 0.0; // should be -inf, we only consider bins where we expect non-zero entries: no expected entries we throw away.
    }
    Double_t poisson = TMath::Poisson(n, nhat);
    if (poisson == 0) return 0.0; // only when model and data are close enough do we care about their values, so if P=0, do nothing, since we are not in the region we care about.
    else              return TMath::Log(poisson);
  }

  double getLogQuality(const TH1D* data, const TH1D* model, int di, double data_bkg = 0.0001, double model_bkg = 0.0001) {

    double q = 0.0;

    // If you calculate over the whole histogram you lose bins at edges: only when dt=0 are all N bins considered, everywhere else 
    // either data or model bins get lost at the edge. So integrate over a smaller range:
    int i_low = data->GetNbinsX() / 4; // Start at a quarter of the hist
    int i_hgh = data->GetNbinsX() / 4 * 3; // End at three quarters
    for (int i = i_low; i <= i_hgh; ++i) {
      double n = data_bkg;
      double nhat = model_bkg;
      if (i >= 1 && i <= data->GetNbinsX()) {
        // Count in bin i for the data
        n = data->GetBinContent(i);
      }
      if (i+di >= 1 && i+di <= model->GetNbinsX()) {
        // Count in bin i+di for the model
        nhat = model->GetBinContent(i+di);
      }
      double logP = logPoison(n, nhat);
      q += logP; // Added 0 for when the background is zero in model and data, but then P(n|0) gives inf! So removed overwritten logPmin
    }
    return q;
  }
}


/**
 * \file
 *
 * Auxiliary program to fit L1dt output data to an L1dt model, both created with JMonitorL1dt
 *
 * \author kmelis, lnauta
 */
int main(int argc, char **argv)
{
  using namespace std;
  using namespace FITL1DTSLICES;
  using namespace JPP;

  string   inputFile;
  string   modelFile;
  string   outputFile;
  string   detectorFile;
  double   dt_fitrange;
  int      rebin;
  double   TMax_ns;
  bool     normaliseMC;
  int      debug;

  try { 

    JParser<> zap("Program to calculate log-likelihood distributions between DOM pairs and fit a poly2 to find the shape, used in FitL1dt to find time offsets");
    
    zap['f'] = make_field(inputFile,    "input file")                                        = "monitor.root";
    zap['m'] = make_field(modelFile,    "model file");
    zap['o'] = make_field(outputFile,   "output file")                                       = "fitslices.root";
    zap['a'] = make_field(detectorFile, "detector file");
    zap['T'] = make_field(TMax_ns,      "scan range around 0 in data-model comparison [ns]") = 50.0;
    zap['t'] = make_field(dt_fitrange,  "fitrange of polynomial to quality [ns]")            = 5.0;
    zap['r'] = make_field(rebin,        "rebin the histograms")                              = 1;
    zap['N'] = make_field(normaliseMC,  "normalize the MC histogram");
    zap['d'] = make_field(debug) = 0;

    if (zap.read(argc, argv) != 0) {
      return 1;
    }
  }
  catch(const exception &error) {
    FATAL(error.what() << endl);
  }


  const double TSignal_ns = 750.0; // The width of the signal for integrating out the background

  JDetector detector;

  try {
    load(detectorFile, detector);
  }
  catch(const JException& error) {
    FATAL(error);
  }


  TFile* in = TFile::Open(inputFile.c_str(), "exist");
  TFile* in_model = TFile::Open(modelFile.c_str(), "exist");

  if (in == NULL || !in->IsOpen()) {
    FATAL("File: " << inputFile << " not opened." << endl);
  }
  if (in_model == NULL || !in_model->IsOpen()) {
    FATAL("File: " << modelFile << " not opened." << endl);
  }

  TFile out(outputFile.c_str(), "recreate");

  TH2D h2A("Aij", "Aij", detector.size(), -0.5, detector.size()-0.5, detector.size(), -0.5, detector.size()-0.5);
  TH2D h2B("Bij", "Bij", detector.size(), -0.5, detector.size()-0.5, detector.size(), -0.5, detector.size()-0.5);
  TH2D h2C("Cij", "Cij", detector.size(), -0.5, detector.size()-0.5, detector.size(), -0.5, detector.size()-0.5);

  // The bare histogram can be used for debugging, it contains the actual time offsets. 
  // Histogram Bij contains Bij/2Aij due to a different representation used in the code compared to the mathematical formulation.
  TH2D h2Bbare("Bij_bare", "Bij_bare", detector.size(), -0.5, detector.size()-0.5, detector.size(), -0.5, detector.size()-0.5);
  for (JDetector::iterator moduleA = detector.begin(); moduleA != detector.end(); ++moduleA) {
    const int idom = distance(detector.begin(), moduleA);
    TString label = Form("%i", moduleA->getID());
    h2A.GetXaxis()->SetBinLabel(idom+1, label);
    h2A.GetYaxis()->SetBinLabel(idom+1, label);
    h2B.GetXaxis()->SetBinLabel(idom+1, label);
    h2B.GetYaxis()->SetBinLabel(idom+1, label);
    h2C.GetXaxis()->SetBinLabel(idom+1, label);
    h2C.GetYaxis()->SetBinLabel(idom+1, label);

    h2Bbare.GetXaxis()->SetBinLabel(idom+1, label);
    h2Bbare.GetYaxis()->SetBinLabel(idom+1, label);
  }

  for (JDetector::iterator moduleA = detector.begin(); moduleA != detector.end(); ++moduleA) {

    DEBUG("Module " << moduleA->getID() << endl);
    const JString title = JString(moduleA->getID());
    const int idom = distance(detector.begin(), moduleA);

    // data
    TH2D* d2s = (TH2D*) in->Get((title + ".2S").c_str());
    TH1D* d1l = (TH1D*) in->Get((title + ".1L").c_str());

    if (d2s == NULL || d1l == NULL) {
      WARNING("No data in data histogram " << title << endl);
      continue;
    }

    // model
    TH2D* m2s = (TH2D*) in_model->Get((title + ".2S").c_str());
    TH1D* m1l = (TH1D*) in_model->Get((title + ".1L").c_str());

    if (m2s == NULL || m1l == NULL) {
      WARNING("No mata in model histogram " << title << endl);
      continue;
    }


    for (JDetector::iterator moduleB = moduleA; moduleB != detector.end(); ++moduleB) {
      const int jdom = distance(detector.begin(), moduleB);

      if (moduleB->getID() == moduleA->getID()) { 
        continue; 
      }
      if (moduleB->getString() != moduleA->getString()) { 
        continue; 
      }

      TH1D* d1s = d2s->ProjectionY((title + Form(".%i.d1S", moduleB->getID())).c_str(), 
          d2s->GetXaxis()->FindBin(TString(Form("%i", moduleB->getID()))), 
          d2s->GetXaxis()->FindBin(TString(Form("%i", moduleB->getID()))), "e");
      TH1D* m1s = m2s->ProjectionY((title + Form(".%i.m1S", moduleB->getID())).c_str(), 
          m2s->GetXaxis()->FindBin(TString(Form("%i", moduleB->getID()))), 
          m2s->GetXaxis()->FindBin(TString(Form("%i", moduleB->getID()))), "e");
     
      if (d1s->Integral() <= 0.0 || m1s->Integral() <= 0.0) { 
        continue; 
      }

      // background fit: data
      double binCenterMaxBin = d1s->GetXaxis()->GetBinCenter(d1s->GetMaximumBin());
      double backgroundrate_s = 0.0;
      int nbins = 0;
      for (int j = 1; j <= d1s->GetXaxis()->GetNbins(); ++j) {
        if (fabs(d1s->GetXaxis()->GetBinCenter(j) - binCenterMaxBin) > TSignal_ns ) { 
          backgroundrate_s += d1s->GetBinContent(j); 
          nbins++; 
        }
      }
      backgroundrate_s /= nbins;

      // background fit: model
      binCenterMaxBin = m1s->GetXaxis()->GetBinCenter(m1s->GetMaximumBin());
      double backgroundrate_m = 0.0;
      nbins = 0;
      for (int j = 1; j <= m1s->GetXaxis()->GetNbins(); ++j) {
        if (fabs(m1s->GetXaxis()->GetBinCenter(j) - binCenterMaxBin) > TSignal_ns ) { 
          backgroundrate_m += m1s->GetBinContent(j); 
          nbins++; 
        }
      }
      backgroundrate_m /= nbins;

      // scale livetimes and remove background
      const double Ld = d1l->GetBinContent(d1l->GetXaxis()->FindBin(TString(Form("%i", moduleB->getID()))));
      const double Lm = m1l->GetBinContent(m1l->GetXaxis()->FindBin(TString(Form("%i", moduleB->getID()))));

      for (int j = 1; j <= d1s->GetXaxis()->GetNbins(); ++j) {
        const double y  = d1s->GetBinContent(j);
        const double dy = d1s->GetBinError(j);

        d1s->SetBinContent(j, y);
        d1s->SetBinError(j, dy);
      }

      // scale the livetime of the model
      for (int j = 1; j <= m1s->GetXaxis()->GetNbins(); ++j) {
        const double y  = m1s->GetBinContent(j);
        const double dy = m1s->GetBinError(j);

        m1s->SetBinContent(j, y * Ld/Lm);
        m1s->SetBinError(j, dy*Ld/Lm);
      }


      if (normaliseMC) {
        const double scalefactor = d1s->Integral()/m1s->Integral();
        m1s->Scale(scalefactor);
      }

      // rebin
      d1s->Rebin(rebin);
      m1s->Rebin(rebin);

      // write background subtracted, normalised slices
      d1s->Write();
      m1s->Write();

      // get quality
      int ddt_min = -0.5*TMax_ns;  // maximum time shift [ns]
      int ddt_max =  0.5*TMax_ns;  // maximum time shift [ns]

      TH1D q1((title + Form(".%i.q1", moduleB->getID())).c_str(), (title + Form(".%i.q1", moduleB->getID())).c_str(), (ddt_max-ddt_min)/rebin+1, ddt_min-0.5*rebin, ddt_max+0.5*rebin);

      for (int di = 1; di <= q1.GetNbinsX(); di++) {
        q1.SetBinContent(di, getLogQuality(d1s, m1s, (int)(q1.GetBinCenter(di)/rebin), 0.0, 0.0));
      }

      // Fit peak with parabola
      const double dt_fitmid = q1.GetBinCenter(q1.GetMaximumBin());

      TF1 q1f((title + Form(".%i.q1f", moduleB->getID())).c_str(), "[0]*x*x + [1]*x + [2]", dt_fitmid - dt_fitrange, dt_fitmid + dt_fitrange);
      q1f.SetParLimits(0, -1e5, 0.0); // Added as to improve fits. Aij should always be l.t. 0 to fit a maximum to the LLH

      string fitoptions = "R";
      if (debug < JEEP::debug_t) {
        fitoptions += " Q";
      }
      q1.Fit(&q1f, fitoptions.c_str());

      // write quality and fit
      q1.Write();

      // Fit results to global fit matrix
      const double Aij = q1f.GetParameter(0);
      const double Bij = q1f.GetParameter(1);
      const double Cij = q1f.GetParameter(2);

      double bareBij = Bij / Aij;
      if (Aij == 0) bareBij = 0;

      if (fabs(-0.5*Bij/Aij - dt_fitmid) > 3 || fabs(-0.5*Bij/Aij) > dt_fitrange) {
        WARNING("  Please check fit of pair " << idom << ", " << jdom << " : " << moduleA->getID() << ", " << moduleB->getID() << endl); 
        WARNING("  DOMpair " << idom << "; " << jdom << "  A=" << Aij << " B=" << Bij << " , max at: " << dt_fitmid << " [ns], best fit at " << -0.5*Bij/Aij << " ns " << endl);
      }
      DEBUG("DOMpair " << idom << "; " << jdom << "  A=" << Aij << " B=" << Bij << " , max at: " << dt_fitmid << " [ns], best fit at " << -0.5*Bij/Aij << " ns " << endl);

      h2A.SetBinContent(idom+1, jdom+1, Aij);
      h2A.SetBinContent(jdom+1, idom+1, Aij);
      h2B.SetBinContent(idom+1, jdom+1, Bij);
      h2B.SetBinContent(jdom+1, idom+1,-Bij);
      h2C.SetBinContent(idom+1, jdom+1, Cij);
      h2C.SetBinContent(jdom+1, idom+1, Cij);

      h2Bbare.SetBinContent(idom+1, jdom+1, bareBij);
      h2Bbare.SetBinContent(jdom+1, idom+1, 0);
    }
  }

  out.cd();

  h2A.Write();
  h2B.Write();
  h2C.Write();

  h2Bbare.Write();

  out.Close();
}