clear;

%% Data

rng(2022);

L = 10;
phiarray = cell(L + 1, 2);
meanarray = zeros(L + 1, 1);
phi = @(x) x;

lambdaVec = [-10, -50, -250];

phiFunc = @(x) x;
truth = exp(lambdaVec);

X0 = 1;

Nlarray = 2.^(0:L)';
hlarray = 1./Nlarray;
Mlarray = (L*hlarray)/(hlarray(end)^2);

for i = 1 : length(lambdaVec)
    
    disp(['Iteration ' num2str(i)]);
    
    lambda = lambdaVec(i);
    mu = sqrt(lambda);
   
    hEM = - (2*lambda + mu^2)/(lambda^2);
    Nlarray = 2.^(0:L)';
    hlarray = 1./Nlarray;
    Mlarray = round((L*hlarray)/(hlarray(end)^2))+1;

    cost = zeros(L,1);
    MCEst = zeros(L,1);
    VarEst = zeros(L,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 = prod((1 + mu*DW)*(1-hl*lambda), 2).*(X0*ones(Ml, 1));
    phiarray{1, 1} = phi(X);
    phiarray{1, 2} = zeros(Ml,1);
    
    meanarray(1) = sum(phiarray{1,1} - phiarray{1,2}, 1)/Ml;
    
    
    for l = 2:L+1
            
            disp(['Level ', num2str(l-1)]);
            
            Ml = Mlarray(l);
            hl = hlarray(l);
            Nl = Nlarray(l);
            
            disp(Ml)
            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);
            phiarray{l, 1} = phi(X);
            
            DW = W(:,3:2:end) - W(:,1:2:end-2);
            X = X0*prod((1 + mu*DW)*(1-hl*lambda), 2);
            phiarray{l, 2} = phi(X);   
            
            meanarray(l) = sum(phiarray{l,1} - phiarray{l,2}, 1)/Ml;
            
            cost(l-1) = cost(l-1) + Ml*Nl;
            disp(Ml*Nl)
    end
   
    barZ = sum(meanarray);

    
    figure;
    loglog(hlarray(2:end), cost, 's-', 'Linewidth', 2, 'Markersize', 10);
    hold on;
    xlabel('$h_L$', 'interpreter', 'latex', 'fontsize', 20);
    ylabel('cost', 'interpreter', 'latex', 'fontsize', 20);
    title(['$\lambda = ' num2str(lambda) '$'], 'interpreter', 'latex', 'fontsize', 20);
    ylim([1, 1e10]);
    set(gca, 'FontSize', 20);
    saveas(gcf, ['figures/exercise3_2_', num2str(i), '.png']);
    
end