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

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

// Function which calls NAG AD Library routines.
template <typename T> void func(T r, T &x);

// Driver with the adjoint calls.
// Finds the zero point f(xv) = 0 for the function
// f(x) =  exp(-x) - r*x.
// Also, computes the gradient dxdr = dxv/dr using a symbolic adjoint of the
// root finder.
void driver(const double &rv, double &xv, double &dxdr);

// Callback that evaluates the non-linear function
template <typename T>
void NAG_CALL f(void *&ad_handle, const T &x, T &z, Integer iuser[], T ruser[]);

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

  // Set problem parameters
  double rv = 2.0;
  // Solution x
  double xv;
  // Derivative of x
  double dxdr;

  // Call driver
  driver(rv, xv, dxdr);

  // sPrint outputs
  std::cout << "\n Derivatives calculated: First order adjoints\n";
  std::cout << " Computational mode    : symbolic (expert mode)\n";

  // Print derivatives
  std::cout.setf(std::ios::scientific, std::ios::floatfield);
  std::cout.precision(12);
  std::cout << "\n Solution:\n";
  std::cout << " x = " << xv << std::endl;
  std::cout << "\n Derivative of solution x w.r.t. parameter r:\n";
  std::cout << " dx/dr(x) = " << dxdr << std::endl;

  return 0;
}

// Driver with the adjoint calls.
// Finds the zero point f(xv) = 0 for the function
// f(x) =  exp(-x) - r*x.
// Also, computes the gradient dxdr = dxv/dr using a symbolic adjoint of the
// root finder
void driver(const double &rv, double &xv, double &dxdr)
{
  using T = dco::ga1s<double>::type;

  // Create the AD tape
  dco::ga1s<double>::global_tape = dco::ga1s<double>::tape_t::create();

  // Function parameter r
  T r = rv;

  // Register variables to differentiate w.r.t.
  dco::ga1s<double>::global_tape->register_variable(r);

  // Variable to differentiate
  T x;

  // Call the NAG AD Lib functions
  func(r, x);
  // Extract the computed solution
  xv = dco::value(x);

  dco::ga1s<double>::global_tape->register_output_variable(x);
  dco::derivative(x) = 1.0;
  dco::ga1s<double>::global_tape->interpret_adjoint();

  // Extract the derivatives
  dxdr = dco::derivative(r);

  // Remove tape
  dco::ga1s<double>::tape_t::remove(dco::ga1s<double>::global_tape);
}

// Function which calls NAG AD Library routines
template <typename T> void func(T r, T &x)
{
  // Active variables
  T a = 0.0, b = 1.0;
  T eps = 1.0e-5, eta = 0.0;
  // Create AD configuration data object
  Integer ifail     = 0;
  void *  ad_handle = 0;
  nag::ad::x10aa(ad_handle, ifail);
  // Set computational mode
  ifail        = 0;
  Integer mode = nagad_symbolic_expert;
  nag::ad::x10ac(ad_handle, mode, ifail);
  // Routine for computing an approximation to a simple zero of the function f
  ifail = 0;
  nag::ad::c05ay(ad_handle, a, b, eps, eta, f<T>, x, 0, nullptr, 1, &r, ifail);
  // Remove computational data object
  ifail = 0;
  nag::ad::x10ab(ad_handle, ifail);
}

// Callback that evaluates the non-linear function
template <class T>
void NAG_CALL f(void *&ad_handle, const T &x, T &z, Integer iuser[], T ruser[])
{
  using value_t = typename dco::mode<T>::value_t;
  T r           = ruser[0];

  // Get the callback mode
  Integer ifail = 0, cb_mode;
  nag::ad::x10bd(ad_handle, cb_mode, ifail);

  // Extract the value of our state (x) and parameters (ruser).
  // This is only needed because we are using c05ay in expert symbolic
  // adjoint mode
  value_t xv = dco::value(x);
  value_t rv = dco::value(r);

  if (cb_mode == nagad_primal)
    {
      // We're in the forward pass of the root finder c05ay only wants
      // the primal evaluation, so compute with values (but
      // assign to active output, since it's the only way to pass
      // value back to caller)
      z = exp(-xv) - xv * rv;
    }
  else if (cb_mode == nagad_dstate)
    {
      // c05ay wants derivatives w.r.t state, so use values of the
      // parameter data ruser; this is required to not propagate
      // adjoints to parameters accidentally
      z = exp(-x) - x * rv;
    }
  else if (cb_mode == nagad_dparam)
    {
      // c05ay wants derivatives w.r.t the parameter data, so
      // use values of the state data x; this is only required for
      // efficiency reasons; values would still be correct if x were
      // used actively.
      z = exp(-xv) - xv * r;
    }
  else
    {
      // In this case cb_mode == nagad_dall c05ay wants all
      // derivatives, so just compute with active data
      z = exp(-x) - x * r;
    }
}