#include <string>
#include <iostream>
#include <iomanip>
#include <map>
#include <cmath>

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

#include "TROOT.h"
#include "TFile.h"
#include "TH1D.h"
#include "TH2D.h"
#include "TF2.h"
#include "TMath.h"
#include "TFitResult.h"
#include "Math/MinimizerOptions.h"

#include "JPhysics/JConstants.hh"
#include "JDetector/JDetector.hh"
#include "JDetector/JDetectorToolkit.hh"
#include "JDetector/JPMTParametersMap.hh"
#include "JDetector/JPMTParametersToolkit.hh"

#include "JROOT/JRootFileWriter.hh"
#include "JROOT/JRootToolkit.hh"
#include "JTools/JRange.hh"
#include "JTools/JQuantile.hh"
#include "JSupport/JMeta.hh"
#include "JSupport/JSingleFileScanner.hh"

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

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

namespace {

  double  STDEV            =  2.0;        //!< number of standard deviations
  double  QE_MIN_VALUE     =  0.01;       //!< minimal QE value (below, set to zero)
  double  QE_MAX_ERROR     =  0.5;        //!< maximal QE error (above, set to zero)
  double  MINIMAL_RATE_HZ  =  1.0e-2;     //!< minimal coincidence rate (bin-by-bin)

  /**
   * Check validity of QE measurement.
   *
   * The QE measurement is valid when:
   *  - fitted value is at least STDEV standard deviations larger than QE_MIN_VALUE; and
   *  - fitted error is smaller than QE_MAX_ERROR.
   *
   * \param  qe_value      QE value
   * \param  qe_error      QE error
   * \return               true if valid; else false
   */
  inline bool is_valid(const double qe_value, const double qe_error)
  {
    return (qe_value > QE_MIN_VALUE  +  STDEV * qe_error  &&
	    qe_error < QE_MAX_ERROR);
  }
}


/**
 * \file
 *
 * Auxiliary program to fit PMT parameters from JMergeCalibrateK40.cc output.\n
 * To force drawing of fit functions bin-by-bin compatible to the corresponding histogram data,  \n
 * the function values will be stored as histograms, rather than as associated function objects. \n
 * See JCalibrateK40.hh for the corresponding extensions of histogram names.
 * \author mdejong
 */
int main(int argc, char **argv)
{
  using namespace std;
  using namespace JPP;
  using namespace KM3NETDAQ;

  typedef JRange<double>      JRange_t;

  const double      epsilon = 1.0e-10;    // minimal error for ROOT plotting
  JFitK40Parameters fitk40;               // setting of internal fit parameters

  string    inputFile;
  string    outputFile;
  string    detectorFile;
  string    pmtFile;
  JTDC_t    TDC;
  bool      reverse;
  bool      overwriteDetector;
  bool      writeFits;
  bool      fitAngDist;
  bool      fitModel;
  JRange_t  X;
  JRange_t  Y;
  string    option;
  int       debug;

  try { 

    JProperties properties;

    properties.insert(gmake_property(STDEV));
    properties.insert(gmake_property(QE_MIN_VALUE));
    properties.insert(gmake_property(QE_MAX_ERROR));
    properties.insert(gmake_property(MINIMAL_RATE_HZ));
    properties.insert(gmake_property(fitk40.Rate_Hz));
    properties.insert(gmake_property(fitk40.p1));
    properties.insert(gmake_property(fitk40.p2));
    properties.insert(gmake_property(fitk40.p3));
    properties.insert(gmake_property(fitk40.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 offset 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 offsets of PMTs 0 and 23 for 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 correct time offsets.");
    zap['w'] = make_field(writeFits,         "write fit results.");
    zap['D'] = make_field(fitAngDist,        "fit angular distribution; fix normalisation.");
    zap['M'] = make_field(fitModel,          "fit angular distribution; fix PMT QEs = 1.0.");
    zap['x'] = make_field(X,                 "ROOT fit range (PMT pairs).")                                    = JRange_t();
    zap['y'] = make_field(Y,                 "ROOT fit range (time residual).")                                = JRange_t();
    zap['O'] = make_field(option,            "ROOT fit option, see TH1::Fit.")                                 = "";
    zap['d'] = make_field(debug,             "debug.")                                                         = 1;

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


  ROOT::Math::MinimizerOptions::SetDefaultMaxFunctionCalls(1000000);


  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);
  }

  TDC.is_valid(true);


  //----------------------------------------------------------
  // load detector file and PMT parameters
  //----------------------------------------------------------

  JDetector detector;

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


  JPMTParametersMap parameters;

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


  if (option.find('O') == string::npos) { option += '0'; }
  if (option.find('S') == string::npos) { option += 'S'; }
  if (option.find('Q') == string::npos && debug < JEEP::debug_t) { option += 'Q'; }


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

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


  TDirectory::AddDirectory(NULL);

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

  // Histograms with global fit results.

  TH1D hc("chi2", NULL, 500, 0.0, 5.0);

  TH1D p1("p1", NULL, 501, -5.0, +5.0);
  TH1D p2("p2", NULL, 501, -5.0, +5.0);
  TH1D p3("p3", NULL, 501, -5.0, +5.0);
  TH1D p4("p4", NULL, 501, -5.0, +5.0);
  TH1D bg("bg", NULL, 501, -0.1, +0.1);
  TH1D cc("cc", NULL, 501, -0.1, +0.1);


  //----------------------------------------------------------
  // loop over modules in detector file to fit histogram
  //----------------------------------------------------------

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

    NOTICE("Module " << setw(10) << module->getID() << ' ' << module->getLocation() << endl);

    if (module->empty()) {
      continue;
    }

    //----------------------------------------------------------
    // get histogram for this module
    //----------------------------------------------------------

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

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

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

      continue;
    }

    //----------------------------------------------------------
    // get constraints
    //----------------------------------------------------------

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

    //----------------------------------------------------------
    // set-up fitting procedure
    //----------------------------------------------------------

    JFitK40 fit(*module,
		h2s->GetXaxis()->GetXmin(),
		h2s->GetXaxis()->GetXmax(),
		h2s->GetYaxis()->GetXmin(),
		h2s->GetYaxis()->GetXmax(),
		range.first == range.second);

    fit.setModelParameters(fitk40.getModelParameters());

    for (JTDC_t::const_iterator i = range.first; i != range.second; ++i) {
      fixParameter(fit, fit.getModelParameter(i->second, &JPMTParameters_t::t0));
    }

    if (fitModel) {
      for (int pmt = 0; pmt != NUMBER_OF_PMTS; ++pmt) {
	fixParameter(fit, fit.getModelParameter(pmt, &JPMTParameters_t::QE));
      }

    } else {

      fixParameter(fit, fit.getModelParameter(&JFitK40::Rate_Hz));
      
      if (!fitAngDist) {
	fixParameter(fit, fit.getModelParameter(&JFitK40::p1));
	fixParameter(fit, fit.getModelParameter(&JFitK40::p2));
	fixParameter(fit, fit.getModelParameter(&JFitK40::p3));
	fixParameter(fit, fit.getModelParameter(&JFitK40::p4));
      }
    }

    //----------------------------------------------------------
    // check validity of data
    //----------------------------------------------------------

    vector<int> buffer(NUMBER_OF_PMTS, 1);

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

      Double_t value = 0.0;
      Double_t error = 0.0;

      for (int iy = 1; iy <= h2s->GetYaxis()->GetNbins(); ++iy) {
	value += h2s->GetBinContent(ix,iy);
	error += h2s->GetBinError(ix,iy) * h2s->GetBinError(ix,iy);
      }

      error = sqrt(error);

      if (value  <  MINIMAL_RATE_HZ + STDEV*error) {

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

	buffer[pair.first]  += 1;
	buffer[pair.second] += 1;
      }
    }

    try {

      for (int pmt = 0; pmt != NUMBER_OF_PMTS; ++pmt) {
      
	if (buffer[pmt] == NUMBER_OF_PMTS) {

	  WARNING("PMT " << setw(10) << module->getID() << ' ' << setw(2) << pmt << " too low a rate; switch off." << endl);

	  fit.disablePMT(pmt);
	}
      }
    }
    catch(const JException& error) {

      NOTICE("Skip fit for module " << module->getID() << "; set corresponding QEs to 0." << endl);

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

      continue;
    }

    // constrain fit to specified range(s)

    if (X != JRange_t()) {
      h2s->GetXaxis()->SetRangeUser(max(X.getLowerLimit(), h2s->GetXaxis()->GetXmin()),
				    min(X.getUpperLimit(), h2s->GetXaxis()->GetXmax()));
    }
    if (Y != JRange_t()) {
      h2s->GetYaxis()->SetRangeUser(max(Y.getLowerLimit(), h2s->GetYaxis()->GetXmin()),
				    min(Y.getUpperLimit(), h2s->GetYaxis()->GetXmax()));
    }

    // 1D-histogram with peak positions

    TH1D h1(MAKE_CSTRING(module->getID() << _1D), NULL, h2s->GetXaxis()->GetNbins(), h2s->GetXaxis()->GetXmin(), h2s->GetXaxis()->GetXmax());

    const Double_t sigma[] = {

      (min(Y.getUpperLimit(), h2s->GetYaxis()->GetXmax()) -
       max(Y.getLowerLimit(), h2s->GetYaxis()->GetXmin())) / sqrt(12.0),

      sqrt(JPMTParameters_t().TTS * JPMTParameters_t().TTS * 2.0  +  fit.getSigmaK40() * fit.getSigmaK40())
    };


    for (Int_t ix = 1; ix <= h2s->GetXaxis()->GetNbins(); ++ix) {

      Double_t y0    = 0.0;
      Double_t zmax  = 0.0;
      Double_t value = 0.0;
      Double_t error = 0.0;

      for (Int_t iy = 1; iy <= h2s->GetYaxis()->GetNbins(); ++iy) {

	const Double_t y = h2s->GetYaxis()->GetBinCenter(iy);
	const Double_t z = h2s->GetBinContent(ix, iy);
	const Double_t w = h2s->GetBinError  (ix, iy);

	if (Y(y)) {

	  if (z > zmax) {
	    y0   = y;
	    zmax = z;
	  }

	  value += z;
	  error += w*w;
	}
      }

      error = sqrt(error);

      h1.SetBinContent(ix, y0);
      h1.SetBinError  (ix, value > MINIMAL_RATE_HZ + STDEV * error ? sigma[1] * error / value : sigma[0]);
    }

    {
      JFitK40_t<TF1> f1(fit);

      TFitResultPtr result = f1(h1, option.c_str());

      if (result.Get() == NULL) {
	FATAL("Invalid TFitResultPtr " << h1.GetName() << endl);
      }

      if (writeFits) {

	TH1D h2(MAKE_CSTRING(module->getID() << _1F), NULL, h1.GetXaxis()->GetNbins(), h1.GetXaxis()->GetXmin(), h1.GetXaxis()->GetXmax());

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

	  const Double_t x = h2.GetXaxis()->GetBinCenter(ix);

	  h2.SetBinContent(ix, f1.Eval(x));
	  h2.SetBinError  (ix, 0.0);
	}

	out << h1 << h2;
      }

      if (debug >= JEEP::notice_t) {

	cout << "Fit result chi2 / NDF " << result->Chi2() << ' ' << result->Ndf() << ' ' << (result->IsValid() ? "" : "failed") << endl;

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

	    cout << ' '
		 << setw(2)    << pmt             << ' '
		 << FIXED(7,2) << f1.getT0 (pmt)  << ' ';
	    cout << (result->IsParameterFixed(f1.getModelParameter(pmt, &JPMTParameters_t::t0))  ? " *** fixed t0  *** " : "");
	    cout << (f1.is_disabled  (pmt) ? " (disabled)" : "");
	    cout << (f1.is_average_t0(pmt) ? " (averaged)" : "");
	    cout << endl;
	}
      }

      fit.setModelParameters(f1.getModelParameters());
    }


    TFitResultPtr result = fit(*h2s, option);


    bool refit = false;

    try {

      const JFitK40Parameters values(fit.GetParameters());
      const JFitK40Parameters errors(fit.GetParErrors());
    
      for (int pmt = 0; pmt != NUMBER_OF_PMTS; ++pmt) {

	if (!result->IsParameterFixed(fit.getModelParameter(pmt, &JPMTParameters_t::QE))  &&  !is_valid(values.getQE(pmt), errors.getQE(pmt))) {

	  WARNING("PMT " << setw(10) << module->getID() << ' ' << setw(2) << pmt << ' '
		  << "measured QE " 
		  << FIXED(5,2) << values.getQE(pmt) << " +/- " 
		  << FIXED(5,2) << errors.getQE(pmt) 
		  << " compatible with zero -> set QE = 0 and t0 = 0; -> refit." << endl);


	  fit.disablePMT(pmt);
	  
	  refit = true;
	}
      }
    }
    catch(const JException& error) {
      
      NOTICE("Skip fit for module " << module->getID() << "; set corresponding QEs to 0." << endl);

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

    
    if (refit) {

      JFitK40_t<TF1> f1(fit);

      f1(h1, option);

      fit.setModelParameters(f1.getModelParameters());

      result = fit(*h2s, option);
    }


    const JFitK40Parameters values(fit.GetParameters());
    const JFitK40Parameters errors(fit.GetParErrors());

    hc.Fill(TMath::Log10(result->Chi2() / result->Ndf()));

    TH2D h2f(MAKE_CSTRING(module->getID() << _2F),
	     NULL,
	     h2s->GetXaxis()->GetNbins(), h2s->GetXaxis()->GetXmin (), h2s->GetXaxis()->GetXmax (),
	     h2s->GetYaxis()->GetNbins(), h2s->GetYaxis()->GetXmin (), h2s->GetYaxis()->GetXmax ());
    
    for (int ix = 1; ix <= h2f.GetXaxis()->GetNbins(); ++ix) {
      for (int iy = 1; iy <= h2f.GetYaxis()->GetNbins(); ++iy) {

	const Double_t x = h2f.GetXaxis()->GetBinCenter(ix);
	const Double_t y = h2f.GetYaxis()->GetBinCenter(iy);

	h2f.SetBinContent(ix, iy, fit.Eval(x,y));
	h2f.SetBinError  (ix, iy, 0.0);
      }
    }
     

    if (debug >= JEEP::notice_t) {

      cout << "Fit result chi2 / NDF " << result->Chi2() << ' ' << result->Ndf() << ' ' << (result->IsValid() ? "" : "failed") << endl;

      cout << "Rate_Hz= " << FIXED(8,4) << values.Rate_Hz << (result->IsParameterFixed(fit.getModelParameter(&JFitK40::Rate_Hz)) ? " *** fixed *** " : "") << endl;

      cout << "p1=      " << FIXED(8,4) << values.p1      << (result->IsParameterFixed(fit.getModelParameter(&JFitK40::p1))      ? " *** fixed *** " : "") << endl;
      cout << "p2=      " << FIXED(8,4) << values.p2      << (result->IsParameterFixed(fit.getModelParameter(&JFitK40::p2))      ? " *** fixed *** " : "") << endl;
      cout << "p3=      " << FIXED(8,4) << values.p3      << (result->IsParameterFixed(fit.getModelParameter(&JFitK40::p2))      ? " *** fixed *** " : "") << endl;
      cout << "p4=      " << FIXED(8,4) << values.p4      << (result->IsParameterFixed(fit.getModelParameter(&JFitK40::p3))      ? " *** fixed *** " : "") << endl;

      cout << "Background (correlated):   " << FIXED(8,5) << values.bg << (result->IsParameterFixed(fit.getModelParameter(&JFitK40::bg)) ? " *** fixed *** " : "") << endl;
      cout << "Background (uncorrelated): " << FIXED(8,5) << values.cc << (result->IsParameterFixed(fit.getModelParameter(&JFitK40::cc)) ? " *** fixed *** " : "") << endl;

      cout << " PMT t0 [ns] TTS [ns]   R" << endl;

      JQuantile Q[3];
      
      for (int pmt = 0; pmt != NUMBER_OF_PMTS; ++pmt) {

	cout << ' ' 
	     << setw(2)    << pmt             << ' '
	     << FIXED(7,2) << fit.getT0 (pmt) << ' '
	     << FIXED(7,2) << fit.getTTS(pmt) << ' '
             << FIXED(7,3) << fit.getQE (pmt);

	cout << (result->IsParameterFixed(fit.getModelParameter(pmt, &JPMTParameters_t::QE))  ? " *** fixed QE  *** " : "");
	cout << (result->IsParameterFixed(fit.getModelParameter(pmt, &JPMTParameters_t::TTS)) ? " *** fixed TTS *** " : "");
	cout << (result->IsParameterFixed(fit.getModelParameter(pmt, &JPMTParameters_t::t0))  ? " *** fixed t0  *** " : "");
	cout << (fit.is_disabled  (pmt) ? " (disabled)" : "");
	cout << (fit.is_average_t0(pmt) ? " (averaged)" : "");
	cout << endl;

	Q[0].put(fit.getT0 (pmt));
	Q[1].put(fit.getTTS(pmt));
	Q[2].put(fit.getQE (pmt));
      }

      cout << "<>  "
	   << FIXED(7,2) << Q[0].getMean()  << ' '
	   << FIXED(7,2) << Q[1].getMean()  << ' '
	   << FIXED(7,3) << Q[2].getMean()  << endl;
      cout << "+/- "
	   << FIXED(7,2) << Q[0].getSTDev() << ' '
	   << FIXED(7,2) << Q[1].getSTDev() << ' '
	   << FIXED(7,3) << Q[2].getSTDev() << endl;
    }

    h2s->Write();
    h2f .Write();

    if (writeFits) {

      // Histograms with fit results.

      p1.Fill(values.p1);
      p2.Fill(values.p2);
      p3.Fill(values.p3);
      p4.Fill(values.p4);
      bg.Fill(values.bg);
      cc.Fill(values.cc);

      TH1D h1t(MAKE_CSTRING(module->getID() << ".1t0"),  NULL, NUMBER_OF_PMTS, -0.5, NUMBER_OF_PMTS - 0.5);
      TH1D h1s(MAKE_CSTRING(module->getID() << ".1TTS"), NULL, NUMBER_OF_PMTS, -0.5, NUMBER_OF_PMTS - 0.5);
      TH1D h1q(MAKE_CSTRING(module->getID() << ".1QE"),  NULL, NUMBER_OF_PMTS, -0.5, NUMBER_OF_PMTS - 0.5);
      
      h1t.Sumw2();
      h1s.Sumw2();
      h1q.Sumw2();

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

	h1t.SetBinContent(pmt + 1, values.getT0 (pmt));
	h1s.SetBinContent(pmt + 1, values.getTTS(pmt));
	h1q.SetBinContent(pmt + 1, values.getQE (pmt));

	h1t.SetBinError  (pmt + 1, max(errors.getT0 (pmt), epsilon));
	h1s.SetBinError  (pmt + 1, max(errors.getTTS(pmt), epsilon));
	h1q.SetBinError  (pmt + 1, max(errors.getQE (pmt), epsilon));
      }
      
      out << h1t << h1s << h1q;
    }


    // save output

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

      // correction to QE

      const double    P    = getSurvivalProbability(data);

      if (P > 0.0) 
	data.QE = values.getQE(pmt) / P;
      else
	data.QE = 0.0;
	
      if (overwriteDetector) {
	module->getPMT(pmt).addT0(fit.getT0(pmt));
      }
    }
  }

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

  if (writeFits) {
    out << hc << p1 << p2 << p3 << p4 << bg << cc;
  }

  out.Close();


  for (JSingleFileScanner<JMeta> in(inputFile); in.hasNext(); ) {
    
    const JMeta meta(*in.next());

    parameters.comment.add(meta);
    detector  .comment.add(meta);
  }

  if (overwriteDetector) {

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

    store(detectorFile, detector);
  }

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

  return 0;
}