#!/usr/bin/env python
import postpic as pp
import numpy as np
import numexpr as ne
import matplotlib as mpl; mpl.use('Agg')
import matplotlib.pyplot as plt
from scitools.std import ndgrid

def fig_imshow(data,ext,cmap,figname,xlim,ylim):
    fig, ax = plt.subplots()
    p = ax.imshow(data.T, interpolation='none', origin='lower', extent=ext, cmap=cmap)
    ax.grid(color='k',linestyle='dashed',linewidth=0.1)
    ax.tick_params(direction='in') #,grid_alpha=0.5)
    ax.set_xlim(xlim); ax.set_ylim(ylim)
    fig.colorbar(p); plt.savefig(figname)
    plt.close(fig)
    return None

def polar_to_xy(Xn,Yn):
    r = np.sqrt(Xn**2+Yn**2)
    theta = np.arctan2(Yn,Xn)
    return r, theta

def T(r,theta):
    x = r*np.cos(theta)
    y = r*np.sin(theta)
    return x, y

x = np.linspace(-8, 8, 800, endpoint=False)
y = np.linspace(-8, 8, 800, endpoint=False)
X,Y = ndgrid(x,y)
Phi = np.arctan2(X,Y)
Rho2 = X**2+Y**2
Rho = np.sqrt(Rho2)

Z3 = ne.evaluate('exp(1j*3*Phi)*exp(-Rho2/9)*(Rho*1.414/3)**3')
ZSum = Z3.real # PIC field data is only real

# (a)
E1 = pp.Field(ZSum,axes=[pp.Axis(grid=x),pp.Axis(grid=y)])
fig_imshow(E1.matrix,E1.extent,'RdYlBu','E1_xoy.png',[-8,8],[-8,8])

# (b)
E2 = E1.topolar()
fig_imshow(E2.matrix,E2.extent,'RdYlBu','E2_polar.png',[-np.pi,np.pi],[0,10])

# (c)
E21 = E1.map_coordinates(E2.axes, T)
fig_imshow(E21.matrix,E21.extent,'RdYlBu','E3_polar1.png',[-3*np.pi,3*np.pi],[0,10])

# (d)
E3 = E2.map_coordinates(E1.axes,polar_to_xy)
fig_imshow(E3.matrix,E3.extent,'RdYlBu','E4_back_xy.png',[-8,8],[-8,8])

# (e)
E3 = E21.map_coordinates(E1.axes,polar_to_xy)
fig_imshow(E3.matrix,E3.extent,'RdYlBu','E4_back1_xy.png',[-8,8],[-8,8])