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


DATA_POINTS = 200


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


def main():
    a = 0.9
    n = np.arange(DATA_POINTS)
    x = n*(a**n)*u(n)

    # Plot samples of x[n]
    plt.figure()
    plt.stem(n, x)
    plt.xlabel("n")
    plt.ylabel(r"$x[n]$")
    
    # Plot the analytical DTFT
    numerator = [0, a]
    denominator = [1, -2*a, a**2]
    omega, h_0 = sp.signal.freqz(numerator, denominator, 256)
    
    plt.figure()
    plt.plot(omega, np.abs(h_0), label="Analytical DTFT")
    plt.xlabel("Frequency")
    plt.ylabel("Magnitude Response")
    
    # Plot the truncated DTFT for K = 3, 10, 20
    for K in (3, 10, 20):
        omega, h = sp.signal.freqz(x[:K], 1, 256)
        plt.plot(omega, np.abs(h), label=f"Truncated DTFT ($K = {K}$)")
    plt.legend()

    # Calculate the maximum error between the truncated DTFT
    # and the analytical DTFT for values of K from 1 to 200
    k = np.arange(DATA_POINTS)
    coeffs = np.zeros_like(k, dtype=np.float32)

    for i_k in k:
        omega, h_k = sp.signal.freqz(x[:i_k], 1, 256)
        err_k = np.abs(h_0 - h_k)
        coeffs[i_k] = np.max(err_k)
    
    # Plot the maximum error previously calculated
    plt.figure()
    plt.stem(k, coeffs)
    plt.xlabel("k")
    plt.ylabel("Supremum coefficients of the error")
    plt.show()

if __name__ == "__main__":
    main()