function [dsTrain,layers,lossFcn,options] = {functionName}(params)
%% Signal Classification by Sweeping Hyperparameters
% This template function defines the training data, network architecture, loss function, and 
% training options for a signal classification 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.
% 
% This example uses 162 ECG recordings from three PhysioNet databases: MIT-BIH 
% Arrhythmia Database [1][2], MIT-BIH Normal Sinus Rhythm Database [1], and The 
% BIDMC Congestive Heart Failure Database [1][3]. In total, there are 96 recordings 
% from persons with arrhythmia, 30 recordings from persons with congestive heart 
% failure, and 36 recordings from persons with normal sinus rhythms. The goal 
% is to train a classifier to distinguish between arrhythmia (ARR), congestive 
% heart failure (CHF), and normal sinus rhythm (NSR).

dataURL = "https://raw.githubusercontent.com/mathworks/physionet_ECG_data/main/ECGData.zip";
datasetFolder = fullfile(tempdir,"PhysionetECGData");
if ~exist(datasetFolder,"dir")
     mkdir(datasetFolder)
     zipFile = websave(tempdir,dataURL);
     unzip(zipFile,datasetFolder)
     delete(zipFile)
end
%% 
% Use a <matlab:doc('signalDatastore') |signalDatastore|> and a <matlab:doc('signalDatastore') 
% arrayDatastore> to create a datastore that combines signals and labels.

ECGData = importdata(fullfile(datasetFolder,"ECGData.mat"));
data = ECGData.Data;
numSignals = size(data,1);
data = mat2cell(data,ones(numSignals,1),size(data,2));
labels = ECGData.Labels;
labels = categorical(labels);
ds = combine(signalDatastore(data,SampleRate=128),arrayDatastore(labels));
%% 
% 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;
idx = splitlabels(labels,dataRatio);
ds = subset(ds,idx{1});
labels = labels(idx{1});
%% 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.

ds = transform(ds,@helperResizeData);
%% Extract Features
% In classification tasks, it is generally better to use extracted time-frequency 
% features rather than raw signals. Features highight important information in 
% data and thus help models do a better job of identifying and learning patterns. 
% Features also reduce noise and dimensionality, which improves model accuracy 
% and efficiency.
% 
% Define the feature extraction function here and choose different features 
% for experimentation.

dsFeature = transform(ds,@(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 to perform a 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.

trainValIdx = splitlabels(labels,[0.8,0.2]);
trainIdx = trainValIdx{1};
validIdx = trainValIdx{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>.

numF = 4;
numClasses = length(unique(labels));
layers = [
    imageInputLayer(featureSize,Normalization="none")

    convolution2dLayer(5,numF,Padding="same")
    batchNormalizationLayer
    reluLayer
    maxPooling2dLayer(3,Stride=2,Padding="same")

    convolution2dLayer(3,2*numF,Padding="same")
    batchNormalizationLayer
    reluLayer
    maxPooling2dLayer(3,Stride=2,Padding="same")

    convolution2dLayer(3,4*numF,Padding="same")
    batchNormalizationLayer
    reluLayer

    convolution2dLayer(3,4*numF,Padding="same")
    batchNormalizationLayer
    reluLayer

    averagePooling2dLayer(2)

    dropoutLayer(0.3)
    fullyConnectedLayer(numClasses)
    softmaxLayer
    ];
%% 
% For classification 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"|.

miniBatchSize = 20;
options = trainingOptions("adam", ...
    InitialLearnRate=params.initialLearnRate, ...
    MaxEpochs=80, ...
    LearnRateSchedule="piecewise", ...
    LearnRateDropFactor=0.9, ...
    LearnRateDropPeriod=10, ...
    MiniBatchSize=miniBatchSize, ...
    Shuffle="every-epoch", ...
    Plots="training-progress", ...
    Metrics="accuracy", ...
    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 categorical label.
% * |outputCell| is a two-element cell array containing the ECG signal features 
% and a categorical label.

function outputCell = helperExtractFeatures(inputCell,Nvargs)
    arguments
        inputCell
        Nvargs.FeatureType = "stft"
    end
    Fs = 128;
    sigs = inputCell(:,1);
    features = cell(size(sigs));
    switch Nvargs.FeatureType
        case "stft"
            % Compute time-frequency maps
            for idx = 1:length(sigs)
                [s,f,~] = stft(sigs{idx},Fs,"FrequencyRange","onesided");
                findices = (f > 0.5) & (f < 40);
                features{idx} = abs(s(findices,:));
            end
        case "cwt"
            for idx = 1:length(sigs)
                [s,f] = cwt(sigs{idx},Fs);
                findices = (f > 0.5) & (f < 40);
                % Resize the feature to avoid excessively large feature sizes
                features{idx} = imresize(abs(s(findices,:)),[64,1024]);
            end
    end
    outputCell = [features inputCell(:,2)];
end
%% 
% *helperResizeData*
% 
% Pad or truncate input ECG signal into segments of target length.
%% 
% * |inputCell| is a two-element cell array containing an ECG signal and a label.
% * |outputCell| is a two-column cell array containing all signal segments with 
% the same length and labels.

function outputCell = helperResizeData(inputCell)
    targetLength = 60000;
    % Pad or truncate sequence data to target length
    sig = padsequences(inputCell(:,1),2, ...
        Length=targetLength, ...
        PaddingValue="symmetric");
    % Normalize data
    sig = normalize(sig);
    outputCell = [{sig} inputCell(2)];
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] Moody, G. B., and R. G. Mark. "The impact of the MIT-BIH Arrhythmia Database."  
% _IEEE Engineering in Medicine and Biology Magazine_. Vol. 20. Number 3, May-June 
% 2001, pp. 45–50. (PMID: 11446209).
% 
% [3 Baim, D. S., W. S. Colucci, E. S. Monrad, H. S. Smith, R. F. Wright, A. 
% Lanoue, D. F. Gauthier, B. J. Ransil, W. Grossman, and E. Braunwald. "Survival 
% of patients with severe congestive heart failure treated with oral milrinone."  
% _Journal of the American College of Cardiology_. Vol. 7, Number 3, 1986, pp. 
% 661–670.
% 
% _Copyright 2023 The MathWorks, Inc._