% Generates simulated data as the sum of outputs from 3 latent networks
clear; clc

SAVEFILE = 'sim.mat';

N_REGIONS = 5;

DECAY1 = 0.98;
MU1 = 5;
DECAY2 = 0.9;
MU2 = 30;

SEED = 1337;
rng(SEED)

N_WINDOWS = 1e4;
WIN_LEN = 5;
FS = 500;

LOOP1DELAY = 10;
LOOP2DELAY = 3;
LOOP1STR = 3e-3;
LOOP2STR = 2e-2;

NOISE_COV1 = 1;
NOISE_COV2 = 1;
NOISE_COV3 = 1;

%% create VAR models to simulate data

% generate intra-region AR coefficient from desired poles
pole1a = DECAY1*exp(1i*MU1/FS*(2*pi));
pole1b = DECAY1*exp(-1i*MU1/FS*(2*pi));

pole2a = DECAY1*exp(1i*MU2/FS*(2*pi));
pole2b = DECAY1*exp(-1i*MU2/FS*(2*pi));

sys1 = zpk([], [pole1a, pole1b], LOOP1STR, 1/FS);
sys1Denom = tf(sys1).Denominator{1};
nCoefs1 = numel(sys1Denom)-1;
nLags1 = max(LOOP1DELAY, nCoefs1);
coef1intra = zeros(nLags1,1);
coef1intra(1:nCoefs1) = -sys1Denom(2:end);

sys2 = zpk([], [pole2a, pole2b], LOOP2STR, 1/FS);
sys2Denom = tf(sys2).Denominator{1};
nCoefs2 = numel(sys2Denom)-1;
nLags2 = max(LOOP2DELAY, nCoefs2);
coef2intra = zeros(nLags2,1);
coef2intra(1:nCoefs2) = -sys2Denom(2:end);

% make inter region AR coefficients simple impulse
coef1inter = zeros(nLags1,1);
coef1inter(LOOP1DELAY) = LOOP1STR;
coef2inter = zeros(nLags1,1);
coef2inter(LOOP2DELAY) = LOOP2STR;

% populate VAR1 parameters
net1 = varm(N_REGIONS, nLags1);
preCov = randn(N_REGIONS)/N_REGIONS;
net1Covariance = NOISE_COV1*((preCov + preCov')/2 + eye(N_REGIONS));
net1.Constant = zeros(N_REGIONS,1);
for k = 1:nLags1
    thisLag = zeros(N_REGIONS);
    
    % the network simulated by this var is 1>2>3>4 with recurrent
    % connections within 1,2,3,&4
    for r = 1:4
       thisLag(r,r) = coef1intra(k); 
    end
    thisLag(2,1) = coef1inter(k);
    thisLag(3,2) = coef1inter(k);
    thisLag(4,3) = coef1inter(k);
    net1.AR{k} = thisLag;
end

% populate VAR2 parameters
net2 = varm(N_REGIONS, nLags2);
preCov = randn(N_REGIONS)/N_REGIONS;
net2Covariance = NOISE_COV2*((preCov + preCov')/2 + eye(N_REGIONS));
net2.Constant = zeros(N_REGIONS,1);
for k = 1:nLags2
    thisLag = zeros(N_REGIONS);
    
    % the network simulated by this var is 1<2<3 with recurrent
    % connections within 1,2,&3
    for r = 1:3
       thisLag(r,r) = coef2intra(k); 
    end
    thisLag(1,2) = coef2inter(k);
    thisLag(2,3) = coef2inter(k);
    net2.AR{k} = thisLag;
end

% populate VAR3 parameters
net3 = varm(N_REGIONS, nLags1);
net3Covariance = NOISE_COV3*((preCov + preCov')/2 + eye(N_REGIONS));
net3.Constant = zeros(N_REGIONS,1);
for k = 1:nLags1
    thisLag = zeros(N_REGIONS);
    
    % the network simulated by this var is 5>1 + recurrent connections within region 1 & 5
    thisLag(1,1) = coef1intra(k); 
    thisLag(5,5) = coef1intra(k);
    thisLag(1,5) = coef1inter(k);
    net3.AR{k} = thisLag;
end


%% Simulate data from VAR model
nObsPerWindow = FS*WIN_LEN;

% initialize network data arrays
data1 = zeros(nObsPerWindow, N_REGIONS, N_WINDOWS);
data2 = zeros(nObsPerWindow, N_REGIONS, N_WINDOWS);
data3 = zeros(nObsPerWindow, N_REGIONS, N_WINDOWS);

% generate scores for each network
scores = rand(3, N_WINDOWS);

% simulate each window w/ scaled innovation term
for w = 1:N_WINDOWS
    net1.Covariance = net1Covariance * scores(1,w);
    net2.Covariance = net2Covariance * scores(2,w);
    net3.Covariance = net3Covariance * scores(3,w);
    
    data1(:,:,w) = simulate(net1, nObsPerWindow);
    data2(:,:,w) = simulate(net2, nObsPerWindow);
    data3(:,:,w) = simulate(net3, nObsPerWindow);
end

% sum network contributions and normalize to get final data
X = data1 + data2 + data3;
X = bsxfun(@rdivide, X, mean(std(X,1),3));

%% create labels structure
labels.fs = FS;
labels.windowLength = WIN_LEN;
labels.area = {'A'; 'B'; 'C'; 'D'; 'E'};


save(SAVEFILE, 'X', 'scores', 'labels',...
    'data1', 'data2', 'data3', 'net1', 'net2', 'net3', '-v7.3')
