classdef OnsetDetector < audioPlugin
    %OnsetDetector Determine onset location of music signal

    %   Copyright 2016-2022 The MathWorks, Inc.
    %#codegen
    
    %-----------------------------------------------------------------------
    % Public Properties - Tunable properties
    %-----------------------------------------------------------------------
    properties (Access = public)
        %OnsetDetectionType Onset detection type
        %
        %   Specify OnsetDetectionType as 'TimeEnvelope', 'SpectralFlux',
        %   'EnergyFlux'.
        %
        %   The default is 'SpectralFlux'. This property is tunable.
        %
        %   Time Envelope - analyzes the signal in 1 second frames
        % with 0% overlap using a rectangular window. The envelope is
        % calculated using a Hilbert transform, and then peaks of the
        % envelope above a relative threshold are declared onsets.
        %
        %   Spectral Flux - analyzes the signal in 0.1 second frames
        % with 75% overlap using a rectangular window. The spectral history
        % is saved for the last one hundred analysis frames and used to
        % determine a mean.  The flux is calculated for three frames in the
        % past (to account for changing mean) and if spectral flux is above
        % a given threshold the frame is declared as containing an onset.
        %
        %   Energy Flux - analyzes the signal in 0.01 second frames with
        % 0% overlap using a Hann window. The spectral history
        % is saved for the last one hundred analysis frames and used to
        % determine a mean.  The flux is calculated for three frames in the
        % past (to account for changing mean) and if spectral flux is above
        % a given threshold the frame is declared as containing an onset.
        %

        OnsetDetectionType = audioexample.OnsetDetectionEnum.SpectralFlux;
    end
    properties (Access = public)
        OnsetHistory = nan(10000,1);
    end
    
    
    %-----------------------------------------------------------------------
    % Private Properties - Used for internal storage
    %-----------------------------------------------------------------------
    properties (Access = private)
        % Used for all methods:
        OnsetCounter = 1;
        Buffer
        
        % Used for spectral methods:
        SpectralHistory = zeros(round(192000/2),1);
        SpectralPowerHistory
        
        % Used for time-domain method:
        MovingAvg
    end
    properties (SetAccess = private, GetAccess = public)
        ElapsedAnalyzedTime = 0;
        pElapsedSamples
    end
    
    %-----------------------------------------------------------------------
    % Constant Properties - Used to define plugin interface
    %-----------------------------------------------------------------------
    properties (Constant)
        PluginInterface = audioPluginInterface( ...
            audioPluginParameter('OnsetDetectionType', ...
                'DisplayName','Onset Detection Method', ...
                'Mapping',{'enum','Time Envelope','Spectral Flux','Energy Flux'}), ...
            'InputChannels',1, ...
            'OutputChannels',1, ...
            'PluginName','Onset Detector');
    end
    
    
    %-----------------------------------------------------------------------
    % Public Methods - End user interacts with these
    %-----------------------------------------------------------------------
    methods
        %-------------------------------------------------------------------
        % Construct objects used in plugin
        %-------------------------------------------------------------------
        function plugin = OnsetDetector(varargin)
            plugin.Buffer            = audiopluginexample.private.AnalysisAndSynthesisBuffer(0,192000);
            plugin.Buffer.FrameTime             = 0.1;
            plugin.Buffer.AnalysisFrameOverlap  = 75;
            plugin.Buffer.WindowType            = audioexample.WindowEnum.Rectangle;

            for i = 1:2:nargin
                plugin.(varargin{i}) = varargin{i+1};
            end

            plugin.MovingAvg            = dsp.MovingAverage(500);
            plugin.SpectralPowerHistory = dsp.AsyncBuffer(100);
        end
        
        %-------------------------------------------------------------------
        % Main processing algorithm
        %-------------------------------------------------------------------
        function out = process(plugin,in)
            % PROCESS Determine onset location of music signal
            %
            % audioOut = PROCESS(oD,audioIn) detects signal onsets.The
            %   audio is passed through the plugin unchanged
            %   (audioIn==audioOut).
            
            % Write audio input to buffer
            writeAnalysisBuffer(plugin.Buffer,in);
            
            % Read from input buffer and process signal
            while (plugin.Buffer.AnalysisNumHopsAvailable > 0)
                % Update elapsed analyzed time
                plugin.ElapsedAnalyzedTime = plugin.ElapsedAnalyzedTime + ...
                    plugin.Buffer.AnalysisSamplesPerHop/getSampleRate(plugin);
                
                % Read from buffer
                analysisFrame = readAnalysisBuffer(plugin.Buffer);
                
                % Get onset location in time
                if plugin.OnsetDetectionType     == audioexample.OnsetDetectionEnum.TimeEnvelope
                    onsets = timeEnvelope(plugin,analysisFrame);
                elseif plugin.OnsetDetectionType == audioexample.OnsetDetectionEnum.SpectralFlux
                    onsets = spectralFlux(plugin,analysisFrame);
                elseif plugin.OnsetDetectionType == audioexample.OnsetDetectionEnum.EnergyFlux
                    onsets = energyFlux(plugin,analysisFrame);
                else
                    onsets = nan;
                end

                % Update onset history
                plugin.OnsetHistory(plugin.OnsetCounter:plugin.OnsetCounter + numel(onsets)-1) = onsets;
                plugin.OnsetCounter = plugin.OnsetCounter + numel(onsets);
            end

            out = in;
        end
        
        %-------------------------------------------------------------------
        % Reset memory states and update samplerate-dependent parameters
        %-------------------------------------------------------------------
        function reset(plugin)
            % RESET resets internal states to initial conditions
            %
            % RESET(bD) clears history of onset and BPM tracking, and
            %   resets all states to initial conditions.
            
            % Clear analysis buffer
            plugin.Buffer.SampleRate = getSampleRate(plugin);
            reset(plugin.Buffer);
            
            % Reset onset history
            plugin.OnsetHistory = nan(10000,1);
            
            % Clear SpectralHistory
            plugin.SpectralHistory = zeros(round(192000/2),1);
            
            % Reset elapsed time
            plugin.ElapsedAnalyzedTime = 0;
        end

        function set.OnsetDetectionType(plugin,type)
            if type == audioexample.OnsetDetectionEnum.SpectralFlux
                plugin.Buffer.FrameTime             = 0.1;%#ok<MCSUP>
                plugin.Buffer.AnalysisFrameOverlap  = 75;%#ok<MCSUP>
                plugin.Buffer.WindowType            = audioexample.WindowEnum.Rectangle;%#ok<MCSUP>
            elseif type == audioexample.OnsetDetectionEnum.EnergyFlux
                plugin.Buffer.FrameTime             = 0.01;%#ok<MCSUP>
                plugin.Buffer.AnalysisFrameOverlap  = 0;%#ok<MCSUP>
                plugin.Buffer.WindowType            = audioexample.WindowEnum.Hann;%#ok<MCSUP>
            elseif type == audioexample.OnsetDetectionEnum.TimeEnvelope
                plugin.Buffer.FrameTime             = 1;%#ok<MCSUP>
                plugin.Buffer.AnalysisFrameOverlap  = 0;%#ok<MCSUP>
                plugin.Buffer.WindowType            = audioexample.WindowEnum.Rectangle;%#ok<MCSUP>
            end
        end
        
        function val = get.OnsetDetectionType(plugin)
            if  plugin.Buffer.FrameTime             == 0.1
                val = audioexample.OnsetDetectionEnum.SpectralFlux;
            elseif plugin.Buffer.FrameTime             == 0.01
                val =  audioexample.OnsetDetectionEnum.EnergyFlux;
            else
                val = audioexample.OnsetDetectionEnum.TimeEnvelope;
            end
        end
        
        %-------------------------------------------------------------------
        % Get click track
        %-------------------------------------------------------------------
        function clickTrack = getClickTrack(plugin)
            % GETCLICKTRACK Returns click track of detected onsets since reset
            %
            % clickTrack = getClickTrack(bD)
            
            % Get click sound
            [click,clickFs] = audioread('Click-16-44p1-mono-0.2secs.wav');
            
            % Resample click sound to the sample rate that the plugin is
            % run at.
            fs    = getSampleRate(plugin);
            click = resample(click,clickFs,fs);
            
            % Convert onset location in time to samples
            a               = round(plugin.OnsetHistory.*fs); % Convert time to samples
            b               = a(a>0);                         % Remove zeros
            onsetsInSamples = b(~isnan(b));                   % Remove NaNs
            
            % Initialize zero vector
            track = zeros(max(onsetsInSamples),1);
            
            % Place marker at location of onsets
            track(onsetsInSamples) = nan;
            
            % At each pulse, place a click
            for i = 1:numel(track)
                if isnan(track(i))
                    track(i:i+numel(click)-1) = click;
                end
            end
            
            % Output click track
            clickTrack = track;
        end
        function s = saveobj(obj)
            s = saveobj@audioPlugin(obj);
            s.Buffer = saveobj(obj.Buffer);
            s.MovingAvg = matlab.System.saveObject(obj.MovingAvg);
            s.SpectralPowerHistory = saveobj(clone(obj.SpectralPowerHistory)); 
            s.OnsetDetectionType = obj.OnsetDetectionType;
            s.OnsetHistory = obj.OnsetHistory;
            s.OnsetCounter = obj.OnsetCounter;
            s.SpectralHistory = obj.SpectralHistory;
            s.ElapsedAnalyzedTime = obj.ElapsedAnalyzedTime;
            s.pElapsedSamples = obj.pElapsedSamples;
        end
        function obj = reload(obj,s)
            obj = reload@audioPlugin(obj,s);
            obj.Buffer = audiopluginexample.private.AnalysisAndSynthesisBuffer.loadobj(s.Buffer);
            obj.MovingAvg = matlab.System.loadObject(s.MovingAvg);
            obj.SpectralPowerHistory = loadobj(s.SpectralPowerHistory);
            obj.OnsetDetectionType = s.OnsetDetectionType;
            obj.OnsetHistory = s.OnsetHistory;
            obj.OnsetCounter = s.OnsetCounter;
            obj.SpectralHistory = s.SpectralHistory;
            obj.ElapsedAnalyzedTime = s.ElapsedAnalyzedTime;
            obj.pElapsedSamples = s.pElapsedSamples;
        end
    end

    %-----------------------------------------------------------------------
    % Private methods - Used for onset detection functions
    %-----------------------------------------------------------------------
    methods (Access = private)
        
        %-------------------------------------------------------------------
        % Onset Detection (Energy flux)
        %-------------------------------------------------------------------
        function onsets = energyFlux(plugin,x)
            % Adapted from:
            % Laroche, Jean. Efficient Tempo and Beat Tracking in Audio
            % Recordings
            
            spectralHistory = plugin.SpectralHistory(1:round(getSampleRate(plugin)/2));
            
            % Magnitude Spectrum
            X     = abs(fft(x,getSampleRate(plugin)));
            half  = round(getSampleRate(plugin)/2);
            Xhalf = (X(1:half)).^0.5;
            
            % Differentiation
            if isnan(spectralHistory)
                Z = Xhalf;
            else
                Z = max(Xhalf - spectralHistory,0);
            end
            
            % Accumulation
            A = sum(Z(100:10000),1);
            
            % Normalization
            write(plugin.SpectralPowerHistory,A);
            spectralPowerHistory = read(plugin.SpectralPowerHistory,100,99);
            meanSpectralPower = mean(spectralPowerHistory(~isnan(spectralPowerHistory)));
            
            energyFlux = max(spectralPowerHistory - meanSpectralPower*1.5,0);
            
            % Locate peaks ------------------------------------------------
            [~,locs] = findpeaks(energyFlux,'MinPeakDistance',5);

            % Look 30 hops in the past for peak detection.
            if any(locs == 70)
                onsets = plugin.ElapsedAnalyzedTime - plugin.Buffer.AnalysisSamplesPerHop*31/getSampleRate(plugin);
            else
                onsets = [];
            end

            % Update spectral history
            plugin.SpectralHistory(1:numel(Xhalf)) = Xhalf;
        end
        
        %-------------------------------------------------------------------
        % Onset Detection - Spectral Flux
        %-------------------------------------------------------------------
        function onsets = spectralFlux(plugin,x)
            spectralHistory = plugin.SpectralHistory(1:round(numel(x)/2));
            
            % Magnitude Spectrum
            X     = abs(fft(x));
            half  = round(numel(X)/2);
            Xhalf = X(1:half);
            
            % Differentiation
            if isnan(spectralHistory)
                Z = Xhalf;
            else
                Z = max(Xhalf - spectralHistory,0);
            end
            
            % Accumulation
            A = sum(Z,1);
            
            % Normalization
            write(plugin.SpectralPowerHistory,A);
            spectralPowerHistory = read(plugin.SpectralPowerHistory,100,99);
            meanSpectralPower = mean(spectralPowerHistory(~isnan(spectralPowerHistory)));

            spectralFlux = max(spectralPowerHistory - meanSpectralPower,0);
            
            % Locate Peaks ------------------------------------------------
            [~,locs] = findpeaks(spectralFlux);
            
            % Look three hops in the past for peak detection.
            if any(locs == 97)
                onsets = plugin.ElapsedAnalyzedTime - plugin.Buffer.AnalysisSamplesPerHop*4/getSampleRate(plugin);
            else
                onsets = [];
            end

            % Update spectral history
            plugin.SpectralHistory(1:numel(Xhalf)) = Xhalf;
        end
        
        %-------------------------------------------------------------------
        % Onset Detection - Time Envelope
        %-------------------------------------------------------------------
        function onsets = timeEnvelope(plugin,in)
            [row,col] = size(in);
            
            % Get envelope of signal
            IN = fft(in);
            h  = zeros(row,1);
            if 2*fix(row/2) == row
                h([1 row/2+1]) = 1;
                h(2:row/2) = 2;
            elseif row>0
                h(1) = 1;
                h(2:(row+1)/2) = 2;
            end
            env = abs(ifft(IN.*h(:,ones(1,col))));

            env = step(plugin.MovingAvg,env);
            
            % Locate peaks ------------------------------------------------
            % Determine peak detection threshold (Get top quarter)
            envSorted     = sort(env);
            element       = round(row*0.65);
            minPeakHeight = envSorted(element);
            
            % Find peaks
            [~,loc] = findpeaks(env,getSampleRate(plugin), ...
                'MinPeakHeight',minPeakHeight, ...
                'MinPeakDistance',0.1);
            
            % Ignore first and last 0.01 seconds
            loc(loc<0.01) = [];
            loc(loc>0.99) = [];

            % Place loc in time relative to last reset
            loc = 1 - loc;
            onsets = plugin.ElapsedAnalyzedTime - loc;
        end
       
    end

    methods(Static)
        function obj = loadobj(s)
            if isstruct(s)
                obj = audiopluginexample.private.OnsetDetector;
                obj = reload(obj,s);
            end
        end     
    end
end