function onset_data = SMIDIBT_soundmod_onsets(data_file,params)
% [onset_data] = SMIDIBT_soundmod_onsets(data_file,params)
%   Extracts the send MIDI onset and offset times and audio onset times
%
%   INPUTS:
%   data_file       The name of a ".wav" file ending in "_trigs" that
%                   contains the MIDI onset and offset triggers. This file
%                   MUST have a complementary audio file ending in "_audio"
%                   to extract the audio onsets.
%
%   params          A file containing the parameters for onset extraction
%                   (e.g., amplitude and timing thresholds). The parameters
%                   used in Schultz (2017, Experiment 4) are set as default.
%
%   OUTPUT:
%   onset_data      The times (in milliseconds) of the MIDI onsets of the
%                   MIDI send (first column), MIDI offsets of the MIDI send
%                   (second column), and the audio onsets (third column).
%
%   This version has dependencies appended at the bottom
%
%   2017-09-28 ben.schultz@maastrichtuniversity.nl
%   Copyright (c) 2017, Benjamin Schultz, Maastricht University.

%   This script is described in more detail in the publication:
%   Schultz, B. G. (submitted). The Schultz MIDI Benchmarking toolbox for
%   MIDI interfaces, percussion pads, and sound cards. Behavior Research 
%   Methods.
%
%   The SMIDIB Toolbox is distributed in the hope that it will be useful, but
%   WITHOUT ANY WARRANTY; without even the implied warranty of
%   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
%   General Public License for more details.
% 
%   You should have received a copy of the GNU General Public License
%   along with the SMIDIB Toolbox; if not, write to the Free Software
%   Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA
%   02110-1301 USA
% 
%   See the file "COPYING" for the text of the license.


%% Check inputs
% check input parameters and set defaults
params = check_params(params);

% check file names
midi_ind = strfind(data_file,'_trigs');
audio_file = strrep(data_file,'_trigs.wav','_audio.wav');

if isempty(midi_ind)
    error('Filename does not have the correct format. Should end in ''_trigs''');
end

if ~exist(data_file,'file') || ~exist(audio_file,'file')
    error('One or more files do not exist in specified location');
end

% load midi triggers and audio
[trig_data,Fs]=audioread(data_file);
trig_data = trig_data/max(abs(trig_data(Fs*3:end-Fs)));

[aud_data,aud_Fs]=audioread(audio_file);
aud_data = aud_data(1:length(trig_data));

if Fs~=aud_Fs
     error('Sample rates do not match. Files may not be synchronous.');
end

% set thresholds
time_thresh_samps = round(params.time_thresh_ms*(Fs/1000));
off_time_thresh_samps = round(params.off_time_thresh_ms*(Fs/1000));
moving_max_wind_samps = round(params.moving_max_wind_ms*(Fs/1000));
trig_time_thresh_samps = round(params.trig_time_thresh_ms*(Fs/1000));

% apply filters to the data (if any)
aud_data = abs(aud_data);
aud_data = getMovingMax(aud_data,moving_max_wind_samps);
aud_data = (aud_data-min(aud_data(Fs*3:end-Fs*3)))/(max(aud_data(Fs*3:end-Fs*3))-min(aud_data(Fs*3:end-Fs*3)));

%% Get onsets
% preset start parameters
prev_offset = 1;
cur_pos = 1;
onset_data = ones(params.n_trigs,3)*NaN;

% get MIDI trigger onsets
cur_onset_temp = find(trig_data(prev_offset:end)>params.trig_on_thresh,1,'first')+prev_offset-1;

while ~isempty(cur_onset_temp)
    
    % get MIDI trigger onset (MIDI send)
    cur_onset = cur_onset_temp;
    
    % go back to onset
    while trig_data(cur_onset-1)<trig_data(cur_onset)
        cur_onset = cur_onset-1;
    end
    cur_onset = cur_onset+1;
    
    % get MIDI trigger offset (MIDI sent)
    cur_offset = find(trig_data(cur_onset_temp+trig_time_thresh_samps:end)<params.trig_off_thresh,1,'first')+cur_onset_temp+trig_time_thresh_samps-1;
    
    % go back to offset
    cur_max = cur_offset;
    
    while trig_data(cur_max-1)> trig_data(cur_max)
        cur_max = cur_max-1;
    end
    
    cur_max = cur_max+1;
    
    % catch if audio file gets cut off
    if isempty(cur_max)
        cur_max = NaN;
    end
    
    % get audio onset
    cur_aud_onset_temp = find(aud_data(cur_onset_temp:cur_onset_temp+time_thresh_samps)>params.aud_on_thresh,1)+cur_onset_temp;
    
    if isempty(cur_aud_onset_temp)
        cur_aud_onset = NaN;
    else
        cur_aud_onset = cur_aud_onset_temp;
    end
    
    % set previous offset
    prev_offset = cur_offset+off_time_thresh_samps;
    
    % add data
    onset_data(cur_pos,1:3) = [cur_onset,cur_max,cur_aud_onset];    
    cur_pos = cur_pos+1;
    
    % get next onset
    cur_onset_temp = find(trig_data(prev_offset:end)>params.trig_on_thresh,1,'first')+prev_offset-1;
    
end

% remove excess NaNs and turn into milliseconds
onset_data = onset_data(~isnan(onset_data(:,1)),:)/(Fs/1000);

%% Save onsets (recommended for large datasets)
if isfield(params,'out_dir')
    [~,data_filename,data_ext] = fileparts(data_file);
    out_filename = strrep([data_filename,data_ext],'_trigs.wav','_onsets.mat');
    save(fullfile(params.out_dir,out_filename),'onset_data','Fs');
end

%% Print plot
if isfield(params,'plot_dir')
    t = (1:length(trig_data))/Fs;
    cur_fig = figure;
    plot(t,trig_data); 
    hold on;
    plot(t,aud_data,'r');
    hold off;
    [~,data_filename,~] = fileparts(data_file);
    saveas(cur_fig,fullfile(params.plot_dir,strrep(data_filename,'_trigs','')),'fig');
    close(cur_fig);    
end

%% FUNCTIONS
function out = getMovingMax(data,wind)
% Gets moving maximum

out = ones(size(data))*NaN;

for i = 1:length(data)
    
    if i <= wind
        cur_data = data(1:i);
    else
        cur_data = data(i-wind:i);
    end
    
    out(i)=max(cur_data);
    
end

function params = check_params(params)

% check parameters and thresholds (set defaults)
if ~isfield(params,'time_thresh_ms');
    params.time_thresh_ms = 20; % time after trig to look for audio onset
end
if ~isfield(params,'trig_time_thresh_ms');
    params.trig_time_thresh_ms = 0.5; % time after trig to look for offset
end
if ~isfield(params,'trig_on_thresh');
    params.trig_on_thresh = 0.01; % threshold to detect MIDI onset
end
if ~isfield(params,'trig_off_thresh');
    params.trig_off_thresh = 0.01; % threshold to detect MIDI offset
end
if ~isfield(params,'aud_on_thresh');
    params.aud_on_thresh = 0.1; % threshold to detect audio onset
end
if ~isfield(params,'moving_max_wind_ms');
    params.moving_max_wind_ms = 1; % window for moving maximum
end
if ~isfield(params,'off_time_thresh_ms');
    params.off_time_thresh_ms = 80; % time to wait for next MIDI onset
end
if ~isfield(params,'n_trigs');
    params.n_trigs = 201; % expected number of onsets
end

