#include "dco.hpp"
/* G02DA_A1W_F C++ Header Example Program.
 *
 * Copyright 2024 Numerical Algorithms Group.
 * Mark 30.0, 2024.
 */

#include <iostream>
#include <nag.h>
#include <nagad.h>
#include <stdio.h>
#include <string>
using namespace std;

int main()
{
  int               exit_status = 0;
  nag::ad::handle_t ad_handle;
  nagad_a1w_w_rtype tol;
  Integer           ifail = 0;

  cout << "G02DA_A1W_F C++ Header Example Program Results\n\n";

  tol = 0.000001;

  // Skip heading in data file
  string mystr;
  getline(cin, mystr);

  // Read problem sizes
  Integer n, m;
  cin >> n;
  cin >> m;

  // Allocate arrays depending on m and n.
  nagad_a1w_w_rtype *x = 0, *y = 0;
  Integer *          isx = 0;
  double *           dy  = 0;
  if (!(x = NAG_ALLOC(m * n, nagad_a1w_w_rtype)) ||
      !(y = NAG_ALLOC(n, nagad_a1w_w_rtype)) || !(dy = NAG_ALLOC(n, double)) ||
      !(isx = NAG_ALLOC(m, Integer)))
  {
    cout << "Allocation failure\n";
    exit_status = -1;
    exit(exit_status);
  }

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

  // Create AD configuration data object
  ifail = 0;

  // Read model data
  double dd;
  for (int i = 0; i < n; ++i)
  {
    for (int j = 0; j < m; ++j)
    {
      cin >> dd;
      int k = i + j * n;
      x[k]  = dd;
    }
    cin >> dd;
    y[i] = dd;
    dco::ga1s<double>::global_tape->register_variable(y[i]);
  }

  // Calculate ip
  Integer ip = 0;
  for (int j = 0; j < m; ++j)
  {
    cin >> isx[j];
    if (isx[j] > 0)
    {
      ip++;
    }
  }
  // Mean = 'M'
  ip++;

  // Allocate arrays depending on ip
  nagad_a1w_w_rtype *b = 0, *cov = 0, *h = 0, *p = 0, *q = 0, *res = 0, *se = 0,
                    *wk = 0;
  double *dbdy;
  Integer lcov = (ip * ip + ip) / 2, lp = ip * (ip + 2), lq = n * (ip + 1);
  Integer lwk = ip * ip + 5 * (ip - 1);
  if (!(b = NAG_ALLOC(ip, nagad_a1w_w_rtype)) ||
      !(cov = NAG_ALLOC(lcov, nagad_a1w_w_rtype)) ||
      !(h = NAG_ALLOC(n, nagad_a1w_w_rtype)) ||
      !(p = NAG_ALLOC(lp, nagad_a1w_w_rtype)) ||
      !(q = NAG_ALLOC(lq, nagad_a1w_w_rtype)) ||
      !(res = NAG_ALLOC(n, nagad_a1w_w_rtype)) ||
      !(se = NAG_ALLOC(ip, nagad_a1w_w_rtype)) ||
      !(wk = NAG_ALLOC(lwk, nagad_a1w_w_rtype)) ||
      !(dbdy = NAG_ALLOC(n * ip, double)))
  {
    cout << "Allocation failure\n";
    exit_status = -2;
    exit(exit_status);
  }

  // Perform Regression
  Integer           idf, irank;
  nagad_a1w_w_rtype rss, wt[1];
  logical           svd;
  ifail = 0;
  nag::ad::g02da(ad_handle, "M", "U", n, x, n, m, isx, ip, y, wt, rss, idf, b,
                 se, cov, res, h, q, n, svd, irank, p, tol, wk, ifail);

  // Display results
  if (svd)
  {
    cout << "Model is not of full rank, rank = " << irank << endl;
  }
  cout << "Residual sum of squares         = " << dco::value(rss);
  cout << "\nDegrees of freedom              = " << idf << endl;
  cout << "\nVariable   Parameter estimate   Standard error\n\n";
  cout.setf(ios::scientific, ios::floatfield);
  cout.precision(2);
  for (int i = 0; i < ip; ++i)
  {
    cout.width(5);
    cout << i << "       ";
    cout.width(12);
    cout << dco::value(b[i]) << "       ";
    cout.width(12);
    cout << dco::value(se[i]) << endl;
  }

  cout << "\n\n Derivatives calculated: First order adjoints\n";
  cout << " Computational mode    : algorithmic\n";
  cout << "\n Derivatives:\n\n";

  // Obtain derivatives
  double inc = 1.0;
  dco::derivative(rss) += inc;
  ifail                                              = 0;
  dco::ga1s<double>::global_tape->sparse_interpret() = true;
  dco::ga1s<double>::global_tape->interpret_adjoint();
  for (int j = 0; j < n; j++)
  {
    dy[j] = dco::derivative(y[j]);
  }
  cout << "    i        d(rss)/dy(i)\n";
  cout.precision(4);
  for (int i = 0; i < n; ++i)
  {
    cout.width(5);
    cout << i << "       ";
    cout.width(12);
    cout << dy[i] << endl;
  }

  for (int i = 0; i < ip; i++)
  {
    // Reset adjoints, initialize derivative, and evaluate adjoint
    dco::ga1s<double>::global_tape->zero_adjoints();
    dco::derivative(b[i]) += inc;
    ifail                                              = 0;
    dco::ga1s<double>::global_tape->sparse_interpret() = true;
    dco::ga1s<double>::global_tape->interpret_adjoint();
    for (int j = 0; j < n; j++)
    {
      int    k  = j + i * n;
      double dd = dco::derivative(y[j]);
      dbdy[k]   = dd;
    }
  }
  // Print matrix routine
  cout << endl;
  NagError fail;
  INIT_FAIL(fail);
  x04cac(Nag_ColMajor, Nag_GeneralMatrix, Nag_NonUnitDiag, n, ip, dbdy, n,
         "  Derivatives db/dy", 0, &fail);

  ifail = 0;

  dco::ga1s<double>::tape_t::remove(dco::ga1s<double>::global_tape);

  return exit_status;
}