% Generate simulated data for mixed effects model analysis
% Inputs and sample values:
% m = 20;                     % number of subjects
% n_i = 20;                   % n observations within-subject. OR: [1 x m] vector of n per subject
% p = 2;                        % Number of model parameters (regressors), including intercept
% s2 = 5;                       % Within-person residual noise variance. OR: [1 x m] vector of s^2 per subject
% U_g = [20 0; 0 10];    % p x p  covariance matrix of true population-level parameters
% b_g = [5 2]';               % p x 1 vector of true fixed effect parameters
% regressortype = 'continuous';     % type of regressor: 'continuous' or 'binary'
% Use these to generate a sample dataset:
% [y, Z_i, x1, sim_data] = sim_generate_mixedfx_data1(m, n_i, p, b_g, s2, U_g, regressortype);
% To fit the dataset, we'll need one of two things:
% 1. these variables:
%         y{i = 1...m}            The DV/outcome
%         Z_i{i = 1...m}         The design matrix for each person
%         CANlab functions like glmfit_multilevel, igls, and rigls can use these cell arrays directly.
%
% 2.  sim_data, a long-format data table
%         For fitlm and fitlme, and lmer in R (etc.), we will want these in stacked, 'long format' tabular form for mixed effects models.  This requires a Grouping (in this case a Subject ID variable) as well. We will also create a column for each predictor and leave out the intercept, as it's added automatically.
%
% To plot data, we want to extract a predictor (x1) and plot it against the dv.  If you enter p > 1, the first predictor (x1) is already extracted in a cell array format (like y). You can extract other predictors like this:
% x1 = cellfun(@(x) x(:, 2), Z_i, 'UniformOutput', false);  % 2 indicates predictor 1 here, as the intercept is first.
%

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

% Expand if scalar choices have been entered
% -----------------------------------------------------
if isscalar(n_i)
    n_i = n_i * ones(m, 1);
end

if isscalar(s2)
    s2 = s2 .* ones(1, m);
end

% Check the sizes
% -----------------------------------------------------
if any(size(b_g) - [p 1])
    error('Population effects b_g is the wrong size')
end

if ~all(size(U_g) == p)
    error('Covariance matrix U_g is the wrong size')
end

Z_i = generate_within_design_matrix_Z_i(n_i, p, regressortype);
% Generate [n_i(i) x p] within-person design matrix Z_i for each group {1....m}


%% Create simulated dataset
% Generate subject-level true parameters
% Assume a normal distribution with covariance  U_g
% beta_i here is an [p x m] matrix where column i is the [p x 1] true beta coefficient vector for subject i.

beta_i = mvnrnd(b_g, U_g, m)';

% Generate data (y)

y = cell(1, m);

for i = 1:m         

    y{i} = Z_i{i} * beta_i(:, i) + sqrt(s2(i)) .* randn(n_i(i), 1);

end

% To fit the dataset, we'll need these variables:
% y{i = 1...m}            The DV/outcome
% Z_i{i = 1...m}         The design matrix for each person
% CANlab functions like glmfit_multilevel, igls, and rigls can use these cell arrays directly.
% For fitlm and fitlme, and lmer in R (etc.), we will want these in stacked, 'long format' tabular form for mixed effects models.  This requires a Grouping (in this case a Subject ID variable) as well. We will also create a column for each predictor and leave out the intercept, as it's added automatically.

% Make a long-format table with simulated data
sim_data = generate_data_table_object(y, m, p, Z_i);

% Extract the first column of x1

if p > 1
    x1 = cellfun(@(x) x(:, 2), Z_i, 'UniformOutput', false);
end



end % main function

%

function Z_i = generate_within_design_matrix_Z_i(n_i, p, regressortype)
% Generate [n_i(i) x p] within-person design matrix Z_i for each group {1....m}

m = length(n_i);
Z_i = cell(1, m);

for i = 1:m         % allows subjects to vary later if needed

    Z_within = [ones(n_i(i), 1)]; % intercept

    Z_tmp = randn(n_i(i), p - 1);

    switch regressortype
        case 'continuous'
            % do nothing
        case {'binary', 'categorical'}
            Z_tmp(Z_tmp >= 0) = 1;
            Z_tmp(Z_tmp <= 0) = -1;
        otherwise
            error('Unknown regressortype')
    end

    Z_within = [Z_within Z_tmp];

    Z_i{i} = Z_within;

end

end % generate_within_design_matrix_Z_i

%


function sim_data = generate_data_table_object(y, m, p, Z_i)
% Make a long-format table with simulated data

sim_data = table(cat(1, y{:}), 'VariableNames', {'Y'});

sid_tmp = cell(1, m);       % Subject ID
var_tmp = cell(1, p - 1);   % predictor variables

for i = 1:m         

    sid_tmp{i} = i .* Z_i{i}(:, 1);  % use intercept in case n's vary

    for j = 1:(p-1)

        var_tmp{j}{i} = Z_i{i}(:, j+1); 

    end

end

sim_data.sid = cat(1, sid_tmp{:});

% add other variables
for j = 1:(p-1)

    varname = sprintf('x%02d', j);
    sim_data.(varname) = cat(1, var_tmp{j}{:});

end

end % generate table function
