#ifndef __JFIT__JFITTOOLKIT__
#define __JFIT__JFITTOOLKIT__

#include <cmath>

#include "JPhysics/JK40Rates.hh"

#include "JMath/JMath.hh"
#include "JMath/JMathSupportkit.hh"

#include "JFit/JVectorNZ.hh"
#include "JFit/JMatrixNZ.hh"


/**
 * \file
 *
 * Auxiliary methods to evaluate Poisson probabilities and chi2.
 * \author mdejong
 */

/**
 * Auxiliary classes and methods for linear and iterative data regression.
 */
namespace JFIT {}
namespace JPP { using namespace JFIT; }

namespace JFIT {

  using JPHYSICS::JK40Rates;


  /**
   * Get Poisson probability to observe a hit or not
   * for given expectation value for the number of hits.
   *
   * \param  expval           expectation value
   * \param  hit              hit
   * \return                  probability
   */
  inline double getP(const double expval, bool hit)
  {
    if (hit)
      return -expm1(-expval);       // 1 - P(0)
    else
      return exp(-expval);          // P(0)
  }

  
  /**
   * Get chi2 corresponding to given probability.
   *
   * \param  P                probability
   * \return                  chi2
   */
  inline double getChi2(const double P)
  {
    return -log(P);
  }


  /**
   * Get chi2 to observe a hit or not for given expectation value for the number of hits.
   *
   * \param  expval           expectation value
   * \param  hit              hit
   * \return                  probability
   */
  inline double getChi2(const double expval, bool hit)
  {
    if (hit)
      return -log1p(-exp(-expval));
    else
      return expval;
  }
  
  
  /**
   * Determine chi2 of a hit for a given model and time resolution.
   *
   * The template argument <tt>JModel_t</tt> refers to a data structure 
   * which should provide for the following method:
   * <pre>
   *   double %getT(const JHit_t& hit) const;    // get expected time of hit
   * </pre>
   *
   * \param  model            model
   * \param  hit              hit
   * \param  sigma            sigma
   * \return                  chi2
   */
  template<class JModel_t, class JHit_t>
  inline double getChi2(const JModel_t& model, const JHit_t& hit, const double sigma)
  {
    const double ds = (hit.getT() - model.getT(hit)) / sigma;

    return ds * ds;
  }

  
  /**
   * Determine chi2 of data for given model and time resolution.
   *
   * The template argument <tt>JModel_t</tt> refers to a data structure 
   * which should provide for the following method:
   * <pre>
   *   double %getT(const value_type& hit) const;    // expected time of hit
   * </pre>
   * where <tt>value_type</tt> corresponds to the value type if the input data.
   *
   * \param  model            model
   * \param  __begin          begin of data
   * \param  __end            end   of data
   * \param  sigma            time resolution [ns]
   * \return                  chi2
   */
  template<class JModel_t, class T>
  inline double getChi2(const JModel_t& model, T __begin, T  __end, const double sigma)
  {
    double chi2 = 0.0;

    for (T hit = __begin; hit != __end; ++hit) {
      chi2 += getChi2(model, *hit, sigma);
    }

    return chi2;
  }


  /**
   * Determine chi2 using full covariance matrix.
   *
   * \param  Y                residuals
   * \param  V                covariance matrix
   * \return                  chi2
   */
  inline double getChi2(const JVectorNZ& Y,
                        const JMatrixNZ& V)
  {
    return V.getDot(Y);
  }


  /**
   * Determine chi2 of data for given track using full covariance matrix.
   *
   * \param  track            track
   * \param  __begin          begin of data
   * \param  __end            end   of data
   * \param  V                covariance matrix
   * \return                  chi2
   */
  template<class T>
  inline double getChi2(const JLine1Z&   track,
                        T                __begin,
                        T                __end,
                        const JMatrixNZ& V)
  {
    const JVectorNZ Y(track, __begin, __end);

    return getChi2(Y, V);
  }


  /**
   * Determine chi2 of data for given track and angular and time resolution.
   *
   * \param  track            track
   * \param  __begin          begin of data
   * \param  __end            end   of data
   * \param  alpha            angular resolution [deg]
   * \param  sigma            time resolution    [ns]
   * \return                  chi2
   */
  template<class T>
  inline double getChi2(const JLine1Z& track, T __begin, T __end, const double alpha, const double sigma)
  {
    JMatrixNZ V(track, __begin, __end, alpha, sigma);

    V.invert();

    return getChi2(track, __begin, __end, V);
  }


  /**
   * Determine difference between chi2 with and without hit using full covariance matrix.
   *
   * \param  Y                residuals
   * \param  V                covariance matrix
   * \param  i                index of excluded hit
   * \return                  chi2
   */
  inline double getChi2(const JVectorNZ& Y,
                        const JMatrixNZ& V,
                        const int        i)
  {
    double chi2 = 0.0;

    for (size_t j = 0; j != V.size(); ++j) {
      chi2 += V(i,j) * Y[j];
    }

    return chi2*chi2 / V(i,i);
  }


  /**
   * Get chi2 of data for given model and fit function.
   *
   * The template argument <tt>JFit_t</tt> refers to a data structure 
   * which should provide for the function object operator:
   * <pre>
   *   double operator()(const JModel_t& model, const value_type&) const;    // chi2
   * </pre>
   * where 
   * <tt>JModel_t</tt> refers to the given model and 
   * <tt>value_type</tt> to the value type if the input data.
   * The return value should correspond to the chi2 of the hit.
   *
   * \param  model          model
   * \param  fit            fit function
   * \param  __begin        begin of data
   * \param  __end          end   of data
   */
  template<class JModel_t, class JFit_t, class T>
  inline double getChi2(const JModel_t& model, const JFit_t& fit, T __begin, T __end)
  {
    double chi2 = 0.0;

    for (T hit = __begin; hit != __end; ++hit) {
      chi2 += fit(model, *hit);
    }

    return chi2;
  }


  /**
   * Get probability due to random background.
   *
   * \param  N                    number of active PMTs
   * \param  M                    multiplicity of hits
   * \param  R_Hz                 rates       [Hz]
   * \param  T_ns                 time window [ns]
   * \return                      probability
   */
  inline double getProbability(const size_t     N,
                               const size_t     M,
                               const JK40Rates& R_Hz,
                               const double     T_ns)
  {
    using namespace JPP;

    if (N == 0) {
      return (M == 0 ? 1.0 : 0.0);
    }

    const double T =  T_ns * 1.0e-9;
    const double p = -expm1(-R_Hz.getSinglesRate() * N * T);    // 1 - P(0) due to singles rate
    
    double P = 0.0;

    for (multiplicity_type i = R_Hz.getLowerL1Multiplicity(); i <= R_Hz.getUpperL1Multiplicity(); ++i) {
      P -= expm1(-R_Hz.getMultiplesRate(i) * T);                // 1 - P(0) due to multiples rates
    }

    P *= (1.0 + p);                                             // 1 - P(0) due to multiples rates M-1 combined with singles rate
    
    return (poisson(M, R_Hz.getSinglesRate() * N * T)  -                 // P(M) due to singles rate
	    expm1(-R_Hz.getMultiplesRate(M)      * T)  -                 // P(M) due to multiples rate
	    expm1(-R_Hz.getMultiplesRate(M-1)    * T) * p) / (1.0 + P);  // P(M-1) * P(1) due to multiples rate combined with singles rate
  }
}

#endif