#ifndef __JCALIBRATE_JFITK40__
#define __JCALIBRATE_JFITK40__

#include "TMath.h"
#include "TString.h"
#include "TF1.h"
#include "TF2.h"
#include "TH1.h"
#include "TH2.h"
#include "TFitResult.h"

#include "km3net-dataformat/online/JDAQ.hh"
#include "JROOT/JRootToolkit.hh"
#include "JDetector/JModule.hh"
#include "JMath/JMathToolkit.hh"
#include "JLang/JException.hh"
#include "JLang/JNullType.hh"
#include "JCalibrate/JCalibrateK40.hh"


/**
 * \author mdejong
 */

namespace JCALIBRATE {}
namespace JPP { using namespace JCALIBRATE; }

namespace JCALIBRATE {

  using JLANG::JIndexOutOfRange;
  using JLANG::JNullType;
  using JROOT::JFitParameter_t;
  using KM3NETDAQ::NUMBER_OF_PMTS;
  using JDETECTOR::JModule;


  static const double FITK40_QE_MIN      =  0.0;   //!< Minimal quantum efficiency [unit]
  static const double FITK40_QE_MAX      = 10.0;   //!< Maximal quantum efficiency [unit]
  static const double FITK40_TTS_MIN_NS  =  0.0;   //!< Minimal transition-time spread [ns] 
  static const double FITK40_TTS_MAX_NS  =  3.5;   //!< Maximal transition-time spread [ns]


  /**
   * Fit parameters for single PMT.
   */
  struct JPMTParameters_t {
    /**
     * Default constructor.
     */
    JPMTParameters_t() :
      QE (1.0),
      TTS(2.0),
      t0 (0.0)
    {}


    Double_t  QE;             //!< quantum efficiency [unit]
    Double_t  TTS;            //!< transition-time spread [ns]
    Double_t  t0;             //!< time offset [ns]
  };


  /**
   * Fit parameters for two-fold coincidence rate due to K40.
   */
  struct JFitK40Parameters {
    /**
     * Default constructor.
     *
     * The default parameter values are set to those obtained from a designated simulation
     * of K40 decays (see http://wiki.km3net.de/index.php/OMGsim_simulations_for_K40_fit).\n
     * If you change these values, you may also want to change the corresponding values in JK40DefaultSimulator.hh.
     */
    JFitK40Parameters() :
      Rate_Hz(18.460546),
      p1     ( 3.0767),
      p2     (-1.2078),
      p3     ( 0.9905),
      p4     ( 0.9379),
      bg     ( 1.0e-3),
      cc     ( 1.0e-3)
    {}


    /**
     * Copy constructor.
     *
     * \param  data         data
     */
    JFitK40Parameters(const Double_t* data) 
    {
      if (data != NULL) {
	setModelParameters(data);
      }
    }


    /** 
     * Get number of model parameters.
     *
     * \return              number of parameters
     */
    static Int_t getNumberOfModelParameters()
    {
      return sizeof(JFitK40Parameters) / sizeof(Double_t);
    }


    /**
     * Get model parameters.
     *
     * \return              pointer to parameters
     */
    const Double_t* getModelParameters() const
    {
      return &Rate_Hz;
    }


    /**
     * Get model parameters.
     *
     * \return              pointer to parameters
     */
    Double_t* getModelParameters()
    {
      return &Rate_Hz;
    }


    /**
     * Set model parameters.
     *
     * \param  data         pointer to parameters
     */
    void setModelParameters(const Double_t* data)
    {
      for (Int_t i = 0; i != getNumberOfModelParameters(); ++i) {
	getModelParameters()[i] = data[i];
      }
    }


    /**
     * Get model parameter.
     *
     * \param  i            parameter index
     * \return              parameter value
     */
    Double_t getModelParameter(Int_t i) const
    {
      return getModelParameters()[i];
    }


    /**
     * Get model parameter.
     *
     * \param  p            pointer to data member
     * \return              parameter index and value
     */
    JFitParameter_t getModelParameter(Double_t JFitK40Parameters::*p) const
    {
      const Int_t i = &(this->*p) - getModelParameters();

      return JFitParameter_t(i, getModelParameter(i));
    }


    /**
     * Get model parameter.
     *
     * \param  pmt          PMT number
     * \param  p            pointer to data member of PMT parameters 
     * \return              parameter index and value
     */
    JFitParameter_t getModelParameter(Int_t pmt, Double_t JPMTParameters_t::*p) const
    {
      const Int_t i = &(parameters[pmt].*p) - getModelParameters();

      return JFitParameter_t(i, getModelParameter(i));
    }


    /**
     * Get QE of given PMT.
     *
     * \param  pmt          pmt address
     * \return              QE
     */
    Double_t getQE(const int pmt) const
    {
      return parameters[pmt].QE;
    }


    /**
     * Set QE of given PMT.
     *
     * \param  pmt          pmt address
     * \param  QE           QE
     */
    void setQE(const int pmt, const Double_t QE)
    {
      parameters[pmt].QE = QE;
    }


    /**
     * Get time resolution of given PMT.
     *
     * \param  pmt          pmt address
     * \return              TTS [ns]
     */
    Double_t getTTS(const int pmt) const
    {
      return parameters[pmt].TTS;
    }


    /**
     * Set time resolution of given PMT.
     *
     * \param  pmt          pmt address
     * \param  TTS          TTS [ns]
     */
    void setTTS(const int pmt, const Double_t TTS)
    {
      parameters[pmt].TTS = TTS;
    }

  
    /**
     * Get time offset of given PMT.
     *
     * \param  pmt          pmt address
     * \return              time offset [ns]
     */
    Double_t getT0(const int pmt) const
    {
      return parameters[pmt].t0;
    }

  
    /**
     * Set time offset of given PMT.
     *
     * \param  pmt          pmt address
     * \param  t0           time offset [ns]
     */
    void setT0(const int pmt, const Double_t t0)
    {
      parameters[pmt].t0 = t0;
    }


    /**
     * Get K40 coincidence rate as a function of cosine angle between PMT axes.
     *
     * \param  ct           cosine angle between PMT axes
     * \return              rate [Hz]
     */
    Double_t getValue(const Double_t ct) const
    {
      return Rate_Hz * TMath::Exp(-(p1+p2+p3+p4)) * TMath::Exp(ct*(p1+ct*(p2+ct*(p3+ct*p4))));
    }
    

    // fit parameters

    Double_t Rate_Hz;            //!< maximal coincidence rate [Hz]
    Double_t p1;                 //!< angle dependence coincidence rate
    Double_t p2;                 //!< angle dependence coincidence rate
    Double_t p3;                 //!< angle dependence coincidence rate
    Double_t p4;                 //!< angle dependence coincidence rate
    Double_t bg;                 //!< remaining constant background
    Double_t cc;                 //!< fraction of signal correlated background

    JPMTParameters_t parameters[NUMBER_OF_PMTS];
  };


  /**
   * Template definition of two-fold coincidence rate due to K40 and other radioactive decays.
   */
  template<class JFunction_t = JNullType>
  struct JFitK40_t;


  /**
   * Template specialisation of two-fold coincidence rate due to K40 and other radioactive decays.
   *
   * This class serves as a base class for the ROOT fit classes.
   */
  template<>
  struct JFitK40_t<JNullType> :
    public JFitK40Parameters,
    public JModule,
    public JCombinatorics
  {
    using  JFitK40Parameters::getValue;
    using  JFitK40Parameters::getT0;
    using  JFitK40Parameters::setT0;

    /**
     * Data structure for a pair of addresses.
     */
    typedef JCombinatorics::pair_type       pair_type;


    /**
     * Get intrinsic K40 arrival time spread.
     *
     * \return              sigma [ns]
     */
    Double_t getSigmaK40() const
    {
      return this->sigmaK40_ns;
    }


    /**
     * Set intrinsic K40 arrival time spread.
     *
     * \param  sigma        sigma [ns]
     */
    void setSigmaK40(const Double_t sigma)
    {
      this->sigmaK40_ns = sigma;
    }


    /**
     * Test PMT status.
     *
     * \param  pmt          pmt address
     * \return              true if given pmt is disabled; else false
     */
    bool is_disabled(int pmt) const
    {
      return disable[pmt];
    }


    /**
     * Test PMT status.
     *
     * \param  pmt          pmt address
     * \return              true if given pmt is used for average t0; else false
     */
    bool is_average_t0(int pmt) const
    {
      return pmt == index_of_average_t0;
    }

  
    /**
     * Get time offset of given PMT.
     *
     * Note that if the average time offset is constraint,
     * the time offset of each PMT is corrected by the average time offset of all PMTs.
     *
     * \param  pmt          pmt address
     * \return              time offset [ns]
     */
    Double_t getT0(const int pmt) const
    {
      if (index_of_average_t0 == -1) {

	return JFitK40Parameters::getT0(pmt);

      } else {

	if (disable[pmt]) {

	  return parameters[pmt].t0;

	} else {	

	  Int_t    n0  =  0;       // number  of non-fixed t0's
	  Double_t t0  =  0.0;     // average of non-fixed t0's
	  
	  for (int i = 0; i != NUMBER_OF_PMTS; ++i) {
	    if (!disable[i] && i != index_of_average_t0) {
	      n0  +=  1;
	      t0  +=  parameters[i].t0;
	    }
	  }
	  
	  t0  /=  (n0 + 1);

	  if (pmt != index_of_average_t0)
	    return parameters[pmt].t0 - t0;
	  else
	    return -t0;
	}
      }
    }
 

    /**
     * Enable PMT.
     *
     * \param  pmt          pmt address
     */
    void enablePMT(const int pmt)
    {
      disable[pmt] = false;
    }


    /**
     * Disable PMT.
     *
     * Note that if the average time offset is constraint to zero,
     * the corresponding index may change.
     *
     * \param  pmt          pmt address
     */
    void disablePMT(const int pmt)
    {
      disable[pmt] = true;

      if (index_of_average_t0 == pmt) {

	for (index_of_average_t0 = 0; index_of_average_t0 != NUMBER_OF_PMTS; ++index_of_average_t0) {
	  if (!disable[index_of_average_t0]) {
	    break;
	  }
	}
	  
	if (index_of_average_t0 == NUMBER_OF_PMTS) {
	  THROW(JIndexOutOfRange, "JFitK40_t::disable PMT: No free index for average t0.");
	}
      }
    }


    /**
     * Get cosine of space angle between PMT axes.
     *
     * \param  pair         pair of addresses
     * \return              cosine
     */
    double getDot(const pair_type& pair) const
    {
      return JMATH::getDot(this->getPMT(pair.first),
			   this->getPMT(pair.second));
    }

  
    /**
     * Get time resolution of given PMT pair.
     *
     * \param  pair         pair of addresses
     * \return              sigma [ns]
     */
    Double_t getSigma(const pair_type& pair) const
    {
      return TMath::Sqrt(getTTS(pair.first)  * getTTS(pair.first)  +
			 getTTS(pair.second) * getTTS(pair.second) +
			 getSigmaK40()       * getSigmaK40());
    }


    /**
     * Get time offset of given PMT pair.
     *
     * \param  pair         pair of addresses
     * \return              time offset [ns]
     */
    Double_t getTimeOffset(const pair_type& pair) const
    {
      if (index_of_average_t0 == -1) {

	return     parameters[pair.first].t0 - parameters[pair.second].t0;

      } else {

	if        (pair.first  == index_of_average_t0) {
	  
	  return  -parameters[pair.second].t0;
	  
	} else if (pair.second == index_of_average_t0) {
	  
	  return   parameters[pair.first].t0;
	
	} else {
	  
	  return   parameters[pair.first].t0 - parameters[pair.second].t0;
	}
      }
    }


    /**
     * Get K40 coincidence rate.
     *
     * \param  pair         pair of addresses
     * \return              rate [Hz]
     */
    Double_t getValue(const pair_type& pair) const
    {
      const Double_t ct = getDot(pair);

      return (getValue(ct)         *
	      getQE(pair.first)    *
	      getQE(pair.second));
    }


    /**
     * Get K40 coincidence rate.
     *
     * \param  pair         pair of addresses
     * \param  dt_ns        time difference [ns]
     * \return              rate [Hz/ns]
     */
    Double_t getValue(const pair_type& pair, const Double_t dt_ns) const
    {
      const Double_t t0    = getTimeOffset(pair);
      const Double_t sigma = getSigma     (pair);

      return getValue(pair) * TMath::Gaus(dt_ns, t0, sigma, kTRUE);
    }


  protected:

    Double_t       sigmaK40_ns;                         //!< intrinsic K40 arrival time spread [ns]
    int            index_of_average_t0;                 //!< index of t0 used for average time offset
    bool           disable[NUMBER_OF_PMTS];             //!< disable PMT from fit


    /**
     * Constructor.
     *
     * \param  module       detector module
     * \param  option       true = fix average t0; false = free average t0
     */
    JFitK40_t(const JModule& module, 
	      const bool     option) :
      sigmaK40_ns(0.54)
    {
      static_cast<JModule&>(*this) = module;
    
      this->configure(module.size());
      
      this->sort(JPairwiseComparator(module));

      index_of_average_t0 = (option ? 0 : -1);

      for (int i = 0; i != NUMBER_OF_PMTS; ++i) {
	disable[i] = false;
      }
    }


    /**
     * Default constructor.
     */
    JFitK40_t()
    {}


    /**
     * Get unique instance of fit object.
     *
     * \return              reference to fit object.
     */
    static JFitK40_t& getInstance()
    {
      static JFitK40_t object;

      return object;
    }
  };


  /**
   * Template specialisation of two-fold coincidence rate due to K40 and other radioactive decays.
   *
   * Note that for use in ROOT fit operations, the member method JFitK40_t::getRate is static.
   */
  template<>
  struct JFitK40_t<TF2> :
    public JFitK40_t<>,
    public TF2
  {
    /**
     * Constructor.
     *
     * \param  module       detector module
     * \param  xmin         minimal x
     * \param  xmax         maximal x
     * \param  ymin         minimal y
     * \param  ymax         maximal y
     * \param  option       true = fix average t0; false = free average t0
     */
    JFitK40_t(const JModule& module, 
	      const Double_t xmin,
	      const Double_t xmax,
	      const Double_t ymin,
	      const Double_t ymax,
	      const bool     option) :
      JFitK40_t<>(module, option),
      TF2("f2",
	  JFitK40_t<TF2>::getRate,
	  xmin, xmax, ymin, ymax,
	  getNumberOfModelParameters())
    {}


    /**
     * Get K40 coincidence rate as a function of the fit parameters.
     *
     * To speed up the calculation, it is assumed that the parameter values do not
     * change unless also the x[0] value changes (i.e.\ the index of the PMT pair).
     *
     * \param  x            pointer to data
     * \param  data         pointer to parameter values
     * \return              rate [Hz/ns]
     */
    static Double_t getRate(const Double_t* x, const Double_t* data)
    {
      static int       id    =  -1;
      static Double_t  t0;
      static Double_t  sigma;
      static Double_t  rate;
      
      const int      ix    = (int) x[0];
      const Double_t dt_ns = x[1];

      if (ix != id) {
	
	getInstance().setModelParameters(data);
	
	const pair_type& pair = getInstance().getPair(ix);

	t0     =  getInstance().getTimeOffset(pair);
	sigma  =  getInstance().getSigma     (pair);
	rate   =  getInstance().getValue     (pair);
	id     =  ix;
      }

      return getInstance().bg  +  rate * (getInstance().cc + TMath::Gaus(dt_ns, t0, sigma, kTRUE));
    }


    /**
     * Fit 2D-histogram.
     *
     * Note that the PMT parameters are partially reset before the fit according the status of each PMT and
     * the obtained fit parameters are copied back to the model parameters after the fit.
     *
     * \param  h2           ROOT histogram
     * \param  option       fit option
     */
    TFitResultPtr operator()(TH2& h2, const std::string& option) 
    {
      for (int pmt = 0; pmt != NUMBER_OF_PMTS; ++pmt) {

	parameters[pmt].QE  = JPMTParameters_t().QE;
	parameters[pmt].TTS = JPMTParameters_t().TTS;

	if (disable[pmt]) {

	  parameters[pmt].QE = 0.0;

	  fixParameter(*this, this->getModelParameter(pmt, &JPMTParameters_t::QE));
	  fixParameter(*this, this->getModelParameter(pmt, &JPMTParameters_t::TTS));
	  fixParameter(*this, this->getModelParameter(pmt, &JPMTParameters_t::t0));

	} else {

	  // Note that setting limits to a parameter will also release it

	  if (!isParameterFixed(*this, this->getModelParameter(pmt, &JPMTParameters_t::QE)))  {
	    setParLimits(*this, this->getModelParameter(pmt, &JPMTParameters_t::QE),  FITK40_QE_MIN,            FITK40_QE_MAX);
	  }
	  if (!isParameterFixed(*this, this->getModelParameter(pmt, &JPMTParameters_t::TTS))) {
	    setParLimits(*this, this->getModelParameter(pmt, &JPMTParameters_t::TTS), FITK40_TTS_MIN_NS,        FITK40_TTS_MAX_NS);
	  }
	  if (!isParameterFixed(*this, this->getModelParameter(pmt, &JPMTParameters_t::t0)))  {
	    setParLimits(*this, this->getModelParameter(pmt, &JPMTParameters_t::t0),  h2.GetYaxis()->GetXmin(), h2.GetYaxis()->GetXmax());
	  }
	}
      }

      if (index_of_average_t0 != -1) {

	this->setT0(index_of_average_t0, 0.0);

	fixParameter(*this, this->getModelParameter(index_of_average_t0, &JPMTParameters_t::t0));
      }

      this->SetParameters(this->getModelParameters());

      getInstance() = *this;

      const TFitResultPtr result = h2.Fit(this, option.c_str());

      this->setModelParameters(this->GetParameters());

      return result;
    }
  };


  /**
   * Template specialisation of two-fold coincidence rate due to K40 and other radioactive decays.
   *
   * Note that for use in ROOT fit operations, the member method JFitK40_t::getTimeOffset is static.
   */
  template<>
  struct JFitK40_t<TF1> :
    public JFitK40_t<>,
    public TF1
  {
    /**
     * Constructor.
     *
     * \param  f2           2D-fit function
     */
    JFitK40_t(const JFitK40_t<TF2>& f2) :
      JFitK40_t<>(f2),
      TF1("f1",
	  JFitK40_t<TF1>::getTimeOffset,
	  f2.GetXmin(), f2.GetXmax(),
	  getNumberOfModelParameters())
    {
      this->SetParameters(f2.GetParameters());

      for (size_t i = 0; i != JFitK40_t<>::getNumberOfModelParameters() - NUMBER_OF_PMTS * sizeof(JPMTParameters_t) / sizeof(Double_t); ++i) {
	fixParameter(*this, JFitParameter_t(i, this->getModelParameter(i)));
      }

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

	fixParameter(*this, this->getModelParameter(pmt, &JPMTParameters_t::QE));
	fixParameter(*this, this->getModelParameter(pmt, &JPMTParameters_t::TTS));

	if (isParameterFixed(f2, f2.getModelParameter(pmt, &JPMTParameters_t::t0)))  {
	  fixParameter(*this, f2.getModelParameter(pmt, &JPMTParameters_t::t0));
	}
      }
    }


    /**
     * Get time offset as a function of the fit parameters.
     *
     * \param  x            pointer to data
     * \param  data         pointer to parameter values
     * \return              time offset [ns]
     */
    static Double_t getTimeOffset(const Double_t* x, const Double_t* data)
    {
      const int ix = (int) x[0];

      getInstance().setModelParameters(data);

      const pair_type& pair = getInstance().getPair(ix);

      return getInstance().getTimeOffset(pair);
    }
 

    /**
     * Fit 1D-histogram.
     *
     * \param  h1           ROOT histogram
     * \param  option       fit option
     */
    TFitResultPtr operator()(TH1& h1, const std::string& option) 
    {
      for (int pmt = 0; pmt != NUMBER_OF_PMTS; ++pmt) {

	if (disable[pmt]) {
	  fixParameter(*this, this->getModelParameter(pmt, &JPMTParameters_t::t0));
	}
      }

      if (index_of_average_t0 != -1) {

	this->setT0(index_of_average_t0, 0.0);

	fixParameter(*this, this->getModelParameter(index_of_average_t0, &JPMTParameters_t::t0));
      }

      this->SetParameters(this->getModelParameters());

      getInstance() = *this;

      const TFitResultPtr result = h1.Fit(this, option.c_str());

      this->setModelParameters(this->GetParameters());

      return result;
    }
  };


  /**
   * Type definition for backward compatibility.
   */
  typedef JFitK40_t<TF2>  JFitK40; 
}

#endif