clear;

%% Data

rng(2022);

lambda = 2;
mu = 1;

T = 1;
h = 0.01;
N = round(T/h);
tVec = 0:h:T;

X0 = 1;

ratio = 10;
h_small = h/ratio;
N_small = round(T/h_small);
tVec_small = 0:h_small:T;

M = 10000;
hs = 2.^(-11:-4);

%% Solution SDE

dW_small = sqrt(h_small)*randn(1, N_small);
W_small = [0, cumsum(dW_small)];
W = W_small(1:ratio:end);
dW = diff(W);

X1 = X0*[1, cumprod(1 + h*lambda + mu*dW)];
X2 = X0*[1, cumprod(1 + h*lambda + mu*dW + (1/2)*mu*dW.*(lambda*h + mu*dW))];
XI = X0*exp((lambda - mu^2/2)*tVec_small + mu*W_small);
XS = X0*exp(lambda*tVec_small + mu*W_small);

%% Figure

figure;
hold on;

plot(tVec, X1, 'bo--');
plot(tVec, X2, 'ro--');
plot(tVec_small, XI, 'b-');
plot(tVec_small, XS, 'r-');

ylim([0.5,9]);
xlabel('t', 'Fontsize', 20);
set(gca, 'FontSize', 20);
legend('Euler-Maruyama', 'Trapezoidal', 'Itô', 'Stratonovich', 'Location', 'best');

saveas(gcf, 'figures/exercise5_method.png');

%% Rate

error_I = zeros(1, length(hs));
error_S = zeros(1, length(hs));

for i = 1:length(hs)
    for m = 1:M
        h = hs(i);
        N = round(T/h);
        dW = sqrt(h)*randn(1, N);
        W = [0, cumsum(dW)];
        X1 = X0*prod(1 + h*lambda + mu*dW);
        X2 = X0*prod(1 + h*lambda + mu*dW + (1/2)*mu*dW.*(lambda*h + mu*dW));
        XI = X0*exp((lambda - mu^2/2)*T + mu*W(end));
        XS = X0*exp(lambda*T + mu*W(end));
        error_I(i) = error_I(i) + abs(X1 - XI);
        error_S(i) = error_S(i) + abs(X2 - XS);
    end
    error_I(i) = error_I(i)/M;
    error_S(i) = error_S(i)/M;
end

figure;
loglog(hs, error_I, '-*', 'Linewidth', 3);
hold on;
loglog(hs, error_S, '-*', 'Linewidth', 3);
loglog(hs, sqrt(hs), '--', 'Linewidth', 3);
loglog(hs, hs, '--', 'Linewidth', 3);
set(gca, 'FontSize', 20);
legend('Euler-Maruyama & Itô', 'Trapezoidal & Stratonovich', 'order 1/2', 'order 1', 'Location', 'best');

saveas(gcf, 'figures/exercise5_rate.png');