/* ---------------------------------------------------------------------------
 * Programmer(s): David J. Gardner @ LLNL
 *                Daniel R. Reynolds @ SMU
 * ---------------------------------------------------------------------------
 * 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
 * ---------------------------------------------------------------------------
 * This is the implementation file for ARKODE's MRIStepCoupling tables.
 * ---------------------------------------------------------------------------*/

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "arkode_impl.h"
#include "arkode_mristep_impl.h"
#include <sundials/sundials_math.h>


/* ===========================================================================
 * Exported Functions
 * ===========================================================================*/


/*---------------------------------------------------------------
  Returns MRIStepCoupling table structure for pre-set MRI methods.

  Input:  imeth -- integer key for the desired method
  ---------------------------------------------------------------*/
MRIStepCoupling MRIStepCoupling_LoadTable(ARKODE_MRITableID method)
{
  switch (method) {
#define ARK_MRI_TABLE(name, coeff) case name: coeff break;
#include "arkode_mri_tables.def"
#undef ARK_MRI_TABLE

    default:
      arkProcessError(NULL, ARK_ILL_INPUT, "ARKODE",
                      "MRIStepCoupling_LoadTable",
                      "Unknown coupling table");
      return NULL;
  }
}


/*---------------------------------------------------------------
  Returns MRIStepCoupling table structure for pre-set MRI methods.

  Input:  method -- string key for the desired method
  ---------------------------------------------------------------*/
MRIStepCoupling MRIStepCoupling_LoadTableByName(const char *method)
{
#define ARK_MRI_TABLE(name, coeff) if (strcmp(#name, method) == 0) coeff
#include "arkode_mri_tables.def"
#undef ARK_MRI_TABLE

  arkProcessError(NULL, ARK_ILL_INPUT, "ARKODE",
	   "MRIStepCoupling_LoadTable",
	   "Unknown coupling table");

  return NULL;
}


/*---------------------------------------------------------------
  Routine to allocate an empty MRIStepCoupling structure
  ---------------------------------------------------------------*/
MRIStepCoupling MRIStepCoupling_Alloc(int nmat, int stages,
                                      MRISTEP_METHOD_TYPE type)
{
  int i, j;
  MRIStepCoupling MRIC = NULL;

  /* Check for legal input values */
  if (nmat < 1 || stages < 1) return(NULL);

  /* ------------------------------------------
   * Allocate and initialize coupling structure
   * ------------------------------------------ */

  MRIC = (MRIStepCoupling) malloc(sizeof(struct MRIStepCouplingMem));
  if (!MRIC) return(NULL);

  MRIC->nmat   = nmat;
  MRIC->stages = stages;
  MRIC->q      = 0;
  MRIC->p      = 0;
  MRIC->c      = NULL;
  MRIC->W      = NULL;
  MRIC->G      = NULL;

  /* --------------------------------------------
   * Allocate abscissae and coupling coefficients
   * -------------------------------------------- */

  MRIC->c = (realtype *) calloc( stages, sizeof(realtype) );
  if (!(MRIC->c)) { MRIStepCoupling_Free(MRIC); return(NULL); }

  if (type == MRISTEP_EXPLICIT || type == MRISTEP_IMEX) {

    /* allocate W matrices */
    MRIC->W = (realtype ***) calloc( nmat, sizeof(realtype**) );
    if (!(MRIC->W)) { MRIStepCoupling_Free(MRIC); return(NULL); }

    /* allocate rows of each matrix in W */
    for (i=0; i<nmat; i++) {
      MRIC->W[i] = NULL;
      MRIC->W[i] = (realtype **) calloc( stages, sizeof(realtype*) );
      if (!(MRIC->W[i])) { MRIStepCoupling_Free(MRIC); return(NULL); }
    }

    /* allocate columns of each matrix in W */
    for (i=0; i<nmat; i++)
      for (j=0; j<stages; j++) {
        MRIC->W[i][j] = NULL;
        MRIC->W[i][j] = (realtype *) calloc( stages, sizeof(realtype) );
        if (!(MRIC->W[i][j])) { MRIStepCoupling_Free(MRIC); return(NULL); }
      }
  }

  if (type == MRISTEP_IMPLICIT || type == MRISTEP_IMEX) {

    /* allocate G matrices */
    MRIC->G = (realtype ***) calloc( nmat, sizeof(realtype**) );
    if (!(MRIC->G)) { MRIStepCoupling_Free(MRIC); return(NULL); }

    /* allocate rows of each matrix in G */
    for (i=0; i<nmat; i++) {
      MRIC->G[i] = NULL;
      MRIC->G[i] = (realtype **) calloc( stages, sizeof(realtype*) );
      if (!(MRIC->G[i])) { MRIStepCoupling_Free(MRIC); return(NULL); }
    }

    /* allocate columns of each matrix in G */
    for (i=0; i<nmat; i++)
      for (j=0; j<stages; j++) {
        MRIC->G[i][j] = NULL;
        MRIC->G[i][j] = (realtype *) calloc( stages, sizeof(realtype) );
        if (!(MRIC->G[i][j])) { MRIStepCoupling_Free(MRIC); return(NULL); }
      }
  }

  return(MRIC);
}


/*---------------------------------------------------------------
  Routine to allocate and fill a MRIStepCoupling structure
  ---------------------------------------------------------------*/
MRIStepCoupling MRIStepCoupling_Create(int nmat, int stages, int q, int p,
                                       realtype *W, realtype *G, realtype *c)
{
  int i, j, k;
  MRISTEP_METHOD_TYPE type;
  MRIStepCoupling MRIC = NULL;

  /* Check for legal inputs */
  if (nmat < 1 || stages < 1 || !c) return(NULL);

  /* Check for method coefficients and set method type */
  if (W && G)
    type = MRISTEP_IMEX;
  else if (W && !G)
    type = MRISTEP_EXPLICIT;
  else if (!W && G)
    type = MRISTEP_IMPLICIT;
  else
    return(NULL);

  /* Allocate MRIStepCoupling structure */
  MRIC = MRIStepCoupling_Alloc(nmat, stages, type);
  if (!MRIC) return(NULL);

  /* -------------------------
   * Copy the inputs into MRIC
   * ------------------------- */

  /* Method and embedding order */
  MRIC->q = q;
  MRIC->p = p;

  /* Abscissae */
  for (i=0; i<stages; i++)
    MRIC->c[i] = c[i];

  /* Coupling coefficients stored as 1D arrays of length nmat * stages * stages,
     with each stages * stages matrix stored in C (row-major) order */
  if (type == MRISTEP_EXPLICIT || type == MRISTEP_IMEX) {
    for (k = 0; k < nmat; k++)
      for (i = 0; i < stages; i++)
        for (j = 0; j < stages; j++)
          MRIC->W[k][i][j] = W[stages * (stages * k + i) + j];
  }
  if (type == MRISTEP_IMPLICIT || type == MRISTEP_IMEX) {
    for (k = 0; k < nmat; k++)
      for (i = 0; i < stages; i++)
        for (j = 0; j < stages; j++)
          MRIC->G[k][i][j] = G[stages * (stages * k + i) + j];
  }

  return(MRIC);
}


/*---------------------------------------------------------------
  Construct the MRI coupling matrix for an MIS method based on
  a given 'slow' Butcher table.
  ---------------------------------------------------------------*/
MRIStepCoupling MRIStepCoupling_MIStoMRI(ARKodeButcherTable B,
                                         int q, int p)
{
  int i, j, stages;
  booleantype padding;
  realtype Asum;
  realtype ***C;
  MRISTEP_METHOD_TYPE type;
  MRIStepCoupling MRIC;

  const realtype tol = RCONST(100.0) * UNIT_ROUNDOFF;

  /* Check that input table is non-NULL */
  if (!B) return(NULL);

  /* -----------------------------------
   * Check that the input table is valid
   * ----------------------------------- */

  /* First stage is just old solution */
  Asum = SUNRabs(B->c[0]);
  for (j=0; j<B->stages; j++)
    Asum += SUNRabs(B->A[0][j]);
  if (Asum > tol) return(NULL);

  /* Last stage exceeds 1 */
  if (B->c[B->stages-1] > ONE+tol) return(NULL);

  /* All stages are sorted */
  for (j=1; j<B->stages; j++)
    if ((B->c[j] - B->c[j-1]) < -tol) return(NULL);

  /* Each stage at most diagonally implicit */
  Asum = ZERO;
  for (i=0; i<B->stages; i++)
    for (j=i+1; j<B->stages; j++)
      Asum += SUNRabs(B->A[i][j]);
  if (Asum > tol) return(NULL);

  /* -----------------------------------------
   * determine whether the table needs padding
   * ----------------------------------------- */

  padding = SUNFALSE;

  /* Last stage time should equal 1 */
  if (SUNRabs(B->c[B->stages-1] - ONE) > tol)
    padding = SUNTRUE;

  /* Last row of A should equal b */
  for (j=0; j<B->stages; j++) {
    if (SUNRabs(B->A[B->stages-1][j] - B->b[j]) > tol)
      padding = SUNTRUE;
  }
  stages = (padding) ? B->stages+1 : B->stages;

  /* -------------------------
   * determine the method type
   * ------------------------- */

  /* Check if the table is strictly lower triangular (explicit) */
  type = MRISTEP_EXPLICIT;

  for (i=0; i<B->stages; i++)
    for (j=i; j<B->stages; j++)
      if (SUNRabs(B->A[i][j]) > tol)
        type = MRISTEP_IMPLICIT;

  /* ----------------------------
   * construct coupling structure
   * ---------------------------- */

  MRIC = MRIStepCoupling_Alloc(1, stages, type);
  if (!MRIC) return(NULL);

  /* Copy method/embedding orders */
  MRIC->q = q;
  MRIC->p = p;

  /* Copy abscissae, padding if needed */
  for (i=0; i<B->stages; i++)
    MRIC->c[i] = B->c[i];

  if (padding)
    MRIC->c[stages-1] = ONE;

  /* Construct the coupling table */
  if (type == MRISTEP_EXPLICIT)
    C = MRIC->W;
  else
    C = MRIC->G;

  /* First row is identically zero */
  for (i=0; i<stages; i++)
    for (j=0; j<stages; j++)
      C[0][i][j] = ZERO;

  /* Remaining rows = A(2:end,:) - A(1:end-1,:) */
  for (i=1; i<B->stages; i++)
    for (j=0; j<B->stages; j++)
      C[0][i][j] = B->A[i][j] - B->A[i-1][j];

  /* Padded row = b(:) - A(end,:) */
  if (padding)
    for (j=0; j<B->stages; j++)
      C[0][stages-1][j] = B->b[j] - B->A[B->stages-1][j];

  return(MRIC);
}


/*---------------------------------------------------------------
  Routine to copy a MRIStepCoupling structure
  ---------------------------------------------------------------*/
MRIStepCoupling MRIStepCoupling_Copy(MRIStepCoupling MRIC)
{
  int i, j, k, nmat, stages;
  MRISTEP_METHOD_TYPE type;
  MRIStepCoupling MRICcopy;

  /* Check for legal input */
  if (!MRIC) return(NULL);

  /* Check for method coefficients and set method type */
  if (MRIC->W && MRIC->G)
    type = MRISTEP_IMEX;
  else if (MRIC->W && !(MRIC->G))
    type = MRISTEP_EXPLICIT;
  else if (!(MRIC->W) && MRIC->G)
    type = MRISTEP_IMPLICIT;
  else
    return(NULL);

  /* Check for stage times */
  if (!(MRIC->c)) return(NULL);

  /* Get the number of coupling matrices and stages */
  nmat   = MRIC->nmat;
  stages = MRIC->stages;

  /* Allocate coupling structure */
  MRICcopy = MRIStepCoupling_Alloc(nmat, stages, type);
  if (!MRICcopy) return(NULL);

  /* Copy method and embedding orders */
  MRICcopy->q = MRIC->q;
  MRICcopy->p = MRIC->p;

  /* Copy abscissae */
  for (i=0; i<stages; i++)
    MRICcopy->c[i] = MRIC->c[i];

  /* Copy explicit coupling matrices W */
  if (MRIC->W)
    for (k = 0; k < nmat; k++)
      for (i = 0; i < stages; i++)
        for (j = 0; j < stages; j++)
          MRICcopy->W[k][i][j] = MRIC->W[k][i][j];

  /* Copy implicit coupling matrices G */
  if (MRIC->G)
    for (k = 0; k < nmat; k++)
      for (i = 0; i < stages; i++)
        for (j = 0; j < stages; j++)
          MRICcopy->G[k][i][j] = MRIC->G[k][i][j];

  return(MRICcopy);
}


/*---------------------------------------------------------------
  Routine to query the MRIStepCoupling structure workspace size
  ---------------------------------------------------------------*/
void MRIStepCoupling_Space(MRIStepCoupling MRIC, sunindextype *liw,
                           sunindextype *lrw)
{
  /* initialize outputs and return if MRIC is not allocated */
  *liw = 0;
  *lrw = 0;
  if (!MRIC) return;

  /* fill outputs based on MRIC */
  *liw = 4;
  if (MRIC->c)
    *lrw += MRIC->stages;
  if (MRIC->W)
    *lrw += MRIC->nmat * MRIC->stages * MRIC->stages;
  if (MRIC->G)
    *lrw += MRIC->nmat * MRIC->stages * MRIC->stages;
}


/*---------------------------------------------------------------
  Routine to free a MRIStepCoupling structure
  ---------------------------------------------------------------*/
void MRIStepCoupling_Free(MRIStepCoupling MRIC)
{
  int k, i;

  /* Free each field within MRIStepCoupling structure, and then
     free structure itself */
  if (MRIC) {

    if (MRIC->c)
      free(MRIC->c);

    if (MRIC->W) {
      for (k=0; k<MRIC->nmat; k++)
        if (MRIC->W[k]) {
          for (i=0; i<MRIC->stages; i++)
            if (MRIC->W[k][i]) {
              free(MRIC->W[k][i]);
              MRIC->W[k][i] = NULL;
            }
          free(MRIC->W[k]);
          MRIC->W[k] = NULL;
        }
      free(MRIC->W);
    }

    if (MRIC->G) {
      for (k=0; k<MRIC->nmat; k++)
        if (MRIC->G[k]) {
          for (i=0; i<MRIC->stages; i++)
            if (MRIC->G[k][i]) {
              free(MRIC->G[k][i]);
              MRIC->G[k][i] = NULL;
            }
          free(MRIC->G[k]);
          MRIC->G[k] = NULL;
        }
      free(MRIC->G);
    }

    free(MRIC);
  }
}


/*---------------------------------------------------------------
  Routine to print a MRIStepCoupling structure
  ---------------------------------------------------------------*/
void MRIStepCoupling_Write(MRIStepCoupling MRIC, FILE *outfile)
{
  int i, j, k;

  /* check for vaild coupling structure */
  if (!MRIC) return;
  if (!(MRIC->G)) return;

  if (MRIC->W) {
    for (i = 0; i < MRIC->nmat; i++) {
      if (!(MRIC->W[i])) return;
      for (j = 0; j < MRIC->stages; j++)
        if (!(MRIC->W[i][j])) return;
    }
  }

  if (MRIC->G) {
    for (i = 0; i < MRIC->nmat; i++) {
      if (!(MRIC->G[i])) return;
      for (j = 0; j < MRIC->stages; j++)
        if (!(MRIC->G[i][j])) return;
    }
  }

  if (!(MRIC->c)) return;

  fprintf(outfile, "  nmat = %i\n", MRIC->nmat);
  fprintf(outfile, "  stages = %i\n", MRIC->stages);
  fprintf(outfile, "  method order (q) = %i\n", MRIC->q);
  fprintf(outfile, "  embedding order (p) = %i\n", MRIC->p);

  fprintf(outfile, "  c = ");
  for (i = 0; i < MRIC->stages; i++)
    fprintf(outfile, "%"RSYM"  ", MRIC->c[i]);
  fprintf(outfile, "\n");

  if (MRIC->W) {
    for (k = 0; k < MRIC->nmat; k++) {
      fprintf(outfile, "  W[%i] = \n", k);
      for (i = 0; i < MRIC->stages; i++){
        fprintf(outfile, "      ");
        for (j = 0; j < MRIC->stages; j++)
          fprintf(outfile, "%"RSYMW"  ", MRIC->W[k][i][j]);
        fprintf(outfile, "\n");
      }
      fprintf(outfile, "\n");
    }
  }

  if (MRIC->G) {
    for (k = 0; k < MRIC->nmat; k++) {
      fprintf(outfile, "  G[%i] = \n", k);
      for (i = 0; i < MRIC->stages; i++) {
        fprintf(outfile, "      ");
        for (j = 0; j < MRIC->stages; j++)
          fprintf(outfile, "%"RSYMW"  ", MRIC->G[k][i][j]);
        fprintf(outfile, "\n");
      }
      fprintf(outfile, "\n");
    }
  }
}


/* ===========================================================================
 * Private Functions
 * ===========================================================================*/


/* ---------------------------------------------------------------------------
 * Stage type identifier: returns one of the constants
 *
 * MRISTAGE_ERK_FAST    -- standard MIS-like stage
 * MRISTAGE_ERK_NOFAST  -- standard ERK stage
 * MRISTAGE_DIRK_NOFAST -- standard DIRK stage
 * MRISTAGE_DIRK_FAST   -- coupled DIRK + MIS-like stage
 *
 * for each nontrivial stage in an MRI-like method. Otherwise (i.e., stage is
 * not in [1,MRIC->stages-1]), returns ARK_INVALID_TABLE (<0).
 *
 * The stage type is determined by 2 factors:
 * (a) Sum |MRIC->G[:][is][is]| (nonzero => DIRK)
 * (b) MRIC->c[is] - MRIC->c[is-1]  (nonzero => fast)
 * ---------------------------------------------------------------------------*/

int mriStepCoupling_GetStageType(MRIStepCoupling MRIC, int is)
{
  int i;
  realtype Gabs, cdiff;
  const realtype tol = RCONST(100.0) * UNIT_ROUNDOFF;

  if ((is < 1) || (is >= MRIC->stages)) return ARK_INVALID_TABLE;

  /* sum of stage diagonal entries across implicit tables */
  Gabs = ZERO;
  if (MRIC->G)
    for (i = 0; i < MRIC->nmat; i++)
      Gabs += SUNRabs(MRIC->G[i][is][is]);

  /* abscissae difference */
  cdiff = MRIC->c[is] - MRIC->c[is-1];

  if (Gabs > tol) {     /* DIRK */
    if (cdiff > tol) {  /* Fast */
      return(MRISTAGE_DIRK_FAST);
    } else {
      return(MRISTAGE_DIRK_NOFAST);
    }
  } else {              /* ERK */
    if (cdiff > tol) {  /* Fast */
      return(MRISTAGE_ERK_FAST);
    } else {
      return(MRISTAGE_ERK_NOFAST);
    }
  }
}


/* ---------------------------------------------------------------------------
 * Computes the stage RHS vector storage maps. With repeated abscissae the
 * first stage of the pair generally corresponds to a column of zeros and so
 * does not need to be computed and stored. The stage_map indicates if the RHS
 * needs to be computed and where to store it i.e., stage_map[i] > -1.
 * ---------------------------------------------------------------------------*/

int mriStepCoupling_GetStageMap(MRIStepCoupling MRIC,
                                int* stage_map,
                                int* nstages_active)
{
  int i, j, k, idx;
  realtype Wsum, Gsum;
  const realtype tol = RCONST(100.0) * UNIT_ROUNDOFF;

  /* ----------------------
   * Check for valid inputs
   * ---------------------- */

  if (!MRIC) return(ARK_ILL_INPUT);
  if (!(MRIC->W) && !(MRIC->G)) return(ARK_ILL_INPUT);
  if (!stage_map || !nstages_active) return(ARK_ILL_INPUT);

  /* -------------------
   * Compute storage map
   * ------------------- */

  /* Number of stage RHS vectors active */
  *nstages_active = 0;

  /* Initial storage index */
  idx = 0;

  /* Check if a stage corresponds to a column of zeros for all coupling
   * matrices by computing the column sums */
  for (j = 0; j < MRIC->stages; j++) {

    Wsum = ZERO;
    Gsum = ZERO;

    if (MRIC->W)
      for (k = 0; k < MRIC->nmat; k++)
        for (i = 0; i < MRIC->stages; i++)
          Wsum += SUNRabs(MRIC->W[k][i][j]);

    if (MRIC->G)
      for (k = 0; k < MRIC->nmat; k++)
        for (i = 0; i < MRIC->stages; i++)
          Gsum += SUNRabs(MRIC->G[k][i][j]);

    if (Wsum > tol || Gsum > tol) {
      stage_map[j] = idx;
      idx++;
    } else {
      stage_map[j] = -1;
    }
  }

  /* Check and set number of stage RHS vectors active */
  if (idx < 1) return(ARK_ILL_INPUT);

  *nstages_active = idx;

  return(ARK_SUCCESS);
}


/*===============================================================
  EOF
  ===============================================================*/
