tempdir = '/mathworks/devel/sandbox/wking';
pathToRecordingsFolder = fullfile(tempdir,'free-spoken-digit-dataset','recordings');
ads = audioDatastore(pathToRecordingsFolder);

ads.Labels = helpergenLabels(ads);
summary(ads.Labels)

rng default;
ads = shuffle(ads);
[adsTrain,adsTest] = splitEachLabel(ads,0.8,0.2);

transTrain = transform(adsTrain,@(x,info)helperReadData(x,info),'IncludeInfo',true);
transTest = transform(adsTest,@(x,info)helperReadData(x,info),'IncludeInfo',true);

numF = 12;
dropoutProb = 0.2;
layers = [
    sequenceInputLayer(1,'Name','input','MinLength',8192,...
         'Normalization',"none")

    convolution1dLayer(5,1,"name","pre-emphasis-filter",...
    "WeightsInitializer",@(sz)kronDelta(sz),"BiasLearnRateFactor",0)  

    stftLayer('Window',hamming(1280),'OverlapLength',900,...
    'OutputMode','spatial','Name','STFT') 
    
    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

    maxPooling2dLayer(3,'Stride',2,'Padding','same')

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

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

    maxPooling2dLayer(2)

    dropoutLayer(dropoutProb)
    fullyConnectedLayer(numel(categories(ads.Labels)))
    softmaxLayer    
    ];
dlnet = dlnetwork(layers);

NumEpochs = 25;
miniBatchSize = 128;
learnRate = 0.001;

mbqTrain = minibatchqueue(transTrain, 2,...
    'MiniBatchSize',miniBatchSize,...
    'MiniBatchFormat', {'CBT', 'CB'}, ... 
    'MiniBatchFcn', @processSpeechMB);

progress = "final-loss";
if progress == "training-progress"
    figure
    lineLossTrain = animatedline;
    ylim([0 inf])
    xlabel("Iteration")
    ylabel("Loss")
    grid on
end

% Initialize some training loop variables.
trailingAvg = [];
trailingAvgSq = [];
iteration = 0;
lossByIteration = 0;

% Loop over epochs. Time the epochs.
start = tic;

for epoch = 1:NumEpochs
    reset(mbqTrain)
    shuffle(mbqTrain)

    % Loop over mini-batches.
    while hasdata(mbqTrain)
        iteration = iteration + 1;
        
        % Get the next minibatch and one-hot coded targets
        [dlX,Y] = next(mbqTrain);
        
        % Evaluate the model gradients and loss 
        [gradients, loss, state] = dlfeval(@modelGradSTFT,dlnet,dlX,Y);
        if progress == "final-loss"
            lossByIteration(iteration) = loss;
        end

        % Update the network state
        dlnet.State = state;
        
        % Update the network parameters using an Adam optimizer.
        [dlnet,trailingAvg,trailingAvgSq] = adamupdate(...
            dlnet, gradients, trailingAvg, trailingAvgSq, iteration, learnRate);        
        
        % Display the training progress.
        D = duration(0,0,toc(start),'Format','hh:mm:ss');
        if progress == "training-progress"
            addpoints(lineLossTrain,iteration,loss)
            title("Epoch: " + epoch + ", Elapsed: " + string(D))
        end
        
    end
    disp("Training loss after epoch " + epoch + ": " + loss); 

end

if progress == "final-loss"
        plot(1:iteration,lossByIteration)
        grid on 
        title('Training Loss by Iteration')
        xlabel("Iteration")
        ylabel("Loss")
end

miniBatchSize = 32;
mbqTest = minibatchqueue(transTest, 2,...
    'MiniBatchSize',miniBatchSize,...
    'MiniBatchFormat', {'CBT', 'CB'}, ... 
    'MiniBatchFcn', @processSpeechMB);

numObservations = numel(adsTest.Files);
classes = string(unique(adsTest.Labels));

predictions = [];

% Loop over mini-batches.
while hasdata(mbqTest)    
    % Read mini-batch of data.
    dlX = next(mbqTest);       
    % Make predictions on the minibatch
    dlYPred = predict(dlnet,dlX,'Acceleration','none');   
    % Determine corresponding classes.
    predBatch = onehotdecode(dlYPred,classes,1);
    predictions = [predictions predBatch];
  
end

accuracy = mean(predictions' == categorical(adsTest.Labels))

FIRFilter = dlnet.Layers(2).Weights;
[H,W] = freqz(FIRFilter,1,[],8000);
delta = kronDelta([5 1 1]);
Hinit = freqz(delta,1,[],4000);
plot(W,20*log10(abs([H Hinit])),'linewidth',2)
grid on
xlabel('Hz')
ylabel('dB')
legend('Learned Filter','Initial Filter','Location','SouthEast')
title('Learned Pre-emphasis Filter')

function Labels = helpergenLabels(ads)
% This function is only for use in the "Learn Pre-Emphasis Filter using
% Deep Learning" example. It may change or be removed in a
% future release.

tmp = cell(numel(ads.Files),1);
expression = "[0-9]+_";
for nf = 1:numel(ads.Files)
    idx = regexp(ads.Files{nf},expression);
    tmp{nf} = ads.Files{nf}(idx);
end
Labels = categorical(tmp);
end

function [out,info] = helperReadData(x,info)
% This function is only for use in the "Learn Pre-Emphasis Filter using
% Deep Learning" example. It may change or be removed in a
% future release.

N = numel(x);
x = single(x);
if N > 8192
    x = x(1:8192);
elseif N < 8192
    pad = 8192-N;
    prepad = floor(pad/2);
    postpad = ceil(pad/2);
    x = [zeros(prepad,1) ; x ; zeros(postpad,1)];
end
x = (x-mean(x))./std(x);
x = x(:)';
out = {x,info.Label};
end

function [dlX,dlY] = processSpeechMB(Xcell,Ycell)
% This function is only for use in the "Learn Pre-Emphasis Filter using
% Deep Learning" example. It may change or be removed in a
% future release.

Xcell = cellfun(@(x)reshape(x,1,1,[]),Xcell,'uni',false);
dlX = cat(2,Xcell{:});
dlY = cat(2,Ycell{:});
dlY = onehotencode(dlY,1);
end

function [grads, loss, state] = modelGradSTFT(net, X, T)
% This function is only for use in the "Learn Pre-Emphasis Filter using
% Deep Learning" example. It may change or be removed in a
% future release.

[y, state] = net.forward(X);
loss = crossentropy(y, T);
grads = dlgradient(loss,net.Learnables);
loss = double(gather(extractdata(loss)));
end

function delta = kronDelta(sz)
% This function is only for use in the "Learn Pre-Emphasis Filter using
% Deep Learning" example. It may change or be removed in a
% future release.

L = sz(1);
delta = zeros(L,sz(2),sz(3),'single');
delta(1) = 1/sqrt(L);

end