clc;
close all;
clear all;

%% === STAGE 0: Load Audio Files ===
[y, Fs] = audioread('sem2.wav');
[y1, Fs1] = audioread('flac.wav');
[y2, Fs2] = audioread('flac2.wav');

% Ensure mono
if size(y,2) > 1, y = mean(y,2); end
if size(y1,2) > 1, y1 = mean(y1,2); end
if size(y2,2) > 1, y2 = mean(y2,2); end

fprintf('Loaded audio files:\n');
fprintf('  sem2.wav: %d samples, Fs = %d Hz\n', length(y), Fs);
fprintf('  flac.wav: %d samples, Fs = %d Hz\n', length(y1), Fs1);
fprintf('  flac2.wav: %d samples, Fs = %d Hz\n', length(y2), Fs2);

%% === STAGE 1: IIR Notch Filter ===
fprintf('\n=== STAGE 1: Notch Filtering ===\n');

N2 = length(y2);
X2 = fft(y2);
halfN2 = floor(N2/2);
freq2 = (0:halfN2-1) * (Fs2/N2);

magX2 = abs(X2(1:halfN2));
[pks, locs] = findpeaks(magX2, 'NPeaks', 10, 'SortStr', 'descend', 'MinPeakDistance', 5);
f_peaks = freq2(locs);
f_peaks = f_peaks(f_peaks > 10);
if length(f_peaks) > 5
    f_peaks = f_peaks(1:5);
end

fprintf('Detected tonal noise frequencies: ');
fprintf('%.1f Hz, ', f_peaks);
fprintf('\n');

r = 0.995;
sos_all = [];
gain_all = 1;

for k = 1:length(f_peaks)
    wo = f_peaks(k) / (Fs2/2);
    if wo > 0 && wo < 1
        bw = 0.005;
        [b_notch, a_notch] = iirnotch(wo, bw);
        [sos, g] = tf2sos(b_notch, a_notch);
        sos_all = [sos_all; sos];
        gain_all = gain_all * g;
    end
end

if ~isempty(sos_all)
    y_notch = filtfilt(sos_all, gain_all, double(y));
    y2_notch = filtfilt(sos_all, gain_all, double(y2));
else
    y_notch = double(y);
    y2_notch = double(y2);
end

fprintf('Notch filtering complete.\n');

%% === STAGE 2: Spectral Subtraction ===
fprintf('\n=== STAGE 2: Spectral Subtraction ===\n');
params.frameLen = 2048;
params.hopLen = 512;
params.nfft = 2048;
params.alpha = 4.0;
params.beta = 0.002;
params.gamma = 1.0;
[y_specsub, ~] = spectral_subtraction(y_notch, y2_notch, Fs, params);
fprintf('Spectral subtraction complete.\n');

%% === STAGE 3: MMSE-LSA Enhancement ===
fprintf('\n=== STAGE 3: MMSE-LSA Enhancement ===\n');
params_mmse.frameLen = 2048;
params_mmse.hopLen = 512;
params_mmse.nfft = 2048;
params_mmse.alphaDD = 0.98;
params_mmse.minGain = 0.1;
params_mmse.noisePowSmooth = 0.7;
[y_mmse, ~] = mmse_lsa_denoise(y_specsub, y2_notch, Fs, params_mmse);
fprintf('MMSE-LSA enhancement complete.\n');

%% === STAGE 4: Post-processing ===
fprintf('\n=== STAGE 4: Post-processing ===\n');
[b_hp, a_hp] = butter(2, 80/(Fs/2), 'high');
y_final = filtfilt(b_hp, a_hp, y_mmse);
y_final = y_final / (max(abs(y_final)) + eps) * 0.95;
fprintf('Post-processing complete.\n');

%% === Save outputs ===
audiowrite('sem2_stage1_notch.wav', y_notch/max(abs(y_notch)+eps)*0.95, Fs);
audiowrite('sem2_stage2_specsub.wav', y_specsub/max(abs(y_specsub)+eps)*0.95, Fs);
audiowrite('sem2_stage3_mmse.wav', y_mmse/max(abs(y_mmse)+eps)*0.95, Fs);
audiowrite('sem2_final_cleaned.wav', y_final, Fs);

y2_cleaned = spectral_subtraction(y2_notch, y2_notch, Fs2, params);
y2_cleaned = mmse_lsa_denoise(y2_cleaned, y2_notch, Fs2, params_mmse);
y2_cleaned = y2_cleaned / (max(abs(y2_cleaned)) + eps) * 0.95;
audiowrite('flac2_final_cleaned.wav', y2_cleaned, Fs2);
fprintf('\n✅ All files saved!\n');

%% === SIMPLE VISUALIZATION (same style as older system) ===
t = (0:length(y)-1)/Fs;

figure;
subplot(3,1,1);
plot(t, y); grid on;
title('Original sem2.wav');
xlabel('Time [s]'); ylabel('Amplitude');

subplot(3,1,2);
plot(t, y_notch); grid on;
title('After Notch Filtering');
xlabel('Time [s]'); ylabel('Amplitude');

subplot(3,1,3);
plot(t, y_final); grid on;
title('Final Cleaned sem2\_final\_cleaned.wav');
xlabel('Time [s]'); ylabel('Amplitude');

figure;
subplot(3,1,1);
plot((0:length(y2)-1)/Fs2, y2); grid on;
title('Original flac2.wav (Noise Sample)');
xlabel('Time [s]'); ylabel('Amplitude');

subplot(3,1,2);
plot((0:length(y2_notch)-1)/Fs2, y2_notch); grid on;
title('After Notch Filtering (Noise)');
xlabel('Time [s]'); ylabel('Amplitude');

subplot(3,1,3);
plot((0:length(y2_cleaned)-1)/Fs2, y2_cleaned); grid on;
title('Cleaned flac2.wav (Noise Reference)');
xlabel('Time [s]'); ylabel('Amplitude');

fprintf('\nProcessing summary:\n');
fprintf(' - Stage 1 removed %d tonal frequencies\n', length(f_peaks));
fprintf(' - Stage 2 applied spectral subtraction (α = %.1f)\n', params.alpha);
fprintf(' - Stage 3 applied MMSE-LSA (α_DD = %.2f)\n', params_mmse.alphaDD);
fprintf(' - Final result saved as sem2_final_cleaned.wav (best quality)\n');


function [y_out, noise_psd] = spectral_subtraction(x, noise_ref, Fs, p)
    x = x(:);
    noise_ref = noise_ref(:);
    frameLen = p.frameLen; hopLen = p.hopLen; nfft = p.nfft;
    alpha = p.alpha; beta = p.beta; gamma = p.gamma;
    win = hamming(frameLen, 'periodic');

    [N_stft, ~, ~] = stft(noise_ref, Fs, 'Window', win, ...
        'OverlapLength', frameLen-hopLen, 'FFTLength', nfft);
    noise_psd = mean(abs(N_stft).^2, 2);
    noise_psd = movmean(noise_psd, 5);

    [X_stft, ~, ~] = stft(x, Fs, 'Window', win, ...
        'OverlapLength', frameLen-hopLen, 'FFTLength', nfft);

    [nFreq, nFrames] = size(X_stft);
    Y_stft = zeros(size(X_stft));

    for k = 1:nFrames
        X_frame = X_stft(:, k);
        X_mag = abs(X_frame);
        X_phase = angle(X_frame);
        X_pow = X_mag.^2;
        snr_post = X_pow ./ (noise_psd + eps);
        alpha_freq = alpha * (1 + 0.5 * (1 - min(snr_post, 4)/4));
        Y_pow = X_pow - alpha_freq .* noise_psd;
        Y_pow = max(Y_pow, beta * X_pow);
        Y_mag = sqrt(max(Y_pow, 0));
        Y_stft(:, k) = Y_mag .* exp(1j * X_phase);
    end

    y_out = istft(Y_stft, Fs, 'Window', win, ...
                  'OverlapLength', frameLen-hopLen, ...
                  'FFTLength', nfft, 'ConjugateSymmetric', true);
    y_out = y_out(1:min(length(y_out), length(x)));
    if length(y_out) < length(x)
        y_out = [y_out; zeros(length(x)-length(y_out),1)];
    end
end


function [y_out, G_all] = mmse_lsa_denoise(x, noise_ref, Fs, p)
    x = x(:); noise_ref = noise_ref(:);
    frameLen = p.frameLen; hopLen = p.hopLen; nfft = p.nfft;
    alphaDD = p.alphaDD; Gmin = p.minGain;
    win = hamming(frameLen, 'periodic');

    [N_stft, ~, ~] = stft(noise_ref, Fs, 'Window', win, ...
        'OverlapLength', frameLen-hopLen, 'FFTLength', nfft);
    lambda_d = mean(abs(N_stft).^2, 2);
    lambda_d = movmean(lambda_d, 5);

    [X_stft, ~, ~] = stft(x, Fs, 'Window', win, ...
        'OverlapLength', frameLen-hopLen, 'FFTLength', nfft);

    [nFreq, nFrames] = size(X_stft);
    Y_stft = zeros(size(X_stft));
    G_all = zeros(nFreq, nFrames);
    xi_prev = ones(nFreq, 1);

    for k = 1:nFrames
        X_frame = X_stft(:, k);
        X_mag = abs(X_frame);
        X_phase = angle(X_frame);
        X_pow = X_mag.^2;
        gamma_k = X_pow ./ (lambda_d + eps);
        xi_ml = max(gamma_k - 1, 0);
        xi = alphaDD * xi_prev + (1 - alphaDD) * xi_ml;
        xi = max(xi, 1e-6);
        v = xi .* gamma_k ./ (1 + xi);
        G = zeros(nFreq,1);
        for i = 1:nFreq
            if v(i) < 0.01
                G(i) = xi(i)/(1+xi(i));
            else
                G(i) = xi(i)/(1+xi(i)) * exp(0.5 * expint_approx(v(i)));
            end
        end
        G = max(G, Gmin);
        G = movmean(G, 3);
        Y_mag = G .* X_mag;
        Y_stft(:,k) = Y_mag .* exp(1j*X_phase);
        G_all(:,k) = G;
        xi_prev = xi;
    end

    y_out = istft(Y_stft, Fs, 'Window', win, ...
                  'OverlapLength', frameLen-hopLen, ...
                  'FFTLength', nfft, 'ConjugateSymmetric', true);
    y_out = y_out(1:min(length(y_out), length(x)));
    if length(y_out) < length(x)
        y_out = [y_out; zeros(length(x)-length(y_out),1)];
    end
end


function y = expint_approx(x)
    if x < 1e-10
        y = -log(1e-10) - 0.5772;
    elseif x < 1
        y = -log(x) - 0.5772156649;
        term = x;
        for n = 1:50
            y = y + term / n;
            term = -term * x / (n + 1);
            if abs(term) < 1e-10, break; end
        end
    else
        y = exp(-x) / x;
        term = 1;
        for n = 1:10
            term = -term * n / x;
            y = y + term * exp(-x) / x;
            if abs(term) < 1e-10, break; end
        end
    end
    y = -y;
end