#ifndef __JPP_SHOWER_PDF__
#define __JPP_SHOWER_PDF__

/*! @brief Wrapper class for using Jpp shower pdfs in the Aanet framework.
    @author Jordan Seneca
    @date June 2020
 */

#include <string>
#include <iostream>
#include <fstream>
#include <iomanip>

#include "util.hh"
#include "utilJpp.hh"
#include "rec2/Params.hh"

#include "Hit.hh"
#include "Trk.hh"
#include "Det.hh"
#include "Evt.hh"
#include "Vec.hh"

#include "Math/Minimizer.h"
#include "Math/Factory.h"
#include "Math/Functor.h"

#include "JTrigger/JGetRiseTime.hh"
#include "Jeep/JMessage.hh"
#include "JAAnet/JAAnetToolkit.hh"

#include "JGeometry3D/JPosition3D.hh"
#include "JGeometry3D/JRotation3D.hh"
#include "JGeometry3D/JDirection3D.hh"
#include "JGeometry3D/JAxis3D.hh"

#include "JPhysics/JGeanz.hh"

#include "rec2/JppGradE.hh"
#include "rec2/JppGrad.hh"

//Author: Jordan Seneca

using namespace JPP;

#ifdef __ROOTCLING__
typedef int JShowerPdf_t;
typedef int JShowerCdf_t;
#else

using namespace JPP;

#endif

template< class pdftype >
class JppShower{
  
protected:  
  
  vector<Hit>          _hits;
  vector<Pmt>        _unhits;
  
  vector<Trk>       evaltrks; /*! To keep track of where the class has been evaluated (obsolete?) */

  double _k40_rate; /*! k40 rate [Hz] */
  
  bool _factorized_fudge;

  uint _verbose; /*! Verbosity */
  uint _debug;

  pair<double, double> _time_window;
  
  JGetRiseTime _jgetrisetime;

 public:
  
  pdftype _gradpdf;;

  JppShower( vector<string> shower_pdfs = { "/pbs/throng/km3net/src/Jpp/master/data/J14p.dat",
	"/pbs/throng/km3net/src/Jpp/master/data/J13p.dat"	},
		   double sigma_blur          = 0, 
		   double R_bg                = 0,
		   double t_start             = -100,
		   double t_end               = 900,
		   uint verbose = 0, 
		   uint debug   = 0,
		   uint sample_count = 0 ):
    _factorized_fudge( true ),
    _verbose( verbose ),
    _debug( debug ),
    _jgetrisetime(),
    _gradpdf( shower_pdfs,
	      sigma_blur,
	      R_bg,
	      t_start,
	      t_end,
	      verbose,
	      debug ) {
    if( sample_count ) _gradpdf.set_sample_count( sample_count );
    cout << "Constructed JppShower." << endl;
  }

  ~JppShower() {};

  /*! Initialize an event. The TimeDomSums are calculated for a given detector file
   *  and a vector of hits. The hits are also filtered when their distance to the vertex is
   *  larger than max_dist.
   */
  void init_event( Det& det,
		   vector<Hit>& hits,
		   Vec vertex     = Vec(),
		   double max_dist = 0,
		   bool incl_unhits = true ){

    if( _verbose ) cout << "JppShower<pdftype>::init_event" << endl;
    
    evaltrks.clear();
    _hits.clear();
    _unhits.clear();

    // set all flags of every PMT to 0
    for ( auto& it : det.pmts ) it.second -> flag = 0;
      
    // flag all hit PMTs with 1
    foreach ( h, hits ) {
      if( det.pmts.find ( h.pmt_id ) == det.pmts.end() ) {
	fatal( "cannot find pmt with id", h.pmt_id );
      }
      det.pmts[ h.pmt_id ] -> flag = 1;
    }
    
    for ( auto& dom_pair: det.doms ) {
      Dom dom = dom_pair.second;
      if( max_dist > 0 ) { // check for DOMs outside the region
	double dist = ( vertex - dom.pos ).len();
	if ( dist > max_dist ) continue;
      }

      foreach( pmt, dom.pmts ){
	if ( pmt.flag ) { // hit PMTs
	  vector< Hit > pmt_hits;
	  foreach( h, hits ) { 
	    if ( h.pmt_id == pmt.id ){
	      pmt_hits.push_back( h );  // store the hit times and ToT
	    }
	  }
	  std::sort( pmt_hits.begin(), pmt_hits.end(), sort_t );
	
	  //hit_pmt.first_hit_corrected_t = get_corrected_dt( hit_pmt.first_hit_t, hit_pmt.first_hit_tot );
	  _hits.push_back( pmt_hits[0] );
	} else if( incl_unhits ) _unhits.push_back( pmt );
      }
    }

    cout << "hits size: "   << _hits.size() << endl;
    cout << "unhits size: " << _unhits.size() << endl;

  };

  void init_event( Det& det,
		   vector<Pmt>& unhits ){
    
    if( _verbose ) cout << "JppShower<pdftype>::init_event unhit version" << endl;

    evaltrks.clear();
    _hits.clear();
    _unhits.clear();
        
    // flag all hit PMTs with 1
    foreach ( unhit, unhits ) {
      if( det.pmts.find ( unhit.id ) == det.pmts.end() ) {
	fatal( "cannot find pmt with id", unhit.id );
      }
      _unhits.push_back( unhit );
    }

    if( _verbose ) cout << "hits size: "   << _hits.size() << endl;
    if( _verbose ) cout << "unhits size: " << _unhits.size() << endl;

  };

  void init_hit(Det& det,
		Hit hit) {
    evaltrks.clear();
    _hits.clear();
    _unhits.clear();
    _hits.push_back( hit );

  };
  
  /*! Evaluate the likelihood for a shower hypothesis using the first hits and the first hit PDF */  
  /*  
  template<uint e>
  double eval_lik( Trk& trk, bool track_evals = false );
  */

  void set_debug( const uint debug ){
    _debug = debug;
    _gradpdf.set_debug( debug );
  }

  double eval_lik( const Trk &trk ){
    return _gradpdf.eval_lik( trk, _hits, _unhits, _factorized_fudge );
  };
  
  double eval_lik( double E, double x, double y, double z,
		   double theta, double phi, double t ){
    Trk trk;
    trk.E = E;
    trk.pos = Vec( x, y, z );
    trk.dir.set_angles( theta, phi );
    trk.t = t;

    //cout << "this trk: " << trk << endl;
    return eval_lik( trk );
  };
  
  vector< double > eval_grad( const Trk &trk ){
    double r;
    vector< double > ds;
    _gradpdf.eval_lik_gradient( r, ds, trk, _hits, _unhits );
    return ds;
  };

  vector< double > eval_grad( double E, double x, double y, double z, double theta, double phi, double t ){
    
    Trk trk;
    trk.E = E;
    trk.pos = Vec( x, y, z );
    trk.dir.set_angles( theta, phi );
    trk.t = t;
    
    return eval_grad( trk );
  };

  std::pair< double, vector< double > > eval_lik_grad( const Trk &trk ){
    double r;
    vector< double > ds;
    _gradpdf.eval_lik_gradient( r, ds, trk, _hits, _unhits );
    return std::pair< double, vector< double > >{ r, ds };
  };

  std::pair< double, vector< double > > eval_lik_grad( double E, double x, double y, double z, double theta, double phi, double t ){
    
    Trk trk;
    trk.E = E;
    trk.pos = Vec( x, y, z );
    trk.dir.set_angles( theta, phi );
    trk.t = t;

    return eval_lik_grad( trk );
  };

  vector< double > num_eval_lik_grad( const Trk &trk, double delta = 1e-8 ){

    const uint N = 7;

    vector< double > ds ( N, 0 );

    // E, x, y, z, dx, dy, dz, dphi, dtheta, t
    for( uint ii = 0; ii != N; ii++ ){	
      const uint ipar = ii >= 4 ? ii + 3 : ii;
      std::pair< Trk, Trk > trks = delta_trks( trk, ipar, delta );
      ds[ii] = ( _gradpdf.eval_lik( trks.second, _hits, _unhits, _factorized_fudge )
	       - _gradpdf.eval_lik( trks.first,  _hits, _unhits, _factorized_fudge ) ) / delta;
      //cout << std::setprecision(16) << "trks.first: " << trks.first << endl;
      //cout << std::setprecision(16) <<  "trks.second: " << trks.second << endl;
      //cout << std::setprecision(16) <<  "ii : " << ii << " ds[ii]: " << ds[ii] << endl;
    }

    return ds;

  }


  double get_time_trk(Trk trk,
		      Vec pmt_pos,
		      Vec pmt_dir,
		      const double rn) const{
    HitParams this_pmt( trk, pmt_pos, pmt_dir );
    return get_time(0,
		    trk.E,
		    this_pmt.D,
		    this_pmt.cd,
		    this_pmt.theta,
		    this_pmt.phi,
		    rn);
  }

  double get_time(const uint   N,
		  const double E,
		  const double D,
		  const double cd,
		  const double theta,
		  const double phi,
		  const double rn) const;

  double get_first_time_trk(Trk trk,
			    Vec pmt_pos,
			    Vec pmt_dir,
			    const double rn) const{
    HitParams this_pmt( trk, pmt_pos, pmt_dir );
    return get_first_time(0,
			  trk.E,
			  this_pmt.D,
			  this_pmt.cd,
			  this_pmt.theta,
			  this_pmt.phi,
			  rn);
  }

  double get_first_time(const uint   N,
			const double E,
			const double D,
			const double cd,
			const double theta,
			const double phi,
			const double rn) const;

  /*
  double add_hitL(Trk& trk,
		  Hit& hit,
		  double& hitL){
    const double mu   = _gradpdf.Ntot( trk, hit );
    const double logp1 = log( 1e-12 + 1 - exp( - mu ) );
    hitL += logp1;
    
    return hitL;
  };
  */
  /*
  double add_unhitL(Trk& trk,
		    Pmt& pmt,
		    double& unhitL){
    
    const double mu = _gradpdf.Ntot( trk, pmt );    
    const double logp0 = -mu;
    unhitL += logp0;

    return unhitL;
  }
  */
  /*! Evaluate the full likelihood for this track assuming a pointlike shower*/

  /*
  double eval_lik( Trk& trk,
		   bool track_evals = false){
    
    if (trk.E > 1e14 ) trk.E = 1e14;
    
    double L=0, unhitL=0, hitL=0;

    foreach( hit, _hits ) { add_hitL( trk, hit, hitL ); }
    L -= hitL;

    foreach( unhit, _unhits ) { add_unhitL( trk, unhit, unhitL ); }
    L -= unhitL;
    
    trk.lik = L;
    
    if( track_evals ) evaltrks.push_back( trk );
    
    return L;
  };
  */
  
  /*! Testing functions */   
  const uint test_params ( Trk& trk ) const;
  const uint test_npe_rad ( const double E ) const;

  /*! Utility functions */

  double get_verbose() const{ return _verbose; };

  /*! Get the k40 rate p.e. for time window [ns] */
  double get_k40_npe_time_window() const;

  //Check if return of integral of pdf and cdf the same
  const uint test_cdf( JppShower& pdf,
			const uint precision = 10,
			const double E = 1e6 ) const;

  template<typename F>
  double inverse( const double rn,
		  F& f_cdf,
		  double x_0,
		  double x_1,
		  const double tol ) const{
    double x = (x_1 + x_0) / 2.;
    if( x_1 - x_0 < tol ) return x;
    else if( f_cdf(x) < rn ) inverse( rn, f_cdf, x, x_1, tol);
    else                     inverse( rn, f_cdf, x_0, x, tol);
  };

};

#include "Math/IFunction.h"
template< class pdftype >
struct JppGradFunc : public ROOT::Math::IGradientFunctionMultiDimTempl<double>
{
  JppShower<pdftype>* pdf;
  double eval_offset;

  virtual IGradientFunctionMultiDimTempl<double> *Clone() const { 
    JppGradFunc *pdf = new JppGradFunc(*this); 
    return pdf;
  }
  
  virtual unsigned int NDim () const { return 7; }

  void init_offset(const double* x )
  {
    eval_offset = 0.0;
    eval_offset = DoEval( x );
    if( pdf->get_verbose() ) cout << "LL offset set to " << eval_offset << endl;
    if( eval_offset != eval_offset ){ 
      cout << "ERROR: NaN value returned by DoEval in init_offset" << endl;
      cin.ignore();
      throw( JException( "NaN value encountered in init_offset" ) );
    }
  };
  

  virtual double DoEval(const double *x) const
  {

    double res = pdf -> eval_lik( pow( 10, x[0] ), x[1], x[2], x[3], x[4], x[5], x[6] ); //use pdf here
    res -= eval_offset;
    //std::pair< double, vector< double > > _res = pdf -> eval_lik_grad( pow( 10, x[0] ), x[1], x[2], x[3], x[4], x[5], x[6] );
    //double res = _res.first;
    if( pdf->get_verbose() ) cout  << "Calling DoEval:   " << res << endl << endl << " x[0]: " << x[0] << " x[1]: " << x[1] << " x[2]: " << x[2] << " x[3]: " << x[3] << " x[4]: " << x[4] << " x[5]: " << x[5] << " x[6]: " << x[6] << endl;

    return res;

  }

  virtual double DoDerivative(const double *x,
			      unsigned int icoord) const
  {
    std::pair< double, vector< double > > res = pdf -> eval_lik_grad( pow( 10, x[0] ), x[1], x[2], x[3], x[4], x[5], x[6] );
    
    res.second[0] *= pow( 10, x[0] ) * log(10); //log(10), dE/dlogE
    if( pdf->get_verbose() ) cout  << "Calling DoDerivative: " << res.second[icoord]  << " x[0]: " << x[0] << " x[1]: " << x[1] << " x[2]: " << x[2] << " x[3]: " << x[3] << " x[4]: " << x[4] << " x[5]: " << x[5] << " x[6]: " << x[6] << endl;

    return res.second[icoord];
  }
  
  virtual void FdF (const double* x, double& f, double* df) const
  { 
    std::pair< double, vector< double > > res = pdf -> eval_lik_grad( pow( 10, x[0] ), x[1], x[2], x[3], x[4], x[5], x[6] );
    cout << "offset: " << eval_offset << endl;
    f = res.first - eval_offset;
    for( uint ii = 0; ii != 7; ii++ ) df[ii] = res.second[ii];
    df[0] *= pow( 10, x[0] ) * log(10); //log(10), dE/dlogE
    if( pdf->get_verbose() ) cout  << "Calling FdF: " << f  << endl << ", df0: " << df[0] << ", df1: " << df[1] << ", df2: " << df[2] << ", df3: " << df[3] << ", df4: " << df[4] << ", df5: " << df[5] << ", df6: " << df[6] << endl << " x[0]: " << x[0] << " x[1]: " << x[1] << " x[2]: " << x[2] << " x[3]: " << x[3] << " x[4]: " << x[4] << " x[5]: " << x[5] << " x[6]: " << x[6] << endl;
  }
  
  virtual void Gradient (const double* x, double* df) const
  {
    std::pair< double, vector< double > > res = pdf -> eval_lik_grad( pow( 10, x[0] ), x[1], x[2], x[3], x[4], x[5], x[6] );
    for( uint ii = 0; ii != 7; ii++ ) df[ii] = res.second[ii];
    df[0] *= pow( 10, x[0] ) * log(10); //log(10), dE/dlogE
    if( pdf->get_verbose() ) cout  << "Calling Gradient: " << res.first - eval_offset  << endl << ", df0: " << df[0] << ", df1: " << df[1] << ", df2: " << df[2] << ", df3: " << df[3] << ", df4: " << df[4] << ", df5: " << df[5] << ", df6: " << df[6] << endl << " x[0]: " << x[0] << " x[1]: " << x[1] << " x[2]: " << x[2] << " x[3]: " << x[3] << " x[4]: " << x[4] << " x[5]: " << x[5] << " x[6]: " << x[6] << endl;
  }

};

inline void use_jppshowerpdfmerged(){ //Needed to make templates
  JppShower<JppGrad>  c1;
  JppShower<JppGradE> c2;
};

#endif