import numpy as np
from scipy.signal import iirnotch, freqz, tf2zpk
import matplotlib.pyplot as plt

# Filter parameters
fs = 44100        # Sampling frequency in Hz
f0 = 2000         # Frequency to be removed from signal (notch frequency)
Q = 1      # Quality factor 

# Design notch filter
b, a = iirnotch(f0, Q, fs)

# Display transfer function
print("Numerator coefficients (b):", b)
print("Denominator coefficients (a):", a)

# Optional: frequency response
w, h = freqz(b, a, fs=fs)
plt.plot(w, 20 * np.log10(abs(h)))
plt.title('Notch Filter Frequency Response')
plt.xlabel('Frequency [Hz]')
plt.ylabel('Amplitude [dB]')
plt.grid()
plt.axvline(f0, color='red', linestyle='--', label='Notch Frequency')
plt.legend()
plt.show()

# Optional: show poles and zeros
z, p, k = tf2zpk(b, a)
print("Zeros:", z)
print("Poles:", p)
print("Gain:", k)