import numpy as np
import matplotlib.pyplot as plt


def delta(n):
    return np.where(n==0, 1, 0)


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


def main():
    # Part a
    n = np.arange(13)
    x = delta(n-5) - 2*delta(n-8) + 6*delta(n-11)
    plt.stem(n, x)
    plt.xlabel("$n$")
    plt.ylabel(r"$x[n] = \delta(n-5) - 2\delta(n-8) + 6\delta(n-11)$")
    plt.show()

    # Part b
    n = np.arange(16)
    x = u(n-3) - 2*u(n-8) + u(n-12)
    plt.stem(n, x)
    plt.xlabel("$n$")
    plt.ylabel("$x[n] = u(n-3) - 2u(n-8) + u(n-12)$")
    plt.show()

    # Part c
    n = np.arange(21)
    x = n*u(n)
    plt.stem(n, x)
    plt.xlabel("$n$")
    plt.ylabel("$x[n] = n u(n)$")
    plt.show()

    # Part d
    n = np.arange(11)
    x = n**2 * u(n)
    plt.stem(n, x)
    plt.xlabel(r"$n$")
    plt.ylabel(r"$x[n] = n^2 u(n)$")
    plt.show()

    # Part e
    n = np.arange(31)
    x = 3 * (n-3)**2 * np.exp(-0.3*n) * np.sin(2*n/3) * u(n)
    plt.stem(n, x)
    plt.xlabel("$n$")
    plt.ylabel(r"$x[n] = 3(n-3)^2 e^{-0.3n} \sin(2n/3) u(n)$")
    plt.show()


if __name__ == "__main__":
    main()