/* nag::ad::d02pt Adjoint Example Program.
 */

#include <dco.hpp>
#include <iostream>
#include <nagad.h>
#include <numeric>

// Function which calls NAG AD routines.
template <typename T> void func(std::vector<T> &y0, std::vector<T> &y);

// Driver with the adjoint calls.
// Solves the equation y''=-y, y0(=y(0))=0, y0'(=y'(0))=1 by solving the ODE
// system:
//    y1'=y2,  y2'=-y1
// over the range [0,2*pi] with initial conditions y1=0, y2=1.
// Also, computes the sum of all Jacobian elements dy/dy0.
void driver(const std::vector<double> &y0v,
            std::vector<double> &      yv,
            double &                   dy_dy0);

int main()
{
  std::cout << " nag::ad::d02pt Adjoint Example Program Results\n";

  // Set initial conditions
  std::vector<double> y0v{0.0, 1.0};
  // Solution y
  std::vector<double> yv;
  // Sum of Jacobian elements
  double dy_dy0;

  // Call driver
  driver(y0v, yv, dy_dy0);

  // Print outputs
  std::cout << "\n Derivatives calculated: First order adjoints\n";
  std::cout << " Computational mode    : algorithmic\n";

  std::cout.setf(std::ios::scientific, std::ios::floatfield);
  std::cout.precision(12);
  std::cout << "\n Solution computed with required tolerance " << 1e-4
            << std::endl;
  for (std::size_t i = 0; i < yv.size(); i++)
  {
    std::cout << " y" << i + 1 << " = " << yv[i] << std::endl;
  }

  // Print derivatives
  std::cout
      << "\n Sum of all Jacobian elements of solution vector y w.r.t. initial values:\n";
  std::cout << " sum_ij [dy/dy0]_ij = " << dy_dy0 << std::endl;

  return 0;
}

// Driver with the adjoint calls.
// Solves the equation y''=-y, y0(=y(0))=0, y0'(=y'(0))=1 by solving the ODE
// system:
//    y1'=y2,  y2'=-y1
// over the range [0,2*pi] with initial conditions y1=0, y2=1.
// Also, computes the sum of all Jacobian elements dy/dy0.
void driver(const std::vector<double> &y0v,
            std::vector<double> &      yv,
            double &                   dy_dy0)
{
  using mode = dco::ga1s<double>;
  using T    = mode::type;

  // Create the AD tape
  mode::global_tape = mode::tape_t::create();

  // Size of system of ODEs
  Integer n = y0v.size();
  // y0 stores the initial conditions of the problem
  std::vector<T> y0(begin(y0v), end(y0v));

  // Register variables to differentiate w.r.t.
  mode::global_tape->register_variable(y0);

  // Solution y, variable to differentiate
  std::vector<T> y(n);

  // Call the NAG AD Lib functions
  func(y0, y);

  // Extract the computed solutions
  yv = dco::passive_value(y);
  // Initialize adjoints of y
  mode::global_tape->register_output_variable(y);
  dco::derivative(y) = std::vector<double>(n, 1.0);
  mode::global_tape->interpret_adjoint();

  // Sum the adjoints of y0
  dy_dy0 = 0.0;
  for (T const &yi : y0)
  {
    dy_dy0 += dco::derivative(yi);
  }

  // Remove tape
  mode::tape_t::remove(mode::global_tape);
}

// function which calls NAG AD Library routines
template <typename T> void func(std::vector<T> &y0, std::vector<T> &y)
{
  // Active variables
  const Integer        n    = y0.size();
  const Integer        npts = 8, liwsav = 130, lrwsav = 350 + 32 * n;
  std::vector<T>       thresh(n, 1e-8), ypgot(n), ymax(n), rwsav(lrwsav);
  std::vector<Integer> iwsav(liwsav);

  // Set parameters for the integrator.
  Integer method = 1;
  T       tol = 1e-4, hstart = 0.0, tend = 2.0 * nag_math_pi, tstart = 0.0;

  // Set control for output
  double tinc = 2.0 * nag_math_pi / (double)(npts);
  // Create AD configuration data object
  Integer           ifail = 0;
  nag::ad::handle_t ad_handle;

  // Initialize Runge-Kutta method for integrating ODE
  ifail = 0;
  nag::ad::d02pq(ad_handle, n, tstart, tend, y0.data(), tol, thresh.data(),
                 method, hstart, iwsav.data(), rwsav.data(), ifail);

  
  auto f = [&](nag::ad::handle_t &ad_handle,
            const T &          t,
            const Integer &    n,
            const T            y[],
            T                  yp[])
          {
            yp[0] = y[1];
            yp[1] = -y[0];
          };

  T tgot, twant = tstart;
  for (int j = 0; j < npts; ++j)
  {
    twant = twant + tinc;
    ifail = 0;
    // Solve an initial value problem for a 1st order system of ODEs
    nag::ad::d02pe(ad_handle, f, n, twant, tgot, y.data(), ypgot.data(),
                   ymax.data(), iwsav.data(), rwsav.data(), ifail);
  }

  T       waste, hnext;
  Integer fevals, stepcost, stepsok;
  ifail = 0;
  // This call has no influence on derivatives; it's just called with active
  // data
  nag::ad::d02pt(ad_handle, fevals, stepcost, waste, stepsok, hnext,
                 iwsav.data(), rwsav.data(), ifail);
  std::cout << "\n Cost of the integration in evaluations of f is " << fevals
            << std::endl;
}