function [dsTrain,layers,lossFcn,options] = {functionName}(params)
%% Signal Segmentation by Sweeping Hyperparameters
% This template function defines the training data, network architecture, loss function, and 
% training options for a waveform segmentation experiment using time-frequency 
% features. Experiment Manager uses the output of this function to start the training. 
% For more information about built-in experiments, see <matlab:helpview('deeplearning','exp-mgr-create-experiment') 
% Configure Built-In Training Experiment>.
% 
% Modify this function as a template to set up your own experiment.
%% Input
%% 
% * |params| is a structure with fields from the Experiment Manager hyperparameter 
% table. This template uses two hyperparameters: The type of feature you extract 
% from the signal and the initial learning rate.
%% Output
%% 
% * |dsTrain| is a datastore that contains the training signals and their 
% corresponding labels.
% * |layers| is a layer array that defines the neural network architecture.
% * |lossFcn| is a string that defines the loss function.
% * |options| is a |trainingOptions| object.
%% Download and Load Training Data
% Download the data files into the temporary directory for the system, as returned 
% by |tempdir|. If you intend to place the data files in a folder different from 
% the temporary directory, replace |tempdir| with your directory name.
% 
% The data set consists of 105 two-channel ECG signals. Each channel was independently 
% labeled by the automated expert system, resulting in a total of 210 ECG signals. 
% The signals, along with their region labels, are stored in 210 MAT-files. [1][2]

dataURL = "https://www.mathworks.com/supportfiles/SPT/data/QTDatabaseECGData1.zip";
datasetFolder = fullfile(tempdir,"QTDataset");
if ~exist(datasetFolder,"dir")
    zipFile = websave(tempdir,dataURL);
    unzip(zipFile,tempdir)
    delete(zipFile)
end
%% 
% Create a <matlab:doc('signalDatastore') |signalDatastore|> object that points 
% to the data set.

ds = signalDatastore(datasetFolder,SignalVariableNames=["ecgSignal","signalRegionLabels"]);
%% 
% Reset the random number generator so the data split is consistent across experiments.
% 
% To limit the runtime of the experiment, use only 40% of the data. To use more 
% data, increase the value of |dataRatio|.

rng("default")
dataRatio = 0.4;
numSignals = numel(ds.Files);
dsSubset = subset(ds,randperm(numSignals,ceil(dataRatio*numSignals)));
numSignals = numel(dsSubset.Files);
%% Format Labels
% In segmentation tasks, the network expects a series of categorical labels 
% with the same length as the signal, also known as a signal mask. If your labels 
% are arranged differently, for example as a region-of-interest (<https://www.mathworks.com/help/signal/ref/signalmask.roimask.html 
% ROI>) table, use a transformed datastore and a helper function to restructure 
% them. 
% 
% If your labels are already in the form of a signal mask, comment out this 
% section.

dsFormat = transform(dsSubset,@helperFormatLabel);
%% Resize Data
% Some models or feature extraction methods require specific input sizes. Use 
% a transformed datastore and a helper function to resize the data to the required 
% signal length.
% 
% If no specific signal length is required, comment out the next line.

dsResize = transform(dsFormat,@helperResizeData);
%% Extract Features
% In segmentation tasks, it is important to create features that have the same 
% time resolution as the input signals. The continuous wavelet transform (<matlab:doc('cwt') 
% CWT>) and the Fourier synchrosqueezed transform (<matlab:doc('fsst') FSST>) 
% both have this characteristic. Define the feature extraction function here and 
% choose different features for experimentation.

dsFeature = transform(dsResize,@(x)helperExtractFeatures(x,FeatureType=params.featureType));
feature = preview(dsFeature);
featureSize = size(feature{1});
%% Split Data
% Split the data set into training and validation sets. 
%% 
% * The training set is used to train the model.Use 80% of the data for the 
% training set.
% * The validation set is used to evaluate the performance of the trained model 
% during hyperparameter tuning. Use the remaining 20% of the data for the validation 
% set.
%% 
% In a comprehensive machine learning workflow, it is common to keep a separate 
% test set for the final evaluation of the model. If you intend to evaluate the 
% generalization performance of your finalized model, set aside some data for 
% testing as well.

[trainIdx,validIdx] = dividerand(numSignals,0.8,0.2);
dsTrain = subset(dsFeature,trainIdx);
dsValid = subset(dsFeature,validIdx);
%% Read Data into Memory
% If your system has enough resources, load all your data into memory to speed 
% up the training process. Preloading lets you transform the data beforehand, 
% which eliminates the need to read from disk and apply transformations during 
% each training iteration. If the data do not fit in memory, set |fitMemory| to 
% |false|.

fitMemory = true;
if fitMemory
    dataTrain = readall(dsTrain);
    XTrain = signalDatastore(dataTrain(:,1));
    TTrain = arrayDatastore(vertcat(dataTrain(:,2)),OutputType="same");
    dsTrain = combine(XTrain,TTrain);
    
    dataValid = readall(dsValid);
    XValid = signalDatastore(dataValid(:,1));
    TValid = arrayDatastore(vertcat(dataValid(:,2)),OutputType="same");
    dsValid = combine(XValid,TValid);
end
%% Define Network Architecture
% Create a small neural network using an array of layers. For more information, 
% see <https://uk.mathworks.com/help/deeplearning/ug/example-deep-learning-networks.html 
% Example Deep Learning Networks Architectures>. Use bidirectional long short-term memory 
% (BiLSTM) layers and a fully connected layer to ensure the output channel count 
% matches the number of classes.

numClasses = 4;
layers = [
    sequenceInputLayer(featureSize(1))
    bilstmLayer(50,OutputMode="sequence")
    fullyConnectedLayer(numClasses)
    softmaxLayer
    ];
%% 
% For segmentation tasks, use cross-entropy loss.
lossFcn = "crossentropy";
%% Define Training Hyperparameters
% Call <matlab:doc('trainingOptions') |trainingOptions|> to set the
% hyperparameters to use when training the network. Access the initial
% learning rate defined in the hyperparameter table with the argument value
% |params.initialLearnRate|. Use a mini-batch size of 20. Specify Adam
% (adaptive moment estimation) as the optimizer. To use the parallel pool
% to read the transformed datastore, set |PreprocessingEnvironment| to
% |"parallel"|. Because the training data has sequences with rows and
% columns corresponding to channels and time steps, respectively, specify
% the input data format |"CTB"|  (channel, time, batch).

miniBatchSize = 20;
options = trainingOptions("adam", ...
    InputDataFormats="CTB", ...
    InitialLearnRate=params.initialLearnRate, ...
    MaxEpochs=15, ...
    LearnRateSchedule="piecewise", ...
    LearnRateDropFactor=0.5, ...
    LearnRateDropPeriod=10, ...
    MiniBatchSize=miniBatchSize, ...
    Shuffle="every-epoch", ...
    Plots="training-progress", ...
    Verbose=false, ...
    ValidationData=dsValid, ...
    ExecutionEnvironment="auto", ...
    PreprocessingEnvironment="serial");

%% 
% End of experiment setup function.

end
%% Supporting Functions
% *helperExtractFeatures* 
% 
% Extract features used for training.
%% 
% * |inputCell| is a two-element cell array containing an ECG signal vector 
% and a table of region labels.
% * |outputCell| is a two-element cell array containing the ECG signal features 
% and a categorical label vector mask of the same length as the signal.

function outputCell = helperExtractFeatures(inputCell,Nvargs)
    arguments
        inputCell
        Nvargs.FeatureType = "fsst"
    end
    Fs = 250;
    sigs = inputCell(:,1);
    features = cell(size(sigs));
    switch Nvargs.FeatureType
        case "fsst"
            % Compute time-frequency maps
            for idx = 1:length(sigs)
                [s,f,~] = fsst(sigs{idx},Fs,kaiser(128));
                f_indices = (f > 0.5) & (f < 40);
                features{idx} = [real(s(f_indices,:)); imag(s(f_indices,:))];
            end
        case "cwt"
            for idx = 1:length(sigs)
                [s,f] = cwt(sigs{idx},Fs);
                f_indices = (f > 0.5) & (f < 40);
                features{idx} = [real(s(f_indices,:)); imag(s(f_indices,:))];
            end

        case "none"
            features = sigs;
    end
    outputCell = [features inputCell(:,2)];
end
%% 
% *helperFormatLabel*
% 
% Convert region labels to a mask of labels of size equal to the size of the 
% input ECG signal.
%% 
% * |inputCell| is a two-element cell array containing an ECG signal vector 
% and a table of region labels.
% * |outputCell| is a two-element cell array containing the ECG signal vector 
% and a categorical label vector mask of the same length as the signal.

function outputCell = helperFormatLabel(inputCell)
    sig = inputCell{1};
    roiTable = inputCell{2};
    L = length(sig);
    M = signalMask(roiTable);
    % Get categorical mask and give priority to QRS regions when there is overlap
    priorityOrder = ["QRS";"P";"T"];
    [~,priorityIdx] = ismember(priorityOrder,M.Categories);
    mask = catmask(M,L,OverlapAction="prioritizeByList",PriorityList=priorityIdx);
    % Set missing values to "n/a"
    mask(ismissing(mask)) = "n/a";
    outputCell = {sig,mask};
end
%% 
% *helperResizeData*
% 
% Break input ECG signal and label mask into segments of target length.
%% 
% * |inputCell| is a two-element cell array containing an ECG signal and a label 
% mask.
% * |outputCell| is a two-column cell array containing all the 5000-sample
% signal segments and label masks that could be generated from the input
% data.

function outputCell = helperResizeData(inputCell)
    targetLength = 5000;
    sig = inputCell{1};
    mask = inputCell{2};
    % Get number of chunks
    numChunks = floor(size(sig,1)/targetLength);
    % Truncate signal and mask to integer number of chunks
    sig = sig(1:(numChunks*targetLength));
    mask = mask(1:(numChunks*targetLength));
    % Create a cell array containing signal chunks
    sigOut = reshape(sig,targetLength,numChunks)';
    sigOut = num2cell(sigOut,2);
    % Create a cell array containing mask chunks
    lblOut = reshape(mask,targetLength,numChunks)';
    lblOut = num2cell(lblOut,2); 
    % Output a two-column cell array with all chunks
    outputCell = [sigOut lblOut];
end
%% *References*
% [1] Goldberger, Ary L., Luis A. N. Amaral, Leon Glass, Jeffery M. Hausdorff, 
% Plamen Ch. Ivanov, Roger G. Mark, Joseph E. Mietus, George B. Moody, Chung-Kang 
% Peng, and H. Eugene Stanley. "PhysioBank, PhysioToolkit, and PhysioNet: Components 
% of a New Research Resource for Complex Physiologic Signals." _Circulation_. Vol. 
% 101, No. 23, 2000, pp. e215–e220. doi: 10.1161/01.CIR.101.23.e215.
% 
% [2] Laguna, Pablo, Roger G. Mark, Ary L. Goldberger, and George B. Moody. 
% "A Database for Evaluation of Algorithms for Measurement of QT and Other Waveform 
% Intervals in the ECG." _Computers in Cardiology_. Vol.24, 1997, pp.
% 673–676.
% 
% _Copyright 2023 The MathWorks, Inc._