% Evan Shenkman
% MMI610
% Project 1
% Swept Distortion Product OAE

% --- Setup simulation params
fs = 48000; % sampling rate
Ts = 1/fs; % sampling period
t = 0:Ts:0.200; % time vec
f1 = [1000 6000]; % f1s
f2 = 1.22 * f1; % f2s
f_dp = 2*f1 - f2;
sweepUp = true;

% -- the rest of the simulation params
spl = 70;
irron = 1;
normalizeRms = 1;
subjectNo = 1;
fc = 'all';

% --- Helper Functions
if sweepUp
f1_at_t = @(t) ((f1(2) - f1(1))/endTime * t + f1(1));
f2_at_t = @(t) ((f2(2) - f2(1))/endTime * t + f2(1));
f_dp_at_t = @(t) ((f_dp(2) - f_dp(1))/endTime * t + f_dp(1));
else
f1_at_t = @(t) ((f1(1) - f1(2))/endTime * t + f1(2)); %#ok<*UNRCH>
f2_at_t = @(t) ((f2(1) - f2(2))/endTime * t + f2(2));
f_dp_at_t = @(t) ((f_dp(1) - f_dp(2))/endTime * t + f_dp(2));
end

% --- Signal Generation
% Make the chirp (sine sweep)
endTime = 0.025;

if sweepUp
s1 = chirp(t(1:endTime*fs), f1(1), endTime, f1(2));
s2 = chirp(t(1:endTime*fs), f2(1), endTime, f2(2));
else
s1 = chirp(t(1:endTime*fs), f1(2), endTime, f1(1));
s2 = chirp(t(1:endTime*fs), f2(2), endTime, f2(1));
end

% Compute the weighted sum
s3 = s1 + s2 * db2mag(-10);

% Zero pad
s3 = [s3 zeros(1, length(t) - length(s3))];

% Smooth the discontinuity
low_fc = 10e3;
[B,A] = butter(1, 2*low_fc/fs);
s3 = filter(B,A,s3);

% Normalize
s3 = s3/max(abs(s3));

% Model input
insig = s3';

% Write the Model input to a wav file
if sweepUp
wavFileName = 'dpoae_sweep_up.wav';
else
wavFileName = 'dpoae_sweep_down.wav';
end
%audiowrite(wavFileName, insig, fs);

% --- Run the model for ALL signals
[~,~,E,~]=verhulst2012(insig,fs,fc,spl,normalizeRms,subjectNo,irron);

%% --- Plot the Chirp, Model Output, and spectrogram of the Output
h = figure(1);
set(h, 'Position', [100 100 1100 900]);

subplot(4,1,1)
plot(1e3*t, insig);
if sweepUp
title('Upwards Chirp');
else
title('Downwards Chirp');
end
xlabel('Time (ms)');
ylabel('Amplitude');
ylim([-2.5 2.5]);

subplot(4,1,2)
plot(1e3*t, E);
if sweepUp
title('Model Output for Upwards Chirp');
else
title('Model Output for Downwards Chirp');
end
xlabel('Time (ms)');
ylabel('Amplitude');
ylim([-2.5 2.5]);

subplot(4,1,[3 4])
sgram(E, fs, 'nocolorbar');
if sweepUp
title('Spectrogram of Model Output for Upwards Chirp');
else
title('Spectrogram of Model Output for Downwards Chirp');
end
%saveas(h, 'dpoae_sweep1.png');

% --- Take the STFT (to easily get window timestamps)
win_len = 128;
window = hamming(win_len);
[~, ~, T] = spectrogram(E, window, win_len/2, 2048, fs);

% --- Plot the spectrogram of the Model Output and overlay the mean
% frequency per window location
h = figure(2);
set(h, 'Position', [100 100 1200 700]);
sgram(E, fs);
hold on;

% Plot the frequency-time locations of where the BPF fc will be per frame
T_sec = T(1:19);
stem(T_sec, f1_at_t(T_sec), 'o', 'Color', 'white', 'LineStyle', 'none');
stem(T_sec, f2_at_t(T_sec), 'o', 'Color', 'white', 'LineStyle', 'none');
stem(T_sec, f_dp_at_t(T_sec), 'o', 'Color', 'white', 'LineStyle', 'none');

title(sprintf('Model Output Spectrogram\nwith overlaid BPF Center Frequencies'));
hold off;

% Save the plot
%saveas(h, 'dpoae_sweep2.png');

% --- Make the High-order BPFs
% We need one filter per frame (1025 by # frames)
filtOrder = 1024;
filters = zeros(filtOrder+1, length(T));

% Make each individual BPF
h = figure(3);
set(h, 'Position', [100 100 700 700]);
for i = 1:length(T)
% -- Design the BPFs
% Check if the chirp is in this frame
% We only need to filter the chirp
if T(i) - T(1) < endTime
hiF = (2*f_dp_at_t(T(i))/3 + f1_at_t(T(i))/3);
loF = f_dp_at_t(T(i)) - (hiF - f_dp_at_t(T(i)));
filters(:,i) = fir1(filtOrder, (2/fs)*[loF hiF], 'bandpass');

% Plot the filter response
freqz(filters(:,i), 1, 512);
title(sprintf('BPF #%d\nf_{c} = %.0fHz', i, f_dp_at_t(T(i))));
subplot(2,1,1)
ylim([-150 20]);
subplot(2,1,2)
ylim([-4000 500]);

% Wait and Save
%waitforbuttonpress();
%saveas(h, sprintf('dpoae_sweep3_%d.png', i));

% Otherwise, just pass through
else
filters(1,i) = 1;
end
end

% --- Apply the BPFs using Convolution and OLA (plotting and animating the
% whole process)
h = figure(4);
set(h, 'Position', [100 100 1200 700]);

subplot(2,1,1)
plot(1e3*t, E);
title('Model Output');
xlabel('Time (ms)');
ylabel('Amplitude');
ylim([-2.5 2.5]);

OAE = zeros(size(E));
for i = 1:length(T)
% Window the Model Output to get the current frame
indStart = round(T(i)*fs-win_len/2+1);
indEnd = indStart + win_len - 1;
windowedSig = E(indStart:indEnd, 1) .* hamming(win_len);

% Overlay the window on top of the Model Output
subplot(2,1,1)
hold on;
winH = plot(1e3*t(indStart:indEnd), hamming(win_len), '--k');
title('Accumulated Filter Output');
xlabel('Time (ms)');
ylabel('Amplitude');
ylim([-2.5 2.5]);
hold off;

% Convolve the frame with the frame's corresponding BPF
convOutput = conv(windowedSig, filters(:,i));
convLen = length(convOutput);

% Check the conv length vs the remaining OAE (accumlation vector) length
if indStart+convLen-1 <= size(OAE,1)
OAE(indStart:indStart+convLen-1, 1) = OAE(indStart:indStart+convLen-1, 1) + convOutput;
else
OAE(indStart:end, 1) = OAE(indStart:end, 1) + convOutput(1:(size(OAE,1)-indStart+1));
end

% Update the OAE (accumulation vector)
subplot(2,1,2)
oaeH = plot(1e3*t, OAE(:,1));

% Stop once we're not longer in the chirp section
if i >= 23
break;

else
% Wait and Save the plot
%waitforbuttonpress();
%saveas(h, sprintf('dpoae_sweep4_%d.png', i));

% Delete the plots that are going to be replotted
if i ~= length(T)
delete(winH);
delete(oaeH);
end
end
end

% --- Plot the spectrogram of the Model Output vs the filtered OAE
h = figure(5);
set(h, 'Position', [100 100 900 700]);
subplot(2,1,1)
sgram(E(:,1), fs);
title('Model Output Spectrogram');
subplot(2,1,2)
sgram(OAE(:,1), fs);
title('Filtered OAE Spectrogram');
%saveas(h, 'dpoae_sweep5.png');