% This function performs generalized least squares (GLS)
% Detailed explanation of this function.

function [stat_table, contrast_table] = my_GLS(X, y, C, W, xnames, cnames)

% Initialize
[C, var_chat, se_chat, t_chat, pval_chat] = deal([]);
contrast_table = table();

xnames = cell(1, size(X, 2))';
cnames = cell(1, size(C, 2))';

% Prepare weights
if ~issymmetric(W)
    % assume they are a vector
    W = diag(W); 
end

if size(W, 1) ~= size(X, 1), error('W is the wrong size. Should be [n x 1] or [n x n]'), end


% Core GLM equations
% -------------------------------------------
% remove nans from X and y
% Note: You need CANlabCore tools on your path for this!
% You can omit this line if you don't have NaNs (missing data)
[wasnan, X, y] = nanremove(X, y);

xtxi = inv(X'*W*X);
bhat = xtxi * X' * W * y;
e = y - X * bhat;

[n, p] = size(X);

sigma_sq = e'* W * e ./ (n - p);  % residual variance
var_bhat = sigma_sq * xtxi;
t = bhat ./ (diag(var_bhat) .^ .5);

dfe = n - p;
pval = 2 * (1 - tcdf(abs(t), dfe));

se = sqrt(diag(var_bhat)); % standard error

% Contrasts
if ~isempty(C)

    chat = C' * bhat;
    var_chat = sigma_sq * C' * xtxi * C;
    se_chat = diag(var_chat) .^ 0.5;
    t_chat = chat ./ se_chat;

    pval_chat = 2 * (1 - tcdf(abs(t_chat), dfe));

end

% Collect and print output tables
% -------------------------------------------

stat_table = table(xnames, bhat, se, t, pval, 'VariableNames', {'Name', 'bhat', 'SE', 't', 'P'});

disp('Model parameters:')
stat_table

if ~isempty(C)

    contrast_table = table(cnames, chat, se_chat, t_chat, pval_chat, 'VariableNames', {'Contrast', 'bhat', 'SE', 't', 'P'});

    disp('Contrasts:')
    contrast_table

end

end % Main function