#include <string>
#include <iostream>
#include <iomanip>
#include <type_traits>
#include <functional>
#include <future>
#include <mutex>
#include <thread>
#include <vector>
#include <set>
#include <memory>
#include <algorithm>

#include "km3net-dataformat/online/JDAQEvent.hh"
#include "km3net-dataformat/definitions/fitparameters.hh"
#include "JDAQ/JDAQEventIO.hh"

#include "JDetector/JDetector.hh"
#include "JDetector/JDetectorToolkit.hh"
#include "JDetector/JModuleRouter.hh"

#include "JDynamics/JDynamics.hh"

#include "JSupport/JMeta.hh"
#include "JSupport/JSupport.hh"
#include "JSupport/JMultipleFileScanner.hh"
#include "JSupport/JParallelFileScanner.hh"
#include "JSupport/JSummaryFileRouter.hh"

#include "JReconstruction/JHitW0.hh"
#include "JReconstruction/JEvt.hh"
#include "JReconstruction/JEvtToolkit.hh"
#include "JReconstruction/JMuonGandalfParameters_t.hh"

#include "JTrigger/JHitL0.hh"
#include "JTrigger/JBuildL0.hh"

#include "JFit/JFitToolkit.hh"
#include "JFit/JLine1Z.hh"
#include "JFit/JLine3Z.hh"
#include "JFit/JModel.hh"
#include "JFit/JGandalf.hh"
#include "JFit/JLine3ZRegressor.hh"
#include "JFit/JGradient.hh"

#include "JLang/JSTDObjectReader.hh"
#include "JLang/JVectorize.hh"

#include "JGeometry3D/JRotation3D.hh"

#include "JMath/JMath.hh"
#include "JMath/JQuantile_t.hh"

#include "JTools/JRange.hh"

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


namespace JRECONSTRUCTION {

  using KM3NETDAQ::JDAQEvent;
  using JFIT::JParameter_t;
  using JFIT::JLine1Z;
  using JFIT::JLine3Z;
  using JDETECTOR::JDetector;
  using JTRIGGER::JHit;

  typedef std::vector<JHitW0>                  buffer_type;  //!< hits
  typedef std::map<int, buffer_type>           map_type;     //!< identifier -> hits

  /**
   * Auxiliary data structure to store data and fit in memory.
   */
  struct event_type {
    map_type  data;
    JLine1Z   value;
  };
    
  typedef std::vector<event_type>                    data_type;
  typedef JLANG::JSTDObjectReader<const event_type>  input_type;


  /**
   * Auxiliary class for editing time offset.
   */
  struct JEditor :
    public JParameter_t
  {
    /**
     * Constructor.
     *
     * \param  data                  data
     * \param  id                    identifier
     */
    JEditor(data_type& data, const int id) :
      data(data),
      id(id),
      t0(0.0)
    {}
    

    /**
     * Apply step.
     *
     * \param  step                  step
     */
    virtual void apply(const double step) override
    {
      for (data_type::iterator evt = data.begin(); evt != data.end(); ++evt) {

	map_type::iterator p = evt->data.find(id);

	if (p != evt->data.end()) {
	  for (buffer_type::iterator hit = p->second.begin(); hit != p->second.end(); ++hit) {
	    static_cast<JHit&>(*hit) = JHit(hit->getT1() + step, hit->getToT());
	  }
	}
      }

      this->t0 += step;
    }

    data_type& data;         //!< data
    int        id;           //!< identifier
    double     t0;           //!< time offset [ns]
  };


  typedef JFIT::JRegressorStorage<JFIT::JLine3Z, JFIT::JGandalf>   JRegressorStorage_t;
  typedef JFIT::JRegressor       <JFIT::JLine3Z, JFIT::JGandalf>   JRegressor_t;


  /**
   * Thread pool for fits to data.
   */
  struct JPerth {
    /**
     * Constructor.
     *
     * \param  storage            storage for PDFs
     * \param  data               data
     * \param  ns                 number of threads
     */
    JPerth(const JRegressorStorage_t& storage,
           const data_type&           data,
	   const size_t               ns) :
      input(data),
      stop(false)
    {
      using namespace std;
      using namespace JPP;

      Q.reset();

      for (size_t i = 0; i < ns; ++i) {

	thread worker([this, storage]() {

	  JRegressor_t regressor(storage);

	  regressor.parameters.resize(5);

	  regressor.parameters[0] = JLine3Z::pX();
	  regressor.parameters[1] = JLine3Z::pY();
	  regressor.parameters[2] = JLine3Z::pT();
	  regressor.parameters[3] = JLine3Z::pDX();
	  regressor.parameters[4] = JLine3Z::pDY();

	  buffer_type  data;
	  JLine3Z      value;
	  
	  for ( ; ; ) {
 
	    {
	      unique_lock<mutex> lock(in);

	      cv.wait(lock, [this]() { return stop || this->input.hasNext(); });

	      if (stop && !this->input.hasNext()) {
		return;
	      }

	      const event_type* evt = this->input.next();

	      // re-organise data

	      data.clear();

	      for (map_type::const_iterator p = evt->data.begin(); p != evt->data.end(); ++p) {
		copy(p->second.begin(), p->second.end(), back_inserter(data));
	      }
	      
	      value = evt->value;
	    }
	    
	    const double chi2 = regressor(value, data.begin(), data.end());

	    {
	      unique_lock<mutex> lock(out);

	      Q.put(chi2);
	    }
	  }
	});

	workers.emplace_back(std::move(worker));
      }
    }


    /**
     * Destructor.
     */   
    ~JPerth()
    {
      using namespace std;

      {
	unique_lock<mutex> lock(in);

	stop = true;
      }
      
      cv.notify_all();

      for (auto& worker : workers) {
	worker.join();
      }
    }
    
    static JMATH::JQuantile_t Q;
    
  private:
    std::vector<std::thread>  workers;
    input_type                input;
    std::mutex                in;
    std::mutex                out;
    std::condition_variable   cv;
    bool                      stop;    
  };

  /**
   */
  JMATH::JQuantile_t JPerth::Q;


  /**
   * Auxiliary data structure for chi2 function object.
   */
  struct JPerth_t {
    /**
     * Get chi2.
     *
     * \param  option                option
     * \return                       chi2/NDF
     */
    double operator()(const int option) const
    {
      JPerth(storage, data, ns);

      return JPerth::Q.getMean();
    }

    const JRegressorStorage_t& storage;
    const data_type&           data;
    const size_t               ns;
  };
}


/**
 * \file
 *
 * Program to determine inter-string time calibration.
 *
 * \author mdejong
 */
int main(int argc, char **argv)
{
  using namespace std;
  using namespace JPP;
  using namespace KM3NETDAQ;

  typedef JParallelFileScanner< JTypeList<JDAQEvent, JFIT::JEvt> > JParallelFileScanner_t;
  typedef JParallelFileScanner_t::multi_pointer_type               multi_pointer_type;
  typedef JTYPELIST<JCOMPASS::JOrientation,
                    JACOUSTICS::JEvt>::typelist                    calibration_types;
  typedef JMultipleFileScanner<calibration_types>                  JCalibration_t;

  JMultipleFileScanner<>   inputFile;
  JLimit_t&                numberOfEvents = inputFile.getLimit();
  string                   detectorFile;
  JCalibration_t           calibrationFile;
  double                   Tmax_s;
  string                   pdfFile;
  JMuonGandalfParameters_t parameters;
  bool                     overwriteDetector;
  set<int>                 strings;                              // string identifiers
  set<int>                 modules;                              // module identifiers
  int                      number_of_iterations    =  1000;
  int                      number_of_extra_steps   =     0;
  double                   epsilon                 =   1.0e-4;
  double                   T_ns                    =   2.5;      // step size
  size_t                   threads;                              // number of threads
  int                      debug;

  const int DEFAULT_ID = -1;

  try { 

    JProperties properties;

    properties.insert(gmake_property(number_of_iterations));
    properties.insert(gmake_property(number_of_extra_steps));
    properties.insert(gmake_property(epsilon));
    properties.insert(gmake_property(T_ns));

    JParser<> zap("Program to determine inter-string time calibration.");
    
    zap['f'] = make_field(inputFile);
    zap['a'] = make_field(detectorFile);
    zap['+'] = make_field(calibrationFile)     = JPARSER::initialised();
    zap['T'] = make_field(Tmax_s)              = 100.0;
    zap['n'] = make_field(numberOfEvents)      = JLimit::max();
    zap['P'] = make_field(pdfFile);
    zap['@'] = make_field(parameters)          = JPARSER::initialised();
    zap['A'] = make_field(overwriteDetector, "overwrite detector file provided through '-a' with fitted time offsets.");
    zap['S'] = make_field(strings,           "string identifiers (" << DEFAULT_ID << " => all)") =  JPARSER::initialised();
    zap['M'] = make_field(modules,           "module identifiers (" << DEFAULT_ID << " => all)") =  JPARSER::initialised();
    zap['%'] = make_field(properties,        "fit parameters")                                   =  JPARSER::initialised();
    zap['N'] = make_field(threads,           "number of threads")                                = 1;
    zap['d'] = make_field(debug)               = 1;
    
    zap(argc, argv);
  }
  catch(const exception& error) {
    FATAL(error.what() << endl);
  }


  if (strings.empty() == modules.empty()) {
    FATAL("Set either strings (option -S) or modules (option -M)." << endl);
  }

  JDetector detector;

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

  unique_ptr<JDynamics> dynamics;

  try {

    dynamics.reset(new JDynamics(detector, Tmax_s));

    dynamics->load(calibrationFile);
  }
  catch(const exception& error) {
    if (!calibrationFile.empty()) {
      FATAL(error.what());
    }
  }

  const JModuleRouter router(dynamics ? dynamics->getDetector() : detector);

  NOTICE("Reading PDFs... " << flush);
  
  const JRegressorStorage_t storage(pdfFile, parameters.TTS_ns);

  NOTICE("OK" << endl);

  JRegressor_t::debug              = debug;
  JRegressor_t::T_ns.setRange(parameters.TMin_ns, parameters.TMax_ns);
  JRegressor_t::Vmax_npe           = parameters.VMax_npe;
  JRegressor_t::MAXIMUM_ITERATIONS = parameters.NMax;


  // preprocess input data
  
  const JBuildL0<JHitL0>  buildL0;

  data_type    data;

  counter_type skip    = inputFile.getLowerLimit();
  counter_type counter = inputFile.getLowerLimit();
  
  for (JMultipleFileScanner<>::const_iterator i = inputFile.begin(); i != inputFile.end(); ++i) {

    JSummaryFileRouter  summary(*i, parameters.R_Hz);

    for (JParallelFileScanner_t in(*i); (skip -= in.skip(skip)) == 0 && in.hasNext() && counter != numberOfEvents; ++counter) {

      STATUS("event: " << setw(10) << counter << '\r'); DEBUG(endl);

      multi_pointer_type ps  = in.next();
    
      const JDAQEvent*   tev = ps;
      const JFIT::JEvt*  evt = ps;

      summary.update(*tev);

      if (dynamics) {
	dynamics->update(*tev);
      }

      if (!evt->empty()) {

	vector<JHitL0>  dataL0;
	      
	buildL0(*tev, router, true, back_inserter(dataL0));

	const JFIT::JFit&  fit = (*evt)[0];
      
	const JRotation3D  R (getDirection(fit));
	const JLine1Z      tz(getPosition (fit).rotate(R), fit.getT());
	JRange<double>     Z_m;

	if (fit.hasW(JSTART_LENGTH_METRES)       &&
	    fit.getW(JSTART_LENGTH_METRES) > 0.0) {
	  Z_m = JZRange(parameters.ZMin_m, parameters.ZMax_m + fit.getW(JSTART_LENGTH_METRES));
	}

	const JFIT::JModel<JLine1Z> match(tz, parameters.roadWidth_m, JRegressor_t::T_ns, Z_m);

	// hit selection based on start value

	buffer_type buffer;

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

          JHitW0 hit(*i, summary.getRate(i->getPMTIdentifier()));

	  hit.rotate(R);

	  if (match(hit)) {
	    buffer.push_back(hit);
	  }
	}

	// select first hit

	sort(buffer.begin(), buffer.end(), JHitL0::compare);

	buffer_type::const_iterator __end = unique(buffer.begin(), buffer.end(), equal_to<JDAQPMTIdentifier>());

	// re-organise data
	
	map_type map;

	for (buffer_type::const_iterator hit = buffer.begin(); hit != __end; ++hit) {
	  
	  const JModule& module = router.getModule(hit->getModuleID());

	  if (!strings.empty()) { map[module.getString()].push_back(*hit); }
	  if (!modules.empty()) { map[module.getID()]    .push_back(*hit); }
	}
	
	data.push_back({map, tz});
      }
    }
  }
  STATUS(endl);


  // fit

  JGradient fit(number_of_iterations, number_of_extra_steps, epsilon, 3);

  if (strings.count(DEFAULT_ID) != 0) { strings = getStringIDs(detector); }
  if (modules.count(DEFAULT_ID) != 0) { modules = getModuleIDs(detector); }

  for (int id : strings) { fit.push_back(JModifier_t(MAKE_STRING("string: " << FILL(4,'0') << id << FILL()), new JEditor(data, id), T_ns)); }
  for (int id : modules) { fit.push_back(JModifier_t(MAKE_STRING("module: " << setw(8)     << id),           new JEditor(data, id), T_ns)); }

  const JPerth_t perth = {storage, data, threads};

  const double   chi2  = fit(perth);

  STATUS("result: " << FIXED(12,6) << chi2 << ' ' << setw(6) << fit.numberOfIterations << endl);

  for (size_t i = 0; i != fit.size(); ++i) {
    {
      const JEditor* p = dynamic_cast<const JEditor*>(fit[i].get());

      if (p != NULL) { STATUS(fit[i].name << ' ' << FIXED(9,3) << p->t0 << " [ns]" << endl); }
    }
  }


  if (overwriteDetector) {

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

    map<int, double> calibration;
 
    for (size_t i = 0; i != fit.size(); ++i) {

      const JEditor* p = dynamic_cast<const JEditor*>(fit[i].get());

      if (p != NULL) {
	calibration[p->id] = p->t0;
      }
    }

    const double t0 = getAverage(get_values(calibration), 0.0);

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

      if (!module->empty()) {

	map<int, double>::const_iterator p = (!strings.empty() ? calibration.find(module->getString()) :
					      !modules.empty() ? calibration.find(module->getID())     :
					      calibration.end());
	if (p != calibration.end()) {
	  for (int pmt = 0; pmt != NUMBER_OF_PMTS; ++pmt) {
	    module->getPMT(pmt).addT0(p->second - t0);
	  }
	}
      }
    }

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

    store(detectorFile, detector);
  }
}