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 = n*a**n * np.heaviside(n-1, 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((0, a), (1, -2*a, a*a))
omega, h = X

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

# Set up plot of actual DTFT of $x[n]$ for comparisson 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, 40):
    # Finite geometric series formula
    #X_K = ((a*np.exp(-1j*omega)) - (a*np.exp(-1j*omega))**(K+3)) / (1 - a*np.exp(-1j*omega))**2 - (K+2)*(a*np.exp(-1j*omega))**(K+2) / (1 - a*np.exp(-1j*omega))
    n_K = np.arange(-K,K+1)
    X_K = np.zeros(omega.shape, np.complex128)
    for n in n_K:
        x_n = n* a**n * np.heaviside(n-1, 1)
        X_K += x_n * np.exp(-1j * n * 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.empty(K_range.shape)
for K in K_range:
    X_K = ((a*np.exp(-1j*omega)) - (a*np.exp(-1j*omega))**(K+3)) / (1 - a*np.exp(-1j*omega))**2 - (K+2)*(a*np.exp(-1j*omega))**(K+2) / (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()