#ifndef __JFIT__JSHOWER3EZREGRESSOR__
#define __JFIT__JSHOWER3EZREGRESSOR__

#include "JPhysics/JPDFTypes.hh"
#include "JPhysics/JPDFTable.hh"
#include "JPhysics/JPDFToolkit.hh"
#include "JPhysics/JNPETable.hh"

#include "JTools/JFunction1D_t.hh"
#include "JTools/JFunctionalMap_t.hh"
#include "JPhysics/JConstants.hh"
#include "JTools/JRange.hh"
#include "JTools/JResult.hh"
#include "JTools/JFunction1D_t.hh"
#include "JTools/JFunctionalMap_t.hh"
#include "JTools/JAbstractHistogram.hh"

#include "JGeometry3D/JVector3D.hh"
#include "JGeometry3D/JVersor3D.hh"
#include "JGeometry3D/JDirection3D.hh"

#include "JMath/JZero.hh"

#include "JFit/JTimeRange.hh"
#include "JFit/JPMTW0.hh"
#include "JFit/JSimplex.hh"
#include "JFit/JGandalf.hh"
#include "JFit/JMEstimator.hh"
#include "JFit/JRegressor.hh"
#include "JFit/JShower3EZ.hh"
#include "JFit/JFitToolkit.hh"

#include "Jeep/JMessage.hh"

/**
 * \file
 * Data regression method for JFIT::JShower3EZ.
 * \author mdejong
 */

namespace JFIT {

  using JPHYSICS::JPDFType_t;
  using JPHYSICS::DIRECT_LIGHT_FROM_EMSHOWER;
  using JPHYSICS::SCATTERED_LIGHT_FROM_EMSHOWER;
  using JTOOLS::get_value;

  /**
   * Regressor function object for JShower3EZ fit using JSimplex minimiser.
   */
  template<>
  struct JRegressor<JShower3EZ, JSimplex> :
    public JAbstractRegressor<JShower3EZ, JSimplex>
  {
    using JAbstractRegressor<JShower3EZ, JSimplex>::operator();

    typedef JTOOLS::JSplineFunction1S_t                                     JFunction1D_t;
    typedef JTOOLS::JMapList<JTOOLS::JPolint0FunctionalMap,
	    JTOOLS::JMapList<JTOOLS::JPolint0FunctionalMap,
	    JTOOLS::JMapList<JTOOLS::JPolint0FunctionalGridMap,
            JTOOLS::JMapList<JTOOLS::JPolint0FunctionalGridMap> > > >        JPDFMaplist_t;
    typedef JPHYSICS::JPDFTable<JFunction1D_t, JPDFMaplist_t>                JPDF_t;
    
    typedef JTOOLS::JMapList<JTOOLS::JPolint0FunctionalMap,
	    JTOOLS::JMapList<JTOOLS::JPolint0FunctionalMap,
	    JTOOLS::JMapList<JTOOLS::JPolint0FunctionalGridMap,
            JTOOLS::JMapList<JTOOLS::JPolint0FunctionalGridMap> > > >        JNPEMaplist_t;
    typedef JPHYSICS::JNPETable<double, double, JNPEMaplist_t>               JNPE_t;
   
    /**
     * Parameterized constructor
     *
     * The PDF file descriptor should contain the wild card character JPHYSICS::WILD_CARD which
     * will be replaced by the corresponding PDF types listed in JRegressor<JShower3Z, JGandalf>::pdf_t.
     *
     * \param  fileDescriptor     PDF file descriptor
     */

    JRegressor(const std::string& fileDescriptor):
      estimator(new JMEstimatorNull())
    {
      using namespace std;
      using namespace JPP;

      const JPDF_t::JSupervisor supervisor(new JPDF_t::JDefaultResult(JMATH::zero));
      
      for (int i = 0; i != NUMBER_OF_PDFS; ++i) {
	
	try {
	    
	  JPDF_t pdf;

	  const string file_name = getFilename(fileDescriptor, pdf_t[i]);
	    
	  NOTICE("loading PDF from file " << file_name << "... " << flush);

	  pdf.load(file_name.c_str());
	  
	  NOTICE("OK" << endl);

	  pdf.setExceptionHandler(supervisor);

	  npe[ i ] = JNPE_t(pdf);
	}
	catch(const JException& error) {
	  FATAL(error.what() << endl);
	}
      }

      // Add PDFs
      for (int i = 1; i < NUMBER_OF_PDFS; i += 2) {
  
	npe[ i ].add(npe[i-1]);

	JNPE_t buffer;
	
	npe[i-1].swap(buffer);
      }    
    }

    /**
     * Fit function.
     * This method is used to determine the chi2 of given PMT with respect to shower hypothesis.
     *
     * \param  shower             shower
     * \param  pmt                pmt
     * \return                    chi2 
     */
    double operator()(const JShower3EZ& shower, const JPMTW0& pmt) const
    {
      using namespace JPP;
      
      JPosition3D D(pmt.getPosition());
      JDirection3D U(pmt.getDirection());

      D.sub(shower.getPosition());

      JVersor3D shower_dir(shower.getDirection());

      const double z  = D.getDot(shower_dir);
      const double x  = D.getX()  -  z * shower.getDX();
      const double y  = D.getY()  -  z * shower.getDY();
      const double cosDelta = z/D.getLength();  // Delta = angle between shower direction and PMT position
 
      U.rotate(JRotation3Z(-atan2(y,x)));                  // rotate PMT axis to x-z plane

      const double theta = U.getTheta();
      const double phi   = fabs(U.getPhi());

      JNPE_t::result_type H0 = getH0(pmt.getR());  // getH0 = Get background hypothesis value for time integrated PDF. 
      JNPE_t::result_type H1 = getH1(D.getLength(), cosDelta, theta, phi, shower.getE());  // getH1 = Get signal hypothesis value for time integrated PDF. 

      if (get_value(H1) >= Vmax_npe) {
        H1 *= Vmax_npe / get_value(H1);
      }

      H1 += H0; // now H1 is signal + background

      const bool hit = pmt.getN() != 0;
      const double u = getChi2(get_value(H1), hit); // -log(lik)
      
      return estimator->getRho(u);
    }

    /**
     * Get background hypothesis value for time integrated PDF.
     *
     * \param  R_Hz               rate [Hz]
     * \return                    hypothesis value
     */
    JNPE_t::result_type getH0(const double R_Hz) const
    {
      return JNPE_t::result_type(R_Hz * 1e-9 * T_ns.getLength());
    }

    /**
     * Get signal hypothesis value for time integrated PDF.
     *
     * \param  D                  PMT distance from shower [m]
     * \param  cosDelta           angle between shower direction and PMT position
     * \param  theta              PMT zenith  angle [deg]
     * \param  phi                PMT azimuth angle [deg]
     * \param  E                  shower energy [GeV]
     * \return                    hypothesis value
     */
    JNPE_t::result_type getH1(const double D,
			      const double cosDelta,
			      const double theta,
			      const double phi,
			      const double E) const
    {
      JNPE_t::result_type h1 = JMATH::zero;
	
      for (int i = 0; i != NUMBER_OF_PDFS; ++i) {

	if (!npe[i].empty() && D <= npe[i].getXmax()) {

	  try {

	    JNPE_t::result_type y1 = E * npe[i](std::max(D, npe[i].getXmin()), cosDelta, theta, phi);

	    // safety measure

	    if(y1 < 0){
	      y1 = 0;  
	    }
	    
	    h1 += y1;
	  
	  }
	  catch(JLANG::JException& error) {
	    ERROR(error << std::endl);
	  }
	}
      }

      return h1;
    }


    static JTimeRange   T_ns;                      //!< Time window with respect to Cherenkov hypothesis [ns]
    static double       Vmax_npe;                  //!< Maximal integral of PDF [npe]

    static const int    NUMBER_OF_PDFS = 2;

    static const JPDFType_t pdf_t[NUMBER_OF_PDFS];

    JNPE_t              npe[NUMBER_OF_PDFS];       //!< PDF

    JLANG::JSharedPointer<JMEstimator> estimator;  //!< M-Estimator function
  };  


  /**
   * Regressor function object for JShower3EZ fit using JGandalf minimiser.
   */
  template<>
  struct JRegressor<JShower3EZ, JGandalf> :
    public JAbstractRegressor<JShower3EZ, JGandalf>
  {
    using JAbstractRegressor<JShower3EZ, JGandalf>::operator();

    typedef JTOOLS::JSplineFunction1S_t                                     JFunction1D_t;
    typedef JTOOLS::JMapList<JTOOLS::JPolint0FunctionalMap,
	    JTOOLS::JMapList<JTOOLS::JPolint0FunctionalMap,
	    JTOOLS::JMapList<JTOOLS::JPolint0FunctionalGridMap,
            JTOOLS::JMapList<JTOOLS::JPolint0FunctionalGridMap> > > >        JPDFMaplist_t;
    typedef JPHYSICS::JPDFTable<JFunction1D_t, JPDFMaplist_t>                JPDF_t;
    
    typedef JTOOLS::JMapList<JTOOLS::JPolint1FunctionalMapH,
	    JTOOLS::JMapList<JTOOLS::JPolint1FunctionalMap,
	    JTOOLS::JMapList<JTOOLS::JPolint1FunctionalGridMap,
            JTOOLS::JMapList<JTOOLS::JPolint1FunctionalGridMap> > > >        JNPEMaplist_t;
    typedef JPHYSICS::JNPETable<double, double, JNPEMaplist_t>               JNPE_t;
   
    /**
     * Parameterized constructor
     *
     * The PDF file descriptor should contain the wild card character JPHYSICS::WILD_CARD which
     * will be replaced by the corresponding PDF types listed in JRegressor<JShower3Z, JGandalf>::pdf_t.
     *
     * \param  fileDescriptor     PDF file descriptor
     */

    JRegressor(const std::string& fileDescriptor):
      estimator(new JMEstimatorNull())
    {
      using namespace std;
      using namespace JPP;

      const JPDF_t::JSupervisor supervisor(new JPDF_t::JDefaultResult(JMATH::zero));
      
      for (int i = 0; i != NUMBER_OF_PDFS; ++i) {
	
	try {
	    
	  JPDF_t pdf;

	  const string file_name = getFilename(fileDescriptor, pdf_t[i]);
	    
	  NOTICE("loading PDF from file " << file_name << "... " << flush);

	  pdf.load(file_name.c_str());
	  
	  NOTICE("OK" << endl);

	  pdf.setExceptionHandler(supervisor);

	  npe[i] = JNPE_t(pdf);
	}
	catch(const JException& error) {
	  FATAL(error.what() << endl);
	}
      }

      // Add PDFs
      for (int i = 1; i < NUMBER_OF_PDFS; i += 2) {

	npe[ i ].add(npe[i-1]);

	JNPE_t buffer;

	npe[i-1].swap(buffer);
      }    
    }

    /**
     * Fit function.
     * This method is used to determine the chi2 of given PMT with respect to shower hypothesis.
     *
     * \param  shower             shower
     * \param  pmt                pmt
     * \return                    chi2 
     */
    result_type operator()(const JShower3EZ& shower, const JPMTW0& pmt) const
    {
      using namespace JPP;
      using namespace std;
      
      JPosition3D D(pmt.getPosition());
      JDirection3D U(pmt.getDirection());

      D.sub(shower.getPosition());

      JVersor3D shower_dir(shower.getDirection());

      const double z  = D.getDot(shower_dir);
      const double x  = D.getX()  -  z * shower.getDX();
      const double y  = D.getY()  -  z * shower.getDY();
      const double R2 = D.getLengthSquared() - z*z;
      const double R = sqrt(R2);
      const double d = D.getLength();  
      const double cos_th0 = z/d;  // theta0 = angle between shower direction and PMT position
 
      U.rotate(JRotation3Z(-atan2(y,x)));                  // rotate PMT axis to x-z plane

      const double theta = U.getTheta();
      const double phi   = fabs(U.getPhi());

      JNPE_t::result_type H0 = getH0(pmt.getR());  // getH0 = Get background hypothesis value for time integrated PDF. 
      JNPE_t::result_type H1 = getH1(d, cos_th0, theta, phi, shower.getE());  // getH1 = Get signal hypothesis value for time integrated PDF. 

      if (get_value(H1) >= Vmax_npe) {
        H1 *= Vmax_npe / get_value(H1);
      }

      H1 += H0; // now H1 is signal + background

      const bool hit = pmt.getN() != 0;
      const double u = H1.getChi2(hit);

      result_type result;      

      result.chi2     = estimator->getRho(u);

      result.gradient = JShower3EZ(JPoint4D(JVector3D(x*z / (d*d*d),       // d(cos_th0)/d(x)
 						      y*z / (d*d*d),       // d(cos_th0)/d(y)
						      -(R*R) / (d*d*d)),          // d(cos_th0)/d(z) 
					    0.0),                                 // d(cos_th0)/d(t) 
				   JVersor3Z(x*z*z / (d*R),                // d(cos_th0)/d(dx)
					     y*z*z / (d*R)),               // d(cos_th0)/d(dy)  
				   0.0);                                          // d(cos_th0)/d(E)
      
      result.gradient.mul(estimator->getPsi(u));
      result.gradient.mul(H1.getDerivativeOfChi2(hit));                    // x d(chi2)/d(cos_th0)
       
      return result;
    }

    /**
     * Get background hypothesis value for time integrated PDF.
     *
     * \param  R_Hz               rate [Hz]
     * \return                    hypothesis value
     */
    JNPE_t::result_type getH0(const double R_Hz) const
    {
      return JNPE_t::result_type(R_Hz * 1e-9 * T_ns.getLength(), 0.0);
    }

    /**
     * Get signal hypothesis value for time integrated PDF.
     *
     * \param  D                  PMT distance from shower [m]
     * \param  cosDelta           angle between shower direction and PMT position
     * \param  theta              PMT zenith  angle [deg]
     * \param  phi                PMT azimuth angle [deg]
     * \param  E                  shower energy [GeV]
     * \return                    hypothesis value
     */
    JNPE_t::result_type getH1(const double D,
			      const double cosDelta,
			      const double theta,
			      const double phi,
			      const double E) const
    {
      JNPE_t::result_type h1 = JMATH::zero;
	
      for (int i = 0; i != NUMBER_OF_PDFS; ++i) {

	if (!npe[i].empty() && D <= npe[i].getXmax()) {

	  try {

	    JNPE_t::result_type y1 = E * npe[i](std::max(D, npe[i].getXmin()), cosDelta, theta, phi);
	  
	    if (get_value(y1) > 0.0) {
	      h1 += y1;
	    }
	  
	  }
	  catch(JLANG::JException& error) {
	    ERROR(error << std::endl);
	  }
	}
      }

      return h1;
    }


    static JTimeRange   T_ns;                      //!< Time window with respect to Cherenkov hypothesis [ns]
    static double       Vmax_npe;                  //!< Maximal integral of PDF [npe]

    static const int    NUMBER_OF_PDFS = 2;

    static const JFIT::JPDFType_t pdf_t[NUMBER_OF_PDFS];
    
    JNPE_t              npe[NUMBER_OF_PDFS];       //!< PDF
    
    JLANG::JSharedPointer<JMEstimator> estimator;  //!< M-Estimator function
  };  

  /**
   * PDF types.
   */
  const JPDFType_t JRegressor<JShower3EZ, JSimplex>::pdf_t[] = { DIRECT_LIGHT_FROM_EMSHOWER, 
								 SCATTERED_LIGHT_FROM_EMSHOWER };

  const JPDFType_t JRegressor<JShower3EZ, JGandalf>::pdf_t[] = { DIRECT_LIGHT_FROM_EMSHOWER, 
								 SCATTERED_LIGHT_FROM_EMSHOWER };

  /**
   * Default values.
   */
  JTimeRange JRegressor<JShower3EZ, JSimplex>::T_ns;
  double     JRegressor<JShower3EZ, JSimplex>::Vmax_npe = std::numeric_limits<double>::max();

  JTimeRange JRegressor<JShower3EZ, JGandalf>::T_ns;
  double     JRegressor<JShower3EZ, JGandalf>::Vmax_npe = std::numeric_limits<double>::max();
}

#endif