function cluster_estimations_code2_randomdata

% Copyright (c) 2023, Lazaro Sanchez-Rodriguez et al., 
% "Personalized whole-brain neural mass models reveal combined Aβ and tau hyperexcitable influences in Alzheimer’s disease" 
% (CommsBio)

%% define Surrogate Optimization options
numInits = 20;
% constraints
lb = [-0.02,-0.02,-0.02];
ub = [ 0.02, 0.02, 0.02];
A = [ 1  1  1;...
     -1 -1 -1;...
      1  1 -1;...
      1 -1 -1;...
      1 -1  1;...
     -1  1  1;...
     -1  1 -1;...
     -1 -1  1;...
    ];
b = 0.02*ones(8,1);
intcon = [];
Aeq = [];
beq = [];

opts = optimoptions('surrogateopt','PlotFcn',[],'Display','iter',...
    'MaxFunctionEvaluations',500);
% Note: We preferred to run the algorithm in steps of 500 iterations each,
% by loading the saved files and defining <<'InitialPoints',trials>> as an
% option. This was done to circumvent the queue on Compute Canada as jobs
% with shorter time requests are generally scheduled to run first. All
% estimations converged after less than 2,000 evaluations or 4 runs of 500

%% load data
load('random_example_data.mat')
Taus = Tausnew;
ABetas = Abetasnew;
falff_reg_obs = falff_reg_obs_new;
C = Cnew;
subIndex = 1;
% times for the BOLD simulation
TR = .681;
tspan = (0:860-1)*TR;

%% parallel runs

% Create a "local" cluster object
local_cluster = parcluster('local');

% % Modify the JobStorageLocation to $SLURM_TMPDIR (Compute Canada)
% comment the next line if running on a PC
%local_cluster.JobStorageLocation = getenv('SLURM_TMPDIR');

% Start the parallel pool
parpool(local_cluster,2);%local_cluster.NumWorkers

sc = parallel.pool.Constant(RandStream('Threefry'));
parfor i = 1:numInits

    % set the environment up
    stream = sc.Value; % for reproducibility, cluster environment
    stream.Substream = i;
    

    myObjective = @(theta) fun_cost(theta,tspan,C,falff_reg_obs,TR,Taus,ABetas);
    [Theta_Est,resnorm,exitflag,output,trials] = surrogateopt(myObjective,lb,ub,intcon,A,b,Aeq,beq,opts);
    
    %save the results for this subject and its surrogateopt run (initial condition)
    parsave(sprintf('Rand_WCestimate_Effectv0ewConstraints_eta2_SurrOpt_S%i_Run%i_Inter.mat',subIndex,i),...
        Theta_Est,resnorm,exitflag,output,trials,ABetas,Taus,falff_reg_obs);

end
delete(gcp('nocreate'))

end


function cost = fun_cost(theta,tspan,C,falff_reg_obs,TR,Taus,ABetas)

%% define interactions and pathological influence model
eta = 1;
Nr = 66; etaNr = (eta/Nr);
kABeta      = theta(1);
kTau        = theta(2);
kABetaTau   = theta(3);

v0e = 2.84 + kABeta*ABetas + kTau*Taus + kABetaTau.*ABetas.*Taus;
% excitability = 1./v0e

%% define times
time2transient = 20;
deltat = 0.001;
times = 0:deltat:100 + time2transient;
transient_index = find(times == time2transient);
times_BOLD = 0:deltat:tspan(end) + time2transient;

%% define parameters
% Neural mass parameters 
%(see Supplementary Table)
ae = 1;
ai = 1;
Cee = 6.4;
Cii = 1.2;
Cie = 4.8;
Cei = 6;
v0i = 4;
P = 0.65; 
tauE = 1e-2;   %[s]
tauI = 2*1e-2; %[s]


% BOLD params
%(see Supplementary Table)
Ae = 1;
Ai = 1;
ke_gluc = 1;
ki_gluc = 1;    
c  = 2.5;
d  = 1.6;
gam= 5;
epsilon= 0.8;        
miu= 0;
kf = 1.7;
k0 = 1;
chi= 0.4;
damp = 1;
%
flow_lff     = 0.01;                             % lowpass frequency of filter
fhigh_lff    = 0.08;                             % highpass
V0 = 0.03;
niu0 = 80.6;
E0 = 0.4;
epsi = 0.24;
r0 = 178;
TE = 32*1e-3;
a1 = 4.3*niu0*E0*TE + epsi*r0*E0*TE;
a2 = epsi*r0*E0*TE + epsi - 1;

%% define initial conditions
%for the neural masses
X0PSP = [0.01*ones(Nr,1);0.075*ones(Nr,1)];
X0 = X0PSP;

%for the BOLD generation variables 
y0               = zeros(10*Nr,1);
y0(2*Nr+1: 3*Nr) = ones(Nr,1); %ge
y0(3*Nr+1: 4*Nr) = zeros(Nr,1);%ze
y0(4*Nr+1: 5*Nr) = ones(Nr,1); %gi
y0(5*Nr+1: 6*Nr) = zeros(Nr,1);%zi
y0(6*Nr+1: 7*Nr) = ones(Nr,1); %ff
y0(7*Nr+1: 8*Nr) = zeros(Nr,1);%zf
y0(8*Nr+1: 9*Nr) = ones(Nr,1); %b
y0(9*Nr+1: 10*Nr)= ones(Nr,1); %q

%% calculate

% initial run
opts = odeset('RelTol',1e-5,'AbsTol',1e-7);
[~,x] = ode45(@DifEq_Neur, times, X0,opts);
Eact1 = x(transient_index:end,Nr+1:2*Nr);
Iact1 = x(transient_index:end,1:Nr);
Pinter_Glob = etaNr*C*x(transient_index:end,Nr+1:2*Nr)';
vi_Glob = Cei*Eact1' - Cii*Iact1';
ve_Glob = Cee*Eact1' - Cie*Iact1' + P + Pinter_Glob;
Sig_i_Glob = 1./(1 + exp(-ai*(vi_Glob - v0i))) - 1./(1 + exp(v0i.*ai));
Sig_e_Glob = 1./(1 + exp(-ae*(ve_Glob - v0e))) - 1./(1 + exp(v0e.*ae));
mean_Sig_e = mean(Sig_e_Glob,2);%baseline (average value) Sotero et al 2008
mean_Sig_i = mean(Sig_i_Glob,2);%baseline (average value) Sotero et al 2008

% full run with generation of BOLD signal 
y0(1:2*Nr)  = [Iact1(1,:)';Eact1(1,:)'];
[~,xBOLD] = ode45(@DifEq_BOLD, times_BOLD, y0,opts);
b = xBOLD(:,8*Nr+1:9*Nr);
q = xBOLD(:,9*Nr+1:10*Nr);
BOLD_sim = V0*(a1*(1-q) - a2*(1-b)); %BOLD_sim = 100*BOLD_sim;

% calculating falff
BOLD_sim = BOLD_sim(transient_index:end,:);
boldsig = BOLD_sim(1:TR/deltat:end,:);
falff_reg_sim = f_alff_localfunction(boldsig, TR, fhigh_lff, flow_lff);

% distance to the subject's BOLD signal indicators:
cost = pdist2(falff_reg_obs,falff_reg_sim,'correlation');

    function f = DifEq_Neur(tt,xx)
        % simulate neuronal activity, initial segment

        f = zeros(2*Nr,1);

        iLoc = xx(     1:  Nr);
        eLoc = xx(  Nr+1:2*Nr);

        Pinter = etaNr*C*(eLoc);
        viLoc = Cei*(eLoc) - Cii*(iLoc);
        veLoc = Cee*(eLoc) - Cie*(iLoc) + P + Pinter;

        Sig_iLoc = 1./(1 + exp(-ai*(viLoc - v0i))) - 1./(1 + exp(v0i.*ai));
        Sig_eLoc = 1./(1 + exp(-ae*(veLoc - v0e))) - 1./(1 + exp(v0e.*ae));

        f(1:Nr)      = (1/tauI).*(-iLoc + Sig_iLoc);
        f(Nr+1:2*Nr) = (1/tauE).*(-eLoc + Sig_eLoc);

    end

    function f = DifEq_BOLD(tt,xx)
        % simulate neuronal activity and BOLD, whole time interval

        f = zeros(6*Nr,1);
        
        iLoc = xx(     1:  Nr);
        eLoc = xx(  Nr+1:2*Nr);

        Pinter = etaNr*C*(eLoc);
        viLoc = Cei*(eLoc) - Cii*(iLoc);
        veLoc = Cee*(eLoc) - Cie*(iLoc) + P + Pinter;

        Sig_iLoc = 1./(1 + exp(-ai*(viLoc - v0i))) - 1./(1 + exp(v0i.*ai));
        Sig_eLoc = 1./(1 + exp(-ae*(veLoc - v0e))) - 1./(1 + exp(v0e.*ae));

        f(1:Nr)      = (1/tauI).*(-iLoc + Sig_iLoc);
        f(Nr+1:2*Nr) = (1/tauE).*(-eLoc + Sig_eLoc);

        input_e = Sig_eLoc;
        input_e = (input_e./mean_Sig_e);%
        input_i = Sig_iLoc;
        input_i = (input_i./mean_Sig_i);%

        ge = xx(2*Nr+1:3*Nr);
        ze = xx(3*Nr+1:4*Nr);
        gi = xx(4*Nr+1:5*Nr);
        zi = xx(5*Nr+1:6*Nr);
        ff = xx(6*Nr+1:7*Nr);
        zz  = xx(7*Nr+1:8*Nr);
        bb  = xx(8*Nr+1:9*Nr);
        qq  = xx(9*Nr+1:10*Nr);
        x0 = 1./(1+exp(c.*(d-1)));
        xfun = 1./(1+exp(c.*(d-ge)));
        me = ge.*(2 - xfun)./(2 - x0);
        mi = gi;
        m  = (gam*me + mi)/(gam + 1);
        fou = bb.^(1/chi);

        f(2*Nr+1:3*Nr)  = ze;
        f(3*Nr+1:4*Nr)  = -2*damp*ze./ke_gluc - (ge - 1)./(ke_gluc.^2) + (Ae./ke_gluc).*(input_e - 1);
        f(4*Nr+1:5*Nr)  = zi;
        f(5*Nr+1:6*Nr)  = -2*damp*zi./ki_gluc - (gi - 1)./(ki_gluc.^2) + (Ai./ki_gluc).*(input_i - 1);
        f(6*Nr+1:7*Nr)  = zz;
        f(7*Nr+1:8*Nr)  = -2*damp*zz./kf  - (ff - 1)./(kf.^2) + epsilon.*(input_e - 1) + miu.*(input_i - 1);
        f(8*Nr+1:9*Nr)  = (1/k0)*(ff - fou);
        f(9*Nr+1:10*Nr) = (1/k0)*(m - fou.*qq./bb);

    end

end

function parsave(fname, Theta_Est,resnorm,exitflag,output,trials,ABetas,Taus,falff_reg_obs)
save(fname, 'Theta_Est','resnorm','exitflag','output','trials','ABetas','Taus',"falff_reg_obs");
end

function fALFF_2D = f_alff_localfunction(Signal, ASamplePeriod, HighCutoff, LowCutoff)

% addapted from Yang et al 2018; please also cite the RESTplus toolkit https://doi.org/10.1016/j.scib.2019.05.008 

nDimTimePoints = size(Signal,1);

% Get the frequency index
sampleFreq 	 = 1/ASamplePeriod;
sampleLength = nDimTimePoints;
paddedLength = sampleLength; 
if (LowCutoff >= sampleFreq/2) % All high included
    idx_LowCutoff = paddedLength/2 + 1;
else % high cut off, such as freq > 0.01 Hz
    idx_LowCutoff = ceil(LowCutoff * paddedLength * ASamplePeriod + 1);
end
if (HighCutoff>=sampleFreq/2)||(HighCutoff==0) % All low pass
    idx_HighCutoff = paddedLength/2 + 1;
else % Low pass, such as freq < 0.08 Hz
    idx_HighCutoff = fix(HighCutoff *paddedLength *ASamplePeriod + 1);
end


Signal_copy = Signal;
for i = 1:size(Signal,2)
    dtsig = detrend(Signal(1:end,i)-mean(Signal(1:end,i)));
    dtsig(find(dtsig>3*std(dtsig))) = 3*std(dtsig);
    dtsig(find(dtsig<-3*std(dtsig))) = -3*std(dtsig);
    Signal_copy(:,i) = dtsig;
end

ftranf = 2*abs(fft(Signal_copy))/sampleLength;

% Generate fALFF
fALFF_2D = sum(ftranf(idx_LowCutoff:idx_HighCutoff,:)) ./ sum(ftranf(2:(paddedLength/2 + 1),:));
fALFF_2D(~isfinite(fALFF_2D))=0;

clearvars -except fALFF_2D 

end