import numpy as np
import random
import math
import pyglet
from pyglet.gl import *

def prep_3iq(n,dx):
    A = np.zeros((n,),dtype=np.complex64)
    for k in range(0,n):
        q = 2*math.pi*k/(dx*n) if k<=n//2 else 2*math.pi*(k-n)/(dx*n)
        A[k] = 3*1j*q
    return A

def prep_iq3(n,dx):
    A = np.zeros((n,),dtype=np.complex64)
    for k in range(0,n):
        q = 2*math.pi*k/(dx*n) if k<=n//2 else 2*math.pi*(k-n)/(dx*n)
        A[k] = 1j*q*q*q
    return A

###
N = 256
L = 20
dx = L/N

OP_3iq = prep_3iq(N,dx)
OP_iq3 = prep_iq3(N,dx)

x = np.linspace(-L/2,L/2,N)
u = np.exp(-x*x)*12

t = 0
dt = 0.00001

SCALE=1
H = 400

IMG_DATA = np.zeros(shape=(H,N*SCALE,3),dtype=np.uint8)
plotImg=None

config = pyglet.gl.Config(double_buffer=True)
win = pyglet.window.Window(N*SCALE,H,config=config,resizable=False)
win.set_caption('KdV Equation')

batch = pyglet.graphics.Batch()

def nextstep():
    global u,t,umem
    U = np.fft.fft(u)
    for cnt in range(100):
        U1 = U * np.exp(OP_iq3 * dt)
        K1 = - OP_3iq * np.fft.fft(u*u)
        U2 = U1 + K1 * dt/2
        u2 = np.fft.ifft(U2).real
        K2 = - OP_3iq * np.fft.fft(u2*u2)
        U = U1 + K2 * dt
        u = np.fft.ifft(U).real                
        t = t + dt

shapes=[ ]
def init_plot():
    line = pyglet.shapes.Line(x=1, y=H-10, x2=1, y2=H-20, width=1, color=(255, 255, 255), batch=batch)
    shapes.append(line)
    line = pyglet.shapes.Line(x=1, y=H-20, x2=N*SCALE, y2=H-20, width=1, color=(255, 255, 255), batch=batch)
    shapes.append(line)    
    line = pyglet.shapes.Line(x=N*SCALE, y=H-20, x2=N*SCALE, y2=H-10, width=1, color=(255, 255, 255), batch=batch)
    shapes.append(line)
    
    str = '{0}'.format(-L/2)
    label_x0 = pyglet.text.Label(str,font_name='Helvetica', color=(255,255,255,255), font_size=10, x=0, y=H-30, batch=batch)
    shapes.append(label_x0)

    str = '{0}'.format(L/2)
    label_x1 = pyglet.text.Label(str,font_name='Helvetica', color=(255,255,255,255), font_size=10, x=N*SCALE-30, y=H-30, batch=batch)
    shapes.append(label_x1)    
        
@win.event
def on_draw():
    win.clear()
    if plotImg != None:
        plotImg.blit(0,0)        
    batch.draw()
    str = 'T={0:.3f}'.format(t)
    label_time = pyglet.text.Label(str,font_name='Helvetica', color=(255,255,255,255), font_size=10, x=N*SCALE/2-30, y=H-12)
    label_time.draw()    

def update(deltat):
    global IMG_DATA,u,plotImg
    nextstep()    
    IMG_DATA = np.roll(IMG_DATA,1,axis=0)
    for i in range(N):
        r = int((math.sin(u[i]*0.1)+1)/2*255)
        g = int((math.sin(2*u[i]*0.1)+1)/2*255)
        b = int((1.0+math.cos(u[i]*0.1))/2*255)
        IMG_DATA[0,i*SCALE:(i+1)*SCALE,0] = r
        IMG_DATA[0,i*SCALE:(i+1)*SCALE,1] = g
        IMG_DATA[0,i*SCALE:(i+1)*SCALE,2] = b
    pixels = IMG_DATA.ravel()
    rawData =  (GLubyte * len(pixels))(*pixels)
    plotImg = pyglet.image.ImageData(N*SCALE, H, 'RGB', rawData)

init_plot()    
pyglet.clock.schedule_interval(update, 1/20.0)
pyglet.app.run()