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

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

std::stringstream filecontent("4 \
 2.0   -1.0    0.0    0.0 \
-1.0    2.0   -1.0    0.0 \
 0.0   -1.0    2.0   -1.0 \
 0.0    0.0   -1.0    2.0 ");

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

// Driver with adjoint calls.
// Computes the Nearest Correlation Matrix X for an input matrix G.
// Also, computes the sum of all Jacobian elements of dX/dG.
void driver(const std::vector<double> &gv,
            std::vector<double> &      xv,
            double &                   dxdg);

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

  Integer n;
  filecontent >> n;

  // Input Matrix G, whose NCM we want to compute
  std::vector<double> gv(n * n);
  for (int i = 0; i < n; ++i)
    {
      for (int j = 0; j < n; ++j)
        {
          filecontent >> gv[i + j * n];
        }
    }

  // Output NCM X
  std::vector<double> xv(n * n);

  // Sum of all Jacobian elements dX/dG
  double dxdg;

  driver(gv, xv, dxdg);

  std::cout.setf(std::ios::scientific, std::ios::floatfield);
  std::cout.precision(5);
  std::cout << "\n Nearest Correlation Matrix:\n\n";
  for (int i = 0; i < n; i++)
    {
      for (int j = 0; j < n; j++)
        {
          std::cout.width(13);
          std::cout << xv[i + j * n];
        }
      std::cout << std::endl;
    }

  std::cout << "\n Derivatives calculated: First order adjoint\n";
  std::cout << " Computational mode    : symbolic\n\n";

  // Print derivatives of NCM X w.r.t. input matrix G
  std::cout << "\n Sum of Jacobian elements of NCM X w.r.t. input matrix G:\n";
  std::cout << " sum_ij [dX/dG]_ij = " << dxdg << std::endl;

  return 0;
}

// Driver with adjoint calls.
// Computes the Nearest Correlation Matrix X for an input matrix G.
// Also, computes the sum of all Jacobian elements of dX/dG.
void driver(const std::vector<double> &gv,
            std::vector<double> &      xv,
            double &                   dxdg)
{
  using T = dco::ga1s<double>::type;
  // Create AD tape
  dco::ga1s<double>::global_tape = dco::ga1s<double>::tape_t::create();

  // AD data types
  std::vector<T> g(gv.size()), g1(gv.size());
  dco::value(g) = gv;
  dco::ga1s<double>::global_tape->register_variable(g);
  // nag::ad::g02aa modifies input matrix g.  Since we seek dx/dg, and g is
  // overwritten, we must save these input nodes in the DAG.  Hence we save g
  // and overwrite a copy of it.
  g1 = g;

  // Variable to differentiate
  std::vector<T> x(xv.size());

  // Call the NAG AD Lib functions
  func(g1, x);

  // Extract the computed NCM
  xv = dco::value(x);

  dco::ga1s<double>::global_tape->register_output_variable(x);
  dco::derivative(x) = std::vector<double>(xv.size(), 1.0);
  dco::ga1s<double>::global_tape->interpret_adjoint();

  // Sum of the adjoints of g
  dxdg = 0.0;
  for (int i = 0; i < g.size(); i++)
    {
      dxdg += dco::derivative(g[i]);
    }

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

// Function which calls NAG AD routines.
template <typename T> void func(std::vector<T> &g, std::vector<T> &x)
{
  Integer n   = sqrt(g.size());
  Integer pdg = n;
  Integer pdx = n;

  // Set up method parameters
  T       errtol = 1.00e-7;
  Integer maxits = 200;
  Integer maxit  = 10;
  // Output variables
  Integer iter, feval;
  T       nrmgrd;
  // Create AD configuration data object
  Integer ifail     = 0;
  nag::ad::handle_t ad_handle;

  // Set computational mode
  ad_handle.set_strategy(nag::ad::symbolic);
  // Routine for computing the NCM of G.
  nag::ad::g02aa(ad_handle, g.data(), pdg, n, errtol, maxits, maxit, x.data(),
                 pdx, iter, feval, nrmgrd, ifail);
}