%% Sleepstudy dataset example
% We'll fit the "sleepstudy" dataset, a standard demo dataset in R, with three different mixed effects models:
% fitlme, Matlab's mixed effects model
% IGLS, CANlab's Iterative Generalized Least Squares 2-level model
% glmfit_multilevel, CANlab's two-stage reweighted summary statistics model
% Table of Contents
% Dataset description
%     Description
%     Format
% Load the dataset
% fitlme
% IGLS: Iterative generalized least squares
% glmfit_multilevel
%% Dataset description
%% Description
% The average reaction time per day (in milliseconds) for subjects in a sleep deprivation study.
% Days 0-1 were adaptation and training (T1/T2), day 2 was baseline (B); sleep deprivation started after day 2.
%% Format
% A data frame with 180 observations on the following 3 variables.
% Reaction: Average reaction time (ms)
% Days: Number of days of sleep deprivation
% Subject: Subject number on which the observation was made.
% Details
% These data are from the study described in Belenky et al. (2003), for the most sleep-deprived group (3 hours time-in-bed) and for the first 10 days of the study, up to the recovery period. The original study analyzed speed (1/(reaction time)) and treated day as a categorical rather than a continuous predictor.
% References
% Gregory Belenky, Nancy J. Wesensten, David R. Thorne, Maria L. Thomas, Helen C. Sing, Daniel P. Redmond, Michael B. Russo and Thomas J. Balkin (2003) Patterns of performance degradation and restoration during sleep restriction and subsequent recovery: a sleep dose-response study. Journal of Sleep Research 12, 1--12.
%
%% Load the dataset
% Before you start, you'll want the dataset on your path:
% basedir = '/Users/f003vz1/Dropbox (Dartmouth College)/COURSES/Courses_Dartmouth/2023_Fall_Computational_Foundations';
% g = genpath(fullfile(basedir));
% addpath(g);
%
%

sleepstudy = readtable('sleepstudy.csv');

%
%% fitlme

lme = fitlme(sleepstudy,'Reaction ~ Days + (Days | Subject)');
disp(lme)

% Note: The degrees of freedom and P-values in the fixed effects table are wrong!
% For valid estimates of the dfe, you must correct the df using the Satterthwaite method, or do model comparisons to test significance levels. I prefer the former because it's straightforward.
%

[~, ~, fixedstats] = fixedEffects(lme, 'DFMethod', 'Satterthwaite');
disp(fixedstats)

% % This is another way to do it:
% anova(lme, 'DFMethod', 'Satterthwaite')
%
%
% The default is maximum likelihood estimation, so let's try it using ReML instead:

lme = fitlme(sleepstudy,'Reaction ~ Days + (Days | Subject)', 'FitMethod','REML');
disp(lme)

[~, ~, fixedstats] = fixedEffects(lme, 'DFMethod', 'Satterthwaite');
disp(fixedstats)

% Let's get the BLUPs, the Best Linear Unbiased Predictors.
% These are estimates of the individual parameters for each person, with shrinkage of high-variance estimates towards the group values.
% They are deviations around the group (population) values, so we have to add the population parameters to get the individual slopes.
% These are what Doug Bates calls conditional means, E[beta|Y=y], evaluated at the estimated parameters

[~, ~, randomstats] = randomEffects(lme);
blups = reshape(randomstats.Estimate, 2, length(unique(randomstats.Level)))';
levelnames = unique(randomstats.Level);
bnames = unique(randomstats.Name, 'stable')
blups = array2table(blups, 'Variablenames', bnames, 'RowNames', levelnames )

% Now create the plot
% ---------------------------
indiv_slopes = table2array(blups) + lme.Coefficients.Estimate';
days = length(unique(sleepstudy.Days));

% This is a local function created to show a simple plot of slopes
plot_indiv_slopes(indiv_slopes, days);

%
%% IGLS: Iterative generalized least squares

% To use IGLS and glmfit_multilevel, we need to reformat the dataset to
% cell arrays

sleepunstacked = unstack(sleepstudy,'Reaction','Days');
sleepunstacked = table2array(sleepunstacked(:, 2:end))';
[n, m] = size(sleepunstacked);  % n = days, m  = subjects
sleep_rt_cell = mat2cell(sleepunstacked, n, ones(1, m));

daycode = ((1:n)') - 1 ;
X1 = repmat({daycode}, 1, m);

%

% IGLS requires CANlab tools
% see also rigls.m for ReML estimation

igls_stats = igls(sleep_rt_cell, X1, 'plot', 'all');

%

% BLUPs 
% Create the same simple plot as before, to compare
indiv_slopes = igls_stats.beta_indiv';

plot_indiv_slopes(indiv_slopes, days);

%
% Here is a different plot, with different color options.
% This plots the naive estimates of individual slopes, not the BLUPs. But you can see that they are very close to the BLUPS in this case.
% This requires CANlab tools)

create_figure('example');
handles = line_plot_multisubject(X1, sleep_rt_cell, 'colors', seaborn_colors(length(X1)));
set(handles.point_handles, 'MarkerSize', 8, 'Marker', 'o');
set(gca, 'FontSize', 18);
xlabel('Days of sleep deprivation'); ylabel('Reaction Time (RT)');

%
%% glmfit_multilevel

glmfit_stats = glmfit_multilevel(sleep_rt_cell, X1, [], 'names', {'Days'}, 'verbose', 'weighted', 'plot');

%

function plot_indiv_slopes(indiv_slopes, days)

figure; hold on;
m = size(indiv_slopes, 1); % number of subjects
colors = seaborn_colors(m);

for i = 1:m
    
    % plot individual line
    plot([0 days], [indiv_slopes(i, 1) indiv_slopes(i, 1) + days * indiv_slopes(i, 2)], 'color', colors{i}, 'LineWidth', 2);
    
end

% plot group
% b = lme.Coefficients.Estimate;
% plot([0 days], [b(1) b(1) + days * b(2)], 'color', 'k', 'LineWidth', 4);

set(gca, 'FontSize', 18);
xlabel('Days of sleep deprivation'); ylabel('Reaction Time (RT)');
end