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

#include "TROOT.h"
#include "TFile.h"
#include "TH1D.h"
#include "TH2D.h"
#include "TF1.h"
#include "TMath.h"
#include "TFitResult.h"

#include "JROOT/JMinimizer.hh"

#include "JLang/JLangToolkit.hh"
#include "JSupport/JMeta.hh"
#include "JTools/JRange.hh"

#include "JDetector/JDetector.hh"
#include "JDetector/JDetectorToolkit.hh"
#include "JDetector/JTimeRange.hh"
#include "JDetector/JStringRouter.hh"

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

namespace {

  /**
   * Name of histogram with fit results from JCalibrateMuon.cc.
   */
  const char* h2_t           =  "h2";

  const char* gauss_t        =  "Gauss";           //!< Gaussian
  const char* landau_t       =  "Landau";          //!< Landau
  const char* emg_t          =  "EMG";             //!< Exponentially Modified Gaussian
  const char* breitwigner_t  =  "BreitWigner";     //!< Breit-Wigner
}


/**
 * \file
 * Program to fit time-residuals histogram in output of JCalibrateMuon.cc.
 * \author mdejong
 */
int main(int argc, char **argv)
{
  using namespace std;
  using namespace JPP;

  vector<string>  inputFile;
  string          outputFile;
  string          detectorFile;
  bool            overwriteDetector;
  string          function;
  JTimeRange      T_ns;
  string          option;
  JTimeRange      E_ns;
  double          WMin;
  int             ID;
  int             debug;

  try {

    JParser<> zap("Program to fit time-residuals histogram in output of JCalibrateMuon.cc.");

    zap['f'] = make_field(inputFile,         "input files (output from JCalibrateMuon).");
    zap['o'] = make_field(outputFile,        "output file.")                                                   = "hobbit.root";
    zap['a'] = make_field(detectorFile,      "detector file.");
    zap['A'] = make_field(overwriteDetector, "overwrite detector file provided through '-a' with correct time offsets.");
    zap['F'] = make_field(function,          "fit function")                                                   = gauss_t, landau_t, emg_t, breitwigner_t;
    zap['T'] = make_field(T_ns,              "ROOT fit range around maximum [ns].")                            = JTimeRange(-7.5,+5.0);
    zap['O'] = make_field(option,            "ROOT fit option, see TH1::Fit.")                                 = "";
    zap['E'] = make_field(E_ns,              "validity range of t0 uncertainty [ns].")                         = JTimeRange( 0.0, 0.5);
    zap['W'] = make_field(WMin,              "minimal histogram contents.")                                    = 100.0;
    zap['M'] = make_field(ID,                "module identifier")                                              = -1;
    zap['d'] = make_field(debug)          = 1;

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


  if (!T_ns.is_valid()) {
    FATAL("Invalid fit range [ns] " << T_ns << endl);
  }

  JDetector detector;

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


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

  TH2D* h2 = NULL;

  for (vector<string>::const_iterator i = inputFile.begin(); i != inputFile.end(); ++i) {

    NOTICE("Processing " << *i << endl);

    TFile in(i->c_str(), "exist");

    if (!in.IsOpen()) {
      FATAL("File " << *i << " not opened." << endl);
    }

    TH2D* p = dynamic_cast<TH2D*>(in.Get(h2_t));
      
    if (p == NULL) {
      FATAL("File " << *i << " has no histogram <" << h2_t << ">." << endl);
    }

    if (h2 == NULL)
      h2 = (TH2D*) p->Clone();
    else
      h2->Add(p);

    h2->SetDirectory(0);

    in.Close();
  }

  if (h2 == NULL) {
    FATAL("No histogram <" << h2_t << ">." << endl);
  }

  // auxiliary data structure for fit result

  struct result_type {

    result_type() :
      value(0.0),
      error(0.0)
    {}

    result_type(double value,
		double error) :
      value(value),
      error(error)
    {}

    double value;
    double error;
  };

  typedef map<int, result_type> map_type;    // module index -> fit result

  map_type zmap;


  // histogram fit results

  const TAxis* x_axis = h2->GetXaxis();
  const TAxis* y_axis = h2->GetYaxis();

  TH1D h0("h0", NULL, x_axis->GetNbins(), -0.5, x_axis->GetNbins() - 0.5);
  TH1D hc("hc", NULL, x_axis->GetNbins(), -0.5, x_axis->GetNbins() - 0.5);
  TH1D hq("hq", NULL, x_axis->GetNbins(), -0.5, x_axis->GetNbins() - 0.5);


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

  size_t counts = 0;
  size_t errors = 0;
  
  for (Int_t ix = 1; ix <= x_axis->GetNbins(); ++ix) {

    const JModule& module = detector[ix - 1];

    DEBUG("Module " << setw(8) << module.getID() << ' ' << getLabel(module.getLocation()) << endl);
    
    if (module.getFloor() == 0) {
      continue;
    }

    if (ID != -1 && ID != module.getID()) {
      continue;
    }

    TH1D h1(MAKE_CSTRING(module.getID() << ".1D"), NULL, y_axis->GetNbins(), y_axis->GetXmin(), y_axis->GetXmax());

    // start values
    
    Double_t ymax   =  0.0;
    Double_t x0     =  0.0;       // [ns]
    Double_t sigma  =  4.0;       // [ns]
    Double_t W      =  0.0;

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

      h1.SetBinContent(iy,      h2->GetBinContent(ix,iy));
      h1.SetBinError  (iy, sqrt(h2->GetBinContent(ix,iy)));

      const Double_t x = h1.GetBinCenter (iy);
      const Double_t y = h1.GetBinContent(iy);
      
      W += y;

      if (y > ymax) {
        ymax = y;
        x0   = x;
      }
    }

    if (W <= WMin) {

      WARNING("Not enough entries for slice " << ix << ' ' << W << "; skip fit." << endl);
      
      continue;
    } 

    ymax *= 0.9;

    const Double_t xmin = x0 + T_ns.getLowerLimit();
    const Double_t xmax = x0 + T_ns.getUpperLimit();


    TF1* f1 = NULL;

    if        (function == gauss_t) {
 
      f1 = new TF1(function.c_str(), "[0]*TMath::Gaus(x,[1],[2])");

      f1->SetParameter(0, 0.8*ymax);
      f1->SetParameter(1, x0);
      f1->SetParameter(2, sigma);

      f1->SetParError(0, sqrt(ymax) * 0.1);
      f1->SetParError(1, 0.01);
      f1->SetParError(2, 0.01);

    } else if (function == landau_t) {

      f1 = new TF1(function.c_str(), "[0]*TMath::Landau(x,[1],[2])");

      f1->SetParameter(0, 0.8*ymax);
      f1->SetParameter(1, x0);
      f1->SetParameter(2, sigma);

      f1->SetParError(0, sqrt(ymax) * 0.1);
      f1->SetParError(1, 0.01);
      f1->SetParError(2, 0.01);
    
    } else if (function == emg_t) {
      
      f1 = new TF1(function.c_str(), "[0]*TMath::Exp(0.5*[3]*(2.0*[1]+[3]*[2]*[2]-2.0*x))*TMath::Erfc(([1]+[3]*[2]*[2]-x)/(TMath::Sqrt(2.0)*[2]))");

      f1->SetParameter(0, 0.5*ymax);
      f1->SetParameter(1, x0 - sigma);
      f1->SetParameter(2, sigma);
      f1->SetParameter(3, 0.06);

      f1->SetParError(0, sqrt(ymax) * 0.1);
      f1->SetParError(1, 0.01);
      f1->SetParError(2, 0.01);
      f1->SetParError(3, 1.0e-4);

    } else if (function == breitwigner_t) {
      
      f1 = new TF1(function.c_str(), "(x <= [1])*[0]*[2]*TMath::BreitWigner(x,[1],[2])+(x > [1])*[0]*[3]*TMath::BreitWigner(x,[1],[3])");

      f1->SetParameter(0, 0.8*ymax);
      f1->SetParameter(1, x0);
      f1->SetParameter(2, 15.0);
      f1->SetParameter(3, 25.0);

      f1->SetParError(0, sqrt(ymax) * 0.1);
      f1->SetParError(1, 0.01);
      f1->SetParError(2, 0.1);
      f1->SetParError(3, 0.1);

    } else {

      FATAL("Unknown fit function " << function << endl);
    }


    DEBUG("Start values:" << endl);
      
    for (int i = 0; i != f1->GetNpar(); ++i) {
      DEBUG(left << setw(12) << f1->GetParName  (i) << ' ' << 
	    SCIENTIFIC(12,5) << f1->GetParameter(i) << endl);
    }

    // fit
      
    TFitResultPtr result = h1.Fit(f1, option.c_str(), "same", xmin, xmax);

    if (debug >= notice_t || !result->IsValid()) {
      cout << "Slice: " 
	   << setw(4)    << ix                    << ' '
	   << setw(16)   << h1.GetName()          << ' '
	   << FIXED(7,3) << f1->GetParameter(1)   << " +/- " 
	   << FIXED(7,3) << f1->GetParError(1)    << ' ' 
	   << FIXED(7,3) << result->Chi2()        << '/'
	   << result->Ndf()                       << ' '
	   << (result->IsValid() ? "" : "failed") << endl;
    }

    counts += 1;

    if (!result->IsValid()) {
      errors += 1;
    }

    if (result->IsValid()) {
      zmap[ix] = result_type(f1->GetParameter(1), f1->GetParError (1)); 
    }

    if (result->Ndf() > 0) {
      hc.SetBinContent(ix, result->Chi2() / result->Ndf());
    }

    hq.SetBinContent(ix, result->IsValid() ? 1.0 : 0.0);

    h1.Write();

    delete f1;
  }


  if (!zmap.empty()) {

    // average time offset

    double t0 = 0.0;
      
    for (map_type::const_iterator i = zmap.begin(); i != zmap.end(); ++i) {
      t0 += i->second.value;
    }
    
    t0 /= zmap.size();

    NOTICE("Average time offset [ns] " << FIXED(7,2) << t0 << endl);
    NOTICE("Number of fits passed/failed " << counts << "/" << errors << endl);

    for (map_type::iterator i = zmap.begin(); i != zmap.end(); ++i) {
      i->second.value -= t0;
    }
    
    for (map_type::const_iterator i = zmap.begin(); i != zmap.end(); ++i) {
      h0.SetBinContent(i->first, i->second.value);
      h0.SetBinError  (i->first, i->second.error);
    }


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

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

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

    for (map_type::const_iterator i = zmap.begin(); i != zmap.end(); ++i) {
      hi.SetBinContent(detector[i->first - 1].getString(), 
		       detector[i->first - 1].getFloor(),
		       i->second.value);
    }

    hi.Write();
    

    if (overwriteDetector) {

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

      detector.comment.add(JMeta(argc, argv));

      for (map_type::const_iterator i = zmap.begin(); i != zmap.end(); ++i) {

	if (E_ns(i->second.error))
	  detector[i->first - 1].sub(i->second.value);
	else
	  ERROR("Slice " << setw(4) << i->first << " fit uncertainty " << FIXED(5,2) << i->second.error << " outside specified range (option -E <E_ns>)" << endl);
      }
      
      store(detectorFile, detector);
    }

  } else {

    NOTICE("No calibration results." << endl);
  }

  h0 .Write();
  hc .Write();
  hq .Write();
  h2->Write();

  out.Close();
}