#include <string>
#include <iostream>
#include <iomanip>
#include <cmath>
#include <set>
#include <vector>
#include <algorithm>

#include "TROOT.h"
#include "TFile.h"
#include "TH1D.h"
#include "TH2D.h"

#include "km3net-dataformat/online/JDAQ.hh"
#include "km3net-dataformat/online/JDAQHeader.hh"

#include "JTools/JConstants.hh"
#include "JDetector/JDetector.hh"
#include "JDetector/JDetectorToolkit.hh"
#include "JDetector/JPMTParametersMap.hh"
#include "JDetector/JStringRouter.hh"

#include "JROOT/JRootFileWriter.hh"
#include "JROOT/JRootToolkit.hh"

#include "JTools/JRange.hh"

#include "JLang/JSTDObjectWriter.hh"

#include "JSupport/JMeta.hh"
#include "JSupport/JSingleFileScanner.hh"

#include "JCalibrate/JCalibrateK40.hh"
#include "JCalibrate/JTDC_t.hh"

#include "JFitK40.hh"

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


namespace {

  double MINIMAL_RATE_HZ  =  1.0e-2;     //!< minimal coincidence rate [Hz] (below, disable PMT)
  double STDEV            =  2.0;        //!< number of standard deviations
  size_t MAXIMAL_COUNTS   = 10;          //!< maximal count of too low rates
  double QE_MIN           =  0.1;        //!< minimal QE                    (below, disable PMT)
}


/**
 * \file
 *
 * Auxiliary program to fit PMT parameters from JMergeCalibrateK40.cc output.\n
 * Note that the contebts of the 2D histograms are converted to standard deviations.
 *
 * \author mdejong
 */
int main(int argc, char **argv)
{
  using namespace std;
  using namespace JPP;
  using namespace KM3NETDAQ;

  typedef JRange<double>   JRange_t;

  JK40Parameters K40 = JK40Parameters::getInstance();

  string       inputFile;
  string       outputFile;
  string       detectorFile;
  string       pmtFile;
  JTDC_t       TDC;
  bool         reverse;
  bool         overwriteDetector;
  JCounter     writeFits;
  bool         fitAngle;
  bool         fitNoise;
  bool         fitModel;
  bool         fixQE;
  JRange_t     X;
  JRange_t     Y;
  int          debug;

  try { 

    JProperties properties;

    properties.insert(gmake_property(MINIMAL_RATE_HZ));
    properties.insert(gmake_property(STDEV));
    properties.insert(gmake_property(MAXIMAL_COUNTS));
    properties.insert(gmake_property(QE_MIN));
    properties.insert(gmake_property(K40.R));
    properties.insert(gmake_property(K40.p1));
    properties.insert(gmake_property(K40.p2));
    properties.insert(gmake_property(K40.p3));
    properties.insert(gmake_property(K40.p4));
 
    JParser<> zap("Auxiliary program to fit PMT parameters from JMergeCalibrateK40 output.");
    
    zap['@'] = make_field(properties)        = JPARSER::initialised();
    zap['f'] = make_field(inputFile,         "input file (output from JMergeCalibrateK40).");
    zap['o'] = make_field(outputFile,        "output file.")                                                   = "fit.root";
    zap['a'] = make_field(detectorFile,      "detector file.");
    zap['P'] = make_field(pmtFile,           "specify PMT file name that can be given to JTriggerEfficiency.") = "";
    zap['!'] = make_field(TDC,
                          "Fix time offset(s) of PMT(s) of certain module(s), e.g."
                          "\n-! \"808969848 0 808982077 23\" will fix time offsets of PMT 0 of module 808969848 and of PMT 23 of module 808982077."
                          "\nSame PMT offset can be fixed for all optical modules, e.g."
                          "\n-! \"-1 0 -1 23\" will fix time offsets of PMTs 0 and 23 of all optical modules.") = JPARSER::initialised();
    zap['r'] = make_field(reverse,           "reverse TDC constraints due to option -! <TDC>.");
    zap['A'] = make_field(overwriteDetector, "overwrite detector file provided through '-a' with fitted time offsets.");
    zap['w'] = make_field(writeFits,         "write fit results to ROOT file; -ww also write fitted TTS to PMT file.");
    zap['D'] = make_field(fitAngle,          "fit angular distribution; fix normalisation.");
    zap['B'] = make_field(fitNoise,          "fit background.");
    zap['M'] = make_field(fitModel,          "fit angular distribution as well as normalisation; fix PMT QEs = 1.0.");
    zap['Q'] = make_field(fixQE,             "fix PMT QEs = 1.0.");
    zap['x'] = make_field(X,                 "fit range (PMT pairs).")                                         = JRange_t();
    zap['y'] = make_field(Y,                 "fit range (time residual [ns]).")                                = JRange_t();
    zap['d'] = make_field(debug,             "debug.")                                                         = 1;

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


  if ((fitModel ? 1 : 0)  +
      (fitAngle ? 1 : 0)  +
      (fitNoise ? 1 : 0)  +
      (fixQE    ? 1 : 0)  >  1) {
    FATAL("Use either option -M, -D, -B or -Q" << endl);
  } 

  const int option = (fitModel ? FIT_MODEL_t                       :
		      fitAngle ? FIT_PMTS_AND_ANGULAR_DEPENDENCE_t :
		      fitNoise ? FIT_PMTS_AND_BACKGROUND_t         :
		      fixQE    ? FIT_PMTS_QE_FIXED_t               :
		      FIT_PMTS_t);

  if (reverse) {
    TDC.reverse();
  }

  for (JTDC_t::const_iterator i = TDC.begin(); i != TDC.end(); ++i) {
    DEBUG("PMT " << setw(10) << i->first << ' ' << setw(2) << i->second << " constrain t0." << endl);
  }

  try {
    TDC.is_valid(true);
  }
  catch(const exception &error) {
    FATAL(error.what() << endl);
  }

  JDetector detector;

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


  JPMTParametersMap parameters;

  if (pmtFile != "") {
    try {
      parameters.load(pmtFile.c_str());
    }
    catch(const exception& error) {}
  }
  
  parameters.comment.add(JMeta(argc, argv));


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

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


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

  TH1D h0("chi2", NULL, 500,  0.0,   5.0);
  TH1D hn("hn",   NULL, 501, -0.5, 500.0);
  TH1D hr("rate", NULL, 500,  0.0,  25.0);
  TH1D h1("p1",   NULL, 500, -5.0,  +5.0);
  TH1D h2("p2",   NULL, 500, -5.0,  +5.0);
  TH1D h3("p3",   NULL, 500, -5.0,  +5.0);
  TH1D h4("p4",   NULL, 500, -5.0,  +5.0);
  TH1D hc("cc",   NULL, 500, -0.1,  +0.1);

  const floor_range   range = getRangeOfFloors(detector);
  const JStringRouter string(detector);

  TH2D H2("detector", NULL,
	  string.size() + 0, -0.5, string.size() - 0.5,
	  range.getUpperLimit(), 1 - 0.5, range.getUpperLimit() + 0.5);

  for (Int_t i = 1; i <= H2.GetXaxis()->GetNbins(); ++i) {
    H2.GetXaxis()->SetBinLabel(i, MAKE_CSTRING(string.at(i-1)));
  }
  for (Int_t i = 1; i <= H2.GetYaxis()->GetNbins(); ++i) {
    H2.GetYaxis()->SetBinLabel(i, MAKE_CSTRING(i));
  }


  JFit fit(0, debug);

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

    if (module->getFloor() == 0) { 
      continue;
    }

    const JTDC_t::range_type range = TDC.equal_range(module->getID());

    NOTICE("Module " << setw(10) << module->getID() << ' ' << getLabel(module->getLocation()) << " !" << distance(range.first, range.second) << endl);

    TH2D* h2d = (TH2D*) in->Get(MAKE_CSTRING(module->getID() << _2R));

    if (h2d == NULL || h2d->GetEntries() == 0) {
      
      NOTICE("No data for module " << module->getID() << " -> set QEs to 0." << endl);

      for (int pmt = 0; pmt != NUMBER_OF_PMTS; ++pmt) {
        parameters[JPMTIdentifier(module->getID(), pmt)].QE = 0.0;
      }

      continue;
    }

    JModel model(*module, K40, range, option);

    data_type data;                                                    // input data

    vector<size_t> count[2] = { vector<size_t>(NUMBER_OF_PMTS, 0),     // counter (exclude self)
				vector<size_t>(NUMBER_OF_PMTS, 1) };   // counter (include self)

    for (int ix = 1; ix <= h2d->GetXaxis()->GetNbins(); ++ix) {

      const pair_type pair = model.getPair(ix - 1);
      
      auto& buffer = data[pair];                                       // storage

      double V = 0.0;                                                  // integrated value
      double W = 0.0;                                                  // integrated error

      for (int iy = 1; iy <= h2d->GetYaxis()->GetNbins(); ++iy) {

	const double x = h2d->GetXaxis()->GetBinCenter(ix);
	const double y = h2d->GetYaxis()->GetBinCenter(iy);

	if (X(x) && Y(y)) {

	  double value = h2d->GetBinContent(ix,iy);
	  double error = h2d->GetBinError  (ix,iy);

	  buffer.push_back(rate_type(y, value, error));

	  double width = h2d->GetYaxis()->GetBinWidth(iy);

	  value *= width;
	  error *= width;

	  V += value;
	  W += error * error;
	}
      }

      W = sqrt(W);

      if (V  <=  0.0             - STDEV*W) {
	count[0][pair.first]  += 1;
	count[0][pair.second] += 1;
      }

      if (V  <=  MINIMAL_RATE_HZ + STDEV*W) {
	count[1][pair.first]  += 1;
	count[1][pair.second] += 1;
      }
    }

    for (int pmt = 0; pmt != NUMBER_OF_PMTS; ++pmt) {

      if (count[0][pmt] >= MAXIMAL_COUNTS) {                           // too many paired rates negative

	WARNING("PMT " << setw(10) << module->getID() << '.' << FILL(2,'0') << pmt << FILL() << " some rates negative -> fit background" << endl);

	if (fit.value.parameters[pmt].status) {
	  model.parameters[pmt].bg.set();
	}
      }
      
      if (count[1][pmt] == NUMBER_OF_PMTS) {                           // all paired rates too low

	WARNING("PMT " << setw(10) << module->getID() << '.' << FILL(2,'0') << pmt << FILL() << " all rates to low -> disable" << endl);

	model.parameters[pmt].disable();
      }
    }

    DEBUG("Start value:" << endl << model << endl);

    try {

      fit.value   = model;                                             // start value

      auto result = fit(data);

      bool refit  = false;

      for (int pmt = 0; pmt != NUMBER_OF_PMTS; ++pmt) {

	if (fit.value.parameters[pmt].QE() < QE_MIN && fit.value.parameters[pmt].status) {

	  WARNING("PMT " << setw(10) << module->getID() << '.' << FILL(2,'0') << pmt << FILL() << ' ' << FIXED(5,3) << fit.value.parameters[pmt].QE() << " < " << FIXED(5,3) << QE_MIN << " -> disable" << (!refit ? " and refit" : "") << endl);

	  fit.value.parameters[pmt].disable();

	  refit = true;
	}
      }

      if (refit) {
	refit  = false;
	result = fit(data);
      }

      NOTICE("Fit result " << setw(10) << module->getID() << " chi2 / NDF  " << FIXED(10,2) << result.chi2 << " / " << setw(5) << result.ndf << ' ' << setw(5) << fit.numberOfIterations << endl);

      DEBUG(fit.value);

      if (writeFits) {

	h0.Fill(result.chi2 / result.ndf);
	hn.Fill(fit.numberOfIterations);
	hr.Fill(fit.value.R );
	h1.Fill(fit.value.p1);
	h2.Fill(fit.value.p2);
	h3.Fill(fit.value.p3);
	h4.Fill(fit.value.p4);
	hc.Fill(fit.value.cc);

	for (int ix = 1; ix <= h2d->GetXaxis()->GetNbins(); ++ix) {

	  const pair_type pair = fit.value.getPair(ix - 1);

	  for (int iy = 1; iy <= h2d->GetYaxis()->GetNbins(); ++iy) {
	  
	    const double dt_ns = h2d->GetYaxis()->GetBinCenter(iy);
	      
	    h2d->SetBinContent(ix, iy, fit.value.getValue(pair, dt_ns));
	    h2d->SetBinError  (ix, iy, 0.0);
	  }
	}

	h2d->SetName(MAKE_CSTRING(module->getID() << _2F));
	h2d->Write();

	H2.Fill((double) string.getIndex(module->getString()), (double) module->getFloor(), result.chi2 / result.ndf);
      }

      const double t0 = (fit.value.hasFixedTimeOffset() ? fit.value.getFixedTimeOffset() : 0.0);

      for (int pmt = 0; pmt != NUMBER_OF_PMTS; ++pmt) {
      
	JPMTParameters& data = parameters[JPMTIdentifier(module->getID(), pmt)];

	const double    P    = getSurvivalProbability(data);
	    
	if (P > 0.0) 
	  data.QE = fit.value.parameters[pmt].QE / P;
	else
	  data.QE = 0.0;

	if (writeFits > 1) {
	  data.TTS_ns = fit.value.parameters[pmt].TTS();
	}
        
	module->getPMT(pmt).addT0(fit.value.parameters[pmt].t0.get() - t0);
      }
    }  
    catch(const exception& error) {

      ERROR("Module " << setw(10) << module->getID() << ' ' << error.what() << " -> set QEs to 0." << endl);

      for (int pmt = 0; pmt != NUMBER_OF_PMTS; ++pmt) {
	parameters[JPMTIdentifier(module->getID(), pmt)].QE = 0.0;
      }
    }
  }


  vector<JMeta> meta(1, JMeta(argc, argv));

  {
    JSingleFileScanner<JMeta> reader(inputFile);
    JSTDObjectWriter  <JMeta> writer(meta);

    writer << reader;
  }

  for (vector<JMeta>::const_reverse_iterator i = meta.rbegin(); i != meta.rend(); ++i) {
    parameters.comment.add(*i);
    detector  .comment.add(*i);
  }

  if (overwriteDetector) {

    NOTICE("Store calibration data on file " << detectorFile << endl);

    store(detectorFile, detector);
  }

  if (pmtFile != "") {
    parameters.store(pmtFile.c_str());
  }


  for (vector<JMeta>::const_iterator i = meta.begin(); i != meta.end(); ++i) {
    putObject(out,*i);
  }

  for (JRootFileReader<JDAQHeader> in(inputFile.c_str()); in.hasNext(); ) {
    putObject(out, *in.next());
  }

  if (writeFits) {
    out << h0 << hn << hr << h1 << h2 << h3 << h4 << hc << H2;
  }

  out.Close();

  return 0;
}