import numpy as N
import netCDF4 as S
import matplotlib.pyplot as plt
import mpl_toolkits.basemap as bm
import math as M

direcInput = "/Users/Shared/DYME/AnnularMode_mslp/" # set working directory
direcOutput = "/Users/Shared/AnalysisUnderworld/" # set directory containing plots of model results

omega = 0.00007292
dy = 1.0
g = 9.81
radius_earth = 6370000.0
pi = M.pi

# resolution is 1.0 deg and 1 month
nx_len = 360  # -180 - +180   ; we use "nx" as index of longitude
ny_len = 181  # SP - NP    ; we use "ny" as index of latitude
nt_len = 432  # 36 (ERA-I) years x 12 months;  we use "nt" as index of time
yr_len = 36
mt_len = 12
year0 = 1979
title = 'mass below 285 K - isentrope'


latitude = N.zeros((ny_len), dtype='d')
for ny in range(ny_len):
 latitude[ny] = (90.0 - (ny * dy)) #latitude in degrees

#open file
a = S.Dataset(direcInput+"ERA-I_mslp_1979-2014_mmofdm_1deg_Global.nc", mode='r') #ERA-Int data of v along a zonal strip between fixed latitudes
b = S.Dataset(direcInput+"ERA-I_ps_1979-2014_mmofdm_1deg_Global.nc", mode='r')
c = S.Dataset(direcInput+"ERA-I_p285K_1979-2014_mmofdm_1deg_Global.nc", mode='r')

#create arrays time, lat and lon for plotting
lat = a.variables["latitude"][:] 
lon = a.variables["longitude"][:] 
time = a.variables["time"][:]
time = mt_len * (time - N.min(time)) / ( N.max(time) - N.min(time)) + 1

# meridional wind (v) and temperature (T)
#indices: time, lat, lon  (time, y, x) 
#axis: 0, 1, 2
mslp = a.variables["msl"][:,:,:]
psurf = b.variables["sp"][:,:,:]
p285 = c.variables["pres"][:,:,:]
pdiff = psurf-p285

#Find latitude of maximum extent of the Underworld
latN0 = N.zeros((nt_len,nx_len), dtype='d')
for nt in range(nt_len):
 for nx in range(nx_len):
  latN0[nt,nx] = 90.0
  flag = 0
  for ny in range(ny_len):
   if nx == 90 and nt == 0: print p285[nt,ny,nx],psurf[nt,ny,nx],pdiff[nt,ny,nx],flag, latN0[nt,nx]
   if pdiff[nt,ny,nx]<200 and flag == 0: latN0[nt,nx] = latitude[ny]
   if pdiff[nt,ny,nx]<200 and flag == 0: flag = 1

latS0 = N.zeros((nt_len,nx_len), dtype='d')
for nt in range(nt_len):
 for nx in range(nx_len):
  latS0[nt,nx] = -90.0
  flag = 0
  for ny in range(ny_len-1,0,-1):
   if pdiff[nt,ny,nx]<200 and flag == 0: latS0[nt,nx] = latitude[ny]
   if pdiff[nt,ny,nx]<200 and flag == 0: flag = 1
   if pdiff[nt,ny,nx]<200: pdiff[nt,ny,nx] = 0.0

year = N.zeros((nt_len), dtype='d')
for nt in range(nt_len):
 year[nt] = 1979.0 + ((nt+0.5)/12.0)

plt.figure(figsize=(12,10))
plt.plot(lon, latN0[0,:],label='Northern hemisphere', color='red')
plt.plot(lon, latS0[0,:],label='Southern hemisphere', color='blue')
plt.axis([0,360,-60,60]) 
plt.legend()   
plt.xlabel('longitude]')
plt.ylabel('Latitude, maximum extent Underworld in January 1979')
plt.grid(True)
plt.savefig(direcOutput+"LatitudeUnderworld_Jan1979"+".png") # save plot as png-file 
plt.show()

#Take zonal mean
latN0_zm = N.mean(latN0,axis=1)
latS0_zm = N.mean(latS0,axis=1)

plt.figure(figsize=(12,10))
plt.plot(year, latN0_zm,label='Northern hemisphere', color='red')
plt.plot(year, -latS0_zm,label='Southern hemisphere', color='blue')
plt.axis([1979,year[nt_len-1],25,65])    
plt.xlabel('time[year]')
plt.ylabel('Zonal mean latitude, maximum extent Underworld')
plt.legend()
plt.title('Zonal mean latitude, maximum extent Underworld='+title,fontsize=16) # Title at top of plot

plt.grid(True)
plt.savefig(direcOutput+"ZonalMeanLatitudeUnderworld"+".png") # save plot as png-file 
plt.show()