#ifndef _JACOBIANS_
#define _JACOBIANS_

#include "Params.hh"

class dX_dTs
{

private:
  double sal2, sth2;
  double cal, sal, cth, sth, sph, D, ND, Nphiu, NphiP, _np_dot_pp_p;
  Vec _n, _np, p, pp, Dphiu;
  double x, y, z, xp, yp, zp, _x, _y, _z;
  double aa, bb;
  Params ps;
  bool _reduced;

public:
  // x, y, z, theta, phi, t, E, _x, _y, _z
  vector<vector<double>> ds;

  dX_dTs() : ps(),
             ds(vector<vector<double>>(6, vector<double>(10, 0)))
  {
    _reduced = false;
  };

  dX_dTs(const Trk &trk, const Vec &pos, const Vec &dir, bool reduced = true) : dX_dTs()
  {
    _reduced = reduced;
    ps = Params(trk, pos, dir);
    init(trk, pos, dir);
  }

  void init(const Trk &trk, const Vec &pos, const Vec &dir, const Params &params)
  {
    ps = params;
    init(trk, pos, dir);
  }

  void init(const Trk &trk, const Vec &pos, const Vec &dir, const Paramst &pst)
  {
    ps = Params(pst.E, pst.d, pst.z, pst.theta, pst.phi); // ugly but ok
    init(trk, pos, dir);
  }

  void init(const Trk &trk, const Vec &pos, const Vec &dir)
  {
    fill(ds.begin(), ds.end(), vector<double>(10, 0)); // set ds to 0 everywhere
    /*
      Member ds contains the jacobians with the change of [ E, D, cd, theta_pmt, phi_pmt, dt ] 
      in the first dimension with respect to a change in [ E, x, y, z, dx, dy, dz, theta, phi, t ]
    */

    // Define frequently used factors as analytically as possible
    cal = ps.z;
    sal2 = 1 / (1 - cal * cal);
    // sth2 = 1 / (1 - cth * cth);
    sal = sqrt(1 - cal * cal); // sin( acos( cal ) );
    cth = trk.dir.dot(dir);    // cos( ps.theta );
    sth = sqrt(1 - cth * cth); // sin( ps.theta );

    if (cth != -1.0 and cth != 1.0)
    {
      const Vec pp_p = pos - trk.pos;
      const double cd_ = pp_p.dot(trk.dir);

      Vec A = pp_p - cd_ * trk.dir;
      Vec B = dir - cth * trk.dir;
    
      double pow_ = pow(A.dot(B) / A.len() / B.len(), 2);
      if (pow_ < 1.0)
        sph = sqrt(1 - pow_); // sin( ps.phi );
      else
        sph = 0.0; // Numerical error close to sin(phi) = 0.0
    }
    else
    {
      sph = 0.0; // Just pick arbitrary angle...
    }

    assert(abs(sal - sin(acos(cal))) < 0.001);
    assert(abs(cth - cos(ps.theta)) < 0.001); // sometimes ps.theta is wrong
    assert(abs(sth - sin(ps.theta)) < 0.001);
    


    D = ps.d;

    _n = trk.dir;

    const double an_p[3] = {_n.x, _n.y, _n.z}; // Direction of the track

    _np = dir;

    const double anp_p[3] = {_np.x, _np.y, _np.z};

    p = trk.pos;

    const double ap_p[3] = {p.x, p.y, p.z};

    pp = pos;

    const double app_p[3] = {pp.x, pp.y, pp.z};

    _np_dot_pp_p = _np.dot(pp - p);

    ds[0][0] = 1; // dE_dE

    // Distance
    if (D != 0.0)
    {
      ND = 1. / D;
      ds[1][1] = ND * (p.x - pp.x);
      ds[1][2] = ND * (p.y - pp.y);
      ds[1][3] = ND * (p.z - pp.z);

      // Cosine opening angle
      ds[2][1] = -ND * (_n.x + cal * ds[1][1]);
      ds[2][2] = -ND * (_n.y + cal * ds[1][2]);
      ds[2][3] = -ND * (_n.z + cal * ds[1][3]);
    }
    else
    {
      for (uint ii = 1; ii != 4; ii++)
      {
        ds[1][ii] = 1.0;
        ds[2][ii] = 0.0;
      }
    }

    // Phi
    double ortho_phi = sph * sal * sth;
    if (ortho_phi != 0.0)
    {
      assert(ortho_phi == ortho_phi); // Being hit
      Nphiu = 1. / ortho_phi;

      if (_reduced)
      {

        Dphiu = -(_np * (-D * cal + cth * (-cth * cal * D + _np_dot_pp_p) / sth / sth) + (pp - p) * (-cth + cal * (-cal * cth + _np_dot_pp_p / D) / (1 - cal * cal))) / D / ortho_phi;
      }
      else
      {
        aa = _np_dot_pp_p - cth * cal * D;
        bb = sal * sth * D;
        Vec dAu = -(pp - p) * cth - D * cal * _np;
        Vec dBu = -sth * cal / sal * (pp - p) - D * sal * cth / sth * _np;
        Dphiu = dPhi_dX(aa, bb, dAu, dBu);
      }

      for (uint ii = 1; ii != 4; ii++)
      {
        const uint _ii = ii - 1;

        if (_reduced)
        {

          ds[4][ii] = ((_np_dot_pp_p / D - cal * cth) * (an_p[_ii] * cal + (ap_p[_ii] - app_p[_ii]) / D) / (1 - cal * cal) - cth * an_p[_ii] + anp_p[_ii]) * Nphiu / D;
        }
        else
        {

          double daa = cth * an_p[_ii] - anp_p[_ii];
          double dbb = (sth / sal) * (an_p[_ii] * cal + (ap_p[_ii] - app_p[_ii]) / D);
          ds[4][ii] = dPhi_dX(aa, bb, daa, dbb);
        };
      };
    }
    else
    {
      // cout << "Hit failsafe 1" << endl; //Gets hit
      for (uint ii = 1; ii != 4; ii++)
        ds[4][ii] = 0.0;
    };

    // Time residual
    ds[5][1] = -ds[1][1] / v_light;
    ds[5][2] = -ds[1][2] / v_light;
    ds[5][3] = -ds[1][3] / v_light;

    ds[5][9] = -1; // ddt_dt
  }

  template <class T>
  inline T dPhi_dX(double A, double B, T dA, T dB) const { return (-dA * B + dB * A) / B / B / sph; }

  vector<Vec> get_unit_vector_jacobians_cartesian(const double dir_x, const double dir_y, const double dir_z) const
  {

    // Transformations from absolute directions to normalized directions
    const double abs_n = dir_x * dir_x + dir_y * dir_y + dir_z * dir_z;
    const double sqrt_abs_n = 1 / sqrt(abs_n);

    const double d_xn_d_x = (1 - dir_x * dir_x / abs_n) * sqrt_abs_n;
    const double d_yn_d_y = (1 - dir_y * dir_y / abs_n) * sqrt_abs_n;
    const double d_zn_d_z = (1 - dir_z * dir_z / abs_n) * sqrt_abs_n;

    const double d_xn_d_y = -dir_x * dir_y / abs_n * sqrt_abs_n;
    const double d_xn_d_z = -dir_x * dir_z / abs_n * sqrt_abs_n;

    const double d_yn_d_x = -dir_y * dir_x / abs_n * sqrt_abs_n;
    const double d_yn_d_z = -dir_y * dir_z / abs_n * sqrt_abs_n;

    const double d_zn_d_x = -dir_z * dir_x / abs_n * sqrt_abs_n;
    const double d_zn_d_y = -dir_z * dir_y / abs_n * sqrt_abs_n;

    Vec dnn_d_x(d_xn_d_x, d_yn_d_x, d_zn_d_x);
    Vec dnn_d_y(d_xn_d_y, d_yn_d_y, d_zn_d_y);
    Vec dnn_d_z(d_xn_d_z, d_yn_d_z, d_zn_d_z);

    return {dnn_d_x, dnn_d_y, dnn_d_z};
  };

  vector<Vec> get_unit_vector_jacobians_spherical(const Vec &trk_dir) const
  {

    double abs_xy, cph_trk, sph_trk;

    // Transformations from spherical to cartesian
    if (trk_dir.x != 0.0 or trk_dir.y != 0.0)
    {
      abs_xy = 1 / sqrt(trk_dir.x * trk_dir.x + trk_dir.y * trk_dir.y);
      assert(abs_xy == abs_xy);
      cph_trk = trk_dir.x * abs_xy; // cos( phi );
      sph_trk = trk_dir.y * abs_xy; // sin( phi );
    }
    else
    {
      // cout << "Hit failsafe 2" << endl;  //Being hit
      cph_trk = cos(trk_dir.phi()); // Revert to regular calculation in this edge case
      sph_trk = sin(trk_dir.phi());
    };

    const double cth_trk = trk_dir.z;                       // cos( theta );
    const double sth_trk = sqrt(1 - trk_dir.z * trk_dir.z); // sin( theta );

    assert(cph_trk == cph_trk);
    assert(sph_trk == sph_trk);
    assert(abs(cth_trk - cos(trk_dir.theta())) < 0.00001);
    assert(abs(cph_trk - cos(trk_dir.phi())) < 0.00001);
    assert(abs(sth_trk - sin(trk_dir.theta())) < 0.00001);
    assert(abs(sph_trk - sin(trk_dir.phi())) < 0.00001);

    Vec dnn_dtheta(cph_trk * cth_trk, sph_trk * cth_trk, -sth_trk);
    Vec dnn_dphi(-sph_trk * sth_trk, cph_trk * sth_trk, 0);

    return {dnn_dtheta, dnn_dphi};
  }

  vector<vector<double>> spherical()
  {

    vector<Vec> dnn_ds = get_unit_vector_jacobians_spherical(_n);
    if (D != 0.0)
    {
      ds[2][7] = ND * dnn_ds[0].dot(pp - p); // Cosine opening angle
      ds[2][8] = ND * dnn_ds[1].dot(pp - p);
    }
    else
    {
      ds[2][7] = 0.0;
      ds[2][8] = 0.0;
    };

    if (sth != 0.0)
    {
      ds[3][7] = -dnn_ds[0].dot(_np) / sth; // Theta
      ds[3][8] = -dnn_ds[1].dot(_np) / sth;

      ds[4][7] = Dphiu.dot(dnn_ds[0]);
      ds[4][8] = Dphiu.dot(dnn_ds[1]);
    }
    else
    {
      // cout << "Hit failsafe 3" << endl; //Gets hit
      ds[3][7] = 0.0; // Theta
      ds[3][8] = 0.0;

      ds[4][7] = 0.0; // Phi
      ds[4][8] = 0.0;
    }

    return ds;
  }

  vector<vector<double>> cartesian()
  {

    // const double dy_dx = -_x / sqrt( 1 - _x * _x - _z * _z );
    // const double dz_dx = -_x / sqrt( 1 - _x * _x - _y * _y );

    vector<Vec> dnn_ds = get_unit_vector_jacobians_cartesian(_n.x, _n.y, _n.z);

    if (D != 0.0)
    {
      ds[2][4] = ND * dnn_ds[0].dot(pp - p); // Cosine opening angle
      ds[2][5] = ND * dnn_ds[1].dot(pp - p);
      ds[2][6] = ND * dnn_ds[2].dot(pp - p);
    }
    else
    {
      for (uint ii = 4; ii != 7; ii++)
        ds[2][ii] = 0.0;
    }

    if (sth == 0.0)
    {
      ds[3][4] = 0.0; // Theta
      ds[3][5] = 0.0;
      ds[3][6] = 0.0;

      ds[4][4] = 0.0; // Phi
      ds[4][5] = 0.0;
      ds[4][6] = 0.0;
    }
    else
    {
      ds[3][4] = -dnn_ds[0].dot(_np) / sth; // Theta
      ds[3][5] = -dnn_ds[1].dot(_np) / sth;
      ds[3][6] = -dnn_ds[2].dot(_np) / sth;

      ds[4][4] = Dphiu.dot(dnn_ds[0]); // Phi
      ds[4][5] = Dphiu.dot(dnn_ds[1]);
      ds[4][6] = Dphiu.dot(dnn_ds[2]);
    }

    return ds;
  }
};

const double FIT_A = 50.0;

// Taken from fit of geanz.GetLength function

const vector<double> FIT_B = {0.203403437,
                              0.240667886,
                              0.261085110,
                              0.275016037,
                              0.286157108,
                              0.295753936,
                              0.304396373,
                              0.312420797,
                              0.320045771,
                              0.327429901,
                              0.334700866,
                              0.341973276,
                              0.349361898,
                              0.356995204,
                              0.365033030,
                              0.37369659,
                              0.38332708,
                              0.39452517,
                              0.40856362,
                              0.42924430,
                              0.46753528};

const vector<double> FIT_C = {-1.22238757,
                              -1.16212686,
                              -1.08353550,
                              -1.00963797,
                              -0.93814128,
                              -0.86731766,
                              -0.79593263,
                              -0.72336607,
                              -0.64843861,
                              -0.57022701,
                              -0.48777126,
                              -0.39955936,
                              -0.30406848,
                              -0.19912564,
                              -0.08136130,
                              0.05342942,
                              0.21507097,
                              0.41646702,
                              0.69028536,
                              1.14773638,
                              2.14882754};

struct dTe_dTs
{

protected:
  double _a;
  double _b;
  double _c;

  // const double A;
  // const vector< double > B;
  // const vector< double > C;

public:
  vector<Trk> _elong_trks;
  vector<vector<double>> _elong_params;
  vector<double> _elong_lens;
  vector<vector<vector<double>>> ds;

  dTe_dTs(const Trk trk, const uint n_samples, const bool optimized = false)
  {

    /*
      The change in elongation sample track parameter with respect to global track hypothesis change.

      ds member contains the jacobians of the change of the sample track [ E, x, y, z, theta, phi, t ] 
      with respect to the hypothesis track [ E, x, y, z, dx, dy, dz, theta, phi, t ]
    */

    _elong_trks.resize(1);
    _elong_params.resize(0);
    _elong_lens.resize(0);
   
    if (n_samples == 0)
    {
      _elong_trks[0] = trk;
      _elong_trks[0].dir = _elong_trks[0].dir.normalize();
      _elong_params.push_back({0.0, 0.0, 0.0});
      _elong_lens.push_back(0.0);
    }
   
    else if (n_samples == 1)
    {
      Trk trk_p = trk;
      trk_p.dir = trk_p.dir.normalize();
      const double len = geanz.getMaximum(trk.E);
      _elong_trks[0] = propagate(trk_p, len);

      const double A_max = 49.957; // From fitting geanz.getMaximum(E)
      const double B_max = 0.3348;
      const double C_max = -0.8504;
      _elong_params.push_back({A_max, B_max, C_max});
      _elong_lens.push_back(len);
    }
    else
    {
      _elong_trks.resize(n_samples);
      for (uint ii = 0; ii < n_samples; ii++)
      {
        _elong_trks[ii] = trk;
        _elong_trks[ii].dir = _elong_trks[ii].dir.normalize();
        const double fraction = (ii + 0.5) / ((double)n_samples);
        const double len = push_back_elong_param(trk.E, fraction, optimized);
        propagate(_elong_trks[ii], len);
      }
    };

    dX_dTs dx_dts;

    vector<Vec> dnn_ds_car = dx_dts.get_unit_vector_jacobians_cartesian(trk.dir.x, trk.dir.y, trk.dir.z);
    vector<Vec> dnn_ds_sph = dx_dts.get_unit_vector_jacobians_spherical(trk.dir);

    ds.assign(_elong_trks.size(), vector<vector<double>>(7, vector<double>(10, 0.0)));

    uint ipar = 0;
    for (uint itrk = 0; itrk != _elong_trks.size(); itrk++, ipar++)
    {

      const double len = _elong_lens[ipar];
      const double dDelta_dE = _elong_params[ipar][1] / trk.E; // _b / E
      const Vec dir = _elong_trks[itrk].dir;

      // E, x, y, z, dx, dy, dz, theta, phi, t
      ds[itrk][0][0] = 1;

      ds[itrk][1][0] = dir.x * dDelta_dE;
      ds[itrk][1][1] = 1;
      ds[itrk][1][4] = len * dnn_ds_car[0].x;
      ds[itrk][1][5] = len * dnn_ds_car[0].y;
      ds[itrk][1][6] = len * dnn_ds_car[0].z;
      ds[itrk][1][7] = len * dnn_ds_sph[0].x;
      ds[itrk][1][8] = len * dnn_ds_sph[1].x;

      ds[itrk][2][0] = dir.y * dDelta_dE;
      ds[itrk][2][2] = 1;
      ds[itrk][2][4] = len * dnn_ds_car[1].x;
      ds[itrk][2][5] = len * dnn_ds_car[1].y;
      ds[itrk][2][6] = len * dnn_ds_car[1].z;
      ds[itrk][2][7] = len * dnn_ds_sph[0].y;
      ds[itrk][2][8] = len * dnn_ds_sph[1].y;

      ds[itrk][3][0] = dir.z * dDelta_dE;
      ds[itrk][3][3] = 1;
      ds[itrk][3][4] = len * dnn_ds_car[2].x;
      ds[itrk][3][5] = len * dnn_ds_car[2].y;
      ds[itrk][3][6] = len * dnn_ds_car[2].z;
      ds[itrk][3][7] = len * dnn_ds_sph[0].z;
      ds[itrk][3][8] = len * dnn_ds_sph[1].z;

      ds[itrk][4][4] = dnn_ds_sph[0].x;
      ds[itrk][4][5] = dnn_ds_sph[0].y;
      ds[itrk][4][6] = dnn_ds_sph[0].z;
      ds[itrk][4][7] = 1.0;

      ds[itrk][5][4] = dnn_ds_sph[1].x;
      ds[itrk][5][5] = dnn_ds_sph[1].y;
      ds[itrk][5][6] = dnn_ds_sph[1].z;
      ds[itrk][5][8] = 1.0;

      ds[itrk][6][0] = dDelta_dE / c_light;
      ds[itrk][6][9] = 1.0;
    }
  };


  void init_elong(const double fraction)
  {

    _a = FIT_A;

    if (fraction == 1.0)
    {
      _b = FIT_B[FIT_B.size() - 1];
      _c = FIT_C[FIT_C.size() - 1];
    }
    else
    {
      const uint fl = floor(fraction * (FIT_B.size() - 1));
      const double fr = fraction * (FIT_B.size() - 1) - (double)fl;

      _b = FIT_B[fl] + fr * (FIT_B[fl + 1] - FIT_B[fl]); // Linear interpolation
      _c = FIT_C[fl] + fr * (FIT_C[fl + 1] - FIT_C[fl]);
    };
  };

  double push_back_elong_param(const double E, const double fraction, const bool optimized = false)
  {

    init_elong(fraction);

    const double len = optimized ? param_len(E) : geanz_len(E, fraction);

    _elong_params.push_back({_a, _b, _c, len});
    _elong_lens.push_back(len);
    return len;
  };

  double param_len(const double E) const
  {
    /*
      Parametrized elongation length
    */
    return log(_a * E) * _b + _c;
  };

  double geanz_len(const double E, const double fraction) const
  {
    /*
      Geanz elongation length
    */

    return geanz.getLength(E, fraction, 1e-6);
  };

  double geanz_max(const double E) const
  {
    /*
      Geanz elongation length
    */

    return geanz.getMaximum(E);
  };
};

#endif