%% UQ Exercise 4: control variates
% 
% Exercise 3: Estimating Pi
clear all; clc; close all;

%% (0) Preliminary defintions
Y_func = @(x,y) 4*(x.^2 + y.^2 <= 1);
Z_func = @(x,y,c) (x+y <= c);
EZ = @(c) 0.5 * c.^2 + (c>1) .* (2*c - c.^2 - 1);

%% (b) Find optimal c
M = 1e6; % number of samples

U = rand(M,2); % uniform samples on unit square
Y = Y_func(U(:,1),U(:,2)); % samples of Y

h = 0.01; % grid size for detecting optimal c
cs = h:h:2-h; % grid for detecting optimal c

% Loop over different cs
for i = 1:length(cs)
    c = cs(i);
    
    % samples of Z
    Z = Z_func(U(:,1),U(:,2),c);
    
    % Factor of improvement 
    rhos(i) = 1-corr(Y,Z)^2;
end

% Plot rhos vs. cs
figure();
plot(cs,rhos,'o-','LineWidth',2)
xlabel('c')
ylabel('Reduction in Variance')
[rho_min,ind] = min(rhos);
c_min = cs(ind)
set(gca,'FontSize',22)

% Plotting the line
figure();
x = 0:1e-3:1;
plot(x, sqrt(1-x.^2),'-k', x, c_min-x,'-r','LineWidth',2)
xlim([0,1])
ylim([0,1])
xlabel('x_1'); ylabel('x_2')
axis equal
set(gca,'FontSize',22)

%% (c) Estimating Pi with optimal c
c = c_min;
M = 1e3; % number of samples

U = rand(M,2); % uniform samples on unit square
Y = Y_func(U(:,1),U(:,2)); % samples of Y
Z = Z_func(U(:,1),U(:,2),c); % samples of Z

% common Monte Carlo estimate
pi_hat = mean(Y);

% Monte Carlo estimate with control variate
b = cov(Y,Z);
b = b(2,1) / var(Z); % Estimating optimal b
Y_tilde = Y - b*(Z - EZ(c_min));
pi_tilde = mean(Y_tilde);

% Errors
err_MCCV = abs(pi_tilde - pi)
err_MC = abs(pi_hat - pi)

%% (d) Estimating Pi with non-optimal c
c_bad = 1; b_bad = -1;

c = c_bad;
M = 1e3; % number of samples

U = rand(M,2); % uniform samples on unit square
Y = Y_func(U(:,1),U(:,2)); % samples of Y
Z = Z_func(U(:,1),U(:,2),c); % samples of Z

% common Monte Carlo estimate
pi_hat = mean(Y);

% Monte Carlo estimate with control variate
Y_tilde = Y - b*(Z - EZ(c));
pi_tilde = mean(Y_tilde);

% Errors
err_MCCV = abs(pi_tilde - pi)
err_MC = abs(pi_hat - pi)


%% (e) Plotting Distribution of Errors
M = 1e3; % number of samples

for i = 1:1e5,
    U = rand(M,2); % uniform samples on unit square
    Y = Y_func(U(:,1),U(:,2)); % samples of Y

    % common Monte Carlo estimate
    pi_hat(i) = mean(Y);

    % Monte Carlo estimate with optimal control variate
    c = c_min; 
    Z = Z_func(U(:,1),U(:,2),c); % samples of Z
    b = cov(Y,Z); b = b(2,1) / var(Z); % Estimating optimal b
    Y_tilde = Y - b*(Z - EZ(c_min));
    pi_tilde_min(i) = mean(Y_tilde);


    % Monte Carlo estimate with bad control variate
    c = c_bad; b = b_bad;
    Z = Z_func(U(:,1),U(:,2),c); % samples of Z
    Y_tilde = Y - b*(Z - EZ(c));
    pi_tilde_bad(i) = mean(Y_tilde);

    % Errors
    err_MC(i) = abs(pi_hat(i) - pi);
    err_MCCV_min(i) = abs(pi_tilde_min(i) - pi);
    err_MCCV_bad(i) = abs(pi_tilde_bad(i) - pi);
end
figure();
[f,xi] = ksdensity(err_MC);
plot(xi,f,'-k','LineWidth',2);
hold on; 
[f,xi] = ksdensity(err_MCCV_min); 
plot(xi,f,'-r','LineWidth',2);
legend('CMC','CVMC optimal')
set(gca,'fontsize',22)
[f,xi] = ksdensity(err_MCCV_bad); 
plot(xi,f,'-b','LineWidth',2);
legend('MC','MCCV optimal','MCCV bad')
set(gca,'fontsize',22)