%% Assessing bias and variance in mixed effects model estimates using simulation
%
% Table of Contents
% Specify input parameters
% Test run
% Run simulation
% Print a report of the output: Bias and variance of estimates
%
%
%% Specify input parameters

n_i = 20;           % len
m = 5;              % sub
p = 2;              % num params
b_g = [1 3]';       % fixed effects: Intercept and slope    
s2 = 25;            % residual var
U_g = [2 0; 0 4];   % Cov matrix of random intercept and slope
regressortype = 'continuous';

% These are params whose estimates we will compare later in the report
fixedint = b_g(1);
fixedslope = b_g(2); 
 
withinerr_std = sqrt(s2);

randintvar = U_g(1, 1);
randslopevar = U_g(2, 2);
 

%
%% Test run
% Generate a dataset, fit the model, and plot the results

[y, Z_i, x, sim_data] = sim_generate_mixedfx_data1(m, n_i, p, b_g, s2, U_g, regressortype);

out = igls(y, x, 'plot', 'all');

%
%% Run simulation


iter = 500;

[fixedfx_hat, rfx_hat] = deal(zeros(2, iter));
s2_within = zeros(1, iter);

tic
for iteration = 1:iter

    % Generate data
    [y, Z_i, x, sim_data] = sim_generate_mixedfx_data1(m, n_i, p, b_g, s2, U_g, regressortype);

    % Fit model
    out = igls(y, x, 'noverbose');  % for igls

    % Save estimates to check
    fixedfx_hat(:, iteration) = out.beta;

    rfx_hat(:, iteration) = out.betastar;

    s2_within(:, iteration) = out.Sigma;

    if mod(iteration, 100) == 0
        toc
    end

end
toc

%
%% Print a report of the output: Bias and variance of estimates

% Calculate means, bias (deviations), and variances for estimates across
% iterations

fixedfxmean = mean(fixedfx_hat, 2);
rfxmean = mean(rfx_hat, 2);
s2mean = mean(s2_within, 2);

fixedfxbias = [fixedint; fixedslope] - fixedfxmean;
rfxbias = [randintvar; randslopevar] - rfxmean;
s2bias = withinerr_std .^ 2 - s2mean;

% Significance test for estimation bias
% Does the model produce the true values on average?
% A significant P-value here indicates systematic bias, which is bad for the model.

devs = (fixedfx_hat - [fixedint; fixedslope])'; % iteration-by-iteration deviations
[~, fixedfx_bias_p, ~, fixedfx_bias_stats] = ttest(devs);

devs = (rfx_hat - [randintvar; randslopevar])'; % iteration-by-iteration deviations
[~, rfx_bias_p, ~, rfx_bias_stats] = ttest(devs);

devs = (withinerr_std .^ 2 - s2_within)'; % iteration-by-iteration deviations
[~, s2_bias_p, ~, s2_bias_stats] = ttest(devs);

% Variance of estimates across iterations
% Lower is better

fixedfx_err_var = var(fixedfx_hat, 0, 2);
rfx_err_var = var(rfx_hat, 0, 2);
s2_err_var = var(s2_within, 0, 2);

% Print output
uline = '_______________________________________________________________';
disp(uline)
disp('Bias and variance of model estimates compared with ground truth')
disp(uline)

t = table([fixedint; fixedslope], fixedfxmean, fixedfx_bias_stats.tstat', fixedfx_bias_p', fixedfx_err_var, 'VariableNames', {'Truth' 'Mean Est.' 'Bias t-stat' 'Bias p-val' 'Var of Est.'}, 'RowNames', {'Intercept' 'Slope'});
disp('Fixed effects'); disp(t);

t = table([randintvar; randslopevar], rfxmean, rfx_bias_stats.tstat', rfx_bias_p', rfx_err_var, 'VariableNames', {'Truth' 'Mean Est.' 'Bias t-stat' 'Bias p-val' 'Var of Est.'}, 'RowNames', {'Intercept' 'Slope'});
disp('Random effects'); disp(t);

t = table(withinerr_std .^ 2, s2mean, s2_bias_stats.tstat', s2_bias_p', s2_err_var, 'VariableNames', {'Truth' 'Mean Est.' 'Bias t-stat' 'Bias p-val' 'Var of Est.'}, 'RowNames', {'Error Sigma^2'});
disp('Residual variance'); disp(t);