clear;

%% Data

rng(2022);

lambda = 2;
mu = 1;
X0 = 1;

T = 1;
nSteps = 2^10;
dt = 1/nSteps;
nTrajec = 10^5;

%% Exact solution

dW = sqrt(dt)*randn(nTrajec, nSteps);
W = [zeros(nTrajec, 1), cumsum(dW, 2)];

X = X0*exp((lambda - (mu^2)/2)*T + mu*W(:, end));
X_mean = X0*exp(lambda*T); 

%% Errors

tRatio_vec = 2.^(1:5);
strongError = nan(1, length(tRatio_vec));
weakError = nan(1, length(tRatio_vec));

for i = 1:length(tRatio_vec)
    
    tRatio = tRatio_vec(i);
    
    DeltaT = tRatio*dt;
    dW_EM = W(:, 1+tRatio:tRatio:end) - W(:, 1:tRatio:end-tRatio);
    X_EM = X0*prod(1 + lambda * DeltaT + mu * dW_EM, 2);
    
    strongError(i) = mean(abs(X_EM - X));
    weakError(i) = abs(mean(X_EM) - X_mean);

end

%% Figures

figure;
loglog(dt*tRatio_vec, strongError, 's:k', 'Linewidth', 2);
hold on;
loglog(dt*tRatio_vec, (dt*tRatio_vec).^(1/2), 'k--', 'Linewidth', 2);
xlabel('{\Delta}t', 'Fontsize', 20);
legend('Strong error', 'Order 1/2', 'Location', 'best', 'Fontsize', 20);
set(gca, 'FontSize', 20);
saveas(gcf, 'figures/exercise3.png');

figure;
loglog(dt*tRatio_vec, weakError, 's:k', 'Linewidth', 2);
hold on;
loglog(dt*tRatio_vec, dt*tRatio_vec, 'k--', 'Linewidth', 2);
xlabel('{\Delta}t', 'Fontsize', 20);
legend('Weak error', 'Order 1', 'Location', 'best', 'Fontsize', 20);
set(gca, 'FontSize', 20);
%saveas(gcf, 'figures/exercise4.png');