%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% HAMR@ISMIR 2015
% Name: BrainBeats_tempograms.m
%  
% Team Brain Beats: Sebastian Stober, Meinard Mueller, Thomas Praetzlich
%
% Description: Illustrates a hacky application of the tempogram toolbox to 
%              EEG data
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
addpath('MATLAB-Tempogram-Toolbox_1.0');
clear
close all

%% parameters
tempoWindow = 12; % window length for tempogram computation
smooth_len = 0.5; % window for local average substraction in novelty curve
sID = '14';       % stimulus    ID
pID = '09';       % participant ID

trialNum = 5;     % total number of trials
channelNum = 64;  % total number of EEG channels 
trialIdx   = [2]; % select a set of trials to accumulute
% trialIdx   = [1:trialNum];    % select trials that will be accumulated    

channelIdx = [11 32 46 47 48];  % subset of channels related to auditory cortex
% channelIdx = [1:channelNum];  % select channels

%% configuration: directories, filenames ...
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
dirData = 'data/';

filenameWav = [ sID '.wav'];
filenameAnn = [ sID '.txt'];

% filename for specific stimulus and participant
filenamePart = [ sID '_P' pID '.mat' ];

%% load beat annotations
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
annoData = csvread( [dirData filenameAnn] );

%% load EEG data 
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% EEG data (tensor of trials x channels x samples)
dataEEG = load( [dirData filenamePart]);
dataEEG = dataEEG.data;

featureRateEEG = 512; 

% accumulate EEG over trialsIdx and channelIdx
tmpTensor = sum( dataEEG(trialIdx,channelIdx,:),2);
accumulatedEEG = squeeze( sum( tmpTensor, 1 ))';


% Consider accumulatedEEG as novelty curve and subtract local average 
% -> results in highpass filtering
parameterSmooth = [];
parameterSmooth.smooth_len = smooth_len;
[ noveltyCurveEEG, local_average ] = novelty_smoothedSubtraction_EEG( accumulatedEEG,parameterSmooth );

% Visualize EEG data
parameterVis = [];
parameterVis.featureRate = featureRateEEG;
parameterVis.plotAnn = annoData;

% visualization of accumulated channels with local average curve
[time_axis] = visualize_noveltyCurve(accumulatedEEG,parameterVis);
title(['EEG Channel ' num2str(channelIdx), ' ID: ', num2str(sID) ])
axis tight
hold on;
plot(time_axis, local_average,'r','LineWidth',2)


% visualization of highpass filtered version of accumulatedEEG
visualize_noveltyCurve(noveltyCurveEEG,parameterVis)
axis tight
title(['EEG Channel ' num2str(channelIdx), ' ID: ', num2str(sID) ])

%% tempogram_fourier (EEG)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
featureRate = featureRateEEG;                  
noveltyCurve = noveltyCurveEEG;

parameterTempogram = [];
parameterTempogram.featureRate = featureRate;
parameterTempogram.tempoWindow = tempoWindow;   % window length in sec
parameterTempogram.BPM = 30:1:240;              % tempo values

[tempogram_fourier, T, BPM] = noveltyCurve_to_tempogram_via_DFT(noveltyCurve,parameterTempogram);
tempogram_fourier = normalizeFeature(tempogram_fourier,2, 0.0001);

visualize_tempogram(tempogram_fourier,T,BPM)
title('Tempogram (Fourier) on EEG data')

% mean on tempogram
figure;
mean_tempogram = mean(abs(tempogram_fourier),2);
stem(BPM, mean_tempogram)
xlabel('Tempo (BPM)')
title('Mean tempogram (EEG)')
%% PLP curve (EEG)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

parameterPLP = [];
parameterPLP.featureRate = featureRate;
parameterPLP.tempoWindow = parameterTempogram.tempoWindow;

[PLP,featureRate] = tempogram_to_PLPcurve(tempogram_fourier, T, BPM, parameterPLP);
PLP = PLP(1:length(noveltyCurve));  % PLP curve will be longer (zero padding)

parameterVis = [];
parameterVis.featureRate = featureRate;
parameterVis.plotAnn = annoData;

visualize_noveltyCurve(PLP,parameterVis)
title('PLP curve')

%% load wav file, automatically converted to Fs = 22050 and mono
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

[audio,sideinfo] = wav_to_audio('',[ dirData] ,filenameWav);
Fs = sideinfo.wav.fs;


figure; plot((0:length(audio)-1)/Fs,audio);
axis tight
xlabel('Time (sec)')
title('Waveform')


%% compute novelty curve (audio)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
parameterNovelty = [];
[noveltyCurve,featureRate] = audio_to_noveltyCurve(audio, Fs, parameterNovelty);

parameterVis = [];
parameterVis.featureRate = featureRate;
parameterVis.plotAnn = annoData;
visualize_noveltyCurve(noveltyCurve,parameterVis)
title('Novelty curve on audio')
axis tight

%% tempogram_fourier (audio)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

parameterTempogram = [];
parameterTempogram.featureRate = featureRate;
parameterTempogram.tempoWindow = tempoWindow; % window length in sec
parameterTempogram.BPM = 30:1:240;            % tempo values

[tempogram_fourier, T, BPM] = noveltyCurve_to_tempogram_via_DFT(noveltyCurve,parameterTempogram);
tempogram_fourier = normalizeFeature(tempogram_fourier,2, 0.0001);

visualize_tempogram(tempogram_fourier,T,BPM)
title('Tempogram (Fourier)')

% mean tempogram (audio)
figure;
mean_tempogram = mean(abs(tempogram_fourier),2);
stem(BPM, mean_tempogram)
title('Mean tempogram (audio)')
xlabel('Tempo (BPM)')
%% PLP curve (audio)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

parameterPLP = [];
parameterPLP.featureRate = featureRate;
parameterPLP.tempoWindow = parameterTempogram.tempoWindow;

[PLP,featureRate] = tempogram_to_PLPcurve(tempogram_fourier, T, BPM, parameterPLP);
PLP = PLP(1:length(noveltyCurve));  % PLP curve will be longer (zero padding)

parameterVis = [];
parameterVis.featureRate = featureRate;
parameterVis.plotAnn = annoData;

visualize_noveltyCurve(PLP,parameterVis)
title('PLP curve')