clear;

%% Data

rng(2022);

LVec = 1:10;
lambdaVec = [-10, -50, -250];

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

X0 = 1;

%% Figures

for i = 1 : length(lambdaVec)
    
    disp(['Iteration ' num2str(i)]);
    
    lambda = lambdaVec(i);
    mu = sqrt(lambda);
    hEM = - (2*lambda + mu^2)/(lambda^2);

    lEM = ceil(abs(log2(hEM)));
    
    cost = zeros(size(LVec));
    MCEst = zeros(size(LVec));
    VarEst = zeros(size(LVec));
    
    for L = LVec
        
        disp(['Level ' num2str(L)]);
        
        if L > lEM
            
            phi = cell(1, L - lEM + 1);
            Ml = zeros(1, L - lEM + 1);
            Ml(1) = (2^(2*L))*(L - lEM);
            l = lEM;
            for j = 2 : L - lEM + 1
                l = l + 1;
                Ml(j) = 2^(2*L - l)*(L - lEM);
            end
            hl = 2.^-(lEM:L);
            Nl = 2.^(lEM:L);
            
            DW = sqrt(hl(end))*randn(Ml(end), Nl(end));
            X = X0*prod(1 + hl(end)*lambda + mu*DW, 2);
            phi{end} = phiFunc(X);
            cost(L) = cost(L) + Ml(end)*Nl(end);
            
            for l = L-lEM:-1:1
                MlDiff = Ml(l) - Ml(l+1);
                DW = [DW(:, 1:2:end); sqrt(hl(l)) * randn(MlDiff, Nl(l))];
                X = X0*prod(1 + hl(l) * lambda + mu * DW, 2);
                phi{l} = phiFunc(X);
                cost(L) = cost(L) + Ml(l)*Nl(l);
            end
            
            MCEst(L) = mean(phi{1});
            VarEst(L) = var(phi{1})/Ml(1);
            for l = 2:L-lEM+1
                MCEst(L) = MCEst(L) + mean(phi{l} - phi{l-1}(1:Ml(l)));
                VarEst(L) = VarEst(L) + var(phi{l} - phi{l-1}(1:Ml(l)))/Ml(l);
            end
            
        else
            
            disp('MONTE CARLO');
            
            h = 2^(-lEM);
            M = ceil(h^(-2));
            N = 1/h;
            DW = sqrt(h)*randn(M, N);
            X = X0*prod(1 + h*lambda + mu*DW, 2);
            cost(L) = cost(L) + M*N;
            
            MCEst(L) = mean(phiFunc(X));
            VarEst(L) = var(phiFunc(X));
            
        end
        
    end
    
    hVec = 1./(2.^LVec);
    figure;
    loglog(hVec, cost, 's-', 'Linewidth', 2, 'Markersize', 10);
    hold on;
    xline(hEM, 'r--', 'Linewidth', 2);
    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_1_', num2str(i), '.png']);
    
end