import numpy as np
import matplotlib.pyplot as plt
import scipy.signal
n = np.arange(0,200)
a = 0.9

# $x[n] = a^n u[n]$
x = a**n * np.heaviside(n, 1)

asums = np.zeros(len(n))

for i in range(len(n)):
    asums[i] = np.sum(x[0:i])

# Plot settings for $x[n]$
plt.subplot(121)
plt.plot(n, x)
plt.xlabel("$n$")
plt.ylabel("$x[n]$", rotation="horizontal")

# Plot settings for sum of $x[n]$
plt.subplot(122)
plt.plot(n, asums)
plt.xlabel("$n$")
plt.ylabel("$\sum_n x[n]$", rotation="horizontal")
plt.show()

# Compute DTFT of $x[n]$
X = scipy.signal.freqz(1, (1, -a))
omega, h = X

# Plot DTFT of $x[n]$
plt.plot(omega, np.abs(h))
plt.ylabel("Amplitude")
plt.xlabel("Frequency [rad/sample]")
plt.show()

# Set up plot for DTFT of $x[n]$ for comparrison with truncated DTFTs
plt.plot(omega, np.abs(h), label="Actual DTFT")
plt.ylabel("Amplitude")
plt.xlabel("Frequency [rad/sample]")

# Calculate the truncated DTFTs of $x[n]$ as a function of $K$, $\sum_{n=-K}^K x[n] e^{-j\omega n}$
for K in (3, 10, 20):
    # Finite geometric series formula
    X_K = (1 - (a**(K+1) * np.exp(-1j * omega * (K+1)))) / (1 - a*np.exp(-1j*omega))
    plt.plot(omega, np.abs(X_K), label=f"Truncated DTFT ($K={K}$)")

plt.legend()
plt.show()

# Frequency of maximum difference between actual and truncated DTFT
K_range = np.arange(1,200+1)
max_diffs = np.zeros(K_range.shape)
for K in K_range:
    X_K = (1 - a**(K+1) * np.exp(-1j * omega * (K+1))) / (1 - a*np.exp(-1j*omega))
    abs_diff = np.abs(X_K - h)
    max_diffs[K-1] = np.max(abs_diff)

plt.plot(K_range, max_diffs)
plt.xlabel("K")
plt.ylabel("Maximum Error")
plt.show()