# -*- coding: utf-8 -*-

import numpy
from schrodinger2d.main import Schrodinger2d
import imageio

opencl=True
rendering=1 # 1=image, 2=surface

px_min = 7
py_min = 6
levels = 3
Lx = 2.0
x0=0.7
y0=0.5
k0 = 250
sigma0 = 0.1
E=k0*k0
solver = Schrodinger2d(px_min,py_min,levels,Lx,normalize=False,colormap=2,gamma=0.7)
if opencl:
    solver.opencl_platforms()
    solver.set_opencl_platform_device(0,0)
    solver.opencl_init()
    
solver.potential_disk(64,32,2,E*1.0)
dt = 2*solver.dx**2
solver.schrodinger(dt)

solver.paquet(x0,y0,k0,sigma0)
solver.init()
if opencl:
    solver.opencl_create_memory()
writer = imageio.get_writer('disque.mp4',fps=10)
t = 0.0
delta_t = dt
while t < 200*dt:
    img = solver.offscreen_rendering(t,t+delta_t,opencl=opencl,width=solver.width,height=solver.height,rendering=rendering,surface_height=0.3,angle_x=45,angle_z=10)
    writer.append_data(img)
    t += delta_t
writer.close()
if opencl:
    solver.opencl_release_memory()
solver.close()