clc;
close all;
clear;

%% === LOAD AUDIO FILES ===
[y1, Fs1] = audioread('flac.wav');
[y2, Fs2] = audioread('flac2.wav');

if size(y1,2) > 1, y1 = mean(y1,2); end
if size(y2,2) > 1, y2 = mean(y2,2); end

t1 = (0:length(y1)-1) / Fs1;
t2 = (0:length(y2)-1) / Fs2;

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

%% === TIME-DOMAIN WAVEFORMS ===
figure;
subplot(2,1,1);
plot(t1, y1);
xlabel('Time [s]');
ylabel('Amplitude');
title('flac.wav – Time Domain');
grid on;

subplot(2,1,2);
plot(t2, y2, 'r');
xlabel('Time [s]');
ylabel('Amplitude');
title('flac2.wav – Time Domain');
grid on;

%% === FREQUENCY-DOMAIN ANALYSIS (FFT) ===
N1 = length(y1);
N2 = length(y2);

X1 = fft(y1);
X2 = fft(y2);

freq_shift1 = (-N1/2 : N1/2 - 1) * (Fs1 / N1);
freq_shift2 = (-N2/2 : N2/2 - 1) * (Fs2 / N2);

% --- Real / Imag components ---
figure;
subplot(2,1,1);
plot(freq_shift2, real(fftshift(X2)));
xlabel('Frequency [Hz]');
ylabel('Real Part');
title('flac2.wav – Real Part of FFT');
grid on;
subplot(2,1,2);
plot(freq_shift2, imag(fftshift(X2)), 'r');
xlabel('Frequency [Hz]');
ylabel('Imag Part');
title('flac2.wav – Imaginary Part of FFT');
grid on;

% --- Magnitude Spectrum ---
figure;
plot(freq_shift2, 20*log10(abs(fftshift(X2)) + eps), 'y');
xlabel('Frequency [Hz]');
ylabel('Magnitude [dB]');
title('flac2.wav – Magnitude Spectrum');
xlim([0 2000]);
grid on;

%% === PEAK DETECTION TO FIND TONAL NOISE FREQUENCIES ===
halfN2 = floor(N2/2);
f2 = (0:halfN2-1)*(Fs2/N2);
mag2 = abs(X2(1:halfN2));

[pks, locs] = findpeaks(mag2, 'MinPeakHeight', max(mag2)*0.1, ...
    'MinPeakDistance', round(10*N2/Fs2));

f_tonal = f2(locs);
f_tonal = f_tonal(f_tonal > 30 & f_tonal < Fs2/2);
f_tonal = sort(f_tonal(1:min(5, numel(f_tonal))));

fprintf('\nDetected tonal noise frequencies (from flac2.wav):\n');
fprintf('  %.1f Hz\n', f_tonal);

%% === NOTCH FILTER DESIGN (Z-domain zeros/poles) ===
r = 0.995;         % pole radius
sos_all = [];
g_all = 1;

for i = 1:length(f_tonal)
    f0 = f_tonal(i);
    w0 = 2*pi*f0/Fs1;                % use Fs1 so it matches flac.wav
    z = [exp(1j*w0) exp(-1j*w0)];    % zeros at ±ω0
    p = r*z;                         % poles slightly inside circle
    b = real(poly(z));
    a = real(poly(p));
    b = b / (sum(b)/sum(a));         % normalize gain
    [sos_i, g_i] = tf2sos(b, a);
    sos_all = [sos_all; sos_i];
    g_all = g_all * g_i;
end

%% === VISUALIZE NOTCH FILTER RESPONSE ===
figure;
freqz(sos_all, 2048, Fs1);
title('Notch Filter Cascade – Based on flac2 Noise Frequencies');

%% === APPLY NOTCH FILTERS TO SIGNALS ===
y1_filtered = filtfilt(sos_all, g_all, y1); % Apply to flac.wav
y2_filtered = filtfilt(sos_all, g_all, y2); % Apply to flac2.wav
y1_filtered = y1_filtered / max(abs(y1_filtered));
y2_filtered = y2_filtered / max(abs(y2_filtered));

audiowrite('filtered_flac.wav', y1_filtered, Fs1);
audiowrite('filtered_flac2.wav', y2_filtered, Fs2);

fprintf('\nSaved:\n');
fprintf('  filtered_flac.wav  (filtered signal)\n');
fprintf('  filtered_flac2.wav (filtered noise sample)\n');

%% === COMPARE ORIGINAL VS FILTERED SPECTRA ===
N_plot = length(y1);
f_plot = (0:N_plot/2-1) * Fs1 / N_plot;
X1_orig = fft(y1);
X1_filt = fft(y1_filtered);

figure;
subplot(2,1,1);
plot(f_plot, 20*log10(abs(X1_orig(1:N_plot/2))+eps));
xlabel('Frequency [Hz]'); ylabel('Magnitude [dB]');
title('flac.wav – Original Spectrum');
xlim([0 2000]); grid on;

subplot(2,1,2);
plot(f_plot, 20*log10(abs(X1_filt(1:N_plot/2))+eps), 'r');
xlabel('Frequency [Hz]'); ylabel('Magnitude [dB]');
title('flac.wav – After Notch Filtering');
xlim([0 2000]); grid on;

%% === APPLY THE SAME FILTER TO sem2.wav ===
fprintf('\n=== APPLYING DESIGNED FILTER TO sem2.wav ===\n');

[y_sem2, Fs_sem2] = audioread('sem2.wav');
if size(y_sem2,2) > 1, y_sem2 = mean(y_sem2,2); end

if Fs_sem2 ~= Fs1
    warning('Sample rates differ – resampling sem2.wav to match notch filter design');
    y_sem2 = resample(y_sem2, Fs1, Fs_sem2);
    Fs_sem2 = Fs1;
end

% --- Step 1: Apply multi‑notch filter from flac2 noise analysis ---
y_sem2_notch = filtfilt(sos_all, g_all, y_sem2);

% --- Step 2 (optional): add band‑pass filter for speech 100‑8000 Hz ---
[z_bp, p_bp, k_bp] = butter(4, [100 8000]/(Fs_sem2/2), 'bandpass');
[sos_bp, g_bp] = zp2sos(z_bp, p_bp, k_bp);
y_sem2_final = filtfilt(sos_bp, g_bp, y_sem2_notch);

% --- Normalize both versions ---
y_sem2_notch = y_sem2_notch / (max(abs(y_sem2_notch)) + eps) * 0.95;
y_sem2_final = y_sem2_final / (max(abs(y_sem2_final)) + eps) * 0.95;

% --- Save results ---
audiowrite('sem2_notch_filtered.wav', y_sem2_notch, Fs_sem2);
audiowrite('sem2_final.wav', y_sem2_final, Fs_sem2);
fprintf('Saved:\n');
fprintf('  sem2_notch_filtered.wav  (notch filters only)\n');
fprintf('  sem2_final.wav           (notch + speech band‑pass)\n');

%% === VISUAL CHECK – BEFORE vs AFTER ===
t_sem2 = (0:length(y_sem2)-1)/Fs_sem2;
N_sem2 = length(y_sem2);
f_sem2 = (0:N_sem2/2-1)*Fs_sem2/N_sem2;
Y_orig = fft(y_sem2);
Y_filt = fft(y_sem2_final);

figure;
subplot(2,1,1);
plot(t_sem2, y_sem2);
xlabel('Time [s]'); ylabel('Amplitude');
title('sem2.wav – Original Time Domain'); grid on;

subplot(2,1,2);
plot(t_sem2, y_sem2_final, 'r');
xlabel('Time [s]'); ylabel('Amplitude');
title('sem2\_final.wav – After Notch + Band‑Pass'); grid on;

figure;
subplot(2,1,1);
plot(f_sem2, 20*log10(abs(Y_orig(1:N_sem2/2))+eps));
xlabel('Frequency [Hz]'); ylabel('Magnitude [dB]');
title('sem2.wav – Original Spectrum');
xlim([0 2000]); grid on;

subplot(2,1,2);
plot(f_sem2, 20*log10(abs(Y_filt(1:N_sem2/2))+eps), 'r');
xlabel('Frequency [Hz]'); ylabel('Magnitude [dB]');
title('sem2\_final.wav – After Filtering');
xlim([0 2000]); grid on;