clear;

%% Data

rng(2022);

a = 0.5;
X0 = 0;

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

%% Exact solution

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

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

X = sin(a*W(:,end) + asin(X0));

%% Errors

strongError_1 = nan(1, length(DeltaTs));
strongError_2 = nan(1, length(DeltaTs));
strongError_3 = nan(1, length(DeltaTs));
strongError_4 = nan(1, length(DeltaTs));

for i = 1:length(DeltaTs)
    
    DeltaT = DeltaTs(i);
    ratio = DeltaT/DeltaT_min;
    nSteps = T/DeltaT;
    
    dW_num = W(:, 1+ratio:ratio:end) - W(:, 1:ratio:end-ratio);
    X_1 = X0*ones(nTrajectories, 1);
    X_2 = X0*ones(nTrajectories, 1);
    X_3 = X0*ones(nTrajectories, 1);
    X_4 = X0*ones(nTrajectories, 1);
    for j = 1:nSteps
        X_1 = X_1 + a*sqrt(1 - X_1.^2).*dW_num(:,j) - (a^2)*X_1.*(dW_num(:,j).^2)/2;
        Z_2 = X_2 - (a^2)*X_2*DeltaT/2 + a*sqrt(1 - X_2.^2)*sqrt(DeltaT);
        X_2 = X_2 - (a^2)*X_2*DeltaT/2 + a*sqrt(1 - X_2.^2).*dW_num(:,j) + (a/sqrt(DeltaT))*(sqrt(1 - Z_2.^2) - sqrt(1 - X_2.^2)).*(dW_num(:,j).^2 - DeltaT)/2;
        Z_3 = X_3 + a*sqrt(1 - X_3.^2)*sqrt(DeltaT);
        X_3 = X_3 - (a^2)*X_3*DeltaT/2 + a*sqrt(1 - X_3.^2).*dW_num(:,j) + (a/sqrt(DeltaT))*(sqrt(1 - Z_3.^2) - sqrt(1 - X_3.^2)).*(dW_num(:,j).^2 - DeltaT)/2;
        X_4 = X_4 + a*sqrt(1 - X_4.^2).*dW_num(:,j) - (a^2)*X_4.*(dW_num(:,j).^2)/2 + (a^4)*X_4.*(DeltaT^2)/8 -  (a^3)*sqrt(1 - X_4.^2).*(dW_num(:,j).^3)/6;
    end
    strongError_1(i) = mean(abs(X_1 - X));
    strongError_2(i) = mean(abs(X_2 - X));
    strongError_3(i) = mean(abs(X_3 - X));
    strongError_4(i) = mean(abs(X_4 - X));
end

%% Figures

figure;
loglog(DeltaTs, strongError_1, 's:', 'Linewidth', 2);
hold on;
loglog(DeltaTs, strongError_2, 's:', 'Linewidth', 2);
loglog(DeltaTs, strongError_3, 's:', 'Linewidth', 2);
loglog(DeltaTs, strongError_4, 's:', 'Linewidth', 2);
loglog(DeltaTs, DeltaTs/5, '--', 'Linewidth', 2);
xlabel('{\Delta}t', 'Fontsize', 20);
xlim([DeltaTs(1)/2, DeltaTs(end)*2]);
ylim([1e-5, 1e-1]);
xticks([1e-2, 1e-1]);
legend('Milstein', 'Milstein free 1', 'Milstein free 2', '1.5 Method', 'Order 1', 'Orientation', 'vertical', 'Fontsize', 20, 'Location', 'SE');
set(gca, 'FontSize', 20);
saveas(gcf, 'figures/exercise6.png');