import numpy as np
import scipy as sp
import matplotlib.pyplot as plt


def u(n):
    return np.heaviside(n, 1)

# Transfer function for an ideal lowpass filter
def ideal_lowpass(omega, omega_c):
    return np.where(np.abs(omega) <= omega_c, 1, 0)


def main():
    f_c = 0.4                                           # Cutoff frequency
    omega_c = np.pi*f_c                                 # Angular cutoff frequency
    omega = np.arange(0, np.pi, 0.05*np.pi)             # Frequency range

    n = np.arange(-50, 51)                              # Sample range (100 samples centered at 0)

    x = f_c*np.sinc(f_c*n)                              # Impulse resonse of ideal LPF
    plt.stem(n,x)
    plt.show()

    h_lpf = ideal_lowpass(omega, omega_c)               # Plot X(e^jw)
    plt.plot(omega, h_lpf)
    plt.xlabel("Frequency")
    plt.ylabel("DTFT of Ideal Lowpass Filter")


    for k in (10, 20, 30):                              # K values for truncated DTFT
        w,h = sp.signal.freqz(x[50-k:50+k], 1, omega)   # Calculate truncated DTFT
        plt.plot(omega, np.abs(h))                      # Plot magnitude of truncated DTFT
    plt.show()

if __name__ == "__main__":
    main()