import numpy as np
import matplotlib
import matplotlib.pyplot as plt

import sys
sys.path.append("../")
from quelea import *

nx = 217
ny = 133

x0 = 0
x1 = 30 # lambdas
y0 = 0
y1 = 20 # lambdas

xs = np.linspace(x0, x1, nx)
ys = np.linspace(y0, y1, ny)

# 2d array of (x, y, z, t)
coords = np.array( [ [x, y, 0, 0] for x in xs for y in ys ] )
# for map_fields function this should be converted from 2D to 1D array
coords = coords.reshape((4 * nx * ny,))

ftype = 1 # plane wave
a0 = 1 # normalized field amplitude
omega = 1 # frequency
fparam = [a0, 1, 0, 0, 0, 1, 0, 0, omega] # parameters of the plane wave

ex, ey, ez, bx, by, bz = map_fields(coords, ftype, fparam)
# now convert to 2d arrays
ex = ex.reshape((nx, ny))
ey = ey.reshape((nx, ny))
ez = ez.reshape((nx, ny))
bx = bx.reshape((nx, ny))
by = by.reshape((nx, ny))
bz = bz.reshape((nx, ny))
ex = ex.transpose()
ey = ey.transpose()
ez = ez.transpose()
bx = bx.transpose()
by = by.transpose()
bz = bz.transpose()

plt.imshow(ey, cmap = 'RdYlBu', origin = 'lower', extent = [x0, x1, y0, y1])
plt.colorbar()
plt.clim(-a0, a0)

plt.savefig("map_fields.pdf")


[(6, 'sys.path.append', 'sys.path.append', False, 'import sys\n'), (17, 'numpy.linspace', 'np.linspace', True, 'import numpy as np\n'), (18, 'numpy.linspace', 'np.linspace', True, 'import numpy as np\n'), (21, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (45, 'matplotlib.pyplot.imshow', 'plt.imshow', True, 'import matplotlib.pyplot as plt\n'), (46, 'matplotlib.pyplot.colorbar', 'plt.colorbar', True, 'import matplotlib.pyplot as plt\n'), (47, 'matplotlib.pyplot.clim', 'plt.clim', True, 'import matplotlib.pyplot as plt\n'), (49, 'matplotlib.pyplot.savefig', 'plt.savefig', True, 'import matplotlib.pyplot as plt\n')]