#ifndef POINTSOURCECOMPONENT_HH
#define POINTSOURCECOMPONENT_HH

#include "Component.hh"
#include "rootutil.hh"


static constexpr double must_update = -12345.6789;

/* abstract class that describes a point source component */
struct PointSourceComponent : public Component
{
  //DefaultDetResponse& detres;
  FullDetResponse detector_response;

  Parameter ra    = { "right assencion of the source (radians)", 0, -pi, pi, 0.1 };
  Parameter dec   = { "declinations of the source (radians)",    0, -pi, pi, 0.1 };

  ScalableFlux *flux;

  double _angle_scale;
  double max_angle;      ///< for events further than max_angle from source, the PSF (and thus likellihood) will be assumed 0,
  double min_theta_rec;  ///< do not generate events with zenith angles below this value

  int max_tries = 10;         ///< see generate_event;


  //vector<double> something;

  std::map<flavor, std::map < channel , TH3D > > HFC; ///< 3d histograms with Erec, dayfraction, reco-error-angle



  /*! return a table with Aeff, Rate etc for all energy bins */
  Table rate_info( flavor flav, channel chan );

  /*! return a big table with rate info for all channes/flavors */
  Table rate_info(); 

  /*! Total number of evens for flav/chan combination */

  double Ntot( flavor flav, channel chan )
  {
    double r = detector_response.Aeff(flav, chan ).GetAreaAtDecl( dec.value ).GetTotalRate( *flux );
    return r;
  }



  /*! Total number of events integrated over channels and flavours. It should at all times be equal to norm.value  */

  double Ntot()
  {
    double r = 0;
    for ( auto f : flavors ) {
	    if (f == muon) cout << "MUONS!!!" << endl;
      for ( auto c : channels ) {
	r += Ntot(f, c);
      }
    }
    return r;
  }

  /* update shall be called by the fitter after parameters have changed */

  
  virtual void update()
  {
    MEMCHECK

    // Even if norm, did not change, the normalisation of the flux may 
    // change, because it's spectral index changed.

    flux -> norm *= norm.value / Ntot();
  }


  /*! For PE generation, generate a random flavor, channel combination */

  void random_flav_chan ( flavor& flav, channel& chan ) {

    static vector< double >                 _probs;
    static vector < pair< flavor, channel>> _typs;
    static double  olddec = -999;


    if ( dec.value != olddec ) // <== todo : check if the flux changed
    {
      // recompute the totals
      double tot = 0;
      unsigned i = 0;

      _probs.clear();
      _typs.clear();

      for ( auto f : flavors ) {
        for ( auto c : channels ) {

          tot += Ntot( f, c );
          _probs.push_back(tot);
          _typs.push_back( std::make_pair( f, c ) );
          i++;
        }
      }

      // normalize
      double p = _probs[i - 1];
      for (unsigned j = 0; j < i; j++) _probs[j] /= p;

      for (unsigned i = 0; i < _probs.size() ; i++ ) {
        print ( "prob:", i, _typs[i].first, _typs[i].second, _probs[i] );
      }
    }


    olddec  = dec.value;
    //oldflux = *flux;

    double r = gRandom->Rndm();
    for (unsigned i = 0; i < _probs.size() ; i++ ) {

      if ( _probs[i] > r ) {
        flav = _typs[i].first;
        chan = _typs[i].second;
        return;
      }
    }
  }



  /*! Return angle bewteen the source and a track reconstructed at theta_rec, phi_rec, t */

  double angle( double theta_rec, double phi_rec, TTimeStamp& t )
  {
    EquatorialCoords eq( ra.value, dec.value );
    Vec dtrue = -eq.track_direction( detector_response.det(), t );
    Vec drec; drec.set_angles( theta_rec, phi_rec );
    double alpha = angle_between ( dtrue, drec );
    return alpha;
  }


  /*! Return angle bewteen source and a track reconstructed at theta_rec, phi_rec, t */

  double angle( SearchEvent& sev )
  {
    return angle( sev.zenith, sev.azimuth, sev.t );
  }


  /* The true zenith angle if source is at ra,dec */

  double true_theta( SearchEvent& sev )
  {
    EquatorialCoords eq( ra.value, dec.value );
    Vec dtrue = -eq.track_direction( detector_response.det(), sev.t );
    return dtrue.theta();
  }




  /*! Return the flux object corresponding to the current hypothesis */

  Flux* getflux()
  {
    print ("getflux");
    update();
    return flux;
  }

  

  /*! Return the likelihood of the event -- summed over possible flavours */

  double dN_dOmega_dlogE(SearchEvent& sev )
  {
    TIMEFUNC

    //MEMCHECK
    //print_parameters();

    if ( abs ( dec.value - sev.coordinates.dec() ) > max_angle ) return 0; // for speed

    double a = sev.coordinates.distance( EquatorialCoords( ra.value, dec.value ) );

    if ( a > max_angle ) return 0;

    return dN_dOmega_dlogErec( channel( sev.chan ) , sev.zenith, sev.azimuth, sev.Eproxy, sev.t );
  }



  /*! Return the likelihood of the event -- summed over possible flavours */

  double dN_dOmega_dlogErec( channel chan,
                             double theta_rec, double phi_rec, double logE_rec, TTimeStamp& t )
  {
    TIMEFUNC
    //MEMCHECK

    double alpha, theta_true;
    Vec dtrue;
    double r;
     
    EquatorialCoords eq( ra.value, dec.value );
    dtrue = -eq.track_direction( detector_response.det(), t );
    
    theta_true = dtrue.theta();
    Vec drec; drec.set_angles( theta_rec, phi_rec );
    alpha = angle_between ( dtrue, drec );
   
    r = dN_dOmega_dlogErec_alpha( chan, alpha, logE_rec, theta_true ); 
    
    return r;
  }



  /*! dN_dOmega_dlogErec at a value of angle alpha to the true source, for
    a source at true theta_true.

    Loops over the flavors, integrates over the true energy.
      Will return zero if angle-to-source > max_angle.

    \f[

    \frac{d N}{d \Omega d E_{\rm rec} } ( \alpha, \theta  ) = \int d E_{\rm true} \Phi(E_{\rm Etrue}) A^{\rm eff}(E_true, \theta) \frac{dP}{dE_{\rm rec}

    \f]


   */


  double dN_dOmega_dlogErec_alpha( channel chan,
                                   double alpha, double logE_rec, double theta_true,
                                   bool use_max_angle = true )
  {
    TIMEFUNC
    //MEMCHECK


    double w=0;

    // if ( theta_true < 0 || theta_true > pi )
    // {
    //   fatal("wrong value of theta_true", theta_true );
    // }


    if ( use_max_angle && alpha > max_angle ) return 0;

    // we are going to use to the flux, so update it.
    // update() ; <- update now called by fitter.
    
    // Define the integrant

    flavor flav;                  // to be captured by the lambda
    bool verb = false;

    auto integrand = [&] ( double logEtrue ) {


      double p = detector_response.dP_dlogErec_dOmega_reco_error_log(flav, chan, logE_rec,
                 log10( alpha * 180 / pi ),
                 logEtrue,
                 cos(theta_true) );

      double f = flux->dNdE( flav , logEtrue ) * 24 * 3600 * 365;
      if ( ::isnan( f ) ) {
        fatal( " flux is nan for ", flav, logEtrue );
      }

      if (verb) print ("check", p, f , flav, chan, logE_rec, log10( alpha * 180 / pi) , logEtrue, cos(theta_true) );

      if (f < 0) f = 0;
      return p * f;
    };


     w = 0;


    for ( auto flav_ : flavors )
    {
      flav = flav_;
      w += detector_response.Aeff( flav_, chan ).integrate( cos(theta_true), integrand );


      if ( ::isnan(w) ) {
        print("nan detected after integrating aeff ", flav, chan, theta_true );
        verb = true;
        double x =  detector_response.Aeff( flav, chan ).integrate( cos(theta_true), integrand );
        print (x);
        fatal("quiting after nan.");
      }

    }

    return w;
  }





  virtual double compute_integral()
  {
    print ("compute_integral() not yet implemented for PointSource Component");
    return 0;
  }


  void print_event( SearchEvent& sev )
  {
    print ("----------- point source likelihood for event -------------- ");
    print (sev.__str__() );
    print ();

    double alpha = angle( sev );
    double truetheta = true_theta( sev );

    print (" component norm           = ", norm.value );
    print (" component ra, dec        = ", ra.value, dec.value );
    print (" angle [radians]          = ", alpha);
    print (" angle [degrees]          = ", alpha * 180 / pi );
    print (" truetheta                = ", truetheta);
    print (" dN_dOmega_dlogErec_alpha = ", dN_dOmega_dlogErec_alpha( channel( sev.chan ), alpha, sev.Eproxy, truetheta ) );
    print (" Energy                   = ", sev.Eproxy);
    print (" ---");

    // if (! CheckIfUpdated() ) NormalizeFlux();
    flux->Print();
  }


  TH3D *H = nullptr ;

  /*! The idea is this
      - for each time of day , the source is at a fixed true zenith angle

      - For each zenith angle, we assume we have a fixed 2-d PDF
        of the anglular eror and Erec.

      - this 2d pdf is described by this->dN_dOmega_dlogErec_alpha


      This function produces a 3d PDF of Erec, dayfraction, alpha, by
      transforming theta -> dayfraction.

      The acceptance at theta for our point source is computed inside
      dN_dOmega_dlogErec_alpha
  */

  // we will need one of these for every channel, flavour combination.

  TH3D& _pdf_erec_dayfraction_loga ( flavor flav, channel chan ); // in cc


  SearchEvent generate_event()
  {
    TIMEFUNC

    SearchEvent r;

    channel chan;
    flavor flav;

    double erec, dayf, loga;


    while ( true )
    {
      //generate_channel_and_flavor( chan, flav );

      random_flav_chan ( flav, chan );

      //print ( "generate a evt with ", flav, chan );
      //print_parameters();

      _pdf_erec_dayfraction_loga( flav, chan ).GetRandom3(erec, dayf, loga);

      static TTimeStamp T0 (2020, 1, 1, 0, 0, 0 );  // todo : fix dupliction in _pdf_erec...

      r.t.SetSec( T0.GetSec() + dayf * 24 * 3600 ) ;
      EquatorialCoords c( ra.value, dec.value );
      Vec dtrue = -c.track_direction( detector_response.det(), r.t );

      bool event_is_good = false;

      for (int i = 0 ; i < max_tries; i++ )
      {
        double a = pow(10, loga) * pi / 180 * _angle_scale;

        //a = 0;
        Mat M; M.set(dtrue);
        Vec tr; tr.set_angles( a, gRandom->Rndm() * 2 * pi );
        Vec rr = M * tr;

        r.coordinates = EquatorialCoords( detector_response.det(), -rr, r.t);

        r.flav    = flav;
        r.chan    = chan;
        r.Eproxy  = erec;
        r.zenith  = rr.theta();
        r.azimuth = rr.phi();

        if ( r.zenith < min_theta_rec ) continue; // above the cut -> try again
        event_is_good = true;
        break;
      }


      if (!event_is_good) {
        print ( "PointSourceComponent generated event with true zenith = ", acos( dtrue.z ) );
        print ( "and was not able to generate a rec. track above the cut of ", min_theta_rec );
        print ( "after ", max_tries, "tries ");
        print ( "-> will generate true zenith");
        continue; // goto while(true)
      }

      if ( r.azimuth < 0 ) r.azimuth += 2 * pi;
      return r;
    }
  }


  // -----boilerplate-----
  PointSourceComponent(string name, FullDetResponse& detres_, double r_ = 1, double d_ = 1): Component(name), detector_response(detres_) {
    ra.value    = r_;
    dec.value   = d_;
    max_angle   = pi;
    _angle_scale = 1;

    flavors  = detres_.flavors;
    flavors.erase(std::remove(flavors.begin(), flavors.end(), muon), flavors.end()); //esare-remove idiom to avoid signal calculation from muon histograms
    channels = detres_.channels;
  };

  virtual PointSourceComponent *clone() const // clone
  {
    return new PointSourceComponent(*this); // return pointer to object
  }

  virtual ~PointSourceComponent() {
    if (H != NULL) {
      H->Delete();
      delete H;
    }
    if (flux != NULL) {
      cout << "Warning flux is not cleaned in the daughter of PointSourceComponent" << endl;
    }
  } // destructor

  ClassDef( PointSourceComponent, 1) // for ROOT
};



struct PowerLawComponent : public PointSourceComponent
{
  Parameter gamma = { "slope of the flux", 2., 1., 3., 0.1 };

  double _old_norm;
  double _old_gamma;
  double _old_ra;
  double _old_decl = -999;;

  // cachee optimisation for when gamma,ra,decl does not change 
  vector<double> _dp_domega_dloge; 
  
  
  PowerLawComponent(string name, FullDetResponse& detres_, double g_ = 2, double r_ = 1, double d_ = 1) :
    PointSourceComponent(name, detres_, r_, d_) {
    gamma.value = g_;
    H           = NULL;
    flux = new PowerLawFlux(1, g_);
    _old_gamma = 0;
  }

  /*! update is called by the fit after it has changed one or more parameters */

  virtual void update() {

    // If only norm has changed, we don't have to do antything: dN_dOmega_dlogE 
    // will just scale the cached result.

    if ( gamma.value != _old_gamma || ra.value != _old_ra || dec.value != _old_decl ) {

      print("update" , gamma.value != _old_gamma , ra.value != _old_ra , dec.value != _old_decl );

      ((PowerLawFlux*) flux) -> gamma = gamma.value;
      PointSourceComponent::update(); // recomputes the flux normalisation

      // invalidate the cache
      for( auto& p : _dp_domega_dloge ) p = must_update;
    }

    else if ( norm.value != _old_norm ) {
        flux -> norm *= norm.value / _old_norm;
    }


    _old_decl  = dec.value;
    _old_ra    = ra.value;
    _old_gamma = gamma.value;
    _old_norm  = norm.value;
  }


  virtual void init( DataSet& dset ){
    _dp_domega_dloge.resize( dset.data.size() );
    _old_gamma = -9999;
    _old_norm  = -9999;
  }


  /*! If gamma does not change, we can skip the expensive integration of E that happens
      in PointSourceComponent */


  double dN_dOmega_dlogE(SearchEvent& sev ) {

    TIMEFUNC;

    double& p = _dp_domega_dloge[ sev.index ];

    if ( p == must_update ) p = PointSourceComponent::dN_dOmega_dlogE( sev ) / norm.value ;

    return norm.value * p ;
  }

  PowerLawComponent *clone() const // clone
  {
    PowerLawComponent *r = new PowerLawComponent(*this); // return pointer to object
    r->H = NULL;
    r->flux = new PowerLawFlux(*((PowerLawFlux*)this->flux));
    return r; // return pointer to object
  }

  virtual ~PowerLawComponent() {
    if (flux != NULL) {
      delete ((PowerLawFlux*) flux);
      flux = NULL;
    }
  }

  ClassDef( PowerLawComponent, 1) // for ROOT
};





struct NuFromGammaComponent : public PointSourceComponent
{
  NuFromGammaComponent(string name, /*DefaultDetResponse*/ FullDetResponse& detres_, string gammafluxeq_ = "", double r_ = 1, double d_ = 1, bool _use_spline = true) : PointSourceComponent(name, detres_, r_, d_) {
    flux = new NuFromGammaFlux(gammafluxeq_, norm.value, _use_spline);
    H = NULL;
  }

  NuFromGammaComponent *clone() const // clone
  {
    print( "cloning NuFromGammaComponent");

    NuFromGammaComponent *toreturn = new NuFromGammaComponent(*this); // return pointer to object
    toreturn->H = NULL;
    toreturn->flux = new NuFromGammaFlux(*((NuFromGammaFlux*)this->flux));
    return toreturn; // return pointer to object
  }

  virtual ~NuFromGammaComponent() {
    if (flux != NULL) {
      delete ((NuFromGammaFlux*) flux);
      flux = NULL;
    }
  }

  ClassDef( NuFromGammaComponent, 1) // for ROOT
};


/*
struct NuExprComponent : public PointSourceComponent
{
  Parameter par0;
  Parameter par1;
  Parameter par2;

  NuExprComponent(string name, DefaultDetResponse& detres_, string nufluxeq_ = "", double r_ = 1, double d_ = 1, unsigned int _npar = 0, double _par0 = 1, double _par1 = 1, double _par2 = 1) : PointSourceComponent(name, detres_, r_, d_) {
    flux = new NuExprFlux(norm.value, nufluxeq_, _npar);
    par0.value = _par0;
    par1.value = _par1;
    par2.value = _par2;
    H = NULL;
    NormalizeFlux();
    cout << "normalization done" << endl;
  }

  virtual void NormalizeFlux() {
    unsigned int npar = ((NuExprFlux*)flux)->nuflux.GetNpar();
    if (npar > 0) ((NuExprFlux*)flux)->nuflux.SetParameter(0, par0.value);
    if (npar > 1) ((NuExprFlux*)flux)->nuflux.SetParameter(1, par1.value);
    if (npar > 2) ((NuExprFlux*)flux)->nuflux.SetParameter(2, par2.value);
    PointSourceComponent::NormalizeFlux();
  }

  NuExprComponent *clone() const // clone
  {
    NuExprComponent *toreturn = new NuExprComponent(*this); // return pointer to object
    toreturn->H = NULL;
    toreturn->flux = new NuExprFlux(*((NuExprFlux*)this->flux));
    return toreturn; // return pointer to object
  }
  virtual ~NuExprComponent() {
    if (flux != NULL) {
      delete ((NuExprFlux*) flux);
      flux = NULL;
    }
  }

  ClassDef( NuExprComponent, 1) // for ROOT
};

*/

#endif