function [dsTrain,layers,lossFcn,options] = {functionName}(params)
%% Signal Regression by Sweeping Hyperparameters
% This template function defines the training data, network architecture, loss function, and 
% training options for a signal regression 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 layer you use in 
% the network 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 template uses the Physionet ECG-ID database [1][2], which has 310 ECG 
% records from 90 subjects. Each record contains a raw noisy ECG signal and a 
% manually filtered clean ground truth version.

dataURL = "https://physionet.org/static/published-projects/ecgiddb/ecg-id-database-1.0.0.zip";
datasetFolder = fullfile(tempdir,"ecg-id-database-1.0.0");
if ~isfolder(datasetFolder)
    zipFile = websave(tempdir,dataURL);
    unzip(zipFile,tempdir)
    delete(zipFile);
end
%% 
% Create a <matlab:doc('signalDatastore') |signalDatastore|> object to manage 
% the data. For files with specific formats or extensions, you can define a 
% customized read function to load the signals and the corresponding ground
% truth.

sds = signalDatastore(datasetFolder, ...
                      IncludeSubfolders=true, ...
                      FileExtensions=".dat", ...
                      ReadFcn = @helperReadSignalData);
%% 
% Reset the random number generator so the data split is consistent across experiments.
% 
% To limit the runtime of the experiment, select 10 different subjects at
% random.

rng("default")
subjectIds = unique(extract(sds.Files,"Person_"+digitsPattern));
idx = find(contains(sds.Files,subjectIds(randperm(numel(subjectIds),10))));
sdsSubset = subset(sds,idx);
numSignals = length(idx);
%% 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 the next line.

sdsResize = transform(sdsSubset,@helperResizeData);
%% Extract Features
% In regression tasks, it is important to create features that have the same 
% time resolution as the input signals. The multiresolution analysis of the maximal 
% overlap discrete wavelet transform (<matlab:doc('modwtmra') MODWTMRA>) has this 
% characteristic.
% 
% You have the option of extracting features using a transformed datastore or 
% using a built-in feature extraction layer. This example uses a modwtLayer to 
% extract features directly inside the deep learning model.

switch params.layerType
    case "conv"
        paramsLayer = convolution1dLayer(16,8,Padding="same");
    case "fullyConnected"
        paramsLayer = fullyConnectedLayer(8);
    case "modwt"
        paramsLayer = [
            modwtLayer('Level',8,IncludeLowpass=true, ...
                SelectedLevels=1:8,Wavelet="sym2")
            flattenLayer];
end
%% 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(sdsResize,trainIdx);
dsValid = subset(sdsResize,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>. Here, we experiment with the first layer, 
% i.e., the |paramsLayer|, by testing multiple candidate options to determine 
% the best architecture.

layers = [
    sequenceInputLayer(1,MinLength = 1024)

    paramsLayer

    convolution1dLayer(16,40,Stride=1,Padding="same")
    eluLayer
    batchNormalizationLayer    
    convolution1dLayer(16,20,Stride=1,Padding="same")
    eluLayer
    batchNormalizationLayer
    convolution1dLayer(16,20,Stride=1,Padding="same")
    eluLayer
    batchNormalizationLayer
    convolution1dLayer(16,20,Stride=1,Padding="same")
    eluLayer
    batchNormalizationLayer
    convolution1dLayer(16,40,Stride=1,Padding="same")
    eluLayer
    batchNormalizationLayer 
    convolution1dLayer(16,1,Stride=1,Padding="same")
    eluLayer
    batchNormalizationLayer

    dropoutLayer(0.2)
    convolution1dLayer(16,1,Stride=1,Padding="same")
    ];
%% 
% For regression tasks, use mean squared error loss.
lossFcn = "mse";
%% 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 50. 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 = 50;
options = trainingOptions("adam", ...
    InputDataFormats="CTB", ...
    InitialLearnRate=params.initialLearnRate, ...
    MaxEpochs=50, ...
    LearnRateSchedule="piecewise", ...
    LearnRateDropFactor=0.9, ...
    LearnRateDropPeriod=10, ...
    MiniBatchSize=miniBatchSize, ...
    Shuffle="every-epoch", ...
    Plots="training-progress", ...
    Metrics="rmse", ...
    Verbose=false, ...
    ValidationData=dsValid, ...
    ExecutionEnvironment="auto", ...
    PreprocessingEnvironment="serial");

%% 
% End of experiment setup function.

end
%% Supporting Functions
% *helperReadSignalData* 
% 
% Read signal data from .dat files.
%% 
% * |filename| is a string with the file location and name.
% * |dataOut| is a cell array in which the first cell holds noisy signals and 
% the other contains clean signals.
% * |infoOut| is a structure array holds information like sample rate.

function [dataOut,infoOut] = helperReadSignalData(filename)
    fid = fopen(filename,"r");
    % 1st row : raw data, 2nd row : filtered data
    data = fread(fid,[2 Inf],"int16=>single");
    fclose(fid);
    fid = fopen(replace(filename,".dat",".hea"),"r");
    header = textscan(fid,"%s%d%d%d",2,"Delimiter"," ");
    fclose(fid);
    gain = single(header{3}(2));
    dataOut{1} = data(1,:)/gain; % Noisy, raw data
    dataOut{2} = data(2,:)/gain; % Filtered, clean data
    infoOut.SampleRate = header{3}(1);
end
 
%% 
% *helperResizeData* 
% 
% Pad or truncate input ECG signal into target length.
%% 
% * |inputCell| is a two-element cell array containing an ECG signal and the 
% ground truth signal.
% * |outputCell| is a two-column cell array containing all noisy signal segments 
% with the same length and the corresponding ground truth signal segments.

function outputCell = helperResizeData(inputCell)
    targetLength = ceil(10000/32)*32;
    % Pad or truncate sequence data to target length
    outputCell = padsequences(inputCell,2, ...
        Length=targetLength, ...
        PaddingValue="symmetric", ...
        UniformOutput=false);
end
%% *References*
% [1] Lugovaya, Tatiana. "Biometric Human Identification Based on Electrocardiogram." 
% Master's thesis, Saint Petersburg Electrotechnical University. 2005.
% 
% [2] 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.
% 
% _Copyright 2023 The MathWorks, Inc._