clear();

%% Data

rng(2022);

lambda = 1;
mu = 0.1;
X0 = 0.1;
T = 1;

phi = @(x) x;
Z = X0*exp(lambda*T);
Levels = 3:11;

%% MSE and computational cost

budget = zeros(length(Levels), 1);
rMSE_barZ = zeros(length(Levels), 1);
rMSE_hatZ = zeros(length(Levels), 1);

for iter = 1:length(Levels)
    
    disp(['Iteration ', num2str(iter)]);
    
    L = Levels(iter);
    Nlarray = 2.^(0:L)';
    hlarray = T./Nlarray;
    Mlarray = (L*hlarray)/(hlarray(end)^2);
    
    budget(iter) = sum(Nlarray.*Mlarray, 1);

    disp('MLMC');
    phiarray = cell(L + 1, 2);
    meanarray = zeros(L + 1, 1);

    disp('Level 0');
    
    Ml = Mlarray(1);  
    hl = hlarray(1);  
    Nl = Nlarray(1);
    
    dW = sqrt(hl)*randn(Ml,Nl);
    W = [zeros(Ml,1) cumsum(dW,2)];
    
    DW = W(:,2:end) - W(:,1:end-1);
    X = X0*prod(1 + hl*lambda + mu*DW, 2);
    phiarrayl = phi(X);
    phiarraylm1 = zeros(Ml,1);
    meanarray(1) = sum(phiarrayl - phiarraylm1, 1)/Ml;

    for l = 2:L+1
        
        disp(['Level ', num2str(l-1)]);

        Ml = Mlarray(l);
        hl = hlarray(l);
        Nl = Nlarray(l);
        
        dW = sqrt(hl)*randn(Ml,Nl);
        W  = [zeros(Ml,1) cumsum(dW,2)];

        DW = W(:,2:end) - W(:,1:end-1);
        X = X0*prod(1 + hl*lambda + mu*DW ,2);
        phiarrayl = phi(X);

        DW = W(:,3:2:end) - W(:,1:2:end-2);
        X = X0*prod(1 + 2*hl*lambda + mu*DW, 2);
        phiarraylm1 = phi(X);   

        meanarray(l) = 1/Ml*sum(phiarrayl-phiarraylm1,1);
        
    end

    barZ = sum(meanarray);
    rMSE_barZ(iter) = realsqrt(mean((Z - barZ)^2, 1));

    disp('Monte Carlo');
    eta = budget(iter);
    N = ceil(T/eta^(-1/3));
    h = T/N;   
    M = ceil(eta^(2/3)); 
    
    dW = sqrt(h)*randn(M, N);
    X = X0*prod(1 + h*lambda + mu*dW, 2);
    
    hatZ = mean(phi(X), 1);
    rMSE_hatZ(iter) = realsqrt(mean((Z - hatZ)^2, 1));   

end 

%% Figure
    
figure();

ax = axes('nextplot', 'add', 'box', 'on');
semilogx(budget, rMSE_hatZ, ':s', 'Linewidth', 2, 'Markersize', 10);
semilogx(budget, rMSE_barZ, ':x', 'Linewidth', 2, 'Markersize', 10);
leg = legend({'$\sqrt{\mathrm{MSE}(\hat{Z})}$ (MC)', '$\sqrt{\mathrm{MSE}(\bar{Z})}$ (MLMC)'});
set(leg, 'box', 'off', 'interpreter', 'latex', 'fontsize', 20);
set(get(gca, 'xlabel'), 'string', 'computational cost', 'interpreter', 'latex', 'fontsize', 20);
set(ax, 'yscale', 'log', 'Xscale', 'log');
set(gca, 'FontSize', 20);
saveas(gcf, 'figures/exercise2_3.png');