from __future__ import division
from builtins import zip
from builtins import range
from past.utils import old_div
import numpy as np
from boututils.file_import import file_import
from boututils.surface_average import surface_average
from boutdata.collect import collect
from pylab import plot, show, xlabel, ylabel, figure, legend, gca


path='./data'

gfile='../cbm18_dens8.grid_nx68ny64.nc'

g = file_import(gfile)

var=collect("P", path=path)

sol=surface_average(var, g)
#sol=np.mean(var,axis=3)

p0av=collect("P0", path=path)

q=np.zeros(sol.shape)

for i in range(sol.shape[1]):
    q[:,i]=sol[:,i]+p0av[:,0]


psixy=g.get('psixy')
psi0=g.get('psi_axis')
psix=g.get('psi_bndry')

xarr = psixy[:,0]
xarr = old_div((xarr - psi0), (-psi0 + psix)) #for this grid


fig=figure()

nt=q.shape[1]

plot(xarr, p0av,'k',label='t=0')
plot(xarr,q[:,nt/4],'r',label='t='+np.str(nt/4))
plot(xarr,q[:,nt/2],'b',label='t='+np.str(nt/2))
plot(xarr,q[:,3*nt/4],'g',label='t='+np.str(3*nt/4))
plot(xarr, q[:,-1],'k',label='t='+np.str(nt))

from collections import OrderedDict
handles, labels = gca().get_legend_handles_labels()
by_label = OrderedDict(list(zip(labels, handles)))
legend(list(by_label.values()), list(by_label.keys()))


xlabel(r"$\psi$",fontsize=25)
ylabel(r"$2 \mu_0 <P^2> / B^2$",fontsize=25)

fig.set_tight_layout(True)

show()
