% Solve the Black-Scholes equation for American put option
% Parameters for determine the price of an American put
K = 50;
rf = 0.1;
sigma = 0.4;
T = 5/12;
q = 0;
% Divide S and t into N intervals
NS = 100;
Nt = 20;
Smax = 100;
dS = Smax/NS;
dt = T/Nt;
% Initial
f = zeros(NS+1,Nt+1);
% Terminal condition for t=T
for i = 1:NS+1
    f(i,end) = max(K-(i-1)*dS,0);
end
% Boundary condition for S=0
f(1,:) = max(K-0,0);
% Boundary condition for S=Smax
f(end,:) = max(K-Smax,0);
% Free boundary
% save the index of free boundary as Sf(t)
Sf = zeros(Nt+1,1);
Sf(Nt+1) = K/dS+1;
% Solve f(S,t) from time T to time 0
a = zeros(NS+1,1);
b = zeros(NS+1,1);
c = zeros(NS+1,1);
g = zeros(NS-1,1);
for j = 1:NS+1
    a(j) = (1/2)*(rf-q)*(j-1)*dt-(1/2)*(sigma^2)*((j-1)^2)*dt;
    b(j) = 1+(sigma^2)*((j-1)^2)*dt+rf*dt;
    c(j) = (-1/2)*(rf-q)*(j-1)*dt-(1/2)*(sigma^2)*((j-1)^2)*dt;
end
for t = Nt:-1:1
    % We solve the linear system for f and check the free boundary
    f_temp = zeros(NS+1,1);
    Sf_temp = Sf(t+1);
    cond = 1;
    while cond > 10E-2
        f_temp(1:Sf_temp-1) = f(1:Sf_temp-1,t+1);
        f_temp(end) = f(end,t+1);
        A = diag(b(Sf_temp:end-1),0)+diag([c(Sf_temp)+1; c(Sf_temp+1:end-2)],1)+diag(a(Sf_temp+1:end-1),-1);
        f_temp(Sf_temp:end-1) = A\[f(Sf_temp,t+1)-(a(Sf_temp)-1)*f_temp(Sf_temp-1)-2*dS; f(Sf_temp+1:end-2,t+1); f(end-1,t+1)-c(end-1)*f_temp(end)];
        cond = f_temp(Sf_temp)-f(Sf_temp,end);
        Sf_temp = Sf_temp-1;
    end
    Sf(t) = Sf_temp;
    f(:,t)=f_temp;
end
for i = 1:Nt+1
    Sf(i) = dS*(Sf(i)-1);
end
figure;
plot(Sf,[0:dt:T],'-');
xlabel('S');
ylabel('t');
xlim([0 Smax]);
ylim([0 T]);
title('The free boundary Sf(t)'); 
figure;
[X,Y] = meshgrid(0:dS:Smax,0:dt:T);
surf(X,Y,f');
xlabel('S');
ylabel('t'); 
zlabel('f(S,t)');
xlim([0 Smax]);
ylim([0 T]);
title('The solution f(S,t)');