import iris
import iris.quickplot as qplt
import matplotlib.pyplot as plt
import numpy as np
import scipy as sp
import cartopy.crs as ccrs
import iris.analysis.cartography
import iris.coord_categorisation as coord_cat
import matplotlib.cm as mpl_cm
from iris.coords import DimCoord
from iris.cube import Cube

my_cmap = mpl_cm.get_cmap('brewer_RdBu_11')
levs=-1.+0.2*np.arange(11)

mslp=iris.load_cube('fig2a.nc')
mslp=mslp.intersection(longitude=(0,360))
latitude = mslp.coord('latitude').points
longitude = mslp.coord('longitude').points
mslp_plot=mslp[np.where((latitude < 30) & (latitude > -30))[0],np.where((longitude < 140) & (longitude > 30))[0]]
mslp_plot.units='K'
ax = plt.axes(projection=ccrs.PlateCarree(central_longitude=85))
ax.coastlines()
qplt.contourf(mslp_plot,levels=levs,cmap=my_cmap, extend="both")

plt.plot([50-85,50-85],[-10,10],color='green',linewidth=3)
plt.plot([70-85,70-85],[-10,10],color='green',linewidth=3)
plt.plot([50-85,70-85],[-10,-10],color='green',linewidth=3)
plt.plot([50-85,70-85],[10,10],color='green',linewidth=3)

plt.plot([90-85,90-85],[-10,0],color='green',linewidth=3)
plt.plot([110-85,110-85],[-10,0],color='green',linewidth=3)
plt.plot([90-85,110-85],[-10,-10],color='green',linewidth=3)
plt.plot([90-85,110-85],[0,0],color='green',linewidth=3)

title = 'Observed SST'
ax.set_title(title)
fig1=plt.gcf()
fig1.savefig('fig2a.pdf')
plt.close()

mslp=iris.load_cube('fig2b.nc')
mslp=mslp.intersection(longitude=(0,360))
latitude = mslp.coord('latitude').points
longitude = mslp.coord('longitude').points
mslp_plot=mslp[np.where((latitude < 30) & (latitude > -30))[0],np.where((longitude < 140) & (longitude > 30))[0]]
mslp_plot.units='2K'
ax = plt.axes(projection=ccrs.PlateCarree(central_longitude=85))
ax.coastlines()
qplt.contourf(mslp_plot,levels=levs,cmap=my_cmap, extend="both")
#Ramp 25S-20S, 20N-25N, 40E-45E, 110E-115E
plt.plot([41-85,41-85],[-24,24],color='green',linewidth=3)
plt.plot([114-85,114-85],[-24,24],color='green',linewidth=3)
plt.plot([41-85,114-85],[-24,-24],color='green',linewidth=3)
plt.plot([41-85,114-85],[24,24],color='green',linewidth=3)
title = 'SST perturbation, Orig - IOD_Neg, DP3'
ax.set_title(title)
fig1=plt.gcf()
fig1.savefig('fig2b.pdf')
plt.close()

dmi_ndj = np.load('fig2c_1.npy')
dmi_ndj_noenso = np.load('fig2c_2.npy')
times=1979+np.arange(41)
zeros = np.ones(41) - 1.
plt.plot(times, zeros, linewidth=1, color = 'black')
plt.plot(times, dmi_ndj, color ='red', linewidth=2, label = "Orig")
plt.plot(times, dmi_ndj_noenso, color ='blue', linewidth=2, label = "No ENSO")
plt.legend(loc='upper right')
ax = plt.gca()
ax.set_xlim([1979,2019])
#ax.set_ylim([-2,2])
ax.set_xlabel('Year')
ax.set_ylabel('Dipole Mode Index (K), Nov')
fig1=plt.gcf()
fig1.savefig('fig2c.pdf')
plt.close()

dmi_2019 = np.load('fig2d_1.npy')
dmi_ts = np.load('fig2d_2.npy')
dmi_ts_elop = np.load('fig2d_3.npy')
zeros = np.ones(6) - 1.
plt.plot(np.arange(6), zeros, linewidth=1, color = 'black')
plt.plot(np.arange(6), dmi_2019, linewidth=2, color = 'red', label = "2019/20")
plt.plot(np.arange(6), dmi_ts, linewidth=2, color = 'blue', label = "DMI > 1 sigma (1979-2019)")
plt.legend(loc='upper right')
ax = plt.gca()
ax.fill_between(np.arange(6), dmi_ts - dmi_ts_elop, dmi_ts + dmi_ts_elop, facecolor='lightblue', alpha=0.5)
ax.set_xlabel('Month')
ax.set_xlim([0,5])
ax.set_ylabel('DMI')
ax.set_xticks([0,1,2,3,4,5])
ax.set_xticklabels(['Sep','Oct','Nov','Dec','Jan','Feb'])
fig1=plt.gcf()
fig1.savefig('fig2d.pdf')
plt.close()

my_cmap = mpl_cm.get_cmap('RdBu_r')
levs=-6.+1.5*np.arange(9)

mslp_plot = iris.load_cube('fig2e.nc')
mslp_plot = mslp_plot.intersection(longitude=(-180,180))
lats = mslp_plot.coord('latitude').points
lons = mslp_plot.coord('longitude').points
mslp_plot=mslp_plot[np.where((lats <= 90) & (lats >= 20))[0],np.where((lons <= 60) & (lons >= -60))[0]]
mslp_plot.units='hPa'
ax = plt.axes(projection=ccrs.PlateCarree())
ax.set_ylim([20,90])
ax.set_xlim([-60,60])
ax.coastlines()
qplt.contourf(mslp_plot,levels=levs,cmap=my_cmap, extend="both")
title = 'MSLP(DJF), High(DMI(N)) - Low(DMI(N)), 1979-2019'
ax.set_title(title)
fig1=plt.gcf()
fig1.savefig('fig2e.pdf')
plt.close()

my_cmap = mpl_cm.get_cmap('RdBu')
levs=-1.5+0.25*np.arange(13)

mslp_plot = iris.load_cube('fig2f.nc')
mslp_plot = mslp_plot.intersection(longitude=(-180,180))
lats = mslp_plot.coord('latitude').points
lons = mslp_plot.coord('longitude').points
mslp_plot=mslp_plot[np.where((lats <= 90) & (lats >= 20))[0],np.where((lons <= 60) & (lons >= -60))[0]]
mslp_plot.units='mm/day'
ax = plt.axes(projection=ccrs.PlateCarree())
ax.coastlines()
qplt.contourf(mslp_plot,levels=levs,cmap=my_cmap, extend="both")
title = 'Precip(DJF), High(DMI(N)) - Low(DMI(N)), 1979-2019'
ax.set_title(title)
fig1=plt.gcf()
fig1.savefig('fig2f.pdf')
plt.close()