from __future__ import division
from builtins import range
from past.utils import old_div
import numpy as np
import matplotlib.pyplot as plt
from boututils.file_import import file_import
from matplotlib.ticker import FixedFormatter, FormatStrFormatter, AutoLocator, AutoMinorLocator

g = file_import("./cbm18_dens8.grid_nx68ny64.nc")

majorLocator   = AutoLocator()
majorFormatter = FormatStrFormatter('%3.0e')
minorLocator   = AutoMinorLocator()
Fm = FixedFormatter(['0','$1 \\times 10^4$','$2 \\times  10^4$','$3 \\times 10^4$','$4 \\times 10^4$'])
Fm2 = FixedFormatter(['0','$2 \\times 10^5$','$4 \\times  10^5$','$6 \\times 10^5$'])

bxy=g.get('Bxy')
p=g.get('pressure')
jpar0=g.get('Jpar0')
psixy=g.get('psixy')
btxy=g.get('Btxy')
shiftangle=g.get('ShiftAngle')

nx=g.get('nx')
ny=g.get('ny')

q = np.zeros((nx, ny))
for i in range(ny):
    q[:,i] = old_div(- shiftangle, (2 * np.pi))




xarr = psixy[:,0]
xarr = old_div((xarr + 0.854856), (0.854856 + 0.0760856))


fig=plt.figure()
plt.plot(xarr,q[:,32])
plt.xlabel('normalized $\psi$', fontsize=25)
plt.ylabel('$q$',rotation='horizontal',fontsize=25)

fig.set_tight_layout(True)

fig, ax1 = plt.subplots()

ax1.plot(xarr, p[:,32], 'r-', markevery=1, linewidth=3)
ax1.set_xlabel('normalized $\psi$',fontsize=25)
# Make the y-axis label and tick labels match the line color.
ax1.set_ylabel('Pressure [Pa]', color='k',fontsize=25)

#set y limit
ax1.set_ylim(0,40000,10000)

#define ticks#
ax1.yaxis.set_ticks(np.arange(0, 40000, 10000))
#ax1.yaxis.set_major_locator(majorLocator)
#ax1.yaxis.set_major_formatter(majorFormatter)
ax1.yaxis.set_major_formatter(Fm)

#for the minor ticks, use no labels; default NullFormatter
ax1.xaxis.set_minor_locator(AutoMinorLocator())
ax1.yaxis.set_minor_locator(AutoMinorLocator(10))

#format tick labels
for tl in ax1.get_yticklabels():
    tl.set_color('k')


ax2 = ax1.twinx()

s2 = -jpar0

ax2.plot(xarr, s2[:,32], 'r-',markevery=1,linewidth=3)
ax2.set_ylabel('$J_\parallel  [A/m^2]$', color='k',fontsize=25)
ax2.set_ylim(0,600000)
ax2.yaxis.set_ticks(np.arange(0, 600000, 200000))
ax2.yaxis.set_major_formatter(Fm2)
for tl in ax2.get_yticklabels():
    tl.set_color('k')


fig.set_tight_layout(True)

plt.show()


#plt.savefig('2d.png', transparent=True)
