import numpy as np
import matplotlib.pyplot as plt
from scipy.linalg import toeplitz
from scipy.sparse.linalg import spsolve
from scipy import sparse

# Spatial discretisation
N = 100
x = np.linspace(0, 1, N)
dx = x[1] - x[0]


# Time discretisation
K = 10000
t = np.linspace(0, 10, K)
dt = t[1] - t[0]

alpha = (1j * dt) / (2 * (dx ** 2))

A = sparse.csc_matrix(toeplitz([1 + 2 * alpha, -alpha, *np.zeros(N-4)]), dtype=np.cfloat)  # 2 less for both boundaries
B = sparse.csc_matrix(toeplitz([1 - 2 * alpha, alpha, *np.zeros(N-4)]), dtype=np.cfloat)

# Initial and boundary conditions (localized gaussian)
psi = np.exp((1j * 50 * x) - (200 * (x - .5) ** 2))
b = B.dot(psi[1:-1])
psi[0], psi[-1] = 0, 0

for index, step in enumerate(t):
    # Within the domain
    psi[1:-1] = spsolve(A, b)

    # Enforce boundaries
    # psi[0], psi[N - 1] = 0, 0

    b = B.dot(psi[1:-1])
    # Square integration to show if it's unitary
    print(np.trapz(np.abs(psi) ** 2))