% contract properties
T = 1;
K = 100;

% model properties
S0 = 100;
sigma = 0.2;
r = 0.02;

% branching properties
N = 100;% number of steps
dt = T / N;

lambda = 2;

% risk-neutral probabilities
q1 =   1/2*(1/lambda + (r-0.5*sigma^2)/sigma*sqrt(dt/lambda) + ((r-0.5*sigma^2)/sigma)^2*dt/lambda);
q2 = 1-1/lambda + ((r-0.5*sigma^2)/sigma)^2*dt/lambda;
q3 = 1/2*(1/lambda - (r-0.5*sigma^2)/sigma*sqrt(dt/lambda) + ((r-0.5*sigma^2)/sigma)^2*dt/lambda);

% set up tree space...
V = NaN(2*N+1, N+1);
Vtop = NaN(1,N+1);
Vbottom = NaN(1, N+1);

S = S0 * exp( sigma*sqrt(lambda*dt) * (N - [0 : 1 : 2*N] )' );
    
% payoff at maturity
V(:,end) = max( K - S(:,end) , 0);
Vtop(end) = 2*V(1,end) - V(2,end);
Vbottom(end) = 2*V(end,end) - V(end-1,end);

% store the boundary 
Bd = NaN(N+1);
Bd(end) = K;

for n = N+1 : -1 : 2
    
    V(:,n-1) = exp(-r*dt)*(   q1 * ([Vtop(n) V([1:end-1], n)']' ) ...
                            + q2 * V(:,n) ...
                            + q3 * ( [V([2:end], n)' Vbottom(n)]') );

    idx = find(V(:,n-1) < (K - S ), 1,'first');
    if ~isempty(idx)
       Bd(n-1) = S(idx); 
    end
    
    V(:,n-1) = max( V(:,n-1), K-S);
    
    Vtop(n-1) = 2*V(1,n-1) - V(2,n-1);
    Vbottom(n-1) = 2*V(end,n-1) - V(end-1,n-1);
    
end

VBS = CallPriceBS(K, T, S, sigma, r) - S + K * exp(-r*T);

close all

figure
plot(S, V(:,1), '-b', S, VBS(:,1),'xr');
axis([0 200 0 100]);
xlabel('Spot Price');
ylabel('Option Price');

figure
plot([0:N]*dt, Bd);
xlabel('Time');
ylabel('Spot Price');
legend('Exercise Bd.');