url = 'https://www2.spsc.tugraz.at/databases/PTDB-TUG/SPEECH_DATA_ZIPPED.zip';
downloadFolder = tempdir;
datasetFolder = fullfile(downloadFolder,'PTDB-TUG');
if ~exist(datasetFolder,'dir')
    disp('Downloading PTDB-TUG (3.9 G) ...')
    unzip(url,datasetFolder)
end

ads = audioDatastore([fullfile(datasetFolder,"SPEECH DATA","FEMALE","MIC"),fullfile(datasetFolder,"SPEECH DATA","MALE","MIC")], ...
                     'IncludeSubfolders',true, ...
                     'FileExtensions','.wav');
fileNames = ads.Files;

[audioIn,audioInfo] = read(ads);
fs = audioInfo.SampleRate;
t = (0:size(audioIn,1)-1)/fs;
sound(audioIn,fs)
plot(t,audioIn)
xlabel('Time (s)')
ylabel('Amplitude')
axis([0 t(end) -1 1])
title('Sample Utterance from Training Set')

speakerIDs = extractBetween(fileNames,'mic_','_');
ads.Labels = categorical(speakerIDs);

developmentLabels = categorical(["M01","M02","M03","M04","M06","M07","M08","M09","F01","F02","F03","F04","F06","F07","F08","F09"]);
evaluationLabels = categorical(["M05","M010","F05","F010"]);
adsTrain = subset(ads,ismember(ads.Labels,developmentLabels));
[adsTrain,adsValidation,adsTest] = splitEachLabel(adsTrain,0.8,0.1,0.1);
adsEvaluate = subset(ads,ismember(ads.Labels,evaluationLabels));
[adsEnroll,adsDET] = splitEachLabel(adsEvaluate,3);

countEachLabel(adsTrain)

countEachLabel(adsValidation)

countEachLabel(adsTest)

countEachLabel(adsEnroll)

countEachLabel(adsDET)

speedUpExample = false;
if speedUpExample
    adsTrain = splitEachLabel(adsTrain,20);
    adsDET = splitEachLabel(adsDET,20);
end

desiredFs = 16e3;

windowDuration = 0.03;
hopDuration = 0.005;
windowSamples = round(windowDuration*desiredFs);
hopSamples = round(hopDuration*desiredFs);
overlapSamples = windowSamples - hopSamples;
numCoeffs = 30;
afe = audioFeatureExtractor( ...
    'SampleRate',desiredFs, ...
    'Window',hann(windowSamples,'periodic'), ...
    'OverlapLength',overlapSamples, ...
    ...
    'mfcc',true, ...
    'pitch',false, ...
    'spectralEntropy',false, ...
    'spectralFlux',false);
setExtractorParams(afe,'mfcc','NumCoeffs',numCoeffs)

adsTrain = transform(adsTrain,@(x)resample(x,desiredFs,fs));
[features,YTrain] = xVectorPreprocessBatch(adsTrain,afe);
featuresMAT = cat(1,features{:});
numFeatures = size(featuresMAT,2);
factors = struct('Mean',mean(featuresMAT,1),'STD',std(featuresMAT,1));

XTrain = cellfun(@(x)(x-factors.Mean)./factors.STD,features,'UniformOutput',false);
XTrain = cellfun(@(x)x-mean(x(:)),XTrain,'UniformOutput',false);

adsValidation = transform(adsValidation,@(x)resample(x,desiredFs,fs));
[XValidation,YValidation] = xVectorPreprocessBatch(adsTrain,afe,'Factors',factors);

classes = unique(YTrain);
numClasses = numel(classes);

type('xvecModel')

numFilters = 128;
[parameters,state] = initializexVecModelLayers(numFeatures,numFilters,numClasses)

dsXTrain = arrayDatastore(XTrain,'OutputType','same');
dsYTrain = arrayDatastore(YTrain','OutputType','cell');

dsTrain = combine(dsXTrain,dsYTrain);

miniBatchSize = 128;
numOutputs = 2;
mbq = minibatchqueue(dsTrain,numOutputs, ...
    'MiniBatchSize',miniBatchSize, ...
    'MiniBatchFormat',{'SCB','CB'}, ...
    'MiniBatchFcn',@preprocessMiniBatch);

numEpochs = 6;

learnRate = 0.001;
gradDecay = 0.5;
sqGradDecay = 0.999;
trailingAvg = [];
trailingAvgSq = [];

LearnRateDropPeriod = 2;
LearnRateDropFactor = 0.1;

ValidationsPerEpoch = 2;

iterationsPerEpoch = floor(numel(XTrain)/miniBatchSize);
iterationsPerValidation = round(iterationsPerEpoch/ValidationsPerEpoch);

ExecutionEnvironment = 'gpu';

if ValidationsPerEpoch ~= 0
    [XValidation,YValidation] = preprocessMiniBatch(XValidation,{YValidation});
    XValidation = dlarray(XValidation,'SCB');
    if strcmp(ExecutionEnvironment,'gpu')
        XValidation = gpuArray(XValidation);
    end
end

pp = progressPlotter(categories(classes));

iteration = 0;
for epoch = 1:numEpochs
    
    % Shuffle mini-batch queue
    shuffle(mbq)
    
    while hasdata(mbq)
        
        % Update iteration counter
        iteration = iteration + 1;
        
        % Get mini-batch from mini-batch queue
        [dlX,Y] = next(mbq);
        
        % Place mini-batch on GPU if requested
        if strcmp(ExecutionEnvironment,'gpu')
            dlX = gpuArray(dlX);
        end

        % Evaluate the model gradients, state, and loss using dlfeval and the modelGradients function
        [gradients,state,loss,predictions] = dlfeval(@modelGradients,dlX,Y,parameters,state);

        % Update the network parameters using the Adam optimizer
        [parameters,trailingAvg,trailingAvgSq] = adamupdate(parameters,gradients, ...
            trailingAvg,trailingAvgSq,iteration,learnRate,gradDecay,sqGradDecay,eps('single'));

        % Update the training progress plot
        updateTrainingProgress(pp,'Epoch',epoch,'Iteration',iteration,'LearnRate',learnRate,'Predictions',predictions,'Targets',Y,'Loss',loss)

        % Update the validation plot
        if ~rem(iteration,iterationsPerValidation)
            
            % Pass validation data through x-vector model
            predictions = xvecModel(XValidation,parameters,state,'DoTraining',false);

            % Update plot
            updateValidation(pp,'Iteration',iteration,'Predictions',predictions,'Targets',YValidation)
        end
    end
    
    % Update learn rate
    if rem(epoch,LearnRateDropPeriod)==0
        learnRate = learnRate*LearnRateDropFactor;
    end
    
end

predictedLabels = classes;
predictedLabels(:) = [];

for sample = 1:numel(adsTest.Files)
    [audioIn,xInfo] = read(adsTest);
    audioIn = resample(audioIn,desiredFs,fs);
    f = xVectorPreprocess(audioIn,afe,'Factors',factors,'MinimumDuration',0);
    predictions = zeros(numel(classes),numel(f));
    for segment = 1:numel(f)
        dlX = dlarray(f{segment},'SCB');
        predictions(:,segment) = extractdata(xvecModel(dlX,parameters,state,'DoTraining',false));
    end 
    predictedLabels(sample) = onehotdecode(mean(predictions,2),categories(classes),1);
end
trueLabels = adsTest.Labels;
accuracy = mean(trueLabels==predictedLabels');

figure('Units','normalized','Position',[0.2 0.2 0.6 0.6]);
confusionchart(trueLabels,predictedLabels', ...
    'ColumnSummary','column-normalized', ...
    'RowSummary','row-normalized', ...
    'Title',sprintf('x-vector Speaker Recognition\nAccuracy = %0.2f%%',accuracy*100))

xvecs = zeros(numFilters,numel(YTrain));
for sample = 1:size(YTrain,2)
    dlX = dlarray(XTrain{sample},'SCB');
    
    embedding = xvecModel(dlX,parameters,state,'DoTraining',false,'OutputLayer',7);
    xvecs(:,sample) = extractdata(embedding);
end

numEigenvectors = 32;
projMat = helperTrainProjectionMatrix(xvecs,YTrain,numEigenvectors);

xvecs = projMat*xvecs;

numIterations = 3;
numDimensions = 32;
plda = helperTrainPLDA(xvecs,YTrain,numIterations,numDimensions);

minDur = (numel(afe.Window)+14*(numel(afe.Window)-afe.OverlapLength)+1)/desiredFs;

xvecs = zeros(numEigenvectors,numel(adsEnroll.Files));
reset(adsEnroll)
for sample = 1:numel(adsEnroll.Files)
    [audioIn,xInfo] = read(adsEnroll);
    audioIn = resample(audioIn,desiredFs,fs);
    f = xVectorPreprocess(audioIn,afe,'Factors',factors,'MinimumDuration',minDur);
    embeddings = zeros(numFilters,numel(f));
    for segment = 1:numel(f)
        dlX = dlarray(f{segment},'SCB');

        embeddings(:,segment) = extractdata(xvecModel(dlX,parameters,state,'DoTraining',false,'OutputLayer',7));
    end 
    xvecs(:,sample) = mean(projMat*embeddings,2);
end

labels = adsEnroll.Labels;
uniqueLabels = unique(labels);
atable = cell2table(cell(0,2),'VariableNames',{'xvector','NumSamples'});
for ii = 1:numel(uniqueLabels)
    idx = uniqueLabels(ii)==labels;
    wLocalMean = mean(xvecs(:,idx),2);
    localTable = table({wLocalMean},(sum(idx)), ...
        'VariableNames',{'xvector','NumSamples'}, ...
        'RowNames',string(uniqueLabels(ii)));
    atable = [atable;localTable]; %#ok<AGROW>
end
enrolledLabels = atable

xvecs = zeros(numEigenvectors,numel(adsDET.Files));
reset(adsDET)
for sample = 1:numel(adsDET.Files)
    [audioIn,xInfo] = read(adsDET);
    audioIn = resample(audioIn,desiredFs,fs);
    f = xVectorPreprocess(audioIn,afe,'Factors',factors,'MinimumDuration',minDur);
    embeddings = zeros(numFilters,numel(f));
    for segment = 1:numel(f)
        dlX = dlarray(f{segment},'SCB');
        embeddings(:,segment) = extractdata(xvecModel(dlX,parameters,state,'DoTraining',false,'OutputLayer',7));
    end 
    xvecs(:,sample) = mean(projMat*embeddings,2);
end
labels = adsDET.Labels;
[detTable,scoreRange] = helperDetectionErrorTradeoff(xvecs,labels,enrolledLabels,plda);

plot(detTable.PLDA.Threshold,detTable.PLDA.FAR, ...
    detTable.PLDA.Threshold,detTable.PLDA.FRR)
eer = helperEqualErrorRate(detTable.PLDA);
title(sprintf('Speaker Verification\nDetection Error Tradeoff\nPLDA Scoring\nEqual Error Rate = %0.2f',eer));
xlabel('Threshold')
ylabel('Error Rate')
legend({'FAR','FRR'})

plot(detTable.CSS.Threshold,detTable.CSS.FAR, ...
    detTable.CSS.Threshold,detTable.CSS.FRR)
eer = helperEqualErrorRate(detTable.CSS);
title(sprintf('Speaker Verification\nDetection Error Tradeoff\nCosine Similarity Scoring\nEqual Error Rate = %0.2f',eer));
xlabel('Threshold')
ylabel('Error Rate')
legend({'FAR','FRR'})

function [parameters,state] = initializexVecModelLayers(numFeatures,numFilters,numClasses)
% This function is only for use in this example. It may be changed or
% removed in a future release.

% Initialize Layer 1 (1-D Convolutional)
filterSize1                      = 5;
numChannels1                     = numFeatures;
numFilters1                      = numFilters;

numIn1                           = filterSize1*numFilters1;
numOut1                          = filterSize1*numFilters1;

parameters.conv1.Weights         = initializeGlorot([filterSize1,numChannels1,numFilters1],numOut1,numIn1);
parameters.conv1.Bias            = dlarray(zeros([numFilters1,1],'single'));
parameters.batchnorm1.Offset     = dlarray(zeros([numFilters1,1],'single'));
parameters.batchnorm1.Scale      = dlarray(ones([numFilters1,1],'single'));
state.batchnorm1.TrainedMean     = zeros(numFilters1,1,'single');
state.batchnorm1.TrainedVariance = ones(numFilters1,1,'single');


% Initialize Layer 2 (1-D Convolutional)
filterSize2                      = 3;
numChannels2                     = numFilters1;
numFilters2                      = numFilters;

numIn2                           = filterSize2*numFilters2;
numOut2                          = filterSize2*numFilters2;

parameters.conv2.Weights         = initializeGlorot([filterSize2,numChannels2,numFilters2],numOut2,numIn2);
parameters.conv2.Bias            = dlarray(zeros([numFilters2,1],'single'));
parameters.batchnorm2.Offset     = dlarray(zeros([numFilters2,1],'single'));
parameters.batchnorm2.Scale      = dlarray(ones([numFilters2,1],'single'));
state.batchnorm2.TrainedMean     = zeros(numFilters2,1,'single');
state.batchnorm2.TrainedVariance = ones(numFilters2,1,'single');


% Initialize Layer 3 (1-D Convolutional)
filterSize3                      = 3;
numChannels3                     = numFilters2;
numFilters3                      = numFilters;

numIn3                           = filterSize3*numFilters3;
numOut3                          = filterSize3*numFilters3;

parameters.conv3.Weights         = initializeGlorot([filterSize3,numChannels3,numFilters3],numOut3,numIn3);
parameters.conv3.Bias            = dlarray(zeros([numFilters3,1],'single'));
parameters.batchnorm3.Offset     = dlarray(zeros([numFilters3,1],'single'));
parameters.batchnorm3.Scale      = dlarray(ones([numFilters3,1],'single'));
state.batchnorm3.TrainedMean     = zeros(numFilters3,1,'single');
state.batchnorm3.TrainedVariance = ones(numFilters3,1,'single');


% Initialize Layer 4 (1-D Convolutional)
filterSize4                      = 1;
numChannels4                     = numFilters3;
numFilters4                      = numFilters;

numIn4                           = filterSize4*numFilters4;
numOut4                          = filterSize4*numFilters4;

parameters.conv4.Weights         = initializeGlorot([filterSize4,numChannels4,numFilters4],numOut4,numIn4);
parameters.conv4.Bias            = dlarray(zeros([numFilters4,1],'single'));
parameters.batchnorm4.Offset     = dlarray(zeros([numFilters4,1],'single'));
parameters.batchnorm4.Scale      = dlarray(ones([numFilters4,1],'single'));
state.batchnorm4.TrainedMean     = zeros(numFilters4,1,'single');
state.batchnorm4.TrainedVariance = ones(numFilters4,1,'single');


% Initialize Layer 5 (1-D Convolutional)
filterSize5                      = 1;
numChannels5                     = numFilters4;
numFilters5                      = 1500;

numOut5                          = filterSize5*numFilters5;
numIn5                           = filterSize5*numFilters5;

parameters.conv5.Weights         = initializeGlorot([filterSize5,numChannels5,numFilters5],numOut5,numIn5);
parameters.conv5.Bias            = dlarray(zeros([numFilters5,1],'single'));
parameters.batchnorm5.Offset     = dlarray(zeros([numFilters5,1],'single'));
parameters.batchnorm5.Scale      = dlarray(ones([numFilters5,1],'single'));
state.batchnorm5.TrainedMean     = zeros(numFilters5,1,'single');
state.batchnorm5.TrainedVariance = ones(numFilters5,1,'single');


% Initialize Layer 6 (Statistical Pooling)
numIn6                           = numOut5;
numOut6                          = 2*numIn6;


% Initialize Layer 7 (Fully Connected)
numIn7                           = numOut6;
numOut7                          = numFilters;

parameters.fc7.Weights           = initializeGlorot([numFilters,numIn7],numOut7,numIn7);
parameters.fc7.Bias              = dlarray(zeros([numOut7,1],'single'));
parameters.batchnorm7.Offset     = dlarray(zeros([numOut7,1],'single'));
parameters.batchnorm7.Scale      = dlarray(ones([numOut7,1],'single'));
state.batchnorm7.TrainedMean     = zeros(numOut7,1,'single');
state.batchnorm7.TrainedVariance = ones(numOut7,1,'single');


% Initialize Layer 8 (Fully Connected)
numIn8                           = numOut7;
numOut8                          = numFilters;

parameters.fc8.Weights           = initializeGlorot([numOut8,numIn8],numOut8,numIn8);
parameters.fc8.Bias              = dlarray(zeros([numOut8,1],'single'));
parameters.batchnorm8.Offset     = dlarray(zeros([numOut8,1],'single'));
parameters.batchnorm8.Scale      = dlarray(ones([numOut8,1],'single'));
state.batchnorm8.TrainedMean     = zeros(numOut8,1,'single');
state.batchnorm8.TrainedVariance = ones(numOut8,1,'single');


% Initialize Layer 9 (Fully Connected)
numIn9                           = numOut8;
numOut9                          = numClasses;

parameters.fc9.Weights           = initializeGlorot([numOut9,numIn9],numOut9,numIn9);
parameters.fc9.Bias              = dlarray(zeros([numOut9,1],'single'));
end

function weights = initializeGlorot(sz,numOut,numIn)
% This function is only for use in this example. It may be changed or
% removed in a future release.
Z = 2*rand(sz,'single') - 1;
bound = sqrt(6 / (numIn + numOut));
weights = bound*Z;
weights = dlarray(weights);
end

function [gradients,state,loss,Y] = modelGradients(X,target,parameters,state)
% This function is only for use in this example. It may be changed or
% removed in a future release.
[Y,state] = xvecModel(X,parameters,state,'DoTraining',true);
loss = crossentropy(Y,target);
gradients = dlgradient(loss,parameters);
end

function [sequences,labels] = preprocessMiniBatch(sequences,labels)
% This function is only for use in this example. It may be changed or
% removed in a future release.
lengths = cellfun(@(x)size(x,1),sequences);
minLength = min(lengths);
sequences = cellfun(@(x)randomTruncate(x,1,minLength),sequences,'UniformOutput',false);
sequences = cat(3,sequences{:});
        
labels = cat(2,labels{:});
labels = onehotencode(labels,1);
labels(isnan(labels)) = 0;
end

function y = randomTruncate(x,dim,minLength)
% This function is only for use in this example. It may be changed or
% removed in a future release.
N = size(x,dim);
if N > minLength
    start = randperm(N-minLength,1);
    if dim==1
        y = x(start:start+minLength-1,:);
    elseif dim ==2
        y = x(:,start:start+minLength-1);
    end
else
    y = x;
end
end

function [features,labels] = xVectorPreprocessBatch(ads,afe,nvargs)
% This function is only for use in this example. It may be changed or
% removed in a future release.
    arguments
        ads
        afe
        nvargs.Factors = []
        nvargs.Segment = true;
    end
    if ~isempty(ver('parallel'))
        pool = gcp;
        numpar = numpartitions(ads,pool);
    else
        numpar = 1;
    end
    labels = [];
    features = [];
    parfor ii = 1:numpar
        adsPart = partition(ads,numpar,ii);
        numFiles = numel(adsPart.UnderlyingDatastores{1}.Files);
        localFeatures = cell(numFiles,1);
        localLabels = [];
        for jj = 1:numFiles
            [audioIn,xInfo] = read(adsPart);
            label = xInfo.Label;
            [f,ns] = xVectorPreprocess(audioIn,afe,'Factors',nvargs.Factors,'Segment',nvargs.Segment); %#ok<PFBNS> 
            localFeatures{jj} = f;
            localLabels = [localLabels,repelem(label,ns)];
        end
        features = [features;localFeatures];
        labels = [labels,localLabels];
    end
    features = cat(1,features{:});
    labels = removecats(labels);
end

function [features,numSegments] = xVectorPreprocess(audioData,afe,nvargs)
% This function is only for use in this example. It may be changed or
% removed in a future release.
arguments
    audioData
    afe
    nvargs.Factors = []
    nvargs.Segment = true;
    nvargs.MinimumDuration = 1;
end
% Scale
audioData = audioData/max(abs(audioData(:)));

% Protect against NaNs
audioData(isnan(audioData)) = 0;

% Determine regions of speech
mergeDur = 0.5; % seconds
idx = detectSpeech(audioData,afe.SampleRate,'MergeDistance',afe.SampleRate*mergeDur);

% If a region is less than MinimumDuration seconds, drop it.
if nvargs.Segment
    idxToRemove = (idx(:,2)-idx(:,1))<afe.SampleRate*nvargs.MinimumDuration;
    idx(idxToRemove,:) = [];
end

% Extract features
numSegments = size(idx,1);
features = cell(numSegments,1);
for ii = 1:numSegments
    features{ii} = single(extract(afe,audioData(idx(ii,1):idx(ii,2))));
end

% Standardize features
if ~isempty(nvargs.Factors)
    features = cellfun(@(x)(x-nvargs.Factors.Mean)./nvargs.Factors.STD,features,'UniformOutput',false);
end

% Cepstral mean subtraction (for channel noise)
if ~isempty(nvargs.Factors)
    fileMean = mean(cat(1,features{:}),'all');
    features = cellfun(@(x)x - fileMean,features,'UniformOutput',false);
end

if ~nvargs.Segment
    features = cat(1,features{:});
end
end