/* -----------------------------------------------------------------------------
 * Programmer(s): David J. Gardner @ LLNL
 * -----------------------------------------------------------------------------
 * SUNDIALS Copyright Start
 * Copyright (c) 2002-2024, 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
 * -----------------------------------------------------------------------------
 * This the implementation file for the CVODE nonlinear solver interface.
 * ---------------------------------------------------------------------------*/

#include "cvode_impl.h"
#include "sundials/sundials_math.h"

/* constant macros */
#define ONE SUN_RCONST(1.0) /* real 1.0 */

/* nonlinear solver constants
     NLS_MAXCOR  maximum no. of corrector iterations for the nonlinear solver
     CRDOWN      constant used in the estimation of the convergence rate (crate)
                 of the iterates for the nonlinear equation
     RDIV        declare divergence if ratio del/delp > RDIV
 */
#define NLS_MAXCOR 3
#define CRDOWN     SUN_RCONST(0.3)
#define RDIV       SUN_RCONST(2.0)

/* private functions */
static int cvNlsResidual(N_Vector ycor, N_Vector res, void* cvode_mem);
static int cvNlsFPFunction(N_Vector ycor, N_Vector res, void* cvode_mem);

static int cvNlsLSetup(sunbooleantype jbad, sunbooleantype* jcur,
                       void* cvode_mem);
static int cvNlsLSolve(N_Vector delta, void* cvode_mem);
static int cvNlsConvTest(SUNNonlinearSolver NLS, N_Vector ycor, N_Vector del,
                         sunrealtype tol, N_Vector ewt, void* cvode_mem);

/* -----------------------------------------------------------------------------
 * Exported functions
 * ---------------------------------------------------------------------------*/

int CVodeSetNonlinearSolver(void* cvode_mem, SUNNonlinearSolver NLS)
{
  CVodeMem cv_mem;
  int retval;

  /* Return immediately if CVode memory is NULL */
  if (cvode_mem == NULL)
  {
    cvProcessError(NULL, CV_MEM_NULL, __LINE__, __func__, __FILE__, MSGCV_NO_MEM);
    return (CV_MEM_NULL);
  }
  cv_mem = (CVodeMem)cvode_mem;

  /* Return immediately if NLS memory is NULL */
  if (NLS == NULL)
  {
    cvProcessError(NULL, CV_ILL_INPUT, __LINE__, __func__, __FILE__,
                   "NLS must be non-NULL");
    return (CV_ILL_INPUT);
  }

  /* check for required nonlinear solver functions */
  if (NLS->ops->gettype == NULL || NLS->ops->solve == NULL ||
      NLS->ops->setsysfn == NULL)
  {
    cvProcessError(cv_mem, CV_ILL_INPUT, __LINE__, __func__, __FILE__,
                   "NLS does not support required operations");
    return (CV_ILL_INPUT);
  }

  /* free any existing nonlinear solver */
  if ((cv_mem->NLS != NULL) && (cv_mem->ownNLS))
  {
    retval = SUNNonlinSolFree(cv_mem->NLS);
  }

  /* set SUNNonlinearSolver pointer */
  cv_mem->NLS = NLS;

  /* Set NLS ownership flag. If this function was called to attach the default
     NLS, CVODE will set the flag to SUNTRUE after this function returns. */
  cv_mem->ownNLS = SUNFALSE;

  /* set the nonlinear system function */
  if (SUNNonlinSolGetType(NLS) == SUNNONLINEARSOLVER_ROOTFIND)
  {
    retval = SUNNonlinSolSetSysFn(cv_mem->NLS, cvNlsResidual);
  }
  else if (SUNNonlinSolGetType(NLS) == SUNNONLINEARSOLVER_FIXEDPOINT)
  {
    retval = SUNNonlinSolSetSysFn(cv_mem->NLS, cvNlsFPFunction);
  }
  else
  {
    cvProcessError(cv_mem, CV_ILL_INPUT, __LINE__, __func__, __FILE__,
                   "Invalid nonlinear solver type");
    return (CV_ILL_INPUT);
  }

  if (retval != CV_SUCCESS)
  {
    cvProcessError(cv_mem, CV_ILL_INPUT, __LINE__, __func__, __FILE__,
                   "Setting nonlinear system function failed");
    return (CV_ILL_INPUT);
  }

  /* set convergence test function */
  retval = SUNNonlinSolSetConvTestFn(cv_mem->NLS, cvNlsConvTest, cvode_mem);
  if (retval != CV_SUCCESS)
  {
    cvProcessError(cv_mem, CV_ILL_INPUT, __LINE__, __func__, __FILE__,
                   "Setting convergence test function failed");
    return (CV_ILL_INPUT);
  }

  /* set max allowed nonlinear iterations */
  retval = SUNNonlinSolSetMaxIters(cv_mem->NLS, NLS_MAXCOR);
  if (retval != CV_SUCCESS)
  {
    cvProcessError(cv_mem, CV_ILL_INPUT, __LINE__, __func__, __FILE__,
                   "Setting maximum number of nonlinear iterations failed");
    return (CV_ILL_INPUT);
  }

  /* Reset the acnrmcur flag to SUNFALSE */
  cv_mem->cv_acnrmcur = SUNFALSE;

  /* Set the nonlinear system RHS function */
  if (!(cv_mem->cv_f))
  {
    cvProcessError(cv_mem, CV_ILL_INPUT, __LINE__, __func__, __FILE__,
                   "The ODE RHS function is NULL");
    return (CV_ILL_INPUT);
  }
  cv_mem->nls_f = cv_mem->cv_f;

  return (CV_SUCCESS);
}

/*---------------------------------------------------------------
  CVodeSetNlsRhsFn:

  This routine sets an alternative user-supplied ODE right-hand
  side function to use in the evaluation of nonlinear system
  functions.
  ---------------------------------------------------------------*/
int CVodeSetNlsRhsFn(void* cvode_mem, CVRhsFn f)
{
  CVodeMem cv_mem;

  if (cvode_mem == NULL)
  {
    cvProcessError(NULL, CV_MEM_NULL, __LINE__, __func__, __FILE__, MSGCV_NO_MEM);
    return (CV_MEM_NULL);
  }

  cv_mem = (CVodeMem)cvode_mem;

  if (f) { cv_mem->nls_f = f; }
  else { cv_mem->nls_f = cv_mem->cv_f; }

  return (CV_SUCCESS);
}

/*---------------------------------------------------------------
  CVodeGetNonlinearSystemData:

  This routine provides access to the relevant data needed to
  compute the nonlinear system function.
  ---------------------------------------------------------------*/
int CVodeGetNonlinearSystemData(void* cvode_mem, sunrealtype* tcur,
                                N_Vector* ypred, N_Vector* yn, N_Vector* fn,
                                sunrealtype* gamma, sunrealtype* rl1,
                                N_Vector* zn1, void** user_data)
{
  CVodeMem cv_mem;

  if (cvode_mem == NULL)
  {
    cvProcessError(NULL, CV_MEM_NULL, __LINE__, __func__, __FILE__, MSGCV_NO_MEM);
    return (CV_MEM_NULL);
  }

  cv_mem = (CVodeMem)cvode_mem;

  *tcur      = cv_mem->cv_tn;
  *ypred     = cv_mem->cv_zn[0];
  *yn        = cv_mem->cv_y;
  *fn        = cv_mem->cv_ftemp;
  *gamma     = cv_mem->cv_gamma;
  *rl1       = cv_mem->cv_rl1;
  *zn1       = cv_mem->cv_zn[1];
  *user_data = cv_mem->cv_user_data;

  return (CV_SUCCESS);
}

/* -----------------------------------------------------------------------------
 * Private functions
 * ---------------------------------------------------------------------------*/

int cvNlsInit(CVodeMem cvode_mem)
{
  int retval;

  /* set the linear solver setup wrapper function */
  if (cvode_mem->cv_lsetup)
  {
    retval = SUNNonlinSolSetLSetupFn(cvode_mem->NLS, cvNlsLSetup);
  }
  else { retval = SUNNonlinSolSetLSetupFn(cvode_mem->NLS, NULL); }

  if (retval != CV_SUCCESS)
  {
    cvProcessError(cvode_mem, CV_ILL_INPUT, __LINE__, __func__, __FILE__,
                   "Setting the linear solver setup function failed");
    return (CV_NLS_INIT_FAIL);
  }

  /* set the linear solver solve wrapper function */
  if (cvode_mem->cv_lsolve)
  {
    retval = SUNNonlinSolSetLSolveFn(cvode_mem->NLS, cvNlsLSolve);
  }
  else { retval = SUNNonlinSolSetLSolveFn(cvode_mem->NLS, NULL); }

  if (retval != CV_SUCCESS)
  {
    cvProcessError(cvode_mem, CV_ILL_INPUT, __LINE__, __func__, __FILE__,
                   "Setting linear solver solve function failed");
    return (CV_NLS_INIT_FAIL);
  }

  /* initialize nonlinear solver */
  retval = SUNNonlinSolInitialize(cvode_mem->NLS);

  if (retval != CV_SUCCESS)
  {
    cvProcessError(cvode_mem, CV_ILL_INPUT, __LINE__, __func__, __FILE__,
                   MSGCV_NLS_INIT_FAIL);
    return (CV_NLS_INIT_FAIL);
  }

  return (CV_SUCCESS);
}

static int cvNlsLSetup(sunbooleantype jbad, sunbooleantype* jcur, void* cvode_mem)
{
  CVodeMem cv_mem;
  int retval;

  if (cvode_mem == NULL)
  {
    cvProcessError(NULL, CV_MEM_NULL, __LINE__, __func__, __FILE__, MSGCV_NO_MEM);
    return (CV_MEM_NULL);
  }
  cv_mem = (CVodeMem)cvode_mem;

  /* if the nonlinear solver marked the Jacobian as bad update convfail */
  if (jbad) { cv_mem->convfail = CV_FAIL_BAD_J; }

  /* setup the linear solver */
  retval = cv_mem->cv_lsetup(cv_mem, cv_mem->convfail, cv_mem->cv_y,
                             cv_mem->cv_ftemp, &(cv_mem->cv_jcur),
                             cv_mem->cv_vtemp1, cv_mem->cv_vtemp2,
                             cv_mem->cv_vtemp3);
  cv_mem->cv_nsetups++;

  /* update Jacobian status */
  *jcur = cv_mem->cv_jcur;

  cv_mem->cv_gamrat = ONE;
  cv_mem->cv_gammap = cv_mem->cv_gamma;
  cv_mem->cv_crate  = ONE;
  cv_mem->cv_nstlp  = cv_mem->cv_nst;

  if (retval < 0) { return (CV_LSETUP_FAIL); }
  if (retval > 0) { return (SUN_NLS_CONV_RECVR); }

  return (CV_SUCCESS);
}

static int cvNlsLSolve(N_Vector delta, void* cvode_mem)
{
  CVodeMem cv_mem;
  int retval;

  if (cvode_mem == NULL)
  {
    cvProcessError(NULL, CV_MEM_NULL, __LINE__, __func__, __FILE__, MSGCV_NO_MEM);
    return (CV_MEM_NULL);
  }
  cv_mem = (CVodeMem)cvode_mem;

  retval = cv_mem->cv_lsolve(cv_mem, delta, cv_mem->cv_ewt, cv_mem->cv_y,
                             cv_mem->cv_ftemp);

  if (retval < 0) { return (CV_LSOLVE_FAIL); }
  if (retval > 0) { return (SUN_NLS_CONV_RECVR); }

  return (CV_SUCCESS);
}

static int cvNlsConvTest(SUNNonlinearSolver NLS, N_Vector ycor, N_Vector delta,
                         sunrealtype tol, N_Vector ewt, void* cvode_mem)
{
  CVodeMem cv_mem;
  int m, retval;
  sunrealtype del;
  sunrealtype dcon;

  if (cvode_mem == NULL)
  {
    cvProcessError(NULL, CV_MEM_NULL, __LINE__, __func__, __FILE__, MSGCV_NO_MEM);
    return (CV_MEM_NULL);
  }
  cv_mem = (CVodeMem)cvode_mem;

  /* compute the norm of the correction */
  del = N_VWrmsNorm(delta, ewt);

  /* get the current nonlinear solver iteration count */
  retval = SUNNonlinSolGetCurIter(NLS, &m);
  if (retval != CV_SUCCESS) { return (CV_MEM_NULL); }

  /* Test for convergence. If m > 0, an estimate of the convergence
     rate constant is stored in crate, and used in the test.        */
  if (m > 0)
  {
    cv_mem->cv_crate = SUNMAX(CRDOWN * cv_mem->cv_crate, del / cv_mem->cv_delp);
  }
  dcon = del * SUNMIN(ONE, cv_mem->cv_crate) / tol;

  if (dcon <= ONE)
  {
    cv_mem->cv_acnrm    = (m == 0) ? del : N_VWrmsNorm(ycor, ewt);
    cv_mem->cv_acnrmcur = SUNTRUE;
    return (CV_SUCCESS); /* Nonlinear system was solved successfully */
  }

  /* check if the iteration seems to be diverging */
  if ((m >= 1) && (del > RDIV * cv_mem->cv_delp))
  {
    return (SUN_NLS_CONV_RECVR);
  }

  /* Save norm of correction and loop again */
  cv_mem->cv_delp = del;

  /* Not yet converged */
  return (SUN_NLS_CONTINUE);
}

static int cvNlsResidual(N_Vector ycor, N_Vector res, void* cvode_mem)
{
  CVodeMem cv_mem;
  int retval;

  if (cvode_mem == NULL)
  {
    cvProcessError(NULL, CV_MEM_NULL, __LINE__, __func__, __FILE__, MSGCV_NO_MEM);
    return (CV_MEM_NULL);
  }
  cv_mem = (CVodeMem)cvode_mem;

  /* update the state based on the current correction */
  N_VLinearSum(ONE, cv_mem->cv_zn[0], ONE, ycor, cv_mem->cv_y);

  /* evaluate the rhs function */
  retval = cv_mem->nls_f(cv_mem->cv_tn, cv_mem->cv_y, cv_mem->cv_ftemp,
                         cv_mem->cv_user_data);
  cv_mem->cv_nfe++;
  if (retval < 0) { return (CV_RHSFUNC_FAIL); }
  if (retval > 0) { return (RHSFUNC_RECVR); }

#ifdef SUNDIALS_BUILD_PACKAGE_FUSED_KERNELS
  if (cv_mem->cv_usefused)
  {
    cvNlsResid_fused(cv_mem->cv_rl1, -cv_mem->cv_gamma, cv_mem->cv_zn[1], ycor,
                     cv_mem->cv_ftemp, res);
  }
  else
#endif
  {
    N_VLinearSum(cv_mem->cv_rl1, cv_mem->cv_zn[1], ONE, ycor, res);
    N_VLinearSum(-cv_mem->cv_gamma, cv_mem->cv_ftemp, ONE, res, res);
  }

  return (CV_SUCCESS);
}

static int cvNlsFPFunction(N_Vector ycor, N_Vector res, void* cvode_mem)
{
  CVodeMem cv_mem;
  int retval;

  if (cvode_mem == NULL)
  {
    cvProcessError(NULL, CV_MEM_NULL, __LINE__, __func__, __FILE__, MSGCV_NO_MEM);
    return (CV_MEM_NULL);
  }
  cv_mem = (CVodeMem)cvode_mem;

  /* update the state based on the current correction */
  N_VLinearSum(ONE, cv_mem->cv_zn[0], ONE, ycor, cv_mem->cv_y);

  /* evaluate the rhs function */
  retval = cv_mem->nls_f(cv_mem->cv_tn, cv_mem->cv_y, res, cv_mem->cv_user_data);
  cv_mem->cv_nfe++;
  if (retval < 0) { return (CV_RHSFUNC_FAIL); }
  if (retval > 0) { return (RHSFUNC_RECVR); }

  N_VLinearSum(cv_mem->cv_h, res, -ONE, cv_mem->cv_zn[1], res);
  N_VScale(cv_mem->cv_rl1, res, res);

  return (CV_SUCCESS);
}
