clear;

%% Data

rng(2022);

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

T = 1;
DeltaTs = 2.^(-10:-5);
nTrajectories = 10^4;

%% Exact solution

DeltaT_min = DeltaTs(1);
nSteps = T/DeltaT_min;

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

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

%% Errors

strongError_EM = nan(1, length(DeltaTs));
weakError_EM = nan(1, length(DeltaTs));

strongError_M = nan(1, length(DeltaTs));
weakError_M = nan(1, length(DeltaTs));

for i = 1:length(DeltaTs)
    
    DeltaT = DeltaTs(i);
    ratio = DeltaT/DeltaT_min;
    
    dW_num = W(:, 1+ratio:ratio:end) - W(:, 1:ratio:end-ratio);
    X_EM = X0*prod(1 + lambda*DeltaT + mu*dW_num, 2);
    X_M = X0*prod(1 + lambda*DeltaT + mu*dW_num + ((mu^2)/2)*(dW_num.^2 - DeltaT), 2);
    
    strongError_EM(i) = mean(abs(X_EM - X));
    weakError_EM(i) = abs(mean(X_EM) - mean(X));
    strongError_M(i) = mean(abs(X_M - X));
    weakError_M(i) = abs(mean(X_M) - mean(X));

end

%% Figures

figure;
loglog(DeltaTs, strongError_EM, 's:', 'Linewidth', 2);
hold on;
loglog(DeltaTs, strongError_M, 's:', 'Linewidth', 2);
loglog(DeltaTs, 10*DeltaTs.^(1/2), '--', 'Linewidth', 2);
loglog(DeltaTs, 10*DeltaTs, '--', 'Linewidth', 2);
title('Strong order');
xlabel('{\Delta}t', 'Fontsize', 20);
%legend('Euler-Maruyama', 'Milstein', 'Order 1/2', 'Order 1', 'Orientation', 'horizontal', 'Fontsize', 20);
set(gca, 'FontSize', 20);
saveas(gcf, 'figures/exercise6_strong.png');

figure;
loglog(DeltaTs, weakError_EM, 's:', 'Linewidth', 2);
hold on;
loglog(DeltaTs, weakError_M, 's:', 'Linewidth', 2);
loglog(DeltaTs, 10*DeltaTs.^(1/2), '--', 'Linewidth', 2);
loglog(DeltaTs, 10*DeltaTs, '--', 'Linewidth', 2);
title('Weak order');
xlabel('{\Delta}t', 'Fontsize', 20);
%legend('Euler-Maruyama', 'Milstein', 'Order 1/2', 'Order 1', 'Location', 'best', 'Fontsize', 20);
set(gca, 'FontSize', 20);
saveas(gcf, 'figures/exercise6_weak.png');