import numpy as np
import matplotlib.pyplot as plt

# Define input signal
n = np.array([0, 1, 2, 3, 4])
x = np.array([1, 2, 3, 2, 1])

# Frequency range from -pi to pi
w = np.linspace(-np.pi, np.pi, 2000)

# Calculate DTFT using mathematical formula
X = np.zeros(len(w), dtype=complex)

for k in range(len(w)):
    X[k] = np.sum(x * np.exp(-1j * w[k] * n))

# Calculate magnitude and phase
magnitude = np.abs(X)
phase = np.angle(X)

# Create figure with three subplots
plt.figure(figsize=(10, 12))

# --------------------------------
# Graph 1: Input Signal x[n]
# --------------------------------

plt.subplot(3, 1, 1)

plt.stem(n, x)

plt.title("Input Signal x[n]")
plt.xlabel("n")
plt.ylabel("Amplitude")

plt.xlim(0, 4)
plt.ylim(0, 5)

plt.xticks([0, 1, 2, 3, 4])
plt.yticks([0, 1, 2, 3])

plt.grid(True)

# --------------------------------
# Graph 2: Magnitude Spectrum
# --------------------------------

plt.subplot(3, 1, 2)

plt.plot(w, magnitude, 'b', linewidth=2)

plt.title("Magnitude Spectrum |X(e^jω)|")
plt.xlabel("Frequency (ω)")
plt.ylabel("Magnitude")

plt.xlim(-np.pi, np.pi)
plt.ylim(0, 9)

plt.xticks(
    [-np.pi, -np.pi/2, 0, np.pi/2, np.pi],
    [r'$-\pi$', r'$-\pi/2$', '0',
     r'$\pi/2$', r'$\pi$']
)

plt.grid(True)

# --------------------------------
# Graph 3: Phase Spectrum
# --------------------------------

plt.subplot(3, 1, 3)

plt.plot(w, phase, 'r', linewidth=2)

plt.title("Phase Spectrum ∠X(e^jω)")
plt.xlabel("Frequency (ω)")
plt.ylabel("Phase (Radians)")

plt.xlim(-np.pi, np.pi)
plt.ylim(-5, 5)

plt.xticks(
    [-np.pi, -np.pi/2, 0, np.pi/2, np.pi],
    [r'$-\pi$', r'$-\pi/2$', '0',
     r'$\pi/2$', r'$\pi$']
)

plt.grid(True)

# Adjust spacing and display
plt.tight_layout()
plt.show()