function output = {functionName}(params, trial)
%% Train a Neural Network

%Load Training Data
[XTrain,YTrain] = digitTrain4DArrayData;
classes = categories(YTrain);
numClasses = numel(classes);

%Define Network
layers = [
    imageInputLayer([28 28 1], 'Name', 'input', 'Mean', mean(XTrain,4))
    convolution2dLayer(5, 20, 'Name', 'conv1')
    reluLayer('Name', 'relu1')
    convolution2dLayer(3, 20, 'Padding', 1, 'Name', 'conv2')
    reluLayer('Name', 'relu2')
    convolution2dLayer(3, 20, 'Padding', 1, 'Name', 'conv3')
    reluLayer('Name', 'relu3')
    fullyConnectedLayer(numClasses, 'Name', 'fc')];
lgraph = layerGraph(layers);

dlnet = dlnetwork(lgraph);

% Specify Training Options
velocity = [];
numEpochs = 20;
miniBatchSize = 128;
numObservations = numel(YTrain);
numIterationsPerEpoch = floor(numObservations./miniBatchSize);
initialLearnRate = params.myInitialLearnRate;
momentum = 0.9;
decay = 0.01;

executionEnvironment = "auto";

% Initialize the training progress plot.
vis_panel = uipanel('Parent', []);
vis_axis = axes(vis_panel, 'Position', [0.1 0.1 0.8 0.8]);
lineLossTrain = animatedline(vis_axis);
xlabel(vis_axis, "Iteration");
ylabel(vis_axis, "Loss");
setTrainingPlot(trial, vis_panel);

% Train Model

iteration = 0;
start = tic;
output.trainingInfo.loss = [];

% Loop over epochs.
for epoch = 1:numEpochs
    % Shuffle data.
    idx = randperm(numel(YTrain));
    XTrain = XTrain(:,:,:,idx);
    YTrain = YTrain(idx);
    % Loop over mini-batches.
    for i = 1:numIterationsPerEpoch
        iteration = iteration + 1;
        
        % Read mini-batch of data and convert the labels to dummy
        % variables.
        idx = (i-1)*miniBatchSize+1:i*miniBatchSize;
        X = XTrain(:,:,:,idx);
        
        Y = zeros(numClasses, miniBatchSize, 'single');
        for c = 1:numClasses
            Y(c,YTrain(idx)==classes(c)) = 1;
        end
        
        % Convert mini-batch of data to dlarray.
        dlX = dlarray(single(X),'SSCB');
        
        % If training on a GPU, then convert data to gpuArray.
        if (executionEnvironment == "auto" && canUseGPU) || executionEnvironment == "gpu"
            dlX = gpuArray(dlX);
        end
        
        % Evaluate the model gradients and loss using dlfeval and the
        % modelGradients function.
        [gradients,loss] = dlfeval(@modelGradients,dlnet,dlX,Y);
        
        % Determine learning rate for time-based decay learning rate schedule.
        learnRate = initialLearnRate/(1 + decay*iteration);
        
        % Update the network parameters using the SGDM optimizer.
        [dlnet.Learnables, velocity] = sgdmupdate(dlnet.Learnables, gradients, velocity, learnRate, momentum);
        
        curLoss = double(gather(extractdata(loss)));
        if shouldUpdateMetrics(trial)
            updateMetrics(trial, 'trainingLoss', curLoss);
        end
        output.trainingInfo.loss = [output.trainingInfo.loss; iteration, curLoss];
        
        % Display the training progress.
        
        D = duration(0,0,toc(start),'Format','hh:mm:ss');
        addpoints(lineLossTrain,iteration, curLoss);
        title(vis_axis, "Epoch: " + epoch + ", Elapsed: " + string(D));
        if shouldUpdateTrainingPlot(trial)
            updateTrainingPlot(trial);
        end
        
    end
    output.trainedNet = dlnet;
    updateProgress(trial, (epoch*100.0)/numEpochs);
    if shouldStop(trial)
        return;
    end
end

output.trainedNet = dlnet;
% End of training function.
end

function [gradients,loss] = modelGradients(dlnet,dlX,Y)

dlYPred = forward(dlnet,dlX);
dlYPred = softmax(dlYPred);

loss = crossentropy(dlYPred,Y);
gradients = dlgradient(loss,dlnet.Learnables);

end
