#ifndef __SHOWER_GRAD__
#define __SHOWER_GRAD__

#include <string>
#include <iostream>
#include <fstream>
#include <iomanip>
#include <set>
#include <map>
#include <cmath>

#include "JTools/JFunction1D_t.hh"
#include "JTools/JFunctionalMap_t.hh"
#include "JTools/JQuadrature.hh"
#include "JPhysics/JPDF.hh"
#include "JPhysics/JPDFTable.hh"
#include "JPhysics/Antares.hh"
#include "JPhysics/KM3NeT.hh"

#include "JTools/JElement.hh"
#include "JTools/JGridCollection.hh"
#include "JTools/JGridMap.hh"
#include "JTools/JPolint.hh"
#include "JTools/JMapList.hh"
#include "JTools/JMultiFunction.hh"
#include "Params.hh"

#include "JPhysics/JPDFTable.hh"
#include "JPhysics/JNPETable.hh"
#include "Math/IFunction.h"

#include "../utilJpp.hh"
#include "rootutil.hh"
#include "GradTools.hh"
#include "Jacobians.hh"

using namespace std;
using namespace JPP;

/**
 * Scaling of absorption and scattering length.
 */

inline double getAbsorptionLength(const double lambda)
{
  double absorptionLengthFactor = 1.0;
  return absorptionLengthFactor * getAbsorptionLength(lambda);
};

inline double getScatteringLength(const double lambda)
{
  double scatteringLengthFactor = 1.0;
  return scatteringLengthFactor * getScatteringLength(lambda);
};

namespace
{

  using namespace JPP;

  typedef JPolintFunction1D<1, JPolintElement2S<double, double>, JCollection, JResultPDF<double>> PDG_Function1D_t;

  typedef JMAPLIST<JPolint1FunctionalMapH,
                   JPolint1FunctionalMapH,
                   JPolint1FunctionalGridMapH,
                   JPolint1FunctionalGridMapH>::maplist JMapListGrad_t;
  typedef JPDFTable<PDG_Function1D_t, JMapListGrad_t> JppShowerGradPdf_t;
  typedef JNPETable<double, double, JMapListGrad_t> JppShowerGradNpe_t;

};

/**
 * \file
 *
 * Program to create interpolation tables of the PDF of the arrival time of the Cherenkov light from a shower with derivatives functionality.
 *
 * The PDFs are tabulated as a function of <tt>(D, cos(theta), theta, phi, t)</tt>, where:
 * - <tt>D</tt> is the distance between the vertex and the PMT;
 * - <tt>cos(theta)</tt> the cosine of the photon emission angle;
 * - <tt>(theta, phi)</tt> the orientation of the PMT; and
 * - <tt>t</tt> the arrival time of the light with respect to the Cherenkov hypothesis.
 *
 * The orientation of the PMT is defined in the coordinate system in which
 * the shower developes along the z-axis and the PMT is located in the x-z plane.
 * \author jseneca
 */
class JppGrad
{

private:
  uint _debug;
  vector<Trk>     _elong_trks;
  vector<double>  _elong_lens;
  vector<Paramst> _paramst;
  vector<Params>  _params;
  uint _sample_count;
  double _background_rate;
  const double _fudge1 = 1e-12; // minimum expected npe per second (same as very small background)
  const double _fudge2 = 60;    // maximum integrated npe

  const uint _N = 7; // E, x, y, z, theta, phi, t

  int _verbose;

  vector<JppShowerGradPdf_t> _pdf_gradfuncs; // individual light contributions - not used
  vector<JppShowerGradNpe_t> _npe_gradfuncs;
  JppShowerGradPdf_t *_pdf_gradfunc; // actual pdf
  JppShowerGradNpe_t *_npe_gradfunc; // actual pdf

  uint jj_jji[7];
  uint kk_kki[7];

  vector<vector<vector<double>>> _dXs;     // Pdf derivatives
  vector<vector<vector<double>>> _dX_dTes; // [elong] [] []  Parameter transform jacobians for each elong track
  vector<vector<vector<double>>> _dTe_dTs; // Elong sample jacobians for each elong track

  double _A;

  vector<JppShowerGradPdf_t::result_type> _pdf_cache;
  vector<JppShowerGradNpe_t::result_type> _npe_cache;

  JppShowerGradPdf_t::result_type _res_pdf;
  JppShowerGradNpe_t::result_type _res_npe;

  dX_dTs _dx_dtes;

public:
  double _window_start;
  double _window_end;

  // debug objects
  vector<double> _dL_a;
  vector<vector<double>> _sum_dnpes_a;
  vector<vector<double>> _sum_dXs_a;
  vector<vector<double>> _sum_dTs_a;
  vector<vector<double>> _sum_dTes_a;
  vector<double> _dL_n;
  vector<vector<double>> _sum_dnpes_n;
  vector<vector<double>> _sum_dXs_n;
  vector<vector<double>> _sum_dTs_n;
  vector<vector<double>> _sum_dTes_n;
  vector<double> _times;
  uint bad_dXs;
  uint bad_int_dXs;
  uint bad_dtdXs;
  uint call_dXs;
  uint call_int_dXs;
  uint call_dtdXs;
  uint pdf_call_count;
  uint npe_call_count;
  uint bad_pdf_call_count;
  uint bad_npe_call_count;

  JppGrad(vector<string> shower_pdfs = {"/pbs/throng/km3net/src/Jpp/master/data/J14p.dat", // Scattered one first
                                        "/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);


  JppGrad(int) {};

  void set_sample_count(const uint N)
  {
    // 0 : do nothing, 1 = move to shower max, > 1 : sample
    _sample_count = N;
  };

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

  inline double get_k40_npe(double time_window) const
  {
    if (time_window > 0)
      return _background_rate * 1e-9 * time_window;
    else
      return 0; // for residuals smaller than -100, we don't add background
  }

  // Base functions - no elongation here
  double dnpe_dt(float E, float D, double cd, double theta, double phi, float t)
  {
    _res_pdf = (*_pdf_gradfunc)(D, cd, theta, phi, t);
    return _res_pdf.f.f.f.f.f * E;
  };

  double npe(float E, float D, double cd, double theta, double phi, float t)
  {
    _res_pdf = (*_pdf_gradfunc)(D, cd, theta, phi, t);
    return _res_pdf.f.f.f.f.v * E;
  };

  double npe(float E, float D, double cd, double theta, double phi)
  {
    _res_npe = (*_npe_gradfunc)(D, cd, theta, phi);
    return _res_npe.f.f.f.f * E;
  };

  double dnpe_dt(Paramst p)
  {
    if (_sample_count == 0)
      return dnpe_dt(p.E, p.d, p.z, p.theta, p.phi, p.t);
    else
    {
      Trk trk;
      trk.E = p.E;
      dTe_dTs dte_dts(trk, _sample_count);
      vector<double> e_lens = dte_dts._elong_lens;
      double res = 0.0;
      for (vector<double>::iterator ilen = e_lens.begin(); ilen != e_lens.end(); ++ilen)
      {
        Paramst pp = p.propagate(*ilen);
        res += dnpe_dt(pp.E, pp.d, pp.z, pp.theta, pp.phi, pp.t);
      }
      return res / (double)e_lens.size();
    };
  };

  double npe(Paramst p)
  {
    if (_sample_count == 0)
      return npe(p.E, p.d, p.z, p.theta, p.phi, p.t);
    else
    {
      Trk trk;
      trk.E = p.E;
      dTe_dTs dte_dts(trk, _sample_count);
      vector<double> e_lens = dte_dts._elong_lens;
      double res = 0.0;
      for (vector<double>::iterator ilen = e_lens.begin(); ilen != e_lens.end(); ++ilen)
      {
        Paramst pp = p.propagate(*ilen);
        res += npe(pp.E, pp.d, pp.z, pp.theta, pp.phi, pp.t);
      }
      return res / (double)e_lens.size();
    };
  };

  double npe(Params p)
  {
    if (_sample_count == 0)
      return npe(p.E, p.d, p.z, p.theta, p.phi);
    else
    {

      Trk trk;
      trk.E = p.E;
      dTe_dTs dte_dts(trk, _sample_count);
      vector<double> e_lens = dte_dts._elong_lens;
      double res = 0.0;
      for (vector<double>::iterator ilen = e_lens.begin(); ilen != e_lens.end(); ++ilen)
      {
        Params pp = p.propagate(*ilen);
        res += npe(pp.E, pp.d, pp.z, pp.theta, pp.phi);
      }
      return res / (double)e_lens.size();
    };
  };

  inline vector<double> dnpe_dtdXs(float E, float D, double cd, double theta, double phi, float t)
  {
    // Per time
    try
    {
      if (_debug)
        call_dtdXs++;
      _res_pdf = (*_pdf_gradfunc)(D, cd, theta, phi, t);
      return {
          _res_pdf.f.f.f.f.f,
          _res_pdf.fp.f.f.f.f * E,
          _res_pdf.f.fp.f.f.f * E,
          _res_pdf.f.f.fp.f.f * E,
          _res_pdf.f.f.f.fp.f * E,
          _res_pdf.f.f.f.f.fp * E,
          _res_pdf.f.f.f.f.f * E,
      };
    }
    catch (JException e)
    {
      if (_verbose)
      {
        cout << e << " returning 0..." << endl;
        // cout << "D " << D << ", cd " << cd << ", theta " << theta << ", phi " << phi << ", t " << t << ", E " << E << endl;
        bad_dtdXs++;
      }
      return {0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0};
    }
  };

  inline vector<double> dnpe_dXs(float E, float D, double cd, double theta, double phi, float t)
  {
    // Integral over time
    try
    {
      if (_debug)
        call_dXs++;
      _res_pdf = (*_pdf_gradfunc)(D, cd, theta, phi, t);
      return {
          _res_pdf.f.f.f.f.v,
          _res_pdf.fp.f.f.f.v * E,
          _res_pdf.f.fp.f.f.v * E,
          _res_pdf.f.f.fp.f.v * E,
          _res_pdf.f.f.f.fp.v * E,
          _res_pdf.f.f.f.f.f * E,
          _res_pdf.f.f.f.f.v * E,
      };
    }
    catch (JException e)
    {
      if (_verbose)
      {
        cout << e << " returning 0..." << endl;
        // cout << "D " << D << ", cd " << cd << ", theta " << theta << ", phi " << phi << ", t " << t << ", E " << E << endl;
        bad_dXs++;
      }
      return {0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0};
    }
  };

  inline vector<double> dnpe_dXs(float E, float D, double cd, double theta, double phi)
  {
    // Integral over all time
    try
    {
      if (_debug)
        call_int_dXs++;
      _res_npe = (*_npe_gradfunc)(D, cd, theta, phi);
      return {
          _res_npe.f.f.f.f,
          _res_npe.fp.f.f.f * E,
          _res_npe.f.fp.f.f * E,
          _res_npe.f.f.fp.f * E,
          _res_npe.f.f.f.fp * E,
          0,
          _res_npe.f.f.f.f * E};
    }
    catch (JException e)
    {
      if (_verbose)
      {
        cout << e << " returning 0..." << endl;
        bad_int_dXs++;
      }
      return {0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0};
    }
  };

  inline vector<double> dnpe_dtdXs(const Trk &trk, const Hit &hit)
  {

    dTe_dTs dte_dts(trk, _sample_count);
    vector<Trk> e_trks = dte_dts._elong_trks;
    vector<double> ds;
    ds.assign(7, 0.0);
    for (vector<Trk>::iterator itrk = e_trks.begin(); itrk != e_trks.end(); ++itrk)
    {
      Paramst pst = Paramst(*itrk, hit);
      vector<double> ds_e = dnpe_dtdXs(pst.E, pst.d, pst.z, pst.theta, pst.phi, pst.t);
      for (uint ii = 0; ii != 7; ++ii)
      {
        ds[ii] += ds_e[ii] / (double)e_trks.size();
      }
    }
    return ds;
  };

  inline vector<double> dnpe_dXs(const Trk &trk, const Hit &hit)
  {

    dTe_dTs dte_dts(trk, _sample_count);
    vector<Trk> e_trks = dte_dts._elong_trks;
    vector<double> ds;
    ds.assign(7, 0.0);
    for (vector<Trk>::iterator itrk = e_trks.begin(); itrk != e_trks.end(); ++itrk)
    {
      Paramst pst = Paramst(*itrk, hit);
      vector<double> ds_e = dnpe_dXs(pst.E, pst.d, pst.z, pst.theta, pst.phi, pst.t);
      for (uint ii = 0; ii != 7; ++ii)
      {
        ds[ii] += ds_e[ii] / (double)e_trks.size();
      }
    }
    return ds;
  };

  inline vector<double> dnpe_dXs(const Trk &trk, const Pmt &unhit)
  {

    dTe_dTs dte_dts(trk, _sample_count);
    vector<Trk> e_trks = dte_dts._elong_trks;
    vector<double> ds;
    ds.assign(7, 0.0);
    for (vector<Trk>::iterator itrk = e_trks.begin(); itrk != e_trks.end(); ++itrk)
    {
      Params ps = Params(*itrk, unhit);
      vector<double> ds_e = dnpe_dXs(ps.E, ps.d, ps.z, ps.theta, ps.phi);
      for (uint ii = 0; ii != 7; ++ii)
      {
        ds[ii] += ds_e[ii] / (double)e_trks.size();
      }
    }
    return ds;
  };

  inline vector<double> dnpe_dXs(const JppShowerGradPdf_t::result_type pdf_res, const Paramst &pst) const
  {
    // Integral over time
    return {
        pdf_res.f.f.f.f.v,
        pdf_res.fp.f.f.f.v * pst.E,
        pdf_res.f.fp.f.f.v * pst.E,
        pdf_res.f.f.fp.f.v * pst.E,
        pdf_res.f.f.f.fp.v * pst.E,
        pdf_res.f.f.f.f.f * pst.E,
        pdf_res.f.f.f.f.v * pst.E,
    };
  };

  inline vector<double> dnpe_dtdXs(const JppShowerGradPdf_t::result_type &pdf_res, const Paramst &pst) const
  {

    return {
        pdf_res.f.f.f.f.f,
        pdf_res.fp.f.f.f.f * pst.E,
        pdf_res.f.fp.f.f.f * pst.E,
        pdf_res.f.f.fp.f.f * pst.E,
        pdf_res.f.f.f.fp.f * pst.E,
        pdf_res.f.f.f.f.fp * pst.E,
        pdf_res.f.f.f.f.f * pst.E,
    };
  };

  inline vector<double> dnpe_dXs(const JppShowerGradNpe_t::result_type npe_res, const Params &ps) const
  {
    // Integral over all time
    return {
        npe_res.f.f.f.f,
        npe_res.fp.f.f.f * ps.E,
        npe_res.f.fp.f.f * ps.E,
        npe_res.f.f.fp.f * ps.E,
        npe_res.f.f.f.fp * ps.E,
        0,
        npe_res.f.f.f.f * ps.E};
  };

  double dnpe_dt(const Trk &trk, const Hit &hit)
  {

    dTe_dTs dte_dts(trk, _sample_count);
    vector<Trk> e_trks = dte_dts._elong_trks;
    double res = 0;
    for (vector<Trk>::iterator itrk = e_trks.begin(); itrk != e_trks.end(); ++itrk)
    {
      Paramst pst(*itrk, hit);
      res += dnpe_dt(pst.E, pst.d, pst.z, pst.theta, pst.phi, pst.t) / (double)e_trks.size();
    };
    return res + get_k40_npe(1.0);
  };

  double npe(const Trk &trk, const Hit &hit)
  {

    dTe_dTs dte_dts(trk, _sample_count);
    vector<Trk> e_trks = dte_dts._elong_trks;
    double res = 0;
    for (vector<Trk>::iterator itrk = e_trks.begin(); itrk != e_trks.end(); ++itrk)
    {
      Paramst pst(*itrk, hit);
      res += npe(pst.E, pst.d, pst.z, pst.theta, pst.phi, pst.t) / (double)e_trks.size();
    };
    return res + get_k40_npe(1.0);
  };

  double npe(const Trk &trk, const Pmt &unhit)
  {

    dTe_dTs dte_dts(trk, _sample_count);
    vector<Trk> e_trks = dte_dts._elong_trks;
    double res = 0;
    for (vector<Trk>::iterator itrk = e_trks.begin(); itrk != e_trks.end(); ++itrk)
    {
      Params ps(*itrk, unhit);
      res += npe(ps.E, ps.d, ps.z, ps.theta, ps.phi) / (double)e_trks.size();
    };
    return res + get_k40_npe(1.0);
  };

  double dnpe_dt(const Hit &hit)
  {

    double res = 0.0;

    vector<Paramst>::iterator ipst = _paramst.begin();
    // loop over elongation
    for (vector<JppShowerGradPdf_t::result_type>::iterator iret = _pdf_cache.begin(); iret != _pdf_cache.end(); ++iret, ++ipst)
    {
      res += (*iret).f.f.f.f.f * ipst->E;
    }
    return res / (double)_elong_trks.size() + get_k40_npe(1.0);
  };

  double npe(const Hit &hit)
  {

    double res = 0;

    vector<Paramst>::iterator ipst = _paramst.begin();
    for (vector<JppShowerGradPdf_t::result_type>::iterator iret = _pdf_cache.begin(); iret != _pdf_cache.end(); ++iret, ++ipst)
    {
      res += (*iret).f.f.f.f.v * ipst->E + get_k40_npe(ipst->t - _window_start);
    }
    return res / (double)_elong_trks.size();
  };

  double npe(const Pmt &unhit)
  {

    double res = 0.0;

    vector<Params>::iterator ips = _params.begin();
    for (vector<JppShowerGradNpe_t::result_type>::iterator iret = _npe_cache.begin(); iret != _npe_cache.end(); ++iret, ips++)
    {
      res += (*iret).f.f.f.f * ips->E;
    }
    return res / (double)_elong_trks.size() + get_k40_npe(_window_end - _window_start);
  };




  inline vector<double> dnpe_dtdTs(const Hit &hit)
  {
    vector<double> ds(7, 0);
    for (uint iTe = 0; iTe != _elong_trks.size(); iTe++)
    {
      dnpe_dtdTs1(ds, _dXs[iTe][0], _paramst[iTe], _pdf_cache[iTe], _dX_dTes[iTe], _dTe_dTs[iTe]); // Recycle dnpe_dtdXs
    }
    for (uint ii = 0; ii < 7; ii++)
      ds[ii] /= (float)_elong_trks.size();
    return ds;
  };

  inline vector<double> dnpe_dTs(const Hit &hit)
  {
    vector<double> ds(7, 0);
    for (uint iTe = 0; iTe != _elong_trks.size(); iTe++)
    {
      dnpe_dTs1(ds, _dXs[iTe][1], _paramst[iTe], _pdf_cache[iTe], _dX_dTes[iTe], _dTe_dTs[iTe]); // Recycle dnpe_dXs
    }
    for (int ii = 0; ii < 7; ii++)
      ds[ii] /= (float)_elong_trks.size();
    return ds;
  };

  vector<double> dnpe_int_dTs(const Pmt &unhit)
  {
    vector<double> ds(7, 0);
    for (uint iTe = 0; iTe != _elong_trks.size(); iTe++)
    {
      dnpe_int_dTs1(ds, _dXs[iTe][2], _params[iTe], _npe_cache[iTe], _dX_dTes[iTe], _dTe_dTs[iTe]);
    }
    for (int ii = 0; ii < 7; ii++)
      ds[ii] /= (float)_elong_trks.size();
    return ds;
  };

  inline vector<double> dnpe_dtdTs1(vector<double> &ds,
                                    vector<double> &dXs,
                                    const Paramst &pst,
                                    const JppShowerGradPdf_t::result_type &pdf_res,
                                    const vector<vector<double>> &dTs,
                                    const vector<vector<double>> &dTes)
  {

    dXs = dnpe_dtdXs(pdf_res, pst);

    // E, x, y, z, dtheta, dphi, t
    for (uint kk = 0; kk != 7; ++kk)
    {
      for (uint jj = 0; jj != 7; ++jj)
      {
        for (uint ii = 0; ii != 6; ++ii)
        {
          ds[kk] += dXs[ii] * dTs[ii][jj_jji[jj]] * dTes[jj][kk_kki[kk]];
        }
      }
    }

    return ds;
  }

  inline vector<double> dnpe_dTs1(vector<double> &ds,
                                  vector<double> &dXs,
                                  const Paramst &pst,
                                  const JppShowerGradPdf_t::result_type &pdf_res,
                                  const vector<vector<double>> &dTs,
                                  const vector<vector<double>> &dTes)
  {

    dXs = dnpe_dXs(pdf_res, pst);

    // E, x, y, z, dtheta, dphi,
    for (uint kk = 0; kk != 7; ++kk)
    {
      for (uint jj = 0; jj != 7; ++jj)
      {
        for (uint ii = 0; ii != 6; ++ii)
        {
          ds[kk] += dXs[ii] * dTs[ii][jj_jji[jj]] * dTes[jj][kk_kki[kk]];
        }
      }
    }

    return ds;
  }

  inline vector<double> dnpe_int_dTs1(vector<double> &ds,
                                      vector<double> &dXs,
                                      const Params &ps,
                                      const JppShowerGradNpe_t::result_type &npe_res,
                                      vector<vector<double>> &dTs,
                                      const vector<vector<double>> &dTes)
  {
    dXs = dnpe_dXs(npe_res, ps);

    for (uint ii = 0; ii != dTs[5].size(); ii++)
      dTs[5][ii] = 0.0; // time derivative is 0

    // E, x, y, z, dtheta, dphi, t
    for (uint kk = 0; kk != 7; ++kk)
    {
      for (uint jj = 0; jj != 7; ++jj)
      {
        for (uint ii = 0; ii != 6; ++ii)
        {
          ds[kk] += dXs[ii] * dTs[ii][jj_jji[jj]] * dTes[jj][kk_kki[kk]];
        }
      }
    }

    return ds;
  };

  void set_elongation(const Trk &trk)
  {
    dTe_dTs dte_dts(trk, _sample_count);
    _elong_trks = dte_dts._elong_trks;
    _elong_lens = dte_dts._elong_lens;
    _dTe_dTs = dte_dts.ds;
    _dXs.resize(_elong_trks.size(), vector<vector<double>>(3, vector<double>(6, 0.0)));
    _dX_dTes.resize(_elong_trks.size());
    _paramst.resize(_elong_trks.size());
    _params.resize(_elong_trks.size());
    _pdf_cache.resize(_elong_trks.size());
    _npe_cache.resize(_elong_trks.size());
  };



  void set_paramst(Hit &hit)
  {
    vector<Trk>::iterator itrk = _elong_trks.begin();
   
    vector<JppShowerGradPdf_t::result_type>::iterator iret = _pdf_cache.begin();
  
  
  
    for (vector<Paramst>::iterator ip = _paramst.begin(); ip != _paramst.end(); ip++, ++itrk, ++iret)
    {
      ip->init_t(*itrk, hit);
   
      try
      {
        if (_debug)
          pdf_call_count++;
        *iret = (*_pdf_gradfunc)(ip->d, ip->z, ip->theta, ip->phi, ip->t);
      }
      catch (JException e)
      {
        bad_pdf_call_count++;
        if (_verbose > 2)
          cout << e << " out of range pdf parameters hit..." << endl;
        hit.pattern_flags = 2; // Flag bad hit
        break;
      }
    };
  };

  void set_params(Pmt &unhit)
  {
    vector<Trk>::iterator itrk = _elong_trks.begin();
    vector<JppShowerGradNpe_t::result_type>::iterator iret = _npe_cache.begin();
    for (vector<Params>::iterator ip = _params.begin(); ip != _params.end(); ip++, ++itrk, ++iret)
    {
      ip->init(*itrk, unhit.pos, unhit.dir);
      try
      {
        if (_debug)
          npe_call_count++;
        *iret = (*_npe_gradfunc)(ip->d, ip->z, ip->theta, ip->phi);
      }
      catch (JException e)
      {
        bad_npe_call_count++;
        if (_verbose > 2)
          cout << e << " out of range pdf parameters for unhit..." << endl;
        unhit.flag = 2; // Flag bad pmt
        break;
      }
    };
  };

  void set_jacobians(const Hit &hit)
  {
    vector<vector<vector<double>>>::iterator idx_dtes = _dX_dTes.begin();
    vector<Paramst>::iterator ipst = _paramst.begin();
    for (vector<Trk>::iterator itrk = _elong_trks.begin(); itrk != _elong_trks.end(); ++itrk, ++idx_dtes, ++ipst)
    {
      _dx_dtes.init(*itrk, hit.pos, hit.dir, *ipst);
      _dx_dtes.spherical();
      *idx_dtes = _dx_dtes.ds;
    }
  };

  void set_jacobians(const Pmt &unhit)
  {
    vector<vector<vector<double>>>::iterator idx_dtes = _dX_dTes.begin();
    vector<Params>::iterator ips = _params.begin();
    for (vector<Trk>::iterator itrk = _elong_trks.begin(); itrk != _elong_trks.end(); ++itrk, ++idx_dtes, ips++)
    {
      _dx_dtes.init(*itrk, unhit.pos, unhit.dir, *ips);
      _dx_dtes.spherical();
      *idx_dtes = _dx_dtes.ds;
      for (uint jj = 0; jj != 10; ++jj)
      {
        (*idx_dtes)[5][jj] = 0.0; // Time does not change for unhits (but this shouldnt matter for LL
      };
    }
  };

  inline double dhit_L(const double &dnpe_dtdT,
                       const double &dnpe_dT,
                       const double &A,
                       const double &dt_e) const
  {

    double res = dnpe_dtdT / (A + _fudge1) - dnpe_dT - get_k40_npe(1.0) * dt_e;

    return res;
  };

  inline double get_hit_L(const double nn,
                          const double NN,
                          bool factorized_fudge = true)
  {
    if (factorized_fudge)
      return log(_fudge1 + nn) - min(NN, _fudge2);
    else
      return log(_fudge1 + (nn)*exp(-NN));
  };

  double eval_lik(const Trk &trk,
                  vector<Hit> hits,
                  vector<Pmt> unhits,
                  bool factorized_fudge = true)
  {

    // if( trk.E > 1e14 ) trk.E = 1e14;

    set_elongation(trk);

    double L = 0;
    foreach (hit, hits)
    {
      set_paramst(hit);
      if (hit.pattern_flags == 2)
        continue; // Hit flagged as bad
      L -= get_hit_L(dnpe_dt(hit), npe(hit), factorized_fudge);
    }

    foreach (unhit, unhits)
    {
      set_params(unhit);
      if (unhit.flag == 2)
        continue; // Hit flagged as bad
      L -= -npe(unhit);
    }

    // trk.lik = L;

    return L;
  };



  void eval_lik_gradient(double &r, vector<double> &ds, const Trk &trk, vector<Hit> hits, vector<Pmt> unhits);


  // Debug

  void fill_components(const Trk &trk, Hit &hit)
  {

    numFuncs num;

    set_paramst(hit); // could remove
    set_jacobians(hit);

    const vector<double> dnpe_dtdTs = this->dnpe_dtdTs(hit);
    const vector<double> dnpe_dTs = this->dnpe_dTs(hit);

    vector<double> dnpe_dtdTs_n = this->num_dnpe_dtdTs(trk, hit, 1e-5);
    vector<double> dnpe_dTs_n = this->num_dnpe_dTs(trk, hit, 1e-5);

    for (uint kk = 0; kk != 7; kk++)
    {

      _sum_dnpes_a[kk][0] += dnpe_dtdTs[kk];
      _sum_dnpes_a[kk][1] += dnpe_dTs[kk];

      _sum_dnpes_n[kk][0] += dnpe_dtdTs_n[kk];
      _sum_dnpes_n[kk][1] += dnpe_dTs_n[kk];
    }

    for (uint iTe = 0; iTe != _elong_trks.size(); iTe++)
    {

      Paramst ps = _paramst[iTe];
      vector<double> dtdXs_n = num.num_dnpe_dtdXs(this, ps.E, ps.d, ps.z, ps.theta, ps.phi, ps.t, 1e-6);
      vector<double> dXs_n = num.num_dnpe_dXs(this, ps.E, ps.d, ps.z, ps.theta, ps.phi, ps.t, 1e-6);

      vector<vector<double>> dTs_n = num_dX_dTs(_elong_trks[iTe], hit, 1e-6);
      vector<vector<double>> dTes_n = num_dTe_dTs(_elong_trks[iTe], _elong_lens[iTe], 1e-6);

      for (uint kk = 0; kk != 7; kk++)
      {
        for (uint ii = 0; ii != 6; ii++)
        {

          _sum_dXs_a[kk][0] += _dXs[iTe][0][ii];
          _sum_dXs_n[kk][0] += dtdXs_n[ii];
          _sum_dXs_a[kk][1] += _dXs[iTe][1][ii];
          _sum_dXs_n[kk][1] += dXs_n[ii];

          for (uint jj = 0; jj != 7; jj++)
          {
            _sum_dTs_a[kk][ii] += _dX_dTes[iTe][ii][jj_jji[jj]];
            // cout << "_dX_dTes[" << iTe << "][" << ii << "][" << jj_jji[jj] << "]: " << _dX_dTes[iTe][ii][jj_jji[jj]] << endl;
            if (ii == 0 && jj == 0 && _dX_dTes[iTe][ii][jj_jji[jj]] != 1.0)
            {
              cout << "dE_dE != 1.0" << endl;
              if (_debug)
                cin.ignore();
            }

            _sum_dTes_a[kk][jj] += _dTe_dTs[iTe][jj][kk_kki[kk]];
            _sum_dTs_n[kk][ii] += dTs_n[ii][jj_jji[jj]];
            if (abs(_sum_dTs_n[kk][ii]) > 1e12)
            {
              cout << "dTs_n[  " << ii << "]"
                   << "[" << jj_jji[jj] << "]" << dTs_n[ii][jj_jji[jj]] << endl;
              cout << "sum_dTs_n[" << kk << "]"
                   << "[" << ii << "]" << _sum_dTs_n[kk][ii] << endl;
              if (_debug)
                cin.ignore();
            }
            _sum_dTes_n[kk][jj] += dTes_n[jj][kk_kki[kk]];
          }
        }
      }
    }
  };



  void fill_components(const Trk &trk, Pmt &unhit)
  {

    numFuncs num;

    set_params(unhit); // not necessary actually...
    set_jacobians(unhit);

    const vector<double> dnpe_dTs = this->dnpe_int_dTs(unhit);
    const vector<double> dnpe_dTs_n = this->num_dnpe_dTs(trk, unhit, 1e-5);

    for (uint kk = 0; kk != 7; kk++)
    {
      _sum_dnpes_n[kk][2] += dnpe_dTs_n[kk];
      _sum_dnpes_a[kk][2] += dnpe_dTs[kk];
    }

    for (uint iTe = 0; iTe != _elong_trks.size(); iTe++)
    {

      Params ps = _params[iTe];
      vector<double> dXs_n = num.num_dnpe_dXs(this, ps.E, ps.d, ps.z, ps.theta, ps.phi, 1e-6);

      vector<vector<double>> dTs_n = num_dX_dTs(_elong_trks[iTe], unhit, 1e-6);
      vector<vector<double>> dTes_n = num_dTe_dTs(_elong_trks[iTe], _elong_lens[iTe], 1e-6);

      for (uint kk = 0; kk != 7; kk++)
      {
        for (uint ii = 0; ii != 6; ii++)
        {

          _sum_dXs_a[kk][2] += _dXs[iTe][2][ii];
          _sum_dXs_n[kk][2] += dXs_n[ii];

          for (uint jj = 0; jj != 7; jj++)
          {
            _sum_dTs_a[kk][ii] += _dX_dTes[iTe][ii][jj_jji[jj]];
            _sum_dTes_a[kk][jj] += _dTe_dTs[iTe][jj][kk_kki[kk]];
            _sum_dTs_n[kk][ii] += dTs_n[ii][jj_jji[jj]];
            _sum_dTes_n[kk][jj] += dTes_n[jj][kk_kki[kk]];
          }
        }
      }
    }
  };

  // Numerical functions

  vector<double> num_dnpe_dtdTs(const Trk &trk, const Hit &hit, double delta = 1e-5)
  {
    vector<double> ds(7, 0.0);
    // E, x, y, z, dx, dy, dz, dtheta, dphi, t
    for (uint kk = 0; kk != 7; kk++)
    {
      std::pair<Trk, Trk> trks = delta_trks(trk, kk_kki[kk], delta);
      double dnpe_dt_l = dnpe_dt(trks.first, hit);
      double dnpe_dt_h = dnpe_dt(trks.second, hit);
      ds[kk] += (dnpe_dt_h - dnpe_dt_l) / delta;
    }
    return ds;
  }

  vector<double> num_dnpe_dTs(const Trk &trk, const Hit &hit, double delta = 1e-5)
  {
    vector<double> ds(7, 0.0);
    // E, x, y, z, dx, dy, dz, dphi, dtheta, t
    for (uint kk = 0; kk != 7; kk++)
    {
      std::pair<Trk, Trk> trks = delta_trks(trk, kk_kki[kk], delta);
      double npe_l = npe(trks.first, hit);
      double npe_h = npe(trks.second, hit);
      ds[kk] += (npe_h - npe_l) / delta;
    }
    return ds;
  }

  vector<double> num_dnpe_dTs(const Trk &trk, const Pmt &unhit, double delta = 1e-5)
  {
    vector<double> ds(7, 0.0);
    // E, x, y, z, dx, dy, dz, dphi, dtheta, t
    for (uint kk = 0; kk != 7; kk++)
    {
      std::pair<Trk, Trk> trks = delta_trks(trk, kk_kki[kk], delta);
      double npe_l = npe(trks.first, unhit);
      double npe_h = npe(trks.second, unhit);
      ds[kk] += (npe_h - npe_l) / delta;
    }
    return ds;
  }
};

#endif