function  nma_HW2_math_228B_problem3()
% solves the FitzHugh-Nagumo on unit square

% Math 228B, UC Davis, winter 2011.
%
% This function solves the FitzHugh-Nagumo equations on unit
% square
% V_t = D ( V_xx+V_yy) + (a-V)(v-1)v - w + I
% w_t = epsilon(V- gamma*w)
%
% subject to Nuemann homogenous boundary conditions.
% The solver splitting method to solve V_t = D ( V_xx+V_yy)
% and then solve V_t = (a-V)(v-1)v - w + I. The ADI scheme is
% used to solve V_t = D ( V_xx+V_yy), backward Euler is used to
% solved the reaction equation V_t = (a-V)(v-1)v - w + I.
% To solve w_t = epsilon(V- gamma*w), forward Euler is used.
%
% See my HW2 report for more information.
% by Nasser M. Abbasi  2/14/2011

%PARAMETERS
a       = 0.1;
gamma   = 2;
epsilon = 0.005;
D       = 5*10^-5;  %Diffusion constant
h       = 0.01;     %space step
I       = 0; %zero current
MODE    = 1;  %set this to 2 to plot the W state variable. 1 to plot V .
GENERATE_GIF=false; %turn to true to generate animated gif

% Select the part(b) or part(c) initial data by commenting
% the part needed

% part(b)
ic_v = @(X,Y) exp(-100*(X.^2+Y.^2));
ic_w = @(X,Y) zeros(size(X));

%part(c)
%ic_v = @(X,Y) 1-2*X;
%ic_w = @(X,Y) 0.05*Y;

% Allocate mesh for plotting
[X,Y] = meshgrid(h/2:h:1-h/2,h/2:h:1-h/2);

% generate initial data for w and voltage V
w = ic_w(X,Y);
v = ic_v(X,Y);

% These below are for plotting only to find good limit for z axis
maxv = 1.5*max(max(v));
minv = -0.5+min(min(v));

% Select time step. Use same as space step to keep second order
% since splitting method is used.
delt = h;

% Select a tolerance to use for Newton root solver to use for solving
% the reaction ODE
tolerance = 10^-7;

% Find number of steps to run the simulation based on number of seconds
number_of_seconds = 300;  % use 600 for part (d)
number_of_steps   = round(number_of_seconds/delt);

% Generate the A,B Matrices to use for ADI solver. THis is done once
% and not changed, since the same grid size is used for each step.
[A,A_rhs]= ...
   nma_generate_A_and_ARHS_for_2D_diffusion_Neumman(round(1/h),D,delt,h);

% Get ready to start main loop. Initialize counter and set the time step
stepA        = delt;
step_number  = 0;
fig_handle   = figure();
set(fig_handle,'Name',....
   'FitzHugh-Nagumo equations simulation. HW2,Math 228B, Nasser M. Abbasi');


frame_number = 0;
[im,map]     = initialize_plot(fig_handle,step_number,X,Y,v,minv,maxv,delt,GENERATE_GIF);

while step_number<number_of_steps
   
   if MODE==1
      [im,frame_number]=update_plot(fig_handle,im,map,frame_number,...
         step_number,X,Y,v,minv,maxv,delt,GENERATE_GIF);
   else
      [im,frame_number]=update_plot(fig_handle,im,map,frame_number,...
             step_number,X,Y,w,minv,maxv,delt,GENERATE_GIF);
   end
   
   %splitting method, first step
   [v,w] = RK4_solver(v,w,I,a,delt,epsilon,gamma);
   [v,~] = nma_solve_2D_diffusion_ADI(A,A_rhs,h,delt,D,delt,v,false);
   
   %splitting method, second step
   
   [v,~] = nma_solve_2D_diffusion_ADI(A,A_rhs,h,delt,D,delt,v,false);
   [v,w] = RK4_solver(v,w,I,a,delt,epsilon,gamma);
   
   step_number = step_number + 1;
end

if MODE==1
   [im,~]=update_plot(fig_handle,im,map,frame_number,...
      step_number,X,Y,v,minv,maxv,delt,GENERATE_GIF);
else
   [im,~]=update_plot(fig_handle,im,map,frame_number,...
      step_number,X,Y,w,minv,maxv,delt,GENERATE_GIF);
end

if GENERATE_GIF
   imwrite(im,map,'HW2_problem2.gif','DelayTime',0,'LoopCount',inf)
end

end
%-------------------------------
function f_value = f(v,w,a,I)
% Used by RK4 solver to update RHS of the reaction ODE
f_value = ( (a-v).*(v-1).*v - w + I);
end

%-------------------------------
function g_value = g(v,w,epsilon,gamma)
% Used by RK4 solver to update RHS of the w ODE
g_value = epsilon*( v - gamma * w);
end

%------------------------------------
function [v,w] = RK4_solver(v,w,I,a,delt,epsilon,gamma)
% Solver for the reaction ODE coupled nonlinear ODE system
% INPUT:
%  current value of v and w
%  delt: time step
%  epsilon, gamma are the PDE parameters
%
% OUTPUT:
%  v and w at n+1 time step.
%
m1 = delt * f(v,w,a,I) ;
k1 = delt * g(v,w,epsilon,gamma);

m2 = delt * f(v+(1/2).*m1,w+(1/2).*k1,a,I);
k2 = delt * g(v+(1/2).*m1,w+(1/2).*k1,epsilon,gamma);

m3 = delt * f(v+(1/2).*m2,w+(1/2).*k2,a,I);
k3 = delt * g(v+(1/2).*m2,w+(1/2).*k2,epsilon,gamma);

m4 = delt * f(v+m3,w+k3,a,I);
k4 = delt * g(v+m3,w+k3,epsilon,gamma);

v = v + (1/6) * (m1+2*m2+2*m3+m4);
w = w + (1/6) * (k1+2*k2+2*k3+k4);

end

%----------------------------
function [im,map] = initialize_plot(h,step_number,X,Y,v,minv,maxv,delt,GENERATE_GIF)

set(0,'CurrentFigure',h);
surf(X,Y,v);
%axis tight
set(gca,'nextplot','replacechildren','visible','on')
colormap cool;
title(sprintf...
   ('FitzHug-Nagumo equations simulation.\nSolution at time = %1.3f, concentration=%3.4f',...
   step_number*delt,sum(sum(abs(v)))));
xlabel('x'); ylabel('y'); zlabel('v(x,y,t)');
view(30,60)
zlim([minv maxv]);
drawnow();
if GENERATE_GIF
   f = getframe(gcf);
   [im,map] = rgb2ind(f.cdata,256,'nodither');
   im(1,1,1,20) = 0;
else
   im=[];
   map=[];
end
end
%-------------------------
function [im,frame_number]=update_plot(h,im,map,frame_number,...
   step_number,X,Y,v,minv,maxv,delt,GENERATE_GIF)

%make plot of the solution. Change the mod() statment
%if you do not want to see each frame, to speed things up
if mod(step_number,5)==0
   set(0,'CurrentFigure',h);
   surf(X,Y,v);
   colormap cool;
   title(sprintf('FitzHug-Nagumo equations simulation.\nSolution at time = %1.3f, concentration=%3.4f',...
      step_number*delt,sum(sum(abs(v)))));
   xlabel('x'); ylabel('y'); zlabel('v(x,y,t)');
   view(30,60)
   zlim([minv maxv]);
   drawnow();
   if GENERATE_GIF
      f = getframe(gcf);
      frame_number = frame_number+1;
      im(:,:,1,frame_number)=rgb2ind(f.cdata,map,'nodither');
   else
       im=[];
   end
end

end