% This program performs FFT based noise reduction
%
% Radu Balan April 12, 2017.
clear all

% Load data
%[x,fs] = audioread('SR1Noise0.1.wav');
[x,fs] = audioread('SA1_A10_Noise0.1.wav');
x = x(:); % Turn it into a column vector
N = size(x,1);

% Define the windows
T1 = 40;
T2 = 160;
T3 = 16;
b = T1+T2;
g1 = [(sin((0:T1-1)*pi/(2*T1))).^2,ones(1,T2),(cos((0:T1-1)*pi/(2*T1))).^2,zeros(1,T3)]';
g2 = ones(1,T1+T2+T1+T3)';
% For perfect reconstruction: 
%    sum_k (g1.*g2)(t-k.b) = 1 , for all t ;
%  where b=timestep

% FFT size
T = length(g1);
F = T;
Kt = floor(N/b); % Number of time steps

% Perform time->time/frequency transform
X = zeros(F,Kt);
disp('Windowed FFT analysis')
k = 0;
for t=1:b:N-T
    k = k+1;

    % Window the signal
    xx = x(t:t+T-1).*g1;

    % FFT
    X(1:F,k) = fft(xx);
end
Kt = k; % Ignore additional memory allocation, if any

% Apply Wiener filter assuming noise variance sigma^2 is known
sigma = 0.1;
alpha = 0.5; % Learning rate
theta = 0.5; % Flooring for Wiener filter
% Learn the noise spectral power from the first few frames
NN = 10;
RnExp = zeros(F,1);
for k=1:NN
    RnExp = RnExp + abs(X(1:F,k)).^2;
end
RnExp = RnExp/NN; % Average -- per frequency
sgmAvg = sum(RnExp)/F;

% Else compute the windowed noise spectral power
RnTh = sigma^2*sum(abs(g1).^2)*ones(F,1); % Adjust for windowing & fft

fprintf(1,'Average noise spectral power: %f  vs. Theoretical value: %f \n',sgmAvg,RnTh(1));

Rn = RnTh;

Rx = zeros(F,1);
Y = zeros(F,Kt);
disp('Signal processing')
for k=1:Kt
    % Update signal spectral power
    Rx=(1-alpha)*Rx+alpha*abs(X(1:F,k)).^2;
    
    % Compute the Wiener Filter: signal/(signal+noise) = SNR/(1+SNR)
    W = 1 - Rn./Rx;
    % Set to theta all values below theta 
    W(W<theta) = theta;
    
    % Filter
    Y(1:F,k) = W .* X(1:F,k);
end


% Reconstruct time/frequency -> time
y= zeros(size(x));
disp('Signal reconstruction')
k = 0;
for t=1:b:N-T
    k = k+1;

    % IFFT
    yy = ifft(Y(1:F,k));
    
    % Overlap-add
    y(t:t+T-1) = y(t:t+T-1) + yy.*g2;
end

% Save waveform
audiowrite('SA1DenoiseFFT.wav',y,fs)

% Analysis

% Plot windows
figure(1)
plot(1:T,g1,1:T,g2)
legend('g1','g2')

figure(2)
plot(1:N,x,1:N,y)
legend('x','y')


figure(3)
plot(1:F,RnExp,1:F,RnTh)
legend('Numerical','Theoretical')
