import numpy as np
import pylab as pl

x_min = -0.5
x_max = 0.5
L = x_max - x_min
N = 69
alpha = 0.2

x_vals = np.linspace(x_min, x_max, 10000)

ms = np.arange(-N // 2, N // 2, 1)
ks = N + ms
amplitudes = np.exp(-alpha * ms ** 2)
amplitudes /= np.sum(amplitudes)
y_vals = np.sum(amplitudes[:, None] * np.cos(ks[:, None] * (2 * np.pi * x_vals / L)), axis=0)
y_vals /= np.sum(y_vals)

expect_m_sqrd = np.dot(ks ** 2, amplitudes)
expect_m = np.dot(ks, amplitudes)
sigma_m = np.sqrt(expect_m_sqrd - expect_m ** 2)
print(f'Sigma M: {sigma_m}')

diff = np.diff(y_vals)
minima_indices = np.where(np.r_[True, diff < 0] & np.r_[diff > 0, True])[0]
maxima_indices = np.where(np.r_[True, diff > 0] & np.r_[diff < 0, True])[0]
extremes = np.concatenate((minima_indices, maxima_indices))
envelope_ys = np.abs(y_vals[extremes]) / np.sum(np.abs(y_vals[extremes]))
envelope_xs = x_vals[extremes]
expect_x_sqrd = np.dot(envelope_xs ** 2, envelope_ys)
expect_x = np.dot(envelope_xs, envelope_ys)
sigma_x = np.sqrt(expect_x_sqrd - expect_x ** 2)
print(f'Sigma X: {sigma_x}')

print(f'Sigma K * Sigma X = {sigma_m * sigma_x * 2 * np.pi}')

pl.plot(x_vals, np.array(y_vals), label="wave packet")
pl.plot(x_vals[minima_indices], y_vals[minima_indices], '.', label="minima")
pl.plot(x_vals[maxima_indices], y_vals[maxima_indices], '.', label="maxima")
pl.vlines(x=-sigma_x, ymin=np.min(y_vals), ymax=np.max(y_vals), color="red", label="std min")
pl.vlines(x=sigma_x, ymin=np.min(y_vals), ymax=np.max(y_vals), color="green", label="std max")
pl.legend()
pl.title("Wave Packet")
pl.xlabel("x")
pl.ylabel("psi")
pl.grid()
pl.show()

pl.plot(ms, amplitudes, label="amplitudes")
pl.vlines(x=-sigma_m, ymin=np.min(amplitudes), ymax=np.max(amplitudes), color="red", label="std min")
pl.vlines(x=sigma_m, ymin=np.min(amplitudes), ymax=np.max(amplitudes), color="green", label="std max")
pl.legend()
pl.title("Fourier Amplitudes")
pl.xlabel("m")
pl.ylabel("amplitude")
pl.grid()
pl.show()