#pragma once

#include <stan/math/rev/arr/functor/coupled_ode_system.hpp>

#include "help.cpp"

namespace stan {
namespace math {

/**
 * The coupled ODE system for known initial values and unknown
 * parameters.
 *
 * <p>If the base ODE state is size N and there are M parameters,
 * the coupled system has N + N * M states.
 * <p>The first N states correspond to the base system's N states:
 * \f$ \frac{d x_n}{dt} \f$
 *
 * <p>The next M states correspond to the sensitivities of the
 * parameters with respect to the first base system equation:
 * \f[
 *   \frac{d x_{N+m}}{dt}
 *   = \frac{d}{dt} \frac{\partial x_1}{\partial \theta_m}
 * \f]
 *
 * <p>The final M states correspond to the sensitivities with respect
 * to the second base system equation, etc.
 *
 * @tparam F type of functor for the base ode system.
 */
// template <> indicates that class 
// coupled_ode_system<x_model_namespace::y_functor__, double, var>
// is fully specialized
// y0 is double theta is var
template <>
struct coupled_ode_system<x_model_namespace::y_functor__, double, var> {
  // typedef links x_model_namespace::y_functor__ with coupled_ode_system
  typedef x_model_namespace::y_functor__ F;
  const F& f_;
  const std::vector<double>& y0_dbl_;
  const std::vector<var>& theta_;
  const std::vector<double> theta_dbl_;
  const std::vector<double>& x_;
  const std::vector<int>& x_int_;
  const size_t N_;
  const size_t M_;
  const size_t size_;
  std::ostream* msgs_;

  /**
   * Construct a coupled ODE system with the specified base
   * ODE system, base initial state, parameters, data, and a
   * message stream.
   *
   * @param[in] f the base ODE system functor.
   * @param[in] y0 the initial state of the base ode.
   * @param[in] theta parameters of the base ode.
   * @param[in] x real data.
   * @param[in] x_int integer data.
   * @param[in, out] msgs stream to which messages are printed.
   */
  coupled_ode_system(const F& f, const std::vector<double>& y0,
                     const std::vector<var>& theta,
                     const std::vector<double>& x,
                     const std::vector<int>& x_int, std::ostream* msgs)
      : f_(f),
        y0_dbl_(y0),
        theta_(theta),
        theta_dbl_(value_of(theta)),
        x_(x),
        x_int_(x_int),
        N_(y0.size()),
        M_(theta.size()),
        size_(N_ + N_ * M_),
        msgs_(msgs) {}

  /**
   * Assign the derivative vector with the system derivatives at
   * the specified state and time.
   *
   * <p>The input state must be of size <code>size()</code>, and
   * the output produced will be of the same size.
   *
   * @param[in] z state of the coupled ode system.
   * @param[out] dz_dt populated with the derivatives of
   * the coupled system at the specified state and time.
   * @param[in]  t time.
   * @throw exception if the system function does not return the
   * same number of derivatives as the state vector size.
   *
   * y is the base ODE system state
   *
   */
  void operator()(const std::vector<double>& z, std::vector<double>& dz_dt,
                  double t) const {
    using std::vector;
    vector<double> y(z.begin(), z.begin() + N_);
    help instHelp;
    vector<double> dy_dt(N_); 

    instHelp.fun(y, theta_dbl_, t, dy_dt); 

    vector<double> Jy(N_ * N_);
    
    instHelp.jacobianY(y, theta_dbl_, t, Jy);

    vector<double> Jtheta(N_ * M_);
 
    instHelp.jacobianTheta(y, theta_dbl_, t, Jtheta);

    for (size_t i = 0; i < N_; i++) {
      dz_dt[i] = dy_dt[i]; 
      // zjm=dyj/dthetam=z[N_ + N_ * m + j]
      // d/dt(znm)=dfn/dthetam+sum(j,zjm*dfn/dyj)
      for (size_t j = 0; j < M_; j++) {
        double temp_deriv = Jtheta[j * N_ + i];

        for (size_t k = 0; k < N_; k++)
          temp_deriv += z[N_ + N_ * j + k] * Jy[k * N_ + i];

        dz_dt[N_ + i + j * N_] = temp_deriv;
      }
    }
  }

  /**
   * Returns the size of the coupled system.
   *
   * @return size of the coupled system.
   */
  size_t size() const { return size_; }

  /**
   * Returns the initial state of the coupled system.  Because the
   * initial values are known, the initial state of the coupled
   * system is the same as the initial state of the base ODE
   * system.
   *
   * <p>This initial state returned is of size <code>size()</code>
   * where the first N (base ODE system size) parameters are the
   * initial conditions of the base ode system and the rest of the
   * initial condition elements are 0.
   *
   * @return the initial condition of the coupled system.
   */
  std::vector<double> initial_state() const {
    std::vector<double> state(size_, 0.0);
    for (size_t n = 0; n < N_; n++)
      state[n] = y0_dbl_[n];
    return state;
  }

  /**
   * Returns the base ODE system state corresponding to the
   * specified coupled system state.
   *
   * @param y coupled states after solving the ode
   */
  std::vector<std::vector<var> > decouple_states(
      const std::vector<std::vector<double> >& y) const {
    std::vector<var> temp_vars(N_);
    std::vector<double> temp_gradients(M_);
    std::vector<std::vector<var> > y_return(y.size());

    for (size_t i = 0; i < y.size(); i++) {
      // iterate over number of equations
      for (size_t j = 0; j < N_; j++) {
        // iterate over parameters for each equation
       for (size_t k = 0; k < M_; k++) {
          temp_gradients[k] = y[i][y0_dbl_.size() + y0_dbl_.size() * k + j];
        }
        temp_vars[j] = precomputed_gradients(y[i][j], theta_, temp_gradients);
      }
      y_return[i] = temp_vars;
    }
    return y_return;
  }
};

}  // namespace math
}  // namespace stan