/*
       Formatted test for TS routines.

          Solves U_t=F(t,u)
          Where:

                  [2*u1+u2
          F(t,u)= [u1+2*u2+u3
                  [   u2+2*u3
       We can compare the solutions from euler, beuler and SUNDIALS to
       see what is the difference.

*/

static char help[] = "Solves a linear ODE. \n\n";

#include <petscts.h>
#include <petscpc.h>

extern PetscErrorCode RHSFunction(TS, PetscReal, Vec, Vec, void *);
extern PetscErrorCode RHSJacobian(TS, PetscReal, Vec, Mat, Mat, void *);
extern PetscErrorCode Monitor(TS, PetscInt, PetscReal, Vec, void *);
extern PetscErrorCode Initial(Vec, void *);
extern PetscErrorCode MyMatMult(Mat, Vec, Vec);

extern PetscReal solx(PetscReal);
extern PetscReal soly(PetscReal);
extern PetscReal solz(PetscReal);

int main(int argc, char **argv) {
  PetscInt  time_steps = 100, steps;
  Vec       global;
  PetscReal dt, ftime;
  TS        ts;
  Mat       A = 0, S;

  PetscFunctionBeginUser;
  PetscCall(PetscInitialize(&argc, &argv, (char *)0, help));
  PetscCall(PetscOptionsGetInt(NULL, NULL, "-time", &time_steps, NULL));

  /* set initial conditions */
  PetscCall(VecCreate(PETSC_COMM_WORLD, &global));
  PetscCall(VecSetSizes(global, PETSC_DECIDE, 3));
  PetscCall(VecSetFromOptions(global));
  PetscCall(Initial(global, NULL));

  /* make timestep context */
  PetscCall(TSCreate(PETSC_COMM_WORLD, &ts));
  PetscCall(TSSetProblemType(ts, TS_NONLINEAR));
  PetscCall(TSMonitorSet(ts, Monitor, NULL, NULL));
  dt = 0.001;

  /*
    The user provides the RHS and Jacobian
  */
  PetscCall(TSSetRHSFunction(ts, NULL, RHSFunction, NULL));
  PetscCall(MatCreate(PETSC_COMM_WORLD, &A));
  PetscCall(MatSetSizes(A, PETSC_DECIDE, PETSC_DECIDE, 3, 3));
  PetscCall(MatSetFromOptions(A));
  PetscCall(MatSetUp(A));
  PetscCall(RHSJacobian(ts, 0.0, global, A, A, NULL));
  PetscCall(TSSetRHSJacobian(ts, A, A, RHSJacobian, NULL));

  PetscCall(MatCreateShell(PETSC_COMM_WORLD, 3, 3, 3, 3, NULL, &S));
  PetscCall(MatShellSetOperation(S, MATOP_MULT, (void (*)(void))MyMatMult));
  PetscCall(TSSetRHSJacobian(ts, S, A, RHSJacobian, NULL));

  PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP));
  PetscCall(TSSetFromOptions(ts));

  PetscCall(TSSetTimeStep(ts, dt));
  PetscCall(TSSetMaxSteps(ts, time_steps));
  PetscCall(TSSetMaxTime(ts, 1));
  PetscCall(TSSetSolution(ts, global));

  PetscCall(TSSolve(ts, global));
  PetscCall(TSGetSolveTime(ts, &ftime));
  PetscCall(TSGetStepNumber(ts, &steps));

  /* free the memories */

  PetscCall(TSDestroy(&ts));
  PetscCall(VecDestroy(&global));
  PetscCall(MatDestroy(&A));
  PetscCall(MatDestroy(&S));

  PetscCall(PetscFinalize());
  return 0;
}

PetscErrorCode MyMatMult(Mat S, Vec x, Vec y) {
  const PetscScalar *inptr;
  PetscScalar       *outptr;

  PetscFunctionBeginUser;
  PetscCall(VecGetArrayRead(x, &inptr));
  PetscCall(VecGetArrayWrite(y, &outptr));

  outptr[0] = 2.0 * inptr[0] + inptr[1];
  outptr[1] = inptr[0] + 2.0 * inptr[1] + inptr[2];
  outptr[2] = inptr[1] + 2.0 * inptr[2];

  PetscCall(VecRestoreArrayRead(x, &inptr));
  PetscCall(VecRestoreArrayWrite(y, &outptr));
  PetscFunctionReturn(0);
}

/* this test problem has initial values (1,1,1).                      */
PetscErrorCode Initial(Vec global, void *ctx) {
  PetscScalar *localptr;
  PetscInt     i, mybase, myend, locsize;

  /* determine starting point of each processor */
  PetscCall(VecGetOwnershipRange(global, &mybase, &myend));
  PetscCall(VecGetLocalSize(global, &locsize));

  /* Initialize the array */
  PetscCall(VecGetArrayWrite(global, &localptr));
  for (i = 0; i < locsize; i++) localptr[i] = 1.0;

  if (mybase == 0) localptr[0] = 1.0;

  PetscCall(VecRestoreArrayWrite(global, &localptr));
  return 0;
}

PetscErrorCode Monitor(TS ts, PetscInt step, PetscReal time, Vec global, void *ctx) {
  VecScatter         scatter;
  IS                 from, to;
  PetscInt           i, n, *idx;
  Vec                tmp_vec;
  const PetscScalar *tmp;

  /* Get the size of the vector */
  PetscCall(VecGetSize(global, &n));

  /* Set the index sets */
  PetscCall(PetscMalloc1(n, &idx));
  for (i = 0; i < n; i++) idx[i] = i;

  /* Create local sequential vectors */
  PetscCall(VecCreateSeq(PETSC_COMM_SELF, n, &tmp_vec));

  /* Create scatter context */
  PetscCall(ISCreateGeneral(PETSC_COMM_SELF, n, idx, PETSC_COPY_VALUES, &from));
  PetscCall(ISCreateGeneral(PETSC_COMM_SELF, n, idx, PETSC_COPY_VALUES, &to));
  PetscCall(VecScatterCreate(global, from, tmp_vec, to, &scatter));
  PetscCall(VecScatterBegin(scatter, global, tmp_vec, INSERT_VALUES, SCATTER_FORWARD));
  PetscCall(VecScatterEnd(scatter, global, tmp_vec, INSERT_VALUES, SCATTER_FORWARD));

  PetscCall(VecGetArrayRead(tmp_vec, &tmp));
  PetscCall(PetscPrintf(PETSC_COMM_WORLD, "At t =%14.6e u = %14.6e  %14.6e  %14.6e \n", (double)time, (double)PetscRealPart(tmp[0]), (double)PetscRealPart(tmp[1]), (double)PetscRealPart(tmp[2])));
  PetscCall(PetscPrintf(PETSC_COMM_WORLD, "At t =%14.6e errors = %14.6e  %14.6e  %14.6e \n", (double)time, (double)PetscRealPart(tmp[0] - solx(time)), (double)PetscRealPart(tmp[1] - soly(time)), (double)PetscRealPart(tmp[2] - solz(time))));
  PetscCall(VecRestoreArrayRead(tmp_vec, &tmp));
  PetscCall(VecScatterDestroy(&scatter));
  PetscCall(ISDestroy(&from));
  PetscCall(ISDestroy(&to));
  PetscCall(PetscFree(idx));
  PetscCall(VecDestroy(&tmp_vec));
  return 0;
}

PetscErrorCode RHSFunction(TS ts, PetscReal t, Vec globalin, Vec globalout, void *ctx) {
  PetscScalar       *outptr;
  const PetscScalar *inptr;
  PetscInt           i, n, *idx;
  IS                 from, to;
  VecScatter         scatter;
  Vec                tmp_in, tmp_out;

  /* Get the length of parallel vector */
  PetscCall(VecGetSize(globalin, &n));

  /* Set the index sets */
  PetscCall(PetscMalloc1(n, &idx));
  for (i = 0; i < n; i++) idx[i] = i;

  /* Create local sequential vectors */
  PetscCall(VecCreateSeq(PETSC_COMM_SELF, n, &tmp_in));
  PetscCall(VecDuplicate(tmp_in, &tmp_out));

  /* Create scatter context */
  PetscCall(ISCreateGeneral(PETSC_COMM_SELF, n, idx, PETSC_COPY_VALUES, &from));
  PetscCall(ISCreateGeneral(PETSC_COMM_SELF, n, idx, PETSC_COPY_VALUES, &to));
  PetscCall(VecScatterCreate(globalin, from, tmp_in, to, &scatter));
  PetscCall(VecScatterBegin(scatter, globalin, tmp_in, INSERT_VALUES, SCATTER_FORWARD));
  PetscCall(VecScatterEnd(scatter, globalin, tmp_in, INSERT_VALUES, SCATTER_FORWARD));
  PetscCall(VecScatterDestroy(&scatter));

  /*Extract income array */
  PetscCall(VecGetArrayRead(tmp_in, &inptr));

  /* Extract outcome array*/
  PetscCall(VecGetArrayWrite(tmp_out, &outptr));

  outptr[0] = 2.0 * inptr[0] + inptr[1];
  outptr[1] = inptr[0] + 2.0 * inptr[1] + inptr[2];
  outptr[2] = inptr[1] + 2.0 * inptr[2];

  PetscCall(VecRestoreArrayRead(tmp_in, &inptr));
  PetscCall(VecRestoreArrayWrite(tmp_out, &outptr));

  PetscCall(VecScatterCreate(tmp_out, from, globalout, to, &scatter));
  PetscCall(VecScatterBegin(scatter, tmp_out, globalout, INSERT_VALUES, SCATTER_FORWARD));
  PetscCall(VecScatterEnd(scatter, tmp_out, globalout, INSERT_VALUES, SCATTER_FORWARD));

  /* Destroy idx aand scatter */
  PetscCall(ISDestroy(&from));
  PetscCall(ISDestroy(&to));
  PetscCall(VecScatterDestroy(&scatter));
  PetscCall(VecDestroy(&tmp_in));
  PetscCall(VecDestroy(&tmp_out));
  PetscCall(PetscFree(idx));
  return 0;
}

PetscErrorCode RHSJacobian(TS ts, PetscReal t, Vec x, Mat A, Mat BB, void *ctx) {
  PetscScalar        v[3];
  const PetscScalar *tmp;
  PetscInt           idx[3], i;

  idx[0] = 0;
  idx[1] = 1;
  idx[2] = 2;
  PetscCall(VecGetArrayRead(x, &tmp));

  i    = 0;
  v[0] = 2.0;
  v[1] = 1.0;
  v[2] = 0.0;
  PetscCall(MatSetValues(BB, 1, &i, 3, idx, v, INSERT_VALUES));

  i    = 1;
  v[0] = 1.0;
  v[1] = 2.0;
  v[2] = 1.0;
  PetscCall(MatSetValues(BB, 1, &i, 3, idx, v, INSERT_VALUES));

  i    = 2;
  v[0] = 0.0;
  v[1] = 1.0;
  v[2] = 2.0;
  PetscCall(MatSetValues(BB, 1, &i, 3, idx, v, INSERT_VALUES));

  PetscCall(MatAssemblyBegin(BB, MAT_FINAL_ASSEMBLY));
  PetscCall(MatAssemblyEnd(BB, MAT_FINAL_ASSEMBLY));

  if (A != BB) {
    PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
    PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
  }
  PetscCall(VecRestoreArrayRead(x, &tmp));

  return 0;
}

/*
      The exact solutions
*/
PetscReal solx(PetscReal t) {
  return PetscExpReal((2.0 - PetscSqrtReal(2.0)) * t) / 2.0 - PetscExpReal((2.0 - PetscSqrtReal(2.0)) * t) / (2.0 * PetscSqrtReal(2.0)) + PetscExpReal((2.0 + PetscSqrtReal(2.0)) * t) / 2.0 + PetscExpReal((2.0 + PetscSqrtReal(2.0)) * t) / (2.0 * PetscSqrtReal(2.0));
}

PetscReal soly(PetscReal t) {
  return PetscExpReal((2.0 - PetscSqrtReal(2.0)) * t) / 2.0 - PetscExpReal((2.0 - PetscSqrtReal(2.0)) * t) / PetscSqrtReal(2.0) + PetscExpReal((2.0 + PetscSqrtReal(2.0)) * t) / 2.0 + PetscExpReal((2.0 + PetscSqrtReal(2.0)) * t) / PetscSqrtReal(2.0);
}

PetscReal solz(PetscReal t) {
  return PetscExpReal((2.0 - PetscSqrtReal(2.0)) * t) / 2.0 - PetscExpReal((2.0 - PetscSqrtReal(2.0)) * t) / (2.0 * PetscSqrtReal(2.0)) + PetscExpReal((2.0 + PetscSqrtReal(2.0)) * t) / 2.0 + PetscExpReal((2.0 + PetscSqrtReal(2.0)) * t) / (2.0 * PetscSqrtReal(2.0));
}

/*TEST

    test:
      suffix: euler
      args: -ts_type euler
      requires: !single

    test:
      suffix: beuler
      args:   -ts_type beuler
      requires: !single

TEST*/