% Multiscale LSTM combining fine-grained short-term dynamics (each 6 h time

% block) and coarst long-term dynsmics (all past recordings up to now).

% Wei-Long Zheng, MGH

% Email: weilonglive@gmail.com

% short term lstms

clear

close all

load('feature_sequences\all_features_sequences.mat')

save_path = ('multi_scale\');

for num_neuron = 30

    preds_all = {};

    labels_all = {};

    pred_probability_all = {};

    net_model_all = {};

    pts_id_all = {};

    for itrial = 1:10

time_step = 6;

time_range = 12:time_step:96;

num_fold = 5;

preds = cell(length(time_range),num_fold);

labels = cell(length(time_range),num_fold);

pred_probability = cell(length(time_range),num_fold);

net_model = cell(length(time_range),num_fold);

for i = length(time_range)-5%1:length(time_range)

    % short sequneces

    bs_x = bs(:,time_range(i)-time_step+1:time_range(i),:);

    bs_x = permute(bs_x,[1,3,2]);

    bs_x = reshape(bs_x,[size(bs_x,1),size(bs_x,2)*size(bs_x,3)]);

    spike_x = spike(:,time_range(i)-time_step+1:time_range(i),:);

    spike_x = permute(spike_x,[1,3,2]);

    spike_x = reshape(spike_x,[size(spike_x,1),size(spike_x,2)*size(spike_x,3)]);

    x = cat(3, bs_x, spike_x);

    for ifea = 1:7

        feature_tmp = features{ifea}(:,time_range(i)-time_step+1:time_range(i),:);

        feature_tmp = permute(feature_tmp,[1,3,2]);

        feature_tmp = reshape(feature_tmp,[size(feature_tmp,1),size(feature_tmp,2)*size(feature_tmp,3)]);

        if ifea==4||ifea==5||ifea==6||ifea==7

            feature_tmp = 10*log10(feature_tmp);

        end

        x = cat(3, x, feature_tmp);

    end

    % long sequences

    bs_x = bs(:,time_range(1)-time_step+1:time_range(i),:);

    bs_x = permute(bs_x,[1,3,2]);

    bs_x = reshape(bs_x,[size(bs_x,1),size(bs_x,2)*size(bs_x,3)]);

    spike_x = spike(:,time_range(1)-time_step+1:time_range(i),:);

    spike_x = permute(spike_x,[1,3,2]);

    spike_x = reshape(spike_x,[size(spike_x,1),size(spike_x,2)*size(spike_x,3)]);

    x_long = cat(3, bs_x, spike_x);

    for ifea = 1:7

        feature_tmp = features{ifea}(:,time_range(1)-time_step+1:time_range(i),:);

        feature_tmp = permute(feature_tmp,[1,3,2]);

        feature_tmp = reshape(feature_tmp,[size(feature_tmp,1),size(feature_tmp,2)*size(feature_tmp,3)]);

        if ifea==4||ifea==5||ifea==6||ifea==7

            feature_tmp = 10*log10(feature_tmp);

        end

        x_long = cat(3, x_long, feature_tmp);

    end

    cpc_scores_binary = cpc_scores;

    pos_index = find(cpc_scores<3);

    neg_index = find(cpc_scores>=3);

    cpc_scores_binary(pos_index) = 1;

    cpc_scores_binary(neg_index) = 0;

    X = {};

    X_long = {};

    Y = {};

    pts_id = {};

    for ipts = 1:size(x,1)

        x_tmp = squeeze(x(ipts,:,:));

        x_tmp = x_tmp';

        x_long_tmp = squeeze(x_long(ipts,:,:));

        x_long_tmp = x_long_tmp';

        for itmp = 1:size(x_tmp,1)

            x_tmp(itmp,isinf(x_tmp(itmp,:))) = nan;

            x_long_tmp(itmp,isinf(x_long_tmp(itmp,:))) = nan;

            x_tmp(itmp,isnan(x_tmp(itmp,:))) = nanmean(x_tmp(itmp,:));

            temp = x_long_tmp(itmp,end-time_step*12+1:end);

            temp(1,isnan(temp(1,:))) = nanmean(temp(1,:));

            x_long_tmp(itmp,end-time_step*12+1:end) = temp;

            nan_index = find(isnan(x_long_tmp(itmp,:)));

            for jnan = length(nan_index):-1:1

                if nan_index(jnan)+time_step*12<=size(x_long_tmp,2)

                    x_long_tmp(itmp,nan_index(jnan)) = nanmean(x_long_tmp(itmp,nan_index(jnan)+1:nan_index(jnan)+time_step*12));

                end

            end

%             index = find(isinf(x_long_tmp(itmp,:)));

%             for idex = 1:length(index)

%                 if ~isinf(x_long_tmp(itmp,index(idex)-1))

%                     x_long_tmp(itmp,index(idex)) = x_long_tmp(itmp,index(idex)-1);

%                 else

%                     x_long_tmp(itmp,index(idex)) = x_long_tmp(itmp,index(idex)+1);

%                 end

%             end

        end

        % reshape inputs of long term lstm with the same dimensions

        reduced_dim = 72*2;

        x_long_tmp_short = zeros(9,reduced_dim);

        for itmp = 1:size(x_long_tmp,1)

            if size(x_long_tmp,2)>reduced_dim

                n = fix(size(x_long_tmp,2)/72/2); % for every n points, generate 1 points

                b = arrayfun(@(i) mean(x_long_tmp(itmp,i:i+n-1),2),sort([1:size(x_long_tmp,2)/72:size(x_long_tmp,2),n:size(x_long_tmp,2)/72:size(x_long_tmp,2)])); % the averaged vector

                x_long_tmp_short(itmp,:) = b;

            end

        end

        if ~isnan(sum(sum(x_tmp)))&&~isinf(sum(sum(x_tmp)))

            X = [X; x_tmp];

            X_long = [X_long; x_long_tmp_short];

            Y = [Y; num2str(cpc_scores_binary(ipts))];

            pts_id = [pts_id; unique_names{ipts}];

        end

    end

    XV = [X{:}];

    mu = mean(XV,2);

    sg = std(XV,[],2);

    X = cellfun(@(x)(x-mu)./sg,X,'UniformOutput',false);

    XV = [X_long{:}];

    mu = mean(XV,2);

    sg = std(XV,[],2);

    X_long = cellfun(@(x)(x-mu)./sg,X_long,'UniformOutput',false);

    num_pts = length(Y);

    num_test = round(num_pts/num_fold);

    idx = randperm(num_pts);

    X = X(idx);

    X_long = X_long(idx);

    Y = Y(idx);

    pts_id = pts_id(idx);

%     Y = categorical(Y);

    for ifold = 1:num_fold

        if ifold~=num_fold

            start_index = (ifold-1)*num_test+1;

            end_index = ifold*num_test;

        else

            start_index = (ifold-1)*num_test+1;

            end_index = num_pts;

        end

        train_index = setdiff(1:num_pts,start_index:end_index);

        test_data_short = X(start_index:end_index);

        test_label = Y(start_index:end_index);

        train_data_short = X(train_index);

        train_label = Y(train_index);

        test_data_long = X_long(start_index:end_index);

        train_data_long = X_long(train_index);

        train_set = [train_data_short,train_data_long,train_label];

        test_set = [test_data_short,test_data_long,test_label];

        pts_id_test = pts_id(start_index:end_index);

        ipath_short = [save_path,num2str(i),'h','_train_short_',num2str(round(rand(1)*10e6)),'\'];

        mkdir(ipath_short);

        for itrain = 1:length(train_data_short)

            train_samples = train_data_short{itrain};

            save([ipath_short,num2str(itrain)],'train_samples');

        end

        ipath_long = [save_path,num2str(i),'h','_train_long_',num2str(round(rand(1)*10e6)),'\'];

        mkdir(ipath_long);

        for itrain = 1:length(train_data_long)

            train_samples = train_data_long{itrain};

            save([ipath_long,num2str(itrain)],'train_samples');

        end

        ipath_label = [save_path,num2str(i),'h','_train_label_',num2str(round(rand(1)*10e6)),'\'];

        mkdir(ipath_label);

        for itrain = 1:length(train_label)

            train_labels = train_label{itrain};

            save([ipath_label,num2str(itrain)],'train_labels');

        end

        fds_short = fileDatastore(ipath_short,'ReadFcn',@load_variable,'FileExtensions','.mat');

        fds_long = fileDatastore(ipath_long,'ReadFcn',@load_variable,'FileExtensions','.mat');

        fds_label = fileDatastore(ipath_label,'ReadFcn',@load_variable,'FileExtensions','.mat');

        train_datastore = combine(fds_short,fds_long);

        %% LSTM

        miniBatchSize = 150;

        maxEpochs = 100;

        layers_short = [ ...

            sequenceInputLayer(9,'Name','InputLayer')

            sequenceFoldingLayer('Name','fold')

            splittingLayer('Splitting-1st','1st')

            bilstmLayer(num_neuron,'OutputMode','sequence','Name','lstm1_short')

            dropoutLayer(0.1,'Name','dropout1_short')

            bilstmLayer(num_neuron,'OutputMode','sequence','Name','lstm2_short')

%             dropoutLayer(0.1)

            bilstmLayer(num_neuron,'OutputMode','sequence','Name','lstm3_short')

%             dropoutLayer(0.1)

            bilstmLayer(num_neuron,'OutputMode','last','Name','lstm4_short')

%             dropoutLayer(0.1)

            fullyConnectedLayer(num_neuron,'Name','fc_short')

            concatenationLayer(1,2,'Name','cat')

%             additionLayer(2,'Name','add')

            fullyConnectedLayer(2,'Name','fc')

            softmaxLayer('Name','softmax_short')

            classificationLayer('Name','classOutput')

            ];

        layers_long = [ ...

%             sequenceInputLayer(9,'Name','input_long')

            splittingLayer('Splitting-2nd','2nd')

            bilstmLayer(num_neuron,'OutputMode','sequence','Name','lstm1_long')

            dropoutLayer(0.1,'Name','dropout1_long')

            bilstmLayer(num_neuron,'OutputMode','sequence','Name','lstm2_long')

%             dropoutLayer(0.1)

            bilstmLayer(num_neuron,'OutputMode','sequence','Name','lstm3_long')

%             dropoutLayer(0.1)

            bilstmLayer(num_neuron,'OutputMode','last','Name','lstm4_long')

%             dropoutLayer(0.1)

            fullyConnectedLayer(num_neuron,'Name','fc_long')

%             softmaxLayer('Name','softmax_long')

            ];

        lgraph = layerGraph(layers_short);

        lgraph = addLayers(lgraph,layers_long);

        lgraph = connectLayers(lgraph,'fc_long','cat/in2');

        layers = connectLayers(lgraph,'InputLayer','Splitting-2nd');

%         lgraph = addLayers(lgraph,sequenceInputLayer(9,'Name','input'));

%         lgraph = connectLayers(lgraph,'input','lstm1_short');

%         lgraph = connectLayers(lgraph,'input','lstm1_long');

        figure,plot(lgraph)

        options = trainingOptions('sgdm', ...%adam sgdm

            'MaxEpochs',maxEpochs, ...

            'MiniBatchSize', miniBatchSize, ...

            'InitialLearnRate', 0.1, ... %0.8

            'ExecutionEnvironment',"cpu",...%'GradientThreshold', 1, ... 

            'Shuffle','never', ... %every-epoch

            'plots','training-progress', ...%training-progress none

            'ValidationData',{test_set(:,1:2),categorical(test_label)},...

            'Verbose',false);%'OutputFcn', @(info)savetrainingplot(info)

        % concatenation or addition

        net = trainNetwork(train_datastore,train_label,lgraph,options);

        [pred,probabilities] = classify(net,test_data);

        preds{i,ifold} = pred;

        labels{i,ifold} = test_label;

        pred_probability{i,ifold} = probabilities;

        net_model{i,ifold} = net;

        pts_id_fold{i,ifold} = pts_id_test;

    end

end

        preds_all{itrial} = preds;

        labels_all{itrial} = labels;

        pred_probability_all{itrial} = pred_probability;

        net_model_all{itrial} = net_model;

        pts_id_all{itrial} = pts_id_fold;

    end

save(['D:\Research\Cardiac_arrest_EEG\Codes\ComaPrognosticanUsingEEG-master\deep_learning_results\four_layers\','bilstm_four_neurons_',num2str(num_neuron),'_epoch_',num2str(maxEpochs)],'preds','labels','layers','options','pred_probability','net_model');

end

%[updatedNet,YPred] = predictAndUpdateState(recNet,sequences)

% function stop=savetrainingplot(info)

% stop=false;  %prevents this function from ending trainNetwork prematurely

% if info.State=='done'   %check if all iterations have completed

% % if true

%       saveas(gca,'training_process.png')  % save figure as .png, you can change this

% 

% end

% end

% options = trainingOptions('sgdm',...

%     'InitialLearnRate',0.003,...

%     'Plots','training-progress', ...

%     'ValidationData',garVal,...

%     'ValidationFrequency',40,...

%     'MaxEpochs',1,...

%     'LearnRateSchedule', 'piecewise',...

%     'LearnRateDropPeriod',3,...

%     'Shuffle','every-epoch',...

%     'ValidationPatience',5,...

%     'OutputFcn',@(info)SaveTrainingPlot(info),...

c%     'Verbose',true);

% % ... Training code ...

% % At the end of the script:

% function stop = SaveTrainingPlot(info)

% stop = false;

% if info.State == "done"

%     currentfig = findall(groot,'Type','Figure');

%     savefig(currentfig,'prova.png')

% end

% end