clear;

%% Data

rng(2022);

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

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

%% Errors

DeltaT_min = DeltaTs(1)/8;
nSteps_min = T/DeltaT_min;
dW = sqrt(DeltaT_min)*randn(nTrajectories, nSteps_min);
W = [zeros(nTrajectories, 1), cumsum(dW, 2)];

weakError_MT = nan(1, length(DeltaTs));
weakError_MTDF = nan(1, length(DeltaTs));

for i = 1:length(DeltaTs)
    
    h = DeltaTs(i);
    N = T/h;
    ratio = h/DeltaT_min;
    nSteps = T/h;
    
    dW_num = W(:, 1+ratio:ratio:end) - W(:, 1:ratio:end-ratio);
    X_MT = X0*prod(1 + h*lambda + mu*dW_num + (mu^2)*(dW_num.^2 - h)/2 + (lambda^2)*(h^2)/2 + lambda*mu*h*dW_num, 2);
    
    rn = randi(6, nTrajectories, N);
    xi = sqrt(3)*(rn==5) - sqrt(3)*(rn==6);
    chi = sqrt(h)*xi;
    pZ1  = 1 + lambda*h + mu*chi;
    pZ2p = 1 + lambda*h + mu*sqrt(h);
    pZ2m = 1 + lambda*h - mu*sqrt(h);
    X_MTDF = X0*prod(1 + (lambda*pZ1 + lambda)*(h/2) + (mu*pZ2p + mu*pZ2m + 2*mu)*(chi/4) + 1/4*(mu*pZ2p - mu*pZ2m)*(chi.^2 - h)/sqrt(h), 2);
    
    weakError_MT(i) = abs(mean(X_MT) - exp(lambda*T));
    weakError_MTDF(i) = abs(mean(X_MTDF) - exp(lambda*T));

end

%% Figures

figure;
loglog(DeltaTs, weakError_MT, 's:', 'Linewidth', 2);
hold on;
loglog(DeltaTs, weakError_MTDF, 's:', 'Linewidth', 2);
loglog(DeltaTs, DeltaTs.^2, '--', 'Linewidth', 2);
xlabel('{\Delta}t', 'Fontsize', 20);
legend('Milstein-Talay', 'Derivative free', 'Order 2', 'Location', 'best', 'Fontsize', 20);
set(gca, 'FontSize', 20);
saveas(gcf, 'figures/exercise3.png');