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

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

// Function which calls NAG AD routines.
template <typename T> void func(const T &eps, T &errmax, T &terrmx);

// Driver with the adjoint calls.
// Solves the problem x''=x/r^3, y''=-y/r^3, with r=sqrt(x^2+y^2) and
// initial conditions x(0)=1-eps, x'(0)=0, y(0)=0, y'(0)=sqrt((1+eps)/(1-eps))
// by solving the ODE system:
//    y1'=y3, y2'=y4, y3'=-y1/r^3, y4'=-y2/r^3
// over the range [0,3*pi] and computes worst global error (errmax) as well as
// respective time (terrmx). Also computes the sum of 1st derivatives of outputs
// errmax and terrmx.
void driver(const double &epsv,
            double &      errmaxv,
            double &      terrmxv,
            double &      d_deps);

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

  // Parameter epsilon
  double epsv = 0.7;
  // Maximum weighted approximate true error taken over all solution components
  // and all steps
  double errmaxv;
  // First value of the independent variable where an approximate true error
  // attains the maximum value
  double terrmxv;
  // Sum of 1st derivatives
  double d_deps;

  // Call driver
  driver(epsv, errmaxv, terrmxv, d_deps);

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

  std::cout << "\n Worst global error observed = " << errmaxv;
  std::cout << "\n at T = " << terrmxv << std::endl;

  // Print derivatives
  std::cout.setf(std::ios::scientific, std::ios::floatfield);
  std::cout.precision(6);
  std::cout
      << "\n Sum of 1st derivatives of errmax and terrmx w.r.t. problem parameter epsilon:\n";
  std::cout << " derrmax/deps + dterrmx/deps = " << d_deps << std::endl;

  return 0;
}

// Driver with the adjoint calls.
// Solves the problem x''=x/r^3, y''=-y/r^3, with r=sqrt(x^2+y^2) and
// initial conditions x(0)=1-eps, x'(0)=0, y(0)=0, y'(0)=sqrt((1+eps)/(1-eps))
// by solving the ODE system:
//    y1'=y3, y2'=y4, y3'=-y1/r^3, y4'=-y2/r^3
// over the range [0,3*pi].
// Also computes the sum of all hessian elements dy/deps.
void driver(const double &epsv,
            double &      errmaxv,
            double &      terrmxv,
            double &      d_deps)
{
  using mode = dco::ga1s<double>;
  using T    = mode::type;

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

  // Variable to differentiate w.r.t.
  T eps = epsv;
  // Register variable to differentiate w.r.t.
  mode::global_tape->register_variable(eps);

  // Variables to differentiate
  T errmax, terrmx;

  // Call the NAG AD Lib functions
  func(eps, errmax, terrmx);

  // Extract the computed solutions
  errmaxv = dco::passive_value(errmax);
  terrmxv = dco::passive_value(terrmx);

  // Initialize adjoints
  mode::global_tape->register_output_variable(errmax);
  mode::global_tape->register_output_variable(terrmx);
  dco::derivative(errmax) = 1.0;
  dco::derivative(terrmx) = 1.0;

  // Interpret the tape
  mode::global_tape->interpret_adjoint();

  // Adjoint of eps
  d_deps = dco::derivative(eps);

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

// function which calls NAG AD Library routines
template <typename T> void func(const T &eps, T &errmax, T &terrmx)
{
  // Active variables
  const Integer n      = 4;
  const Integer liwsav = 130, lrwsav = 350 + 32 * n;

  std::vector<T> thresh(n, 1e-10), ypgot(n), ymax(n), rwsav(lrwsav), rmserr(n);
  std::vector<Integer> iwsav(liwsav);

  // Set parameters for the integrator.
  Integer method = 3;
  T       tol = 1e-6, hstart = 0.0, tend = 3.0 * nag_math_pi, tstart = 0.0;
  // Set initial conditions
  std::vector<T> y{1.0 - eps, 0.0, 0.0, sqrt((1.0 + eps) / (1.0 - eps))};

  // Create AD configuration data object
  nag::ad::handle_t ad_handle;

  // Initialize Runge-Kutta method for integrating ODE
  Integer ifail = 0;
  nag::ad::d02pq(ad_handle, n, tstart, tend, y.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[])
          {
            T r   = 1.0 / sqrt(y[0] * y[0] + y[1] * y[1]);
            r     = r * r * r;
            yp[0] = y[2];
            yp[1] = y[3];
            yp[2] = -y[0] * r;
            yp[3] = -y[1] * r;
          };

  T tnow = tstart;
  ifail  = 0;
  nag::ad::d02pe(ad_handle, f, n, tend, tnow, y.data(), ypgot.data(),
                 ymax.data(), iwsav.data(), rwsav.data(), ifail);

  // Get Error estimates
  ifail = 0;
  nag::ad::d02pu(ad_handle, n, rmserr.data(), errmax, terrmx, iwsav.data(),
                 rwsav.data(), ifail);
}