import numpy as np
import matplotlib.pyplot as plt


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


def main():
    # Part a
    a = np.pi/4
    n = np.arange(51)
    x = 5*np.cos(a**2 * n)

    plt.stem(n, x)
    plt.xlabel("$n$")
    plt.ylabel(r"$x[n] = 5\cos(a^2 n)$")
    plt.show()

    # Part b
    A = 5
    b = (1 + 1j)/2
    x = A*np.abs(b)**n * u(n)

    plt.stem(n, x)
    plt.xlabel("$n$")
    plt.ylabel("$x[n] = A|b|^n u(n)$")
    plt.show()


if __name__ == "__main__":
    main()