%% MLMC for Dynamical system with random initial value
%
% 03.05.2016

clear all; close all; clc;

%% (0) Definitions

% Random model for inital value
SampleInitial = @(m) repmat([0.5; 2],1,m) + 0.2*(2*rand(2,m)-1);

% Final time
T = 6;

% Quantity of interest
Q_func = @(u) u(1,end);

% Forward map given discretization
F = @(xi,h) Q_func(DynSys(xi, T, h));

%% (00) Deterministic convergence w.r.t. h
u0 = SampleInitial(1);
h = 2^(-12);
Q = F(u0, h);

h0 = 2^(-1); clear hs; clear err;
for i = 1:9,
    hs(i) = 2^(-i) * h0;
    err(i) = abs(Q - F(u0,hs(i)));
end
h1 = loglog(hs, err,'o-','LineWidth',2);
xlabel('$h$', 'Interpreter', 'latex', 'FontSize', 10);
ylabel('$\mathrm{Error}$', 'Interpreter', 'latex', 'FontSize', 10);

alpha = polyfit(log(hs), log(err), 1);
hold on;
h2 = loglog(hs, exp(alpha(2) - 1 + alpha(1) * log(hs)), '--','LineWidth',2);
hl  = legend([h1; h2], '$\epsilon\left(h_l\right)$', ['$\alpha=' sprintf('%.4f$', alpha(1))], ...
    'Location', 'NorthEast');
set(hl, 'Interpreter','latex');
title('Decay of discretization error w.r.t. h')
set(gca,'fontsize',20)

%% (a1) Estimate average run times for different discretization lengths
h0  = 2^(-3);
hs = 2.^(-(0:7)') * h0;
Cs  = zeros(length(hs), 1);
M = 1e3;
tic;
for i = 1:length(hs)
    t = 0;
    h = hs(i);
    for m = 1:M,
        u0 = SampleInitial(1);

        t0  = tic;
        Q   = F(u0,h);
        t   = t + toc(t0);
    end
    Cs(i) = t / M; 
end
toc;

% Plotting and estimating decay (estimation excludes very coarse levels as outliers)
figure();
h1 = loglog(hs, Cs, '+-','LineWidth',2);
xlabel('$h$', 'Interpreter', 'latex', 'FontSize', 10);
ylabel('$\mathrm{Runtime}$', 'Interpreter', 'latex', 'FontSize', 10);

gamma = polyfit(log(1./hs(end-4:end)), log(Cs(end-4:end)), 1);
hold on;
h2 = loglog(hs, exp(gamma(2)  + gamma(1) * log(1./hs)), '--','LineWidth',2);
hl  = legend([h1; h2], '$\mathcal{C}\left(h_l\right)$', ['$\gamma=' sprintf('%.4f$', gamma(1))], ...
    'Location', 'NorthEast');
set(hl, 'Interpreter','latex');
title('Cost growth w.r.t. h')
set(gca,'fontsize',22)

%% (a2) Compute reference value
% 
% % Reference solution (MC)
% h	= 2^(-12);
% M   = 3e6;
% Qs  = zeros(M, 1);
% tic;
% parfor m = 1:M,
%     u0  = SampleInitial(1);
%     Qs(m)   = F(u0,h);
% end
% toc;
% % Reference value for mean of Q (obtained by many more samples on finer grid)
% trueQ = mean(Qs)
% std(Qs)

trueQ = 1.493005355935809;
%% (a2) Estimate decay of bias

% Number of samples for estimating mean
M       = 1e4;

% Levels
h0 = 2^(-3);
hs     = 2.^(-(0:7)') * h0;

meanQ	= zeros(length(hs), 1);
stdQ	= zeros(length(hs), 1);
errQ	= zeros(length(hs), 1);

tic;
for i = 1:length(hs)   
    Q = zeros(M, 1);
    h = hs(i);
    parfor m = 1:M
        u0  = SampleInitial(1);
        Q(m) = F(u0,h);
    end
    
    meanQ(i)	= mean(Q);
    var_MC(i)   = var(Q);

    % Bias on level l
    errQ(i)     = abs(meanQ(i) - trueQ);
end
toc;
% Plotting and estimating decay rate
figure();
h1 = loglog(hs, errQ, '-+','LineWidth',2);
xlabel('$h$', 'Interpreter', 'latex', 'FontSize', 20);
ylabel('$\mathrm{Bias}$', 'Interpreter', 'latex', 'FontSize', 20);
title('$\mathrm{Decay\ of\ Bias}$', 'Interpreter', 'latex', 'FontSize', 20);

alpha = polyfit(log(hs), log(errQ),1);
hold on;
h2 = loglog(hs, exp(alpha(2) - 1 + alpha(1) * log(hs)), '--','LineWidth',2);
hold off;
hl  = legend([h1; h2], '$|\mathrm{E}\left(Q-Q_{h_l}\right)|$', ['$\alpha=' sprintf('%.4f$', alpha(1))], ...
    'Location', 'NorthEast');
set(hl, 'Interpreter', 'latex', 'FontSize', 20);
set(gca,'fontsize',22)

%% (a3) Estimate Decay of Variance of Differences

% loop over coarse levels
h0 = 2^(-3);
hs = 2.^(-(0:7)') * h0;
M = 1e3;
tic;
VarY = zeros(length(hs),1);
for i = 1:length(hs),
    h_coarse  = hs(i);
    h_fine    = 0.5 * h_coarse;
    
    Y = zeros(M, 1);
   
    parfor m = 1:M
        % make sure it's the same realization
        u0  = SampleInitial(1);
       
        Q_coarse(m) = F(u0,h_coarse);
        Q_fine(m) = F(u0,h_fine);
       
        % Difference of Quantities
        Y(m) = Q_fine(m) - Q_coarse(m);       
    end
    
    % Estimate variance and quantify relative sampling error for variance
    VarY(i) = var(Y);
end
toc;
% Plotting and estimating decay rate
figure()
h1 = loglog(hs, VarY, '-+','LineWidth',2);
xlabel('$h$', 'Interpreter', 'latex', 'FontSize', 10);
ylabel('$\mathrm{Variance\ of\ Y_h}$', 'Interpreter', 'latex', 'FontSize', 10);
title('$\mathrm{Decay\ of\ Variance\ of\ difference\ on\ consecutive\ levels}$', ...
    'Interpreter', 'latex', 'FontSize', 10);

beta = polyfit(log(hs), log(VarY), 1);
hold on;
h2 = loglog(hs, exp(beta(2)-1 + beta(1) * log(hs)), '--','LineWidth',2);
hl  = legend([h1; h2], '$\mathrm{Var}\left(Q_{h_l}-Q_{h_{l-1}}\right)$', ['$\beta=' sprintf('%.4f$', beta(1))], ...
    'Location', 'NorthEast');
set(hl, 'Interpreter', 'latex', 'FontSize', 10);
set(gca,'fontsize',12)


%% (c) MC Simulation

% Levels
h0  = 2^(-5);
hs = 2.^(-(0:5)) * h0;
hs_MC = hs;

% fix costs
C_fix   = 3;

% Vectors for MC sample size, MC estimates and MC confidence intervals
MC_M    = zeros(length(hs), 1);
MC_Est  = zeros(length(hs), 1);
MC_Var	= zeros(length(hs), 1);
MC_Conf	= zeros(length(hs), 2);

% MC computation, loop over levels
for i = 1:length(hs)
    h = hs(i);
  
    % compute allowed sample size for fixed costs
    M(i) = round(C_fix / Cs(i+2));
    
    Q   = zeros(M(i), 1);
    tic;
    for m = 1:M(i),
        u0  = SampleInitial(1);
        Q(m) = F(u0,h);
    end
    toc;
    
    % MC estimate and confidence interval
    MC_Est(i)	= mean(Q);
    MC_Var(i)   = var(Q);
    conf_len    = 2 * sqrt(MC_Var(i)) / sqrt(M(i));
    MC_Conf(i, :) = [MC_Est(i) - conf_len, MC_Est(i) + conf_len];
end

% Plotting the confidence intervals
figure();
for i = 1:length(hs)
    plot([hs(i), hs(i)], [MC_Conf(i, 1), MC_Conf(i, 2)], '*-r','LineWidth',2);
    hold on;
end
plot(hs, MC_Est, 'sk','LineWidth',2);
plot(hs, repmat(trueQ,1,length(hs)), '--','LineWidth',2);
xlabel('$h$', 'Interpreter', 'latex', 'FontSize', 20);
ylabel('$Q\left(\omega\right)$', 'Interpreter', 'latex', 'FontSize', 20);
title(['$\mathrm{Confidence\ intervals\ for\ fixed\ costs}\ ' sprintf('C=%0.1fs$', C_fix)], ...
    'Interpreter', 'latex', 'FontSize', 20);
hold off;
set(gca,'FontSize',22)

%% (d) Develop a 2-level-MC algorithm (L = 2)

% Multilevel-Factor for relating sample sizes M_l by previously determined rates
fac = 2^(-0.5 * round(beta(1) + gamma(1)));
L   = 2;

% fic costs
C_fix   = 3;

% vector hs of finest levels h_L
hs = 2.^(-(5:10)');
hs_MLMC = hs;

% Vector for MLMC estimates and confidence interval
MLMC_Est    = zeros(length(hs), 1);
MLMC_Var    = zeros(length(hs), 1);
MLMC_Conf   = zeros(length(hs), 2);

% Loop over coarsest level
for i = 1:length(hs)
    
    % Determine M0 according to fix total cost
    h0 = hs(i) * 2^L;
    M0 = round(C_fix / (Cs(i) + fac*(Cs(i + 1) + Cs(i)) + ...
            fac^2*(Cs(i + 2) + Cs(i + 1))));
    
    tic;

    Q0   = zeros(M0, 1);
    for m = 1:M0,
        u0  = SampleInitial(1);
        Q0(m) = F(u0,h0);    
    end
    % Mean and variance of MC estimator on coarsest level 
    meanQ0   = mean(Q0);
    varQ0    = var(Q0);
    
    % Loop over 2 levels for computing Y_l
    Ms      = round(fac.^((1:L)') * M0);   % number of samples for Y_l
    meanY   = zeros(L, 1);
    varY    = zeros(L, 1);  
    
    for k = 1:L
        h_coarse  = 2^(-(k-1)) * h0;
        h_fine    = 2^(-k) * h0;
        
        Y = zeros(Ms(k), 1);
        
        % Loop over realizations
        for m = 1:Ms(k),
            % make sure it's the same realization
            u0  = SampleInitial(1);
            Q_coarse	= F(u0,h_coarse);
            Q_fine      = F(u0,h_fine);

            % Difference of Quantities
            Y(m) = Q_fine - Q_coarse;       
        end
        meanY(k)= mean(Y);
        varY(k)	= var(Y);
    end
    toc;
    
    MLMC_Est(i)     = meanQ0 + sum(meanY);
    MLMC_Var(i)     = varQ0 / M0 + sum(varY ./ Ms);
    conf_len        = 2 * sqrt(MLMC_Var(i));
    MLMC_Conf(i, :) = [MLMC_Est(i) - conf_len, MLMC_Est(i) + conf_len];
end

% Plotting of confidence intervals
MLMC_plot = figure();
hold on;
for i = 1:(length(hs))
    loglog([hs(i), hs(i)], [MLMC_Conf(i, 1), MLMC_Conf(i, 2)], '*-b','LineWidth',2);
end
loglog(hs, MLMC_Est, '*k','LineWidth',2);
loglog(hs, repmat(trueQ,1,length(hs)), '--','LineWidth',2);
xlabel('$l$', 'Interpreter', 'latex', 'FontSize', 20);
ylabel('$Q\left(\omega\right)$', 'Interpreter', 'latex', 'FontSize', 20);
title(['$\mathrm{Confidence\ intervals\ of\ MLMC\ simulation\ for\ fixed\ costs}\ ' sprintf('C=%0.1fs$', C_fix)], ...
    'Interpreter', 'latex', 'FontSize', 22);
hold off;
 
%% Comparing both MC and MLMC
% MC
figure(MLMC_plot);
hold on;
for i = 1:length(hs_MC)
    loglog([hs_MC(i), hs_MC(i)], [MC_Conf(i, 1), MC_Conf(i, 2)], '*-r');
end
loglog(hs_MC, MC_Est, 'sk');
