%% Sample dataset: Back pain clinical trial
% We'll load and examine a sample dataset. This dataset is from:
% Ashar et al. 2022, JAMA Psychiatry
% This study is a randomized clinical trial (RCT) of patients with chronic back pain (>= 4/10 pain for at least 3 months, average duration ~10 years). The trial compares Pain Reprocessing Therapy, Placebo (a single injection into the skin of the back), and Usual Care treatments, with pre-treatment and post-treatment (1 month after treatment initiation), and follow-up to 1 year.
% Table of Contents
%     Dependencies and setup
%     Load the dataset
%     Examine some of the key variables
% Visualize the data
% Code hygiene
% Analysis 1: Custom GLM
% Analysis 2: Scaling
%     Explore the relationship between beta-hats and group means
% Analysis 3: Contrasts
% Prep and Examine the dataset
%% Dependencies and setup
% Do you have the relevant files on your Matlab path?

which Ashar2022_outcomes_demographics.csv

% If you get an error, you don't!
% You also need the Statistics and Machine Learning Toolbox in Matlab. Try this to check:

if isempty(ver('stats')), warning('You do not have the statistics toolbox on your Matlab path. You may/will get errors'); end

%
%% Load the dataset
% We'll read the dataset from a Comma Separated Value (csv) file into a table object. Table is a Matlab class that is very convenient, as it:
% (1) stores metadata like table names and descriptions, and
% (2) has associated methods for manipulating the dataset (e.g. selecting cases, etc.)
% In addition, some statistics functions (like fitlme, for mixed effect models) can take table objects directly, allowing you to reference variables by name
% % Sometimes, it may be helpful to change to a directory and add it to your path.
% % Then Matlab can find files
%
% cd('/Users/f003vz1/Dropbox (Dartmouth College)/COURSES/Courses_Dartmouth/2023_Fall_Computational_Foundations/Datasets/Ashar_2022_PRT_backpain')
% addpath(pwd)

backpain_all = readtable('Ashar2022_outcomes_demographics.csv');

backpain_all  % print (some of) the table to screen

%
%% Examine some of the key variables
% It's always critical to look at your data!  Check that the distributions and values make sense.
% We'll work with a few key variables, though there are many more interesting ones to examine!
% backpain_all.id;                             Subject ID
% backpain_all.time;                         Assessment time (1 = pre-treatment, 2 = post)
% backpain_all.group;                      Treatment group (1 = PRT, 2 = Placebo, 3 = Usual Care,
%                                                      NaN = Healthy (nonpatient, not randomized)
% backpain_all.bpi_intensity;           Primary outcome, pain intensity
% backpain_all.gender;                    Sex/gender (not diff here; 1 = male, 2 = fem, 3 = other)
% logical(backpain_all.is_patient);   1 = patient, 0 = healthy control (not randomized/treated)
%
% How many subjects?

length(unique(backpain_all.id))

% How many patients?

is_patient = logical(backpain_all.is_patient); % 1 = patient, 0 = healthy control (not randomized/treated)
length(unique(backpain_all.id(is_patient)))

% Calculate descriptive stats for all variables, and view stats for "group"

descriptive_stats = summary(backpain_all);
descriptive_stats.group

% How many subjects assessed per group and time point?

% Tabulate two variables
tab = crosstab(backpain_all.group, backpain_all.time);

% Print the results with labels (and save table object in tab)
tab = table(tab(:, 1), tab(:, 2), 'VariableNames',{'T1' 'T2'}, 'RowNames', {'PRT' 'Placebo' 'Control'})

%
% Create a smaller, reduced dataset with only a few variables we need
% This is called backpain_stacked because the dataset is "stacked" in "long form", with multiple rows per participant.  This is a repeated measures design, and participants thus have a row for each visit ("Time 1" pre-treatment and "Time 2" post-treatment).
% Long-format is convenient for storing a dataset in a single frame and for mixed effects models, which keep track of which observations are collected on the same participants (and other dependencies) internally.
% This has patients only, because the controls have NaN values for group.

backpain_stacked = backpain_all(:, {'id' 'time' 'group' 'bpi_intensity' 'gender' 'is_patient'})

%
% Unstack the dataset to match outcome values within-person
% For standard single-level GLMs and repeated measures models (paired t-tests, repeated measures ANOVA), the "long format" dataset is inappropriate. We want a "wide format" dataset, with one row per person, and repeated measures from the person in different variables (columns).

backpain = unstack(backpain_stacked,'bpi_intensity','time', 'NewDataVariableNames', {'BPI_Time1', 'BPI_Time2'})

% Now, rows correspond to participants, and columns correspond to measured variables.
% Missing data are coded with NaN (Not a Number), so participants who are missing a time point will have NaN values for that time point.
% Let's add a column for the change in pain (improvement), where negative values mean more improvement:

backpain.BPI_change = backpain.BPI_Time2 - backpain.BPI_Time1;

%
% Create sub-tables with only participants in each group
% We do this for convenience, so we can extract data and plot it easily.

prt = backpain(backpain.group == 1, :);
placebo = backpain(backpain.group == 2, :);
control = backpain(backpain.group == 3, :);

%
% Save simple tables
% We can save simple versions of the tables for later, so that we can download and work with them in any software package.
% writetable(backpain, 'Ashar_backpain_simple_wide.csv');
% writetable(backpain_stacked, 'Ashar_backpain_simple_long.csv');
%
%
%% Visualize the data
% Let's create some plots of the improvement in BPI (Brief Pain Inventory) by group


data_to_plot = {prt.BPI_change, placebo.BPI_change, control.BPI_change};
colors = seaborn_colors(3);

create_figure('Violins');
barplot_columns(data_to_plot, 'Improvement by group', 'colors', colors, 'names', {'PRT', 'Placebo' 'Control'});
ylabel('BPI change, T2-T1')
xlabel('Group');

%


create_figure('Lines');
line1 = lineplot_columns([prt.BPI_Time1 prt.BPI_Time2], 'w', 3, 'color', colors{1}, 'markerfacecolor', colors{1}./1.5);
line2 = lineplot_columns([placebo.BPI_Time1 placebo.BPI_Time2], 'w', 3, 'color', colors{2}, 'markerfacecolor', colors{2}./1.5);
line3 = lineplot_columns([control.BPI_Time1 control.BPI_Time2], 'w', 3, 'color', colors{3}, 'markerfacecolor', colors{3}./1.5);
ylabel('BPI change, T2-T1')
set(gca, 'XLim', [0.5 2.5], 'XTick', 1:2, 'XTickLabel', {'Time 1' 'Time 2'})

%
%% Code hygiene
% To keep the code clean, here are some things I attempted to do:
% Use descriptive names for variables, so that I intuitively know what they mean
% Comment the code to explain key steps, so I can ideally read the code "like a book"
% Print key output to check each step
% Include text to label and explain the output and its interpretation
% To keep the code clean, here are some things I did NOT do:
% I did NOT re-use variable names so that the same variable refers to different content at different points in the script
% I did NOT print a bunch of extraneous output
% I did NOT include long code blocks with calculations in the main script (I would create a subfunction for these).
%
%
%% Analysis 1: Custom GLM
% Steps:
% Code the core GLM equations
% Create a design matrix(X) that codes for differences among the 3 groups
% Estimate the model and report output
% Validate by comparing with an established GLM function
% Turn your code into a function that you can run on any inputs X, y, and regressor names

% Create design matrix for group diffs, and outcome
% -------------------------------------------

% Create X, code for groups
indic = double([backpain.group == 1 backpain.group == 2 backpain.group == 3]);

% We need two orthogonal regressors
C = [1 1 1; 1 -1 0; -0.5 -0.5 1]';
X = indic * C;

cnames = {'Intercept' 'PRT-Pla' 'Treat-Control'}'; % Regressor names

% X = [indic(:, 1) - indic(:, 2)  -indic(:, 1) - indic(:, 2) + 2 * indic(:, 3)]
y = backpain.BPI_change;

% Core GLM equations
% -------------------------------------------
% remove nans from X and y
[wasnan, X, y] = nanremove(X, y);

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

[n, p] = size(X);

var_bhat = e'*e ./ (n - p) * xtxi;
t = bhat ./ (diag(var_bhat) .^ .5);

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

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

% F-test for all model parameters
syy = (y - mean(y))' * (y - mean(y));

F = ( (syy - e'*e) ./ (p - 1) ) ./ (e'*e ./ (n - p));
pval_full_model = 1 - fcdf(F, p - 1, n - p);

R2 = 1 - (e'*e ./ syy);

% -------------------------------------------

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

disp('Analysis of BPI change (negative = improvement)')
stat_table

fprintf('Full model R2 = %3.2f, F(%3.0f, %3.0f) = %3.2f, P = %3.6f\n', R2, p - 1, dfe, F, pval_full_model);

%

% Compare with Matlab's glmfit
[b, dev, stat] = glmfit(X(:, 2:3), y);

disp('Matlab glmfit output:')
glm_table(stat,cnames(2:3)', b)

% for F-test, compare with regress()
[B,BINT,R,RINT,STATS] = regress(y, X);
F_regress = STATS(2)
F_our_function = F

F_P_regress = STATS(3)
F_P_our_function = pval_full_model

R2_regress = STATS(1)

R2_our_function = R2
disp('Squared correlation between predicted and outcome')
corr(y, X*bhat) .^ 2

%
% Turn this into a function
% I copied and pasted the core equations to the end of the script, then used the "refactor" option in live scripts to save time and do it efficiently.

stat_table = my_GLM(X, y, cnames);

%
%% Analysis 2: Scaling
%% Explore the relationship between beta-hats and group means
% What is the relationship between beta-hats in your model and group means for the 3 groups?
% Why?

% mean(y(logical(X(:, 1)))) - mean(y(logical(X(:, 2))))
yy = backpain.BPI_change;

means_by_group = [nanmean(yy(logical(indic(:, 1)))) nanmean(yy(logical(indic(:, 2)))) nanmean(yy(logical(indic(:, 3))))]

group_mean_diffs = means_by_group * C  % Intercept(sum), PRT-Pla, Ctrl-mean(PRT, Pla)

% Betas are not equal to the group diffs, but they are related!

% Because we calculate X'X, and our X regressors are indicators for group means * C,
% we apply C to the group means twice (multiplicatively). 
diag(C'*C) 

% This is the Gram matrix of C (sums of squares/cross-products), norm of columns ^ 2
norm(C(:, 2)).^2

% If we multiply a column of X by a constant a, we multiply the beta-hat by
% the inverse, 1/a. 
% So beta-hats ii should be the group differences scaled by 1/(C'C_ii)
% Multiply by diag(C'*C) to recover the group mean differences
bhat .* diag(C'*C)

%
% So, I should be able to code X regressors to reflect the group mean differences by normalizing C
% We divide each contrast in C by its Gram matrix
% Check that the contrasts still sum to 1!

C = [1 1 1; 1 -1 0; -0.5 -0.5 1]';  % original

Cnormed = C ./ (diag(C'*C)')

%
%%

% Now re-define X with this new scaling and re-run the model

X = indic * Cnormed;
stat_table = my_GLM(X, yy, cnames);

% Compare with group mean diffs
group_mean_diffs

% They match!

% Rule of thumb: To avoid changing scale, sum of contrast weights^2 should be 1 for each contrast. 

%
%% Analysis 3: Contrasts
% So far, we have coded the group comparisons into the design matrix. But we can also apply contrasts after parameter estimation to test the group differences of interest.
% Extend your GLM function to include contrasts for group means
% Run the GLM with the indicator matrix for groups as X, passing in the contrast matrix C
% Compare the output to our previous group difference-coded design matrix
% We have fit a model without an intercept. Is this OK? Why or why not?

xnames = {'Group 1' 'Group 2' 'Group 3'}';
cnames = {'Intercept' 'PRT-Plac' 'Ctrl-mean(PRT, Plac'}';

[stat_table, contrast_table] = my_GLM_with_contrasts(indic, yy, C, xnames, cnames);


%

function stat_table = my_GLM(X, y, cnames)

% 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'*X);
bhat = xtxi * X' * y;
e = y - X * bhat;

[n, p] = size(X);

var_bhat = e'*e ./ (n - p) * xtxi;
t = bhat ./ (diag(var_bhat) .^ .5);

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

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

% -------------------------------------------

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

disp('Analysis:')
stat_table

end

%%
%%
function [stat_table, contrast_table] = my_GLM_with_contrasts(X, y, C, xnames, cnames)

% 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'*X);
bhat = xtxi * X' * y;
e = y - X * bhat;

[n, p] = size(X);

sigma_sq = e'*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
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));


% -------------------------------------------

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

disp('Model parameters:')
stat_table

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

disp('Contrasts:')
contrast_table

end