import iris
import iris.coord_categorisation
import matplotlib.pyplot as plt
import numpy as np
import iris.quickplot as qplt
import netCDF4
import datetime
import scipy
import scipy.signal
import glob
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from scipy.stats import t
import pickle
import os
import glob
import iris.util
from matplotlib.colors import BoundaryNorm
from matplotlib.ticker import MaxNLocator
from mpl_toolkits.axes_grid1 import make_axes_locatable
import iris.plot as iplt
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from matplotlib.transforms import offset_copy
from cartopy.mpl.gridliner import LONGITUDE_FORMATTER, LATITUDE_FORMATTER


file_4_mask = '/data/NAS-geo01/jm953/2021_scripts_outputs/output/bt_st_maxyrly/IPSL-CM6A-LR_119_GBRMPA_austsummerMEAN_st_bt_DIFF.nc'
bt_anomaly_file4timeseries = '/data/NAS-geo01/jm953/2021_scripts_outputs/output/bt_st_maxyrly/bottomtemperature_ERA5_sum_months_anomaly_1980_1999_clim_v19Jan2024.nc'
depth_mask = iris.load_cube(file_4_mask)
depth_mask.data = np.ma.masked_where(depth_mask.data == 9.96920997e+36, depth_mask.data)
depth_mask.data = np.ma.masked_where(depth_mask.data == 1.0000000e+20, depth_mask.data)

bt_anomaly_cube = iris.load_cube(bt_anomaly_file4timeseries)
bt_anomaly_cube.data = np.ma.masked_where(bt_anomaly_cube.data == 9.96920997e+36, bt_anomaly_cube.data)

#keep only negative bottom temp anomalies 
bt_anomaly_cube.data[np.where(bt_anomaly_cube.data > 0.0)] = 0.0
bt_anomaly_cube.data[np.where(bt_anomaly_cube.data != 0.0)] = 1.0

####mask the surface temperature data 
st_anomaly_file4timeseries = '/data/NAS-geo01/jm953/2021_scripts_outputs/output/bt_st_maxyrly/surfacetemperature_ERA5_sum_months_anomaly_1980_1999_clim_v19Jan2024.nc'
st_anomaly_cube = iris.load_cube(st_anomaly_file4timeseries)
st_anomaly_cube.data = np.ma.masked_where(st_anomaly_cube.data == 9.96920997e+36, st_anomaly_cube.data)

#keep only positive st anomalies, this would show where the bt is negative and st is positive
st_anomaly_cube.data[np.where(st_anomaly_cube.data < 0.0)] = 0.0
st_anomaly_cube.data[np.where(st_anomaly_cube.data != 0.0)] = 1.0

bt_st_sum = bt_anomaly_cube + st_anomaly_cube

bt_st_sum.data[np.where(bt_st_sum.data != 2.0)] = 0.0

####change back to 1 so it isn't double counted!!!
bt_st_sum.data[np.where(bt_st_sum.data == 2.0)] = 1.0

bt_st_sum.data.mask = depth_mask.data.mask
cube2 = bt_st_sum.copy()

anomaly_cube = bt_st_sum.collapsed(['time'], iris.analysis.SUM) 

print str('ERA5') + str(np.unique(anomaly_cube.data)) + str('count') + str(np.count_nonzero(anomaly_cube.data))

anomaly_cube.data = (anomaly_cube.data/20)*100
print str('ERA5') + str(np.unique(anomaly_cube.data)) + str('count') + str(np.count_nonzero(anomaly_cube.data))

cube_ERA5 = anomaly_cube.copy()


####CMIP6 
mod_cube = '/data/NAS-geo01/jm953/2021_scripts_outputs/output/bt_st_maxyrly/cooling_locations_models_2000_2019_v19Jan2024.nc'
all_cubes = iris.load_cube(mod_cube)
#in one cell there are 20 years times 4 scenarios times 5 models resulting in 400 possible combinations - make this more exact! 
all_cubes.data = (all_cubes.data/400)*100

cube_CMIP6 = all_cubes.copy()

combo = cube_ERA5 + cube_CMIP6

print str(np.unique(combo.data))


#############ERA5 agreement
fig = plt.figure(figsize=(14, 9))
ax1 = fig.add_subplot(131, projection=ccrs.PlateCarree())
ax1.set_extent([142, 153.5, -25, -10.5], crs=ccrs.PlateCarree())
mesh1 = iplt.pcolormesh(cube_ERA5, cmap='viridis_r',alpha=0.5)
levels = MaxNLocator(nbins=3).tick_values(0, 20)
#Add land and costlines
ax1.coastlines('10m', lw=1, alpha=0.5)
ax1.add_feature(cfeature.LAND, facecolor="gray")

#Aff the red circles and text labels
PlateCarree_transform = ccrs.PlateCarree()._as_mpl_transform(ax1)
text_transform = offset_copy(PlateCarree_transform, units='dots', x=-10)

ax1.plot(145.754120,-16.92549, marker='o', color='red', markersize=6,
            alpha=0.7, transform=ccrs.PlateCarree())
ax1.text(145.754120,-16.92549,'Cairns',
            verticalalignment='center', horizontalalignment='right',
            transform=text_transform)

ax1.plot(146.816956,-19.258965, marker='o', color='red', markersize=6,
            alpha=0.7, transform=ccrs.PlateCarree())
ax1.text(146.816956,-19.258965,'Townsville',
            verticalalignment='center', horizontalalignment='right',
            transform=text_transform)

ax1.plot(149.186813, -21.144337, marker='o', color='red', markersize=6,
            alpha=0.7, transform=ccrs.PlateCarree())
ax1.text(149.186813, -21.144337,'Mackay',
            verticalalignment='center', horizontalalignment='right',
            transform=text_transform)

ax1.plot(151.268356,	-23.843138, marker='o', color='red', markersize=6,
            alpha=0.7, transform=ccrs.PlateCarree())
ax1.text(151.268356,	-23.843138,'Gladstone',
            verticalalignment='center', horizontalalignment='right',
            transform=text_transform)

#Add drig lines and labels
gl = ax1.gridlines(draw_labels=True, color='gray', alpha=0.1)
gl.ylabels_right = False
gl.xlabels_bottom = False

##########CMIP6 strat
ax2 = fig.add_subplot(132, projection=ccrs.PlateCarree())
ax2.set_extent([142, 153.5, -25, -10.5], crs=ccrs.PlateCarree())
mesh2 = iplt.pcolormesh(cube_CMIP6, cmap='viridis_r', alpha=0.5)
levels = MaxNLocator(nbins=3).tick_values(0, 20)
#Add land and costlines
ax2.coastlines('10m', lw=1, alpha=0.5)
ax2.add_feature(cfeature.LAND, facecolor="gray")

#Aff the red circles and text labels
PlateCarree_transform = ccrs.PlateCarree()._as_mpl_transform(ax2)
text_transform = offset_copy(PlateCarree_transform, units='dots', x=-10)

ax2.plot(145.754120,-16.92549, marker='o', color='red', markersize=6,
            alpha=0.7, transform=ccrs.PlateCarree())
ax2.text(145.754120,-16.92549,'Cairns',
            verticalalignment='center', horizontalalignment='right',
            transform=text_transform)

ax2.plot(146.816956,-19.258965, marker='o', color='red', markersize=6,
            alpha=0.7, transform=ccrs.PlateCarree())
ax2.text(146.816956,-19.258965,'Townsville',
            verticalalignment='center', horizontalalignment='right',
            transform=text_transform)

ax2.plot(149.186813, -21.144337, marker='o', color='red', markersize=6,
            alpha=0.7, transform=ccrs.PlateCarree())
ax2.text(149.186813, -21.144337,'Mackay',
            verticalalignment='center', horizontalalignment='right',
            transform=text_transform)

ax2.plot(151.268356,	-23.843138, marker='o', color='red', markersize=6,
            alpha=0.7, transform=ccrs.PlateCarree())
ax2.text(151.268356,	-23.843138,'Gladstone',
            verticalalignment='center', horizontalalignment='right',
            transform=text_transform)

#Add drig lines and labels
gm = ax2.gridlines(draw_labels=True, color='gray', alpha=0.1)
gm.ylabels_right = False
gm.ylabels_left = False
gm.xlabels_bottom = False

#############fancy plot for agreement between era5 and cmip6 
ax3 = fig.add_subplot(133, projection=ccrs.PlateCarree())
ax3.set_extent([142, 153.5, -25, -10.5], crs=ccrs.PlateCarree())
mesh3 = iplt.pcolormesh(combo, cmap='viridis_r',alpha=0.5)
levels = MaxNLocator(nbins=3).tick_values(0, 20)
#Add land and costlines
ax3.coastlines('10m', lw=1, alpha=0.5)
ax3.add_feature(cfeature.LAND, facecolor="gray")

#Aff the red circles and text labels
PlateCarree_transform = ccrs.PlateCarree()._as_mpl_transform(ax3)
text_transform = offset_copy(PlateCarree_transform, units='dots', x=-10)

ax3.plot(145.754120,-16.92549, marker='o', color='red', markersize=6,
            alpha=0.7, transform=ccrs.PlateCarree())
ax3.text(145.754120,-16.92549,'Cairns',
            verticalalignment='center', horizontalalignment='right',
            transform=text_transform)

ax3.plot(146.816956,-19.258965, marker='o', color='red', markersize=6,
            alpha=0.7, transform=ccrs.PlateCarree())
ax3.text(146.816956,-19.258965,'Townsville',
            verticalalignment='center', horizontalalignment='right',
            transform=text_transform)

ax3.plot(149.186813, -21.144337, marker='o', color='red', markersize=6,
            alpha=0.7, transform=ccrs.PlateCarree())
ax3.text(149.186813, -21.144337,'Mackay',
            verticalalignment='center', horizontalalignment='right',
            transform=text_transform)

ax3.plot(151.268356,	-23.843138, marker='o', color='red', markersize=6,
            alpha=0.7, transform=ccrs.PlateCarree())
ax3.text(151.268356,	-23.843138,'Gladstone',
            verticalalignment='center', horizontalalignment='right',
            transform=text_transform)

#Add drig lines and labels
gn = ax3.gridlines(draw_labels=True, color='gray', alpha=0.1)
gn.ylabels_right = False
gn.ylabels_left = False
gn.xlabels_bottom = False

fig.subplots_adjust(bottom=0.10)
cbar_ax = fig.add_axes([0.17, 0.04, 0.65, 0.04])
fig.colorbar(mesh3, cax=cbar_ax, orientation="horizontal")
#bar.set_ticks(levels)
#bar.set_ticklabels(levels)
#bar.ax.set_xticklabels(levels, rotation=45, fontsize=8)
#bar.set_label('Model agreement, ERA5 and CMIP6', fontsize=8)
#plt.savefig('/data/NAS-geo01/jm953/2021_scripts_outputs/plots/refugia_locations/Climate_Model_refugia_20p_v1.pdf', dpi=300)
plt.savefig('/data/NAS-geo01/jm953/bottom_temp_PNAS/2024_plots/ERA5_CMIP6_stbt_v19Jan2024_ALL3_NOSD.png', dpi=600)
plt.show()