function wc = corenhance(wn,fs,ns,SNR)
% sample noise suppression program.
% procedure:    estimate the addative noise and the SNR using speechdetector 
%               generate noisy speech samples of prestored clean ones
%               find the cortcal representation of noisy speech
%               find the mapping from noisy cortical responses to clean cortical
%               apply the mapping to clean a new sample of speech
%               wn: noisy signal
%               fs: sampling frequency (Hz)
%               ns: noise sample signal (optional)
%               SNR: Signal to Noise Ratio (optional) (dB)

% Version 1.0, Nov 2008, Nima Mesgarani
% Version 2.0, Apr 2009, Majid Mirbagheri adding nonlinearity, Mapping from Cor to
% Cor, Additive noise version
% Version 2.1  July 2009, Unknown noise
%% Set parameters
pathstr = fileparts(which('corenhance'));
addpath([pathstr filesep 'SpeechDetectionPackage'],'-begin');
addpath([pathstr filesep 'nsltools'],'-begin');
rv = [1 2 4 8 16 32];
sv = [.5 1  2 4 8];
figure;loadload;close;
%% Estmation of Noise and SNR
if nargin < 3 || isempty(ns)
    disp('extracting nonspeech segments ...')
    [~, ~, ns] = speechdetector (wn, fs);
end
if nargin < 4
    movingPower = conv(abs(ns(:)).^2,ones(fs/10,1))/(fs/10);
    [vals,pws] = hist(movingPower,100);
    [~,pwidx] = max(vals);
    nPower = pws(pwidx);
    Power = sum(abs(wn(:)).^2)/length(wn(:));
    sPower = Power - nPower;
    SNR = 10*log10(sPower/nPower);
    disp(['Estmiated SNR: ' num2str(SNR)]);
end
%% Generate noisy samples by manually adding detected nonspeech segments to
% clean samples
disp 'Finding noisy/clean sample representations in cortical domain...'
%powers of noisy output to be considered for projection (pws=1 :linear
%model, pws = 1:2 a*xn^2+b*xn+c
pws = 1; 
rr = zeros(60*length(pws)+1,60*length(pws)+1,128);
rs = zeros(60*length(pws)+1,60,128);
list = dir([pathstr filesep 'Sounds' filesep '*.wav']);
inds = randperm(length(list));
for cnt1 = 1:30
    disp([num2str(cnt1) ' out of 30 samples loaded']);
    [cwav,fss] = wavread([pathstr filesep 'Sounds' filesep list(inds(cnt1)).name]);
    if fss~=fs
        cwav = resample(cwav,fs,fss);
    end
    CleanResp = wav2nlcor(cwav,1,rv,sv,1,'/tmp/tmptmp');
    % noisy representation
    nwav = addnoise(cwav,ns,SNR);
    NoisyResp = wav2nlcor(nwav,1,rv,sv,pws,'/tmp/tmptmp');
    % full complex numbers
    NoisyResp = cat(1,NoisyResp,ones(1,size(NoisyResp,2),size(NoisyResp,3)));
    for cnt5 = 1:128
        temp = squeeze(NoisyResp(:,cnt5,:));
        rr(:,:,cnt5) = rr(:,:,cnt5) + temp*temp';
        rs(:,:,cnt5) = rs(:,:,cnt5) + temp*squeeze(CleanResp(:,cnt5,:))';
    end
    clear NoisyResp
end
%% Find mappings
disp 'Finding mappings ...'
g = zeros(60,60*length(pws)+1,128);
for cnt5 = 1:128
    g(:,:,cnt5) = rs(:,:,cnt5)'*pinv(rr(:,:,cnt5));
end
%% Denoising sample test
disp 'applying the mappings on noisy signal ...'
NoisyResp = wav2nlcor(wn,1,rv,sv,pws,'/tmp/tmptmp',fs);
NoisyResp = cat(1,NoisyResp,ones(1,size(NoisyResp,2),size(NoisyResp,3)));
CleanedCor = zeros(60,size(NoisyResp,3),128);
for cnt6 = 1:128
    temp = squeeze(NoisyResp(:,cnt6,:));
    CleanedCor(:,:,cnt6) = g(:,:,cnt6)*temp;
end
clear NoisyResp
Cleaned = reshape(CleanedCor,[5 12 size(CleanedCor,2) 128]);
clear CleanedCor 
rat  = cor2aud('/tmp/tmptmp',Cleaned,[.99 0 0 0]);
ra1  = aud_fix(rat);
ra1  = ra1.^2;
rae  = energyEQ(ra1,.02);
wc = aud2wav(rae,wn,[10 10 -2 0 10 0 0]);
end

function wn = addnoise(w,n,SNR)
% add noise n to signal w with the given SNR
nrep = n;
while length(n)<length(w)
    n = cat(find(size(n)>1),n,nrep);
end
nstart = randi(length(n)-length(w)+1);
n = n(nstart:nstart+length(w)-1);
sigPower = sum(abs(w(:)).^2)/length(w(:));
sigPower = 10*log10(sigPower);
noisePower = sum(abs(n(:)).^2)/length(n(:));
noisePower = 10*log10(noisePower);
curSNR = sigPower-noisePower;
noiseGain = sqrt(10^((SNR-curSNR)/10));
wn = w + n/noiseGain;
end