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

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

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

// Driver with the adjoint calls.
// Solves the ODE for a projectile
//   y1' = tan(y3)
//   y2' = alpha * tan(y3) / y2 - beta * y2 / cos(y3)
//   y3' = alpha / y2^2
// until it hits the ground. x is the hit point.
// Also computes the sum of all Jacobian elements d x/d[y0, alpha, beta].
void driver(const std::vector<double> &y0,
            const double &             alpha,
            const double &             beta,
            double &                   x,
            double &                   dx);

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

  // Parameters
  double alpha = -0.032, beta = -0.02;
  // Initial condition
  std::vector<double> y0{0.5, 0.5, 0.2 * nag_math_pi};
  // Hit point
  double x;
  // Sum of Jacobian elements
  double dx;

  // Call driver
  driver(y0, alpha, beta, x, dx);

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

  std::cout.setf(std::ios::scientific, std::ios::floatfield);
  std::cout.precision(6);
  std::cout << "\n Hit point = " << x << std::endl;

  // Print derivatives
  std::cout
      << "\n Sum of all Jacobian elements of hit point x w.r.t. initial condition and parameters (alpha, beta) :\n";
  std::cout << " sum_j [dx / d(y0, alpha, beta)]_j = " << dx << std::endl;

  return 0;
}

// Driver with the adjoint calls.
// Solves the ODE for a projectile
//   y1' = tan(y3)
//   y2' = alpha * tan(y3) / y2 - beta * y2 / cos(y3)
//   y3' = alpha / y2^2
// until it hits the ground. x is the hit point.
// Also computes the sum of all Jacobian elements d x/d[y0, alpha, beta].
void driver(const std::vector<double> &y0v,
            const double &             alphav,
            const double &             betav,
            double &                   xv,
            double &                   dx)
{
  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.
  std::vector<T> y0(begin(y0v), end(y0v));
  T              alpha(alphav), beta(betav);

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

  // Hit point x, variable to differentiate
  T x;

  // Call the NAG AD Lib functions
  func(y0, alpha, beta, x);

  // Extract the computed solutions
  xv = dco::passive_value(x);
  // Initialize adjoints of y
  mode::global_tape->register_output_variable(x);
  dco::derivative(x) = 1.0;

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

  dx = 0.0;
  // Adjoint of y0 and r
  for (T const &y0i : y0)
    dx += dco::derivative(y0i);
  dx += dco::derivative(alpha) + dco::derivative(beta);

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

// function which calls NAG AD Library routines
template <typename T>
void func(const std::vector<T> &y0, const T &alpha, const T &beta, T &x)
{
  Integer n  = y0.size();
  Integer iw = 20 * n;
  // Active variables
  std::vector<T> w(iw), ruser{alpha, beta};
  T              xinit = 0.0, xend = 10.0, tol = 1.0e-5;

  auto fcn = [&](nag::ad::handle_t &ad_handle,
                 const T &          x,
                 const T            y[],
                 T                  f[])
             {
               T alpha, beta;
               alpha = ruser[0];
               beta  = ruser[1];
               
               f[0] = tan(y[2]);
               f[1] = alpha * tan(y[2]) / y[1] + beta * y[1] / cos(y[2]);
               f[2] = alpha / (y[1] * y[1]);
             };

  auto g = [](nag::ad::handle_t&, T const&, T const y[], T & z)
           {
             z = y[0];
           };

  // Create AD configuration data object
  Integer           ifail = 0;
  nag::ad::handle_t ad_handle;

  x                = xinit;
  std::vector<T> y = y0;
  ifail            = 0;
  nag::ad::d02bj(ad_handle, x, xend, n, y.data(), fcn, tol, "D", nullptr, g,
                 w.data(), ifail);
}