function [f0,locs,hr] = HelperPitchTracker(x,fs,varargin)
%HelperPitchTracker Track the fundamental frequency of audio signal
%   f0 = HelperPitchTracker(audioIn,fs) returns an estimate of the
%   fundamental frequency contour for the audio input. Columns of the
%   input are treated as individual channels. The HelperPitchTracker
%   function uses multiple pitch detection algorithms to generate pitch
%   candidates, and uses octave smoothing and a Hidden Markov Model to
%   return an estimate of the fundamental frequency.
%
%   f0 = HelperPitchTracker(...,'HopLength',HOPLENGTH) specifies the number
%   of samples in each hop. The pitch estimate is updated every hop.
%   Specify HOPLENGTH as a scalar integer. If unspecified, HOPLENGTH
%   defaults to round(0.01*fs).
%
%   f0 = HelperPitchTracker(...,'OctaveSmoothing',TF) specifies whether or
%   not to apply octave smoothing. Specify as true or false. If
%   unspecified, TF defaults to true.
%
%   f0 = HelperPitchTracker(...,'EmissionMatrix',EMISSIONMATRIX) specifies
%   the emission matrix used for the HMM during the forward pass. The
%   default emission matrix was trained on the Pitch Tracking Database from
%   Graz University of Technology. The database consists of 4720 speech
%   segments with corresponding pitch trajectories derived from
%   laryngograph signals. The emission matrix corresponds to the
%   probability that a speaker leaves one pitch state to another, in the
%   range [50, 400] Hz. Specify the emission matrix such that rows
%   correspond to the current state, columns correspond to the possible
%   future state, and the values of the matrix correspond to the
%   probability of moving from the current state to the future state. If
%   you specify your own emission matrix, specify its corresponding
%   EMISSIONMATRIXRANGE. EMISSIONMATRIX must be a real N-by-N matrix of
%   integers.
%
%   f0 = HelperPitchTracker(...,'EmissionMatrixRange',EMISSIONMATRIXRANGE)
%   specifies how the EMISSIONMATRIX corresponds to Hz. If unspecified,
%   EMISSIONMATRIXRANGE defaults to 50:400.
%
%   [f0,loc] = HelperPitchTracker(...) returns the locations associated
%   with each pitch decision. The locations correspond to the ceiling of
%   the center of the analysis frames.
%
%   [f0,loc,hr] = HelperPitchTracker(...) returns the harmonic ratio
%   associated with each pitch decision.
%
% See also PITCH, VOICEACTIVITYDETECTOR

% Copyright 2018 The MathWorks, Inc.

% Parse and validate inputs
validateRequiredInputs(x,fs)
params = parseInputs(x,fs,varargin{:});

% -------------------------------------------------------------------------
% Buffer audio into all required frame sizes
x_SRH = buffer(x,round(0.025*fs),round(0.02*fs),'nodelay');
x_PEF = buffer(x,round(0.06*fs),round(0.06*fs-params.HopLength),'nodelay');
x_Backup = buffer(x,round(0.04*fs),round(0.04*fs-params.HopLength),'nodelay');
x_Backup = x_Backup(:,2:end-1); % Account for different frame sizes.
% -------------------------------------------------------------------------

% -------------------------------------------------------------------------
% Generate pitch candidates
params_SRH = struct('Method','SRH', ...
    'Range',params.EmissionMatrixRange, ...
    'WindowLength',round(fs*0.06), ...
    'OverlapLength',round(fs*0.06-params.HopLength), ...
    'SampleRate',fs, ...
    'NumChannels',size(x,2), ...
    'SamplesPerChannel',size(x,1));
multiCandidate_params_SRH = struct('NumCandidates',5,'MinPeakDistance',1);
[SRHcandidates,SRHconfidence] = audio.internal.pitch.SRH(x_SRH, ...
    params_SRH, ...
    multiCandidate_params_SRH);

params_PEF = struct('Method','PEF', ...
    'Range',params.EmissionMatrixRange, ...
    'WindowLength',round(fs*0.06), ...
    'OverlapLength',round(fs*0.06-params.HopLength), ...
    'SampleRate',fs, ...
    'NumChannels',size(x,2), ...
    'SamplesPerChannel',size(x,1));
multiCandidate_params_PEF = struct('NumCandidates',5,'MinPeakDistance',5);
[PEFcandidates,PEFconfidence] = audio.internal.pitch.PEF(x_PEF, ...
    params_PEF, ...
    multiCandidate_params_PEF);

params_NCF = struct('Method','NCF', ...
    'Range',params.EmissionMatrixRange, ...
    'WindowLength',round(fs*0.04), ...
    'OverlapLength',round(fs*0.04-params.HopLength), ...
    'SampleRate',fs, ...
    'NumChannels',size(x,2), ...
    'SamplesPerChannel',size(x,1));
multiCandidate_params_NCF = struct('NumCandidates',5,'MinPeakDistance',1);
bp1 = audio.internal.pitch.NCF(x_Backup, ...
    params_NCF, ...
    multiCandidate_params_NCF);

params_CEP = struct('Method','CEP', ...
    'Range',params.EmissionMatrixRange, ...
    'WindowLength',round(fs*0.04), ...
    'OverlapLength',round(fs*0.04-params.HopLength), ...
    'SampleRate',fs, ...
    'NumChannels',size(x,2), ...
    'SamplesPerChannel',size(x,1));
multiCandidate_params_CEP = struct('NumCandidates',5,'MinPeakDistance',1);
bp2 = audio.internal.pitch.CEP(x_Backup, ...
    params_CEP, ...
    multiCandidate_params_CEP);

BackupCandidates = [bp1(:,1),bp2(:,1)];
% -------------------------------------------------------------------------

% -------------------------------------------------------------------------
% Determine the locations of the pitch decisions
numHops = size(x_PEF,2);
locs = 0:numHops*params.HopLength + round(fs*0.03);
% -------------------------------------------------------------------------

% -------------------------------------------------------------------------
% Determine a new search range based on the thresholded median of the most
% confident PEF and SRH estimates.
win = hamming(size(x_PEF,1));
hr = harmonicRatio(x_PEF,fs,'Window',win,'OverlapLength',0);

idxToKeep      = logical(movmedian(hr>((3/4)*max(hr)),3));
candidatesTemp = [PEFcandidates(:,1),SRHcandidates(:,1)];
candidatesTemp = candidatesTemp(idxToKeep(:),:);
longTermMedian = median(candidatesTemp(:));
lower          = max((2/3)*longTermMedian,params.EmissionMatrixRange(1));
upper          = min((4/3)*longTermMedian,params.EmissionMatrixRange(2));
% -------------------------------------------------------------------------

% -------------------------------------------------------------------------
% Make the first candidate the most confidenct candidate in the range
% [lower,upper].
[PEFcandidates,PEFconfidence] = bubbleUp(PEFcandidates, ...
                                         PEFconfidence, ...
                                         lower,upper);
[SRHcandidates,SRHconfidence] = bubbleUp(SRHcandidates, ...
                                         SRHconfidence, ...
                                         lower,upper);
% -------------------------------------------------------------------------

% -------------------------------------------------------------------------
% Trim candidates based on confidence threshold and combine into matrix
numPEFcan = 2;
numSRHcan = 2;
candidates = [PEFcandidates(:,1:numPEFcan),SRHcandidates(:,1:numSRHcan)];
confidence = [PEFconfidence(:,1:numPEFcan),SRHconfidence(:,1:numSRHcan)];
% -------------------------------------------------------------------------

% -------------------------------------------------------------------------
% Make Distinctive
span                    = 5;
confidenceFactor        = 1;
[candidates,confidence] = makeDistinctive(candidates, ...
                                          confidence, ...
                                          span, ...
                                          confidenceFactor);
% -------------------------------------------------------------------------

% -------------------------------------------------------------------------
% Define number of paths to track through trellis, and initialize a logbook
% to track the paths over time.
numPaths = 4;
logbook = zeros(numHops,3,numPaths);
% -------------------------------------------------------------------------

% -------------------------------------------------------------------------
% Forward Iteration with Octave Smoothing
for hopNumber = 1:numHops
    nowCandidates = candidates(hopNumber,:);
    nowConfidence = confidence(hopNumber,:);
    
    % Remove octave jumps
    if hopNumber>100 && params.OctaveSmoothingFlag
        [nowConfidence,nowCandidates] = octaveSmoothing(logbook, ...
            nowCandidates,nowConfidence,hopNumber,longTermMedian, ...
            BackupCandidates,params.EmissionMatrixRange);
    end
    
    % Create confidence matrix
    confidenceMatrix = createConfidenceMatrix(logbook,hopNumber, ...
        nowCandidates,nowConfidence, ...
        params.EmissionMatrix,params.EmissionMatrixRange);
    
    % Assign an estimate for each path
    logbook = assignPath(confidenceMatrix,logbook,nowCandidates,hopNumber);
end
% -------------------------------------------------------------------------

% -------------------------------------------------------------------------
% Traceback
f0 = traceBack(logbook);
% -------------------------------------------------------------------------

% -------------------------------------------------------------------------
% Smoothe
f0 = movmedian(f0,3);
% -------------------------------------------------------------------------

end

% -------------------------------------------------------------------------
% Bubble up
function [candidates,confidence] = bubbleUp(candidates,confidence,lower,upper)
invalid = candidates<lower | candidates>upper;
confidence(invalid) = confidence(invalid)/10;
[confidence,order] = sort(confidence,2,'descend');
for i = 1:size(candidates,1)
    candidates(i,:) = candidates(i,order(i,:));
end
end
% -------------------------------------------------------------------------

% -------------------------------------------------------------------------
% Octave Smoothing
function [nowConfidence,nowCandidates] = octaveSmoothing(logbook, ...
    nowCandidates,nowConfidence,hopNumber,longTermMedian,BackupCandidates, ...
    emissionMatrixRange)

% Weighted short term median
medianWindowLength  = 50;
aTemp               = logbook(max(hopNumber-min(hopNumber,medianWindowLength),1):hopNumber-1,1,:);
shortTermMedian     = median(aTemp(:));
previousM           = mean([longTermMedian,shortTermMedian]);
lowerTight          = max((2/3)*previousM,emissionMatrixRange(1));
upperTight          = min((4/3)*previousM,emissionMatrixRange(2));
numCandidateOutside = sum([nowCandidates < lowerTight, nowCandidates > upperTight]);
    
if numCandidateOutside > 0
    % Apply the backup pitch estimators
    nowCandidates = [nowCandidates,BackupCandidates(hopNumber,:)];
    nowConfidence = [nowConfidence,repmat(mean(nowConfidence),1,2)];
    
    span = 10;
    confidenceFactor = 1.2;
    [nowCandidates,nowConfidence] = makeDistinctive(nowCandidates,nowConfidence,span,confidenceFactor);
end
end
% -------------------------------------------------------------------------

% -------------------------------------------------------------------------
% Make Distinctive
function [candidates,confidence] = ...
    makeDistinctive(candidates,confidence,span,confidenceFactor)
for r = 1:size(candidates,1)
    for c = 1:size(candidates,2)
        tf = abs(candidates(r,c)-candidates(r,:)) < span;
        candidates(r,c) = mean(candidates(r,tf));
        confidence(r,c) = sum(confidence(r,tf))*confidenceFactor;
    end    
end
end
% -------------------------------------------------------------------------

% -------------------------------------------------------------------------
% Create Confidence Matrix
function confidenceMatrix = createConfidenceMatrix(logbook,hopNumber, ...
    nowCandidates,nowConfidence,emissionMatrix,emissionMatrixRange)
confidenceMatrix = zeros(numel(nowCandidates),size(logbook,3));
for pageIdx = 1:size(logbook,3)
    if hopNumber ~= 1
        pastPitch = floor(logbook(hopNumber-1,1,pageIdx)) - emissionMatrixRange(1) + 1;
    else
        pastPitch = nan;
    end
    for candidateNumber = 1:numel(nowCandidates)
        if hopNumber ~= 1
            % Get the current pitch and convert to an index in the range
            currentPitch = floor(nowCandidates(candidateNumber)) - emissionMatrixRange(1) + 1;
            confidenceMatrix(candidateNumber,pageIdx) = ...
                emissionMatrix(currentPitch,pastPitch)*logbook(hopNumber-1,2,pageIdx)*nowConfidence(candidateNumber);
        else
            confidenceMatrix(candidateNumber,pageIdx) = 1;
        end
    end
end
end
% -------------------------------------------------------------------------

% -------------------------------------------------------------------------
% Assign Path
function logbook = assignPath(confidenceMatrix,logbook,nowCandidates,hopNumber)
for pageIdx = 1:size(logbook,3)
    % Determine most confident transition from past to current pitch
    [~,idx] = max(confidenceMatrix(:));
    
    % Convert confidence matrix index to pitch and logbook page
    [chosenPitch,pastPitchIdx] = ind2sub([numel(nowCandidates),size(logbook,3)],idx);
    
    logbook(hopNumber,:,pageIdx) = ...
        [nowCandidates(chosenPitch), ...
        confidenceMatrix(chosenPitch,pastPitchIdx), ...
        pastPitchIdx];
    
    % Remove the chosen current pitch from the confidence matrix
    confidenceMatrix(chosenPitch,:) = nan;
end
% Normalize confidence
logbook(hopNumber,2,:) = logbook(hopNumber,2,:)/sum(logbook(hopNumber,2,:));
end
% -------------------------------------------------------------------------

% -------------------------------------------------------------------------
% Traceback
function f0 = traceBack(logbook)
numHops     = size(logbook,1);
keepLooking = true;
index       = numHops+1;

while keepLooking
   index = index - 1;
   if max(logbook(index,2,:)) - min(logbook(index,1,:)) ~= 0
      keepLooking = false; 
   end
end

[~,bestPathIdx]    = max(logbook(index,2,:));
bestIndices        = zeros(numHops,1);
bestIndices(index) = bestPathIdx;

for ii = index:-1:2
    bestIndices(ii-1) = logbook(ii,3,bestIndices(ii));
end

bestIndices(bestIndices==0) = 1;
f0 = zeros(numHops,1);
for ii = (numHops):-1:2
    f0(ii) = logbook(ii,1,bestIndices(ii));
end

end
% -------------------------------------------------------------------------

% -------------------------------------------------------------------------
% Parse inputs
function userInput = parseInputs(x,fs,varargin)
parser                 = inputParser;
parser.PartialMatching = false;
parser.FunctionName    = 'HelperPitchTracker';

addRequired(parser,'x');
addRequired(parser,'fs');
[hL,os,em,emr] = getDefaultParameters(fs);
addOptional(parser,'HopLength',hL);
addOptional(parser,'OctaveSmoothingFlag',os);
addOptional(parser,'EmissionMatrix',em);
addOptional(parser,'EmissionMatrixRange',emr);

parse(parser,x,fs,varargin{:});
userInput = parser.Results;
end
% -------------------------------------------------------------------------

% -------------------------------------------------------------------------
% Default parameter values
function [hopLength,octaveSmoothingFlag,em,emr] = getDefaultParameters(fs)
octaveSmoothingFlag = true;
hopLength = round(fs.*0.01);
[em,emr] = getEmissionMatrix();
end
% -------------------------------------------------------------------------

% -------------------------------------------------------------------------
% Get emission matrix
function [emissionMatrix,emissionMatrixRange] = getEmissionMatrix()
load('EmissionMatrix.mat','emissionMatrix','emissionMatrixRange')
end
% -------------------------------------------------------------------------

% -------------------------------------------------------------------------
% Validate required inputs
function validateRequiredInputs(x,fs)
validateattributes(x,{'single','double'}, ...
    {'nonempty','column','real','nonnan','finite'}, ...
    'HelperPitchTracker','audioIn')
validateattributes(fs,{'single','double'}, ...
    {'nonempty','positive','scalar','real','nonnan','finite','integer'}, ...
    'HelperPitchTracker','fs')
end
% -------------------------------------------------------------------------