import numpy as np


def time_integrate(K,N):
    T   = 0.2
    tau = T/N
    x_K = np.linspace(-np.pi,np.pi,K,endpoint=False)
    U0  = np.exp(-10*x_K**2+10*1j*x_K)
    U = 1j*np.zeros((K,N+1))
    U[:,0] = U0
    D_K = -np.diag(np.arange(-K/2,K/2)**2)
    V_K = np.diag(100*(1-np.cos(x_K)))
    for n in range(N):
        DK_FK_f = D_K.dot(np.fft.fftshift(np.fft.fft(U[:,n])))
        lap_approx_Un   = np.fft.ifft(np.fft.ifftshift(DK_FK_f))
        U[:,n+1] = U[:,n] +tau*1j*lap_approx_Un -tau*1j*V_K.dot(U[:,n])
    return U


    
def make_images(K,U):
    plt.semilogy(1.0/np.sqrt(K)*np.linalg.norm(U,axis=0))
    plt.ylim([10**(-1),10**(0)])
    plt.savefig('norms_over_time')
    plt.close()
    plt.imshow(np.abs(U)**2,interpolation='nearest',aspect='auto')
    plt.colorbar()
    plt.savefig('solution')
import matplotlib.pyplot as plt
K = 32
N = 10000
U = time_integrate(K,N)
print("Finished")
#plt.semilogy(np.linalg.norm(U,axis = 0))
#plt.savefig("norms")

make_images(K,U)