clear; close all;

% this code solves the diffusion equation using Crank-Nicholson

N=61;
N2=ceil(N/2); % for i and j at N2 and beyond T=0
L=1.74/2;
hx=2*L/(N-1);
hy=2*hx;
x=0:hx:2*L;
y=2*x;

% define initial conditions
sigma = .01;%initial distribution witdth
beampos = .3;%distance from the collecting electrode in mm
D=.275;%diffusivity
tfinal=18;%simulation time
cutpos_bot_i = floor(.25/hx);%position of the film discontinuity
cutpos_top_i = floor(.25/hx)+1;
cutpos_left_j = find(y>.3,1);
cutpos_right_j = find(y>1.75,1);

for i=1:N
    for j=1:N
        T(i,j)=1e5*exp(-((x(i)-beampos)^2)/(2*sigma^2))*exp(-((y(j)-1)^2)/(2*sigma^2));
    end
end

% set the mask (area to ignore for diffusion)
mask=zeros(N);
mask(cutpos_bot_i:cutpos_top_i,cutpos_left_j:cutpos_right_j)=1;

%Calculate the timestep to use
tau = .8*hx^2/max(D);

% set the number of time steps to take
nsteps=tfinal/tau;

% define a useful constant
coeffx=tau/hx^2;

% load the tridiagonal matrix A with the left side of Crank-Nicholson
% operator and B-, B0, and B+, assuming that the x and y grids are
% identical in size
Ax=zeros(N);
Bmx=D*.5*coeffx;
B0x=1-D*coeffx;
Bpx=Bmx;

% put the outer boundary conditions in for x:
Ax(1,1)=1;
Ax(N,N)=-1.5/hx;Ax(N,N-1)=2/hx;Ax(N,N-2)=-.5/hx;

for j=2:N-1
    Ax(j,j-1) = -D*.5*coeffx;
    Ax(j,j) = (1+D*coeffx);
    Ax(j,j+1) = -D*.5*coeffx;
end

% put the outer bc's in for y
Ay=zeros(N);
Bmy=D*.5*coeffx;
B0y=1-D*coeffx;
Bpy=Bmy;

% put the outer boundary conditions in for y:
Ay(1,1)=-1.5/hy;Ay(1,2)=2/hy;Ay(1,3)=-.5/hy;
Ay(N,N)=-1.5/hy;Ay(N,N-1)=2/hy;Ay(N,N-2)=-.5/hy;

for j=2:N-1
    Ay(j,j-1) = -D*.5*coeffx;
    Ay(j,j) = (1+D*coeffx);
    Ay(j,j+1) = -D*.5*coeffx;
end

% this is the time advance loop

% define the right-hand side vector.
r=zeros(N,1);

for m=1:nsteps
    
    % Crank-Nicholson advance
    % first do the fully implicit step in x by
    % doing a solve in i(x) at each j(y)
    
    for j=2:N-1 % start of j-loop
        
        % load the right-hand side vector and the
        % mask-modified operator Ap
        Ap=Ax;
        for i=2:N-1
            if mask(i,j) == 1
                r(i)=0;
                if i==cutpos_bot_i
                    Ap(i,i-2)=-.5/hx;
                    Ap(i,i-1)=2/hx;
                    Ap(i,i)=-1.5/hx;
                    Ap(i,i+1)=0;
                elseif i==cutpos_top_i
                    Ap(i,i+2)=-.5/hx;
                    Ap(i,i+1)=2/hx;
                    Ap(i,i)=-1.5/hx;
                    Ap(i,i-1)=0;
                end
            else
                r(i)=Bmx*T(i,j-1)+B0x*T(i,j)+Bpx*T(i,j+1);
            end
        end
        
        % set the boundary conditions in i
        r(1)=0;
        r(N)=0;
        
        
        % do the solve to find Thalf
        tmp=Ap\r;
        
        Thalf(:,j)=tmp;
        
    end % end of j-loop
    
    % set the boundary conditions in j
    Thalf(:,1)=4/3*Thalf(:,2)-1/3*Thalf(:,3);
    Thalf(:,N)=4/3*Thalf(:,N-1)-1/3*Thalf(:,N-2);
    
    % now do the second implicit step in y by
    % doing a solve in j(y) at each i(x)
    
    
    for i=2:N-1 % start i-loop
        
        % load r and Ap, the mask modified A
        Ap=Ay;
        for j=2:N-1
            if mask(i,j) == 1
                r(j)=0.;
                if j==cutpos_right_j
                    Ap(j,j-2)=-.5/hy;
                    Ap(j,j-1)=2/hy;
                    Ap(j,j)=-1.5/hy;
                    Ap(j,j+1)=0;
                elseif j==cutpos_left_j
                    Ap(j,j+2)=-.5/hy;
                    Ap(j,j+1)=2/hy;
                    Ap(j,j)=-1.5/hy;
                    Ap(j,j-1)=0;
                end
            else
                r(j)=Bmy*Thalf(i-1,j)+B0y*Thalf(i,j)+Bpy*Thalf(i+1,j);
            end
            
        end
        r(1)=0;r(N)=0; % set the boundary conditions in j
        % for the masked region
        
        % do the final solve
        tmp=Ap\r;
        
        T(i,:)=tmp';
        
    end % end of i-loop
    
    % set the boundary conditions in i
    T(1,:)=0;
    T(N,:)=4/3*T(N-1,:)-1/3*T(N-2,:);
    
    t=m*tau; % set the time
    if mod(m,5)==0
        surf(y,x,T)
        title(num2str(m*tau))
        %axis equal
        axis([0 4*L 0 2*L 0 max([max(T) 50])])
        %view(2)%view(41,41)
        drawnow
    end
    current(m) = sum(T(2,:))/hx;
end