#include "petscts.h"
/*The problem is from
 * chemical kinetics, and consists of the following three rate
 * equations:         
 *    dy1/dt = -.04*y1 + 1.e4*y2*y3
 *    dy2/dt = .04*y1 - 1.e4*y2*y3 - 3.e7*(y2)^2
 *    dy3/dt = 3.e7*(y2)^2
 * on the interval from t = 0.0 to t = 4.e10, with initial
 * conditions: y1 = 1.0, y2 = y3 = 0.
 */
typedef struct {
	PetscInt m;
}AppCtx;

extern PetscErrorCode  FormJacobian(TS,PetscReal,Vec,Mat*,Mat*,MatStructure*,void*),
     FormFunction(TS,PetscReal,Vec,Vec,void*);

int main(int argc,char **argv)
{
	TS ts;
	Vec u,r;
	PetscScalar    *u_localptr;
	Mat J;
	AppCtx user;
	PetscInt its,m;
	PetscReal dt,ftime;
	PetscErrorCode ierr;
	PetscInitialize(&argc,&argv,PETSC_NULL,PETSC_NULL);
	m=3;
	dt=0.1;
  ierr = PetscOptionsGetInt(PETSC_NULL,"-m",&m,PETSC_NULL);CHKERRQ(ierr);
  user.m=m;
  u_localptr[0]=1.0;
  u_localptr[1]=0.0;
  u_localptr[2]=0.0;
  ierr = VecRestoreArray(u,&u_localptr);CHKERRQ(ierr);    
  
  ierr = VecCreateSeq(PETSC_COMM_SELF,m,&u);CHKERRQ(ierr);

  ierr = TSCreate(PETSC_COMM_SELF,&ts);CHKERRQ(ierr);
  ierr = TSSetProblemType(ts,TS_NONLINEAR);CHKERRQ(ierr);
  ierr = TSSetSolution(ts,u);CHKERRQ(ierr);
  ierr = TSSetRHSFunction(ts,FormFunction,&user);CHKERRQ(ierr);
  ierr = TSSetRHSJacobian(ts,J,J,FormJacobian,&user);CHKERRQ(ierr);

  ierr = TSSetType(ts,TSPSEUDO);CHKERRQ(ierr);
  ierr = TSSetInitialTimeStep(ts,0.0,dt);CHKERRQ(ierr);
  ierr = TSPseudoSetTimeStep(ts,TSPseudoDefaultTimeStep,0);CHKERRQ(ierr);
  ierr = TSSetFromOptions(ts);CHKERRQ(ierr);

  ierr = TSSetUp(ts);CHKERRQ(ierr);
  ierr = TSStep(ts,&its,&ftime);CHKERRQ(ierr);
  
  printf("Number of pseudo timesteps = %d final time %4.2e\n",(int)its,ftime);
  ierr = VecDestroy(u);CHKERRQ(ierr);
  ierr = VecDestroy(r);CHKERRQ(ierr);
  ierr = MatDestroy(J);CHKERRQ(ierr);
  ierr = TSDestroy(ts);CHKERRQ(ierr);
  ierr = PetscFinalize();CHKERRQ(ierr);

  return 0;
}



PetscErrorCode FormFunction(TS ts,PetscReal t,Vec X,Vec F,void *ptr)
{

	PetscErrorCode ierr;
	PetscScalar *x,*f;
	PetscInt n;
	n=3;
 	ierr = VecGetArray(X,&x);CHKERRQ(ierr);
  ierr = VecGetArray(F,&f);CHKERRQ(ierr);
	f[0]=-0.04*x[0]+1.0e4*x[1]*x[2];
	f[1]=0.04-1.0e4*x[1]*x[2]-3*1.0e7*x[1]*x[1];
	f[2]=3*1.0e7*x[1]*x[1];
  ierr = VecRestoreArray(X,&x);CHKERRQ(ierr);
  ierr = VecRestoreArray(F,&f);CHKERRQ(ierr);
  return 0;
  
}
PetscErrorCode FormJacobian(TS ts,PetscReal t,Vec X,Mat *J,Mat *B,MatStructure *flag,void *ptr)
{
	
		Mat jac=*B;
		PetscScalar v[3],*x;
		PetscInt row,col;
		PetscErrorCode ierr;
	   ierr = VecGetArray(X,&x);CHKERRQ(ierr);
	   v[0]=-0.04;
	   v[1]=0.04;
	   v[2]=0;
	   row=1;
	   col=1;
	  ierr = MatSetValues(jac,1,&row,1,&col,v,INSERT_VALUES);CHKERRQ(ierr);
   v[0]=1.0e4*x[2];
	   v[1]=-1.0e4*x[2]-6.0e7*x[1];
	   v[2]=6.0e7*x[1];
	   col=2;
	  ierr = MatSetValues(jac,1,&row,1,&col,v,INSERT_VALUES);CHKERRQ(ierr);
	   v[0]=1.0e4*x[1];
	   v[1]=-1.0e4*x[1];
	   v[2]=0;
	   col=3;
	  ierr = MatSetValues(jac,1,&row,1,&col,v,INSERT_VALUES);CHKERRQ(ierr);
	ierr = MatAssemblyBegin(jac,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
  ierr = VecRestoreArray(X,&x);CHKERRQ(ierr);
  ierr = MatAssemblyEnd(jac,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
  return 0;

}