/* -----------------------------------------------------------------------------
 * Programmer(s): David J. Gardner @ LLNL
 * -----------------------------------------------------------------------------
 * SUNDIALS Copyright Start
 * Copyright (c) 2002-2022, Lawrence Livermore National Security
 * and Southern Methodist University.
 * All rights reserved.
 *
 * See the top-level LICENSE and NOTICE files for details.
 *
 * SPDX-License-Identifier: BSD-3-Clause
 * SUNDIALS Copyright End
 * -----------------------------------------------------------------------------
 * Kvaerno-Prothero-Robinson ODE test problem:
 *
 *   [u]' = [ a  b ] [ (-1 + u^2 - r(t)) / (2u) ] + [ r'(t) / (2u) ]
 *   [v]    [ c  d ] [ (-2 + v^2 - s(t)) / (2v) ]   [ s'(t) / (2v) ]
 *
 * This problem has analytical solution given by
 *
 *   u(t) = sqrt(1 + r(t))
 *   v(t) = sqrt(2 + s(t))
 *
 * where, in this test, we use the functions
 *
 *   r(t) = 0.5 * cos(t)
 *   s(t) = cos(2t)
 * ---------------------------------------------------------------------------*/

#include "idas_test_kpr.hpp"

int main(int argc, char* argv[])
{
  // SUNDIALS context object for this simulation
  sundials::Context sunctx;

  // Read input options
  TestOptions opts;

  vector<string> args(argv + 1, argv + argc);

  int flag = ReadInputs(args, opts, sunctx);
  if (check_flag(flag, "ReadInputs")) return 1;

  // Create initial condition
  N_Vector y = N_VNew_Serial(2, sunctx);
  if (check_ptr(y, "N_VNew_Serial")) return 1;

  realtype utrue, vtrue;
  flag = true_sol(ZERO, &utrue, &vtrue);
  if (check_flag(flag, "true_sol")) return 1;

  realtype* ydata = N_VGetArrayPointer(y);
  ydata[0] = utrue;
  ydata[1] = vtrue;

  N_Vector yp = N_VNew_Serial(2, sunctx);
  if (check_ptr(y, "N_VNew_Serial")) return 1;

  realtype uptrue, vptrue;
  flag = true_sol_p(ZERO, &uptrue, &vptrue);
  if (check_flag(flag, "true_sol")) return 1;

  realtype* ypdata = N_VGetArrayPointer(yp);
  ypdata[0] = uptrue;
  ypdata[1] = vptrue;

  // Create IDAS memory structure
  void* ida_mem = IDACreate(sunctx);
  if (check_ptr(ida_mem, "IDACreate")) return 1;

  flag = IDAInit(ida_mem, res, ZERO, y, yp);
  if (check_flag(flag, "IDAInit")) return 1;

  flag = IDASStolerances(ida_mem, opts.rtol, opts.atol);
  if (check_flag(flag, "IDASStolerances")) return 1;

  SUNMatrix A = SUNDenseMatrix(2, 2, sunctx);
  if (check_ptr(A, "SUNDenseMatrix")) return 1;

  SUNLinearSolver LS = SUNLinSol_Dense(y, A, sunctx);
  if (check_ptr(LS, "SUNLinSol_Dense")) return 1;

  flag = IDASetLinearSolver(ida_mem, LS, A);
  if (check_flag(flag, "IDASetLinearSolver")) return 1;

  flag = IDASetJacFn(ida_mem, J);
  if (check_flag(flag, "IDASetJacFn")) return 1;

  realtype udata[4] = {-TWO, HALF, HALF, -ONE};
  flag = IDASetUserData(ida_mem, udata);
  if (check_flag(flag, "IDASetUserData")) return 1;

  flag = IDASetEtaFixedStepBounds(ida_mem,
                                  opts.eta_min_fx, opts.eta_max_fx);
  if (check_flag(flag, "IDASetEtaFixeStepBounds")) return 1;

  flag = IDASetEtaMax(ida_mem, opts.eta_max);
  if (check_flag(flag, "IDASetEtaMax")) return 1;

  flag = IDASetEtaMin(ida_mem, opts.eta_min);
  if (check_flag(flag, "IDASetEtaMin")) return 1;

  flag = IDASetEtaMinErrFail(ida_mem, opts.eta_min_ef);
  if (check_flag(flag, "IDASetEtaMinErrFail")) return 1;

  flag = IDASetEtaConvFail(ida_mem, opts.eta_cf);
  if (check_flag(flag, "IDASetEtaConvFail")) return 1;

  flag = IDASetDeltaCjLSetup(ida_mem, opts.dcj);
  if (check_flag(flag, "IDASetDeltaCjLSetup")) return 1;

  // Initial time and fist output time
  realtype tret  = ZERO;
  realtype tout  = tret + opts.dtout;

  // Output initial contion
  cout << scientific;
  cout << setprecision(numeric_limits<realtype>::digits10);
  cout << "           t              ";
  cout << "          u              ";
  cout << "          v              ";
  cout << "        u err            ";
  cout << "        v err            " << endl;
  for (int i = 0; i < 9; i++)
    cout << "--------------";
  cout << endl;

  cout << setw(22) << tret
       << setw(25) << ydata[0]
       << setw(25) << ydata[1]
       << setw(25) << abs(ydata[0] - utrue)
       << setw(25) << abs(ydata[1] - vtrue) << endl;

  // Advance in time
  for (int i = 0; i < opts.nout; i++)
  {
    flag = IDASolve(ida_mem, tout, &tret, y, yp, IDA_NORMAL);
    if (check_flag(flag, "IDA")) return 1;

    flag = true_sol(tret, &utrue, &vtrue);
    if (check_flag(flag, "true_sol")) return 1;

    cout << setw(22) << tret
         << setw(25) << ydata[0]
         << setw(25) << ydata[1]
         << setw(25) << abs(ydata[0] - utrue)
         << setw(25) << abs(ydata[1] - vtrue) << endl;

    // update output time
    tout += opts.dtout;
  }
  for (int i = 0; i < 9; i++)
    cout << "--------------";
  cout << endl;

  // Print some final statistics
  flag = IDAPrintAllStats(ida_mem, stdout, SUN_OUTPUTFORMAT_TABLE);
  if (check_flag(flag, "IDAPrintAllStats")) return 1;

  // Clean up and return with successful completion
  N_VDestroy(y);
  N_VDestroy(yp);
  SUNMatDestroy(A);
  SUNLinSolFree(LS);
  IDAFree(&ida_mem);

  return 0;
}

/* -----------------------------------------------------------------------------
 * Compute the DAE residual function:
 *   ru = [a  b] * [ (-1 + u^2 - r(t)) / (2*u) ] + [ r'(t) / (2u) ] - [u']
 *   rv = [c  d]   [ (-2 + v^2 - s(t)) / (2*v) ]   [ s'(t) / (2v) ] - [v']
 * ---------------------------------------------------------------------------*/
int res(realtype t, N_Vector y, N_Vector yp, N_Vector rr, void *user_data)
{
  realtype* udata = (realtype *) user_data;
  const realtype a = udata[0];
  const realtype b = udata[1];
  const realtype c = udata[2];
  const realtype d = udata[3];

  realtype* ydata = N_VGetArrayPointer(y);
  const realtype u = ydata[0];
  const realtype v = ydata[1];

  realtype* ypdata = N_VGetArrayPointer(yp);
  const realtype up = ypdata[0];
  const realtype vp = ypdata[1];

  const realtype tmp1 = (-ONE + u * u - r(t)) / (TWO * u);
  const realtype tmp2 = (-TWO + v * v - s(t)) / (TWO * v);

  realtype* rdata = N_VGetArrayPointer(rr);
  rdata[0] = (a * tmp1 + b * tmp2 + rdot(t) / (TWO * u)) - up;
  rdata[1] = (c * tmp1 + d * tmp2 + sdot(t) / (TWO * v)) - vp;

  return 0;
}

/* -----------------------------------------------------------------------------
 * Compute the ODE RHS Jacobin:
 *   [a/2 + (a(1+r(t))-r'(t))/(2u^2) - cj  b/2 + b*(2+s(t))/(2*v^2)            ]
 *   [c/2 + c(1+r(t))/(2u^2)               d/2 + (d(2+s(t))-s'(t))/(2u^2) - cj ]
 * ---------------------------------------------------------------------------*/
int J(realtype t,  realtype cj, N_Vector y, N_Vector yp, N_Vector rr,
      SUNMatrix J, void *user_data, N_Vector tempv1, N_Vector tempv2,
      N_Vector tempv3)
{
  realtype* udata = (realtype *) user_data;
  const realtype a = udata[0];
  const realtype b = udata[1];
  const realtype c = udata[2];
  const realtype d = udata[3];

  realtype* ydata = N_VGetArrayPointer(y);
  realtype* Jdata = SUNDenseMatrix_Data(J);

  const realtype u = ydata[0];
  const realtype v = ydata[1];

  Jdata[0] = (a / TWO + (a * (ONE + r(t)) - rdot(t)) / (TWO * u * u)) - cj;
  Jdata[1] = c / TWO +  c * (ONE + r(t)) / (TWO * u * u);
  Jdata[2] = b / TWO +  b * (TWO + s(t)) / (TWO * v * v);
  Jdata[3] = (d / TWO + (d * (TWO + s(t)) - sdot(t)) / (TWO * v * v)) - cj;

  return 0;
}

/*---- end of file ----*/
