// model for mutation-advection-selection, Chisholm, Lorenzi et al. 2015
// with NIPG
// IMEX scheme: modified Crank-Nicholson-Adams-Bashforth

// domain
int n = 40; // size
mesh Th = square(n,n,flags=1);
Th = adaptmesh(Th, hmax=1./60.);
// FE space - discontinuous Galerkin
fespace Vh(Th,P2dc);
Vh u,v, uold,uold1;

// model parameters
// drug dose
real m = 0.05,vvv=1.2e-2;
// function beta(x,y)
func bb = .03 + .25*y+.05*(1-x);
func dd = m*(.15 + 1.3*(1-x)) + .0;

func vv = -vvv*m*(x<=.4); // adv. field

// define operators: Gradient
macro Grad(u) [ dx(u),dy(u) ] //
// normal times adv. field
macro n() (N.y*vv) //
macro uup(u) (mean(u)-(n<0)*jump(u)/2) //


// parameters
real alpha = 5e-6;
real I0,I1;
// integration, time-step
real dt = Th.hmin/(m*vvv), t=0.; // CFL condition
dt = fmin(1,.1*dt);

cout << "dt=" << dt << endl;

// weak formulation ======================
// =======================================
//
// computing stiffness matrix & vector
// weak formulation parameters: for DG penalisation terms
real sigma0 = 1.2, epsilon = 1, b0 = 1;
// variational formulations
varf a1(u,v) = int2d(Th)(u*v/dt); // time derivative

//  Laplace operator (Riviere ch. 4)
varf a2(u,v)= int2d(Th)( alpha*(dx(u)*dx(v) + dy(u)*dy(v)) )
	- intalledges(Th) ( (1-nTonEdge)*sigma0/(lenEdge^b0) *jump(u)*jump(v) )
	+intalledges(Th)( (1-nTonEdge)*( alpha* mean( N.x*dx(u)+N.y*dy(u) ) )*jump(v) )
	- intalledges(Th) ( epsilon*((1-nTonEdge)*( alpha* mean( N.x*dx(v)+N.y*dy(v) )  )*jump(u)) ) ;

// advection operator (Riviere ch. 4)
varf b(u,v) = int2d(Th)( (vv*dy(v) )*u ) 
	+ intalledges(Th) ( (1-nTonEdge) * n *uup(u)*jump(v) ) 
	- int1d(Th) ( (n>0)*n*u*v );
// reaction terms
varf ll1(u,v) = int2d(Th) (bb *u*v) ;
varf ll2(u,v) = int2d(Th) (dd *u*v) ;

// build stiffness matrices:  bilinear part
matrix A1 =a1(Vh,Vh),A2=a2(Vh,Vh);
// build load vectors: linear part
matrix L1 = ll1(Vh,Vh), L2 = ll2(Vh,Vh), BB = b(Vh,Vh);


Vh rhs,rhs0 ,rhs1, err;
int iter =700;
//real tol = 1.e-4, dtol=1;

// initial data (centering at (x0,y0)
real x0=.2,y0=.8,c1=200.;
uold = c1*exp(-(x-x0)^2/(2*.01^2))*exp(-(y-y0)^2/(2*.01^2))*((x-x0)^2+(y-y0)^2<.1^2); uold=abs(uold); // u{1}
//
// need an auto-starting method for time-integration
// find u{2} by IMEX-theta-scheme
// compute integral
I0 = int2d(Th)(uold);
matrix A = A1+2./3.*A2;
set (A,init=1,solver=UMFPACK);
rhs[] = A1*uold[];
rhs1[] = A2*uold[]; rhs1 = -1./3.*rhs1;
rhs[] += rhs1[];
rhs1[] = L1*uold[];
rhs1 = (1-I0)*rhs1;
rhs[] += rhs1[];
rhs1[]= L2*uold[]; 
rhs[] -=rhs1[];
rhs1[]= BB*uold[]; 
rhs[] +=rhs1[];
u[]=A^-1*rhs[];

A = A1+9./16.*A2;
// initialise solver
set (A,init=1,solver=UMFPACK);
// loop
for (int citer = 0; citer <= iter; citer++)
{
	t +=dt;
// initialise RHS of the variational problem
// compute integrals
	I0 = int2d(Th)(u);
	I1 = int2d(Th)(uold);
//	cout << "Int(u) = "<<I0<< endl;
// add summands
	rhs[]= A1*u[];
	rhs1[] = L1*uold[];
	rhs1 = -.5*(1-I1)*rhs1; // u{n-1}
	rhs[] +=rhs1[];
	rhs0[] = L1*u[];
	rhs0 = 1.5*(1-I0)*rhs0; // u{n}
	rhs[] +=rhs0[];
	rhs0[]= L2*u[]; rhs0 = -1.5*rhs0; // u{n}
	rhs[] +=rhs0[];
	rhs1[]= L2*uold[]; rhs1 = .5*rhs1; // u{n-1}
	rhs[] +=rhs1[];
	rhs0[]= BB*u[]; rhs0= 1.5*rhs0;
	rhs[] +=rhs0[];
	rhs1[]= BB*uold[]; rhs1 =-.5*rhs1;
	rhs[] +=rhs1[];
	rhs0[]=A2*u[]; rhs0 = -3./8.*rhs0; // u{n}
	rhs[]+=rhs0[];
	rhs1[]=A2*uold[]; rhs1 = -1./16.*rhs1; // u{n-1}
	rhs[]+=rhs1[];
// solve explicit part
	uold = u;
	u[]=A^-1*rhs[];
//	err[] = u[] - uold[];
	plot(u,value=1,fill=1,cmm=" t="+t );
//	dtol = err[].linfty;
//	cout << "dtol = "<<dtol<< endl;
}
plot(u,value=1,cmm=" t="+t );