/*
 * ex_vdp.c
 *
 *  Created on: Jun 1, 2012
 *      Author: Hong Zhang
 */
static char help[] = "Solves the van der Pol equation. \n Input parameters include:\n";

/*
 * This program solves the van der Pol equation
 * y' = z                               (1)
 * z' = (((1-y^2)*z-y)/eps              (2)
 * on the domain 0<=x<=0.5, with the initial conditions
 * y(0) = 2,
 * z(0) = -2/3 + 10/81*eps - 292/2187*eps^2-1814/19683*eps^3
 * IMEX schemes are applied to the splitted equation
 * [y'] = [z]  + [0                 ]
 * [z']   [0]    [(((1-y^2)*z-y)/eps]
 *
 * F(x)= [z]
 *       [0]
 *
 * G(x)= [y'] -   [0                 ]
 *       [z']     [(((1-y^2)*z-y)/eps]
 *
 * JG(x) =  G_x + a G_xdot
 */

#include <petscdmda.h>
#include <petscts.h>

typedef struct _User *User;
struct _User {
  PetscReal mu;  /*stiffness control coefficient: epsilon*/
};

static PetscErrorCode RHSFunction(TS,PetscReal,Vec,Vec,void*);
static PetscErrorCode IFunction(TS,PetscReal,Vec,Vec,Vec,void*);
static PetscErrorCode IJacobian(TS,PetscReal,Vec,Vec,PetscReal,Mat,Mat,void*);

int main(int argc, char **argv)
{
  TS                ts;
  Vec               x; /*solution vector*/
  Mat               A; /*Jacobian*/
  PetscInt          steps,mx,eimex_rowcol[2],two;
  PetscErrorCode    ierr;
  PetscScalar       *x_ptr;
  PetscReal         ftime,dt,norm;
  Vec               ref;
  struct _User      user;       /* user-defined work context */
  PetscViewer       viewer;

  PetscCall(PetscInitialize(&argc,&argv,NULL,help));
  /* Initialize user application context */
  ierr = PetscOptionsBegin(PETSC_COMM_WORLD,NULL,"van der Pol options","");PetscCall(ierr);
  user.mu      = 1e0;
  PetscCall(PetscOptionsReal("-eps","Stiffness controller","",user.mu,&user.mu,NULL));
  ierr = PetscOptionsEnd();PetscCall(ierr);

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
   Set runtime options
   - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  /*
   PetscCall(PetscOptionsGetBool(NULL,NULL,"-monitor",&monitor,NULL));
   */

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
   Create necessary matrix and vectors, solve same ODE on every process
   - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  PetscCall(MatCreate(PETSC_COMM_WORLD,&A));
  PetscCall(MatSetSizes(A,PETSC_DECIDE,PETSC_DECIDE,2,2));
  PetscCall(MatSetFromOptions(A));
  PetscCall(MatSetUp(A));
  PetscCall(MatCreateVecs(A,&x,NULL));

  PetscCall(MatCreateVecs(A,&ref,NULL));
  PetscCall(VecGetArray(ref,&x_ptr));
  /*
   * [0,1], mu=10^-3
   */
  /*
   x_ptr[0] = -1.8881254106283;
   x_ptr[1] =  0.7359074233370;*/

  /*
   * [0,0.5],mu=10^-3
   */
  /*
   x_ptr[0] = 1.596980778659137;
   x_ptr[1] = -1.029103015879544;
   */
  /*
   * [0,0.5],mu=1
   */
  x_ptr[0] = 1.619084329683235;
  x_ptr[1] = -0.803530465176385;

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
   Create timestepping solver context
   - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  PetscCall(TSCreate(PETSC_COMM_WORLD,&ts));
  PetscCall(TSSetType(ts,TSEIMEX));
  PetscCall(TSSetRHSFunction(ts,NULL,RHSFunction,&user));
  PetscCall(TSSetIFunction(ts,NULL,IFunction,&user));
  PetscCall(TSSetIJacobian(ts,A,A,IJacobian,&user));

  dt    = 0.00001;
  ftime = 1.1;
  PetscCall(TSSetTimeStep(ts,dt));
  PetscCall(TSSetMaxTime(ts,ftime));
  PetscCall(TSSetExactFinalTime(ts,TS_EXACTFINALTIME_STEPOVER));
  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
   Set initial conditions
   - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  PetscCall(VecGetArray(x,&x_ptr));
  x_ptr[0] = 2.;
  x_ptr[1] = -2./3. + 10./81.*(user.mu) - 292./2187.* (user.mu) * (user.mu)
    -1814./19683.*(user.mu)*(user.mu)*(user.mu);
  PetscCall(TSSetSolution(ts,x));
  PetscCall(VecGetSize(x,&mx));

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
   Set runtime options
   - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  PetscCall(TSSetFromOptions(ts));

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
   Solve nonlinear system
   - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  PetscCall(TSSolve(ts,x));
  PetscCall(TSGetTime(ts,&ftime));
  PetscCall(TSGetStepNumber(ts,&steps));

  PetscCall(VecAXPY(x,-1.0,ref));
  PetscCall(VecNorm(x,NORM_2,&norm));
  PetscCall(TSGetTimeStep(ts,&dt));

  eimex_rowcol[0] = 0; eimex_rowcol[1] = 0; two = 2;
  PetscCall(PetscOptionsGetIntArray(NULL,NULL,"-ts_eimex_row_col",eimex_rowcol,&two,NULL));
  PetscCall(PetscPrintf(PETSC_COMM_WORLD,"order %11s %18s %37s\n","dt","norm","final solution components 0 and 1"));
  PetscCall(VecGetArray(x,&x_ptr));
  PetscCall(PetscPrintf(PETSC_COMM_WORLD,"(%D,%D) %10.8f %18.15f %18.15f %18.15f\n",eimex_rowcol[0],eimex_rowcol[1],(double)dt,(double)norm,(double)PetscRealPart(x_ptr[0]),(double)PetscRealPart(x_ptr[1])));
  PetscCall(VecRestoreArray(x,&x_ptr));

  /* Write line in convergence log */
  PetscCall(PetscViewerCreate(PETSC_COMM_WORLD,&viewer));
  PetscCall(PetscViewerSetType(viewer,PETSCVIEWERASCII));
  PetscCall(PetscViewerFileSetMode(viewer,FILE_MODE_APPEND));
  PetscCall(PetscViewerFileSetName(viewer,"eimex_nonstiff_vdp.txt"));
  PetscCall(PetscViewerASCIIPrintf(viewer,"%D %D %10.8f %18.15f\n",eimex_rowcol[0],eimex_rowcol[1],(double)dt,(double)norm));
  PetscCall(PetscViewerDestroy(&viewer));

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
   Free work space.
   - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  PetscCall(MatDestroy(&A));
  PetscCall(VecDestroy(&x));
  PetscCall(VecDestroy(&ref));
  PetscCall(TSDestroy(&ts));
  PetscCall(PetscFinalize());
  return 0;
}

static PetscErrorCode RHSFunction(TS ts,PetscReal t,Vec X,Vec F,void *ptr)
{
  PetscScalar       *f;
  const PetscScalar *x;

  PetscFunctionBegin;
  PetscCall(VecGetArrayRead(X,&x));
  PetscCall(VecGetArray(F,&f));
  f[0] = x[1];
  f[1] = 0.0;
  PetscCall(VecRestoreArrayRead(X,&x));
  PetscCall(VecRestoreArray(F,&f));
  PetscFunctionReturn(0);
}

static PetscErrorCode IFunction(TS ts,PetscReal t,Vec X,Vec Xdot,Vec F,void *ptr)
{
  User              user = (User)ptr;
  PetscScalar       *f;
  const PetscScalar *x,*xdot;

  PetscFunctionBegin;
  PetscCall(VecGetArrayRead(X,&x));
  PetscCall(VecGetArrayRead(Xdot,&xdot));
  PetscCall(VecGetArray(F,&f));
  f[0] = xdot[0];
  f[1] = xdot[1]-((1.-x[0]*x[0])*x[1]-x[0])/user->mu;
  PetscCall(VecRestoreArrayRead(X,&x));
  PetscCall(VecRestoreArrayRead(Xdot,&xdot));
  PetscCall(VecRestoreArray(F,&f));
  PetscFunctionReturn(0);
}

static PetscErrorCode IJacobian(TS  ts,PetscReal t,Vec X,Vec Xdot,PetscReal a,Mat A,Mat B,void *ptr)
{
  User              user = (User)ptr;
  PetscReal         mu = user->mu;
  PetscInt          rowcol[] = {0,1};
  PetscScalar       J[2][2];
  const PetscScalar *x;

  PetscFunctionBegin;
  PetscCall(VecGetArrayRead(X,&x));
  J[0][0] = a;
  J[0][1] = 0;
  J[1][0] = (2.*x[0]*x[1]+1.)/mu;
  J[1][1] = a - (1. - x[0]*x[0])/mu;
  PetscCall(MatSetValues(B,2,rowcol,2,rowcol,&J[0][0],INSERT_VALUES));
  PetscCall(VecRestoreArrayRead(X,&x));

  PetscCall(MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY));
  PetscCall(MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY));
  if (A != B) {
    PetscCall(MatAssemblyBegin(B,MAT_FINAL_ASSEMBLY));
    PetscCall(MatAssemblyEnd(B,MAT_FINAL_ASSEMBLY));
  }
  PetscFunctionReturn(0);
}

/*TEST

   test:
     args: -ts_type eimex -ts_adapt_type none  -pc_type lu -ts_dt 0.01 -ts_max_time 10 -ts_eimex_row_col 3,3 -ts_monitor_lg_solution
     requires: x

   test:
     suffix: adapt
     args: -ts_type eimex -ts_adapt_type none  -pc_type lu -ts_dt 0.01 -ts_max_time 10 -ts_eimex_order_adapt -ts_eimex_max_rows 7 -ts_monitor_lg_solution
     requires: x

   test:
     suffix: loop
     args: -ts_type eimex  -ts_adapt_type none  -pc_type lu -ts_dt {{0.005 0.001 0.0005}separate output} -ts_max_steps {{100 500 1000}separate output} -ts_eimex_row_col {{1,1 2,1 3,1 2,2 3,2 3,3}separate output}
     requires: x

 TEST*/