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

import numpy
from schrodinger2d.main import Schrodinger2d
import imageio


opencl=False
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()
dt = 2*solver.dx**2
solver.schrodinger(dt)
lignes_nulles_ouverture = [[64,0,30,"u"],[64,34,30,"u"]]
lignes_nulles_fentesA = [[64,0,29,"u"],[64,31,2,"u"],[64,35,29,"u"]]
lignes_nulles_fentesB = [[64,0,28,"u"],[64,30,4,"u"],[64,36,28,"u"]]
lignes_nulles_obstacle = [[64,30,4,"u"]]
lignes_nulles_reseau = []
for i in range(16):
    lignes_nulles_reseau.append([64,i*4,2,"u"])
solver.zero_line_list(lignes_nulles_fentesA)
solver.paquet(x0,y0,k0,sigma0)
solver.init()
if opencl:
    solver.opencl_create_memory()
writer = imageio.get_writer('fentes.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()