import numpy
from schrodinger2d.main import Schrodinger2d
from matplotlib.pyplot import *
import imageio
            

px_min = 7
py_min = 6
levels = 3
Lx = 2.0
solver = Schrodinger2d(px_min,py_min,levels,Lx,normalize=False,colormap=2,gamma=0.6)
             

solver.opencl_platforms()
solver.set_opencl_platform_device(0,0)
solver.opencl_init()
             

dt = 2*solver.dx**2
solver.schrodinger(dt)
             

lignes_nulles_fentesA = [[64,0,29,"u"],[64,31,2,"u"],[64,35,29,"u"]]
solver.zero_line_list(lignes_nulles_fentesA)
             

x0=0.7
y0=0.5
k0 = 250
sigma0 = 0.1
E=k0*k0
solver.paquet(x0,y0,k0,sigma0)
             

solver.init()
solver.opencl_create_memory()
             

ti = 0.0
tf = 50*dt
solver.opencl_iterations(ti,tf)
             

img = solver.get_proba_colors()
figure()
imshow(img)
             

writer = imageio.get_writer('../../../../figures/numerique/diffusion/pyschrodinger2d/fentes-img.mp4',fps=10)
t = 0
solver.paquet(x0,y0,k0,sigma0)
solver.opencl_release_memory()
solver.opencl_create_memory()
delta_t = dt
while t < 200*dt:
    img = solver.offscreen_rendering(t,t+delta_t,opencl=1,width=solver.width,height=solver.height,rendering=1)
    t+=delta_t
    writer.append_data(img) 
writer.close()
              

writer = imageio.get_writer('../../../../figures/numerique/diffusion/pyschrodinger2d/fentes-surf.mp4',fps=10)
t = 0
solver.paquet(x0,y0,k0,sigma0)
solver.opencl_release_memory()
solver.opencl_create_memory()
delta_t = dt
while t < 200*dt:
    img = solver.offscreen_rendering(t,t+delta_t,opencl=1,width=solver.width,height=solver.height,rendering=2,surface_height=0.2,angle_x=45,angle_z=10)
    t+=delta_t
    writer.append_data(img)
writer.close()
              

solver.opencl_release_memory()
solver.close()