#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Aug 17 23:47:10 2021

@author: hdaher
"""

#plot ssh 
import pickle as pic
import matplotlib.pyplot as plt
import numpy as np
import netCDF4 as nc

def find_nearest(array, value):
    array = np.asarray(array)
    idx = (np.abs(array - value)).argmin()
    return idx

path = '/Users/hdaher/Agulhas/ghg_ozone/code/picklefiles/ts/'

lrc07ts_mean, lrc08ts_mean, lrc20ts_mean, hrc07ts_mean, hrc08ts_mean, hrc20ts_mean = pic.load(open(path+'ts_mean.p', 'rb'))
 
lrc_ts_co2 = abs(lrc08ts_mean-lrc07ts_mean)
hrc_ts_co2 = abs(hrc08ts_mean-hrc07ts_mean)
lrc_ts_o3 = abs(lrc20ts_mean-lrc08ts_mean)
hrc_ts_o3 = abs(hrc20ts_mean-hrc08ts_mean)


path = '/Users/hdaher/Agulhas/ghg_ozone/code/picklefiles/ssh/'

#lrc07
lrc07_lon, lrc07_sh_lat, lrc07_ssh_data_annual_mean, lrc07_ssh_data_djf_mean, lrc07_ssh_data_mam_mean, \
    lrc07_ssh_data_jja_mean, lrc07_ssh_data_son_mean = pic.load(open(path+'lrc07_ssh.p' , 'rb'))

#lrc08
lrc08_lon, lrc08_sh_lat, lrc08_ssh_data_annual_mean, lrc08_ssh_data_djf_mean, lrc08_ssh_data_mam_mean, \
    lrc08_ssh_data_jja_mean, lrc08_ssh_data_son_mean = pic.load(open(path+'lrc08_ssh.p' , 'rb'))

#lrc20
lrc20_lon, lrc20_sh_lat, lrc20_ssh_data_annual_mean, lrc20_ssh_data_djf_mean, lrc20_ssh_data_mam_mean, \
    lrc20_ssh_data_jja_mean, lrc20_ssh_data_son_mean = pic.load(open(path+'lrc20_ssh.p' , 'rb'))

#hrc07
hrc07_lon, hrc07_sh_lat, hrc07_ssh_data_annual_mean, hrc07_ssh_data_djf_mean, hrc07_ssh_data_mam_mean, \
    hrc07_ssh_data_jja_mean, hrc07_ssh_data_son_mean = pic.load(open(path+'hrc07_ssh.p' , 'rb'))

#hrc08
hrc08_lon, hrc08_sh_lat, hrc08_ssh_data_annual_mean, hrc08_ssh_data_djf_mean, hrc08_ssh_data_mam_mean, \
    hrc08_ssh_data_jja_mean, hrc08_ssh_data_son_mean = pic.load(open(path+'hrc08_ssh.p' , 'rb'))

#hrc20
hrc20_lon, hrc20_sh_lat, hrc20_ssh_data_annual_mean, hrc20_ssh_data_djf_mean, hrc20_ssh_data_mam_mean, \
    hrc20_ssh_data_jja_mean, hrc20_ssh_data_son_mean = pic.load(open(path+'hrc20_ssh.p' , 'rb'))


#take differences 
#co2 effect: 08-07
lrc_annual_co2 = lrc08_ssh_data_annual_mean - lrc07_ssh_data_annual_mean
lrc_djf_co2 = lrc08_ssh_data_djf_mean - lrc07_ssh_data_djf_mean
lrc_mam_co2 = lrc08_ssh_data_mam_mean - lrc07_ssh_data_mam_mean
lrc_jja_co2 = lrc08_ssh_data_jja_mean - lrc07_ssh_data_jja_mean
lrc_son_co2 = lrc08_ssh_data_son_mean - lrc07_ssh_data_son_mean

hrc_annual_co2 = hrc08_ssh_data_annual_mean - hrc07_ssh_data_annual_mean
hrc_djf_co2 = hrc08_ssh_data_djf_mean - hrc07_ssh_data_djf_mean
hrc_mam_co2 = hrc08_ssh_data_mam_mean - hrc07_ssh_data_mam_mean
hrc_jja_co2 = hrc08_ssh_data_jja_mean - hrc07_ssh_data_jja_mean
hrc_son_co2 = hrc08_ssh_data_son_mean - hrc07_ssh_data_son_mean


#o3 effect: 20-08
lrc_annual_o3 = lrc20_ssh_data_annual_mean - lrc08_ssh_data_annual_mean
lrc_djf_o3 = lrc20_ssh_data_djf_mean - lrc08_ssh_data_djf_mean
lrc_mam_o3 = lrc20_ssh_data_mam_mean - lrc08_ssh_data_mam_mean
lrc_jja_o3 = lrc20_ssh_data_jja_mean - lrc08_ssh_data_jja_mean
lrc_son_o3 = lrc20_ssh_data_son_mean - lrc08_ssh_data_son_mean

hrc_annual_o3 = hrc20_ssh_data_annual_mean - hrc08_ssh_data_annual_mean
hrc_djf_o3 = hrc20_ssh_data_djf_mean - hrc08_ssh_data_djf_mean
hrc_mam_o3 = hrc20_ssh_data_mam_mean - hrc08_ssh_data_mam_mean
hrc_jja_o3 = hrc20_ssh_data_jja_mean - hrc08_ssh_data_jja_mean
hrc_son_o3 = hrc20_ssh_data_son_mean - hrc08_ssh_data_son_mean

lrc_annual_co2 = lrc_annual_co2/lrc_ts_co2
hrc_annual_co2 = hrc_annual_co2/hrc_ts_co2
lrc_annual_o3  = lrc_annual_o3/lrc_ts_o3
hrc_annual_o3 = hrc_annual_o3/hrc_ts_o3 


lrc_annual = lrc_annual_co2+lrc_annual_o3 
hrc_annual = hrc_annual_co2+hrc_annual_o3

#topo
path = '/Users/hdaher/Agulhas/ozoneResponse/'
#plot 
#read in etopo1 file
ncin = nc.Dataset(path+'ETOPO1_Bed_c_gmt4.grd')
topolons = ncin.variables['x'][:]
topolats = ncin.variables['y'][:]
topodepths = ncin.variables['z'][:]
x = find_nearest(topolons, 0)
newtopolons = np.zeros(len(topolons))
newtopolons[:x] = topolons[x:]
newtopolons[x:] = topolons[:x]
for i in range(0, len(newtopolons)):
    if newtopolons[i] < 0:
        newtopolons[i] = newtopolons[i] + 360
newtopodepths = np.zeros((topodepths.shape[0], topodepths.shape[1]))
newtopodepths[:,:x] = topodepths[:, x:]
newtopodepths[:,x:] = topodepths[:, :x]

skip = 10
newtopolons = newtopolons[::skip]
topolats = topolats[::skip]
newtopodepths = newtopodepths[::skip, ::skip]
topolats = topolats[0:540]
newtopodepths = newtopodepths[0:540, :]


fig = plt.figure()
ax = fig.add_subplot(2,3,1)
cmax = 0.30
cmin = -cmax
step = cmax/2
plt.contourf(newtopolons, topolats, newtopodepths, [0, 5000], linewidths = 1.0, colors = ['grey'], zorder = 4)
plt.contourf(newtopolons, topolats, newtopodepths, [-2000, -1000, 0], linewidths = 1.0, colors = ['lightgrey', 'darkgrey'])
plt.pcolor(lrc07_lon, lrc07_sh_lat, lrc_annual_co2/100,  cmap='bwr')
plt.clim(cmin, cmax)
ax.set_yticks([-75,-60, -45, -30, -15])
ax.set_yticklabels(['75$^\circ$S', '60$^\circ$S', '45$^\circ$S', '30$^\circ$S','15$^\circ$S'])
ax.set_xticks([60, 180, 300])
ax.set_xticklabels(['60$^\circ$E', '180$^\circ$', '60$^\circ$W'])
ax.set_xticklabels('')
plt.ylabel('Low resolution')
plt.title('CO$_2$ (08-07)')
plt.suptitle('Annual SSH/TS difference')

ax = fig.add_subplot(2,3,2)
plt.contourf(newtopolons, topolats, newtopodepths, [0, 5000], linewidths = 1.0, colors = ['grey'], zorder = 4)
plt.contourf(newtopolons, topolats, newtopodepths, [-2000, -1000, 0], linewidths = 1.0, colors = ['lightgrey', 'darkgrey'])
plt.pcolor(lrc07_lon, lrc07_sh_lat, lrc_annual_o3/100,  cmap='bwr')
plt.clim(cmin, cmax)
ax.set_yticks([-75,-60, -45, -30, -15])
ax.set_yticklabels(['75$^\circ$S', '60$^\circ$S', '45$^\circ$S', '30$^\circ$S','15$^\circ$S'])
ax.set_xticks([60, 180, 300])
ax.set_xticklabels(['60$^\circ$E', '180$^\circ$', '60$^\circ$W'])
ax.set_xticklabels('')
ax.set_yticklabels('')
plt.title('O$_3$ (20-08)')

ax = fig.add_subplot(2,3,3)
plt.contourf(newtopolons, topolats, newtopodepths, [0, 15000], linewidths = 1.0, colors = ['grey'], zorder = 4)
plt.contourf(newtopolons, topolats, newtopodepths, [-2000, -1000, 0], linewidths = 1.0, colors = ['lightgrey', 'darkgrey'])
plt.pcolor(lrc07_lon, lrc07_sh_lat, lrc_annual/100,  cmap='bwr')
plt.clim(cmin, cmax)
ax.set_yticks([-75,-60, -45, -30, -15])
ax.set_yticklabels(['75$^\circ$S', '60$^\circ$S', '45$^\circ$S', '30$^\circ$S','15$^\circ$S'])
ax.set_xticks([60, 180, 300])
ax.set_xticklabels(['60$^\circ$E', '180$^\circ$', '60$^\circ$W'])
ax.set_xticklabels('')
ax.set_yticklabels('')
plt.title('Total (CO$_2$+O$_3$)')

ax = fig.add_subplot(2,3,4)
plt.contourf(newtopolons, topolats, newtopodepths, [0, 5000], linewidths = 1.0, colors = ['grey'], zorder = 4)
plt.contourf(newtopolons, topolats, newtopodepths, [-2000, -1000, 0], linewidths = 1.0, colors = ['lightgrey', 'darkgrey'])
plt.pcolor(hrc07_lon, hrc07_sh_lat, hrc_annual_co2/100, cmap='bwr')
plt.clim(cmin, cmax)
ax.set_yticks([-75,-60, -45, -30, -15])
ax.set_yticklabels(['75$^\circ$S', '60$^\circ$S', '45$^\circ$S', '30$^\circ$S','15$^\circ$S'])
ax.set_xticks([60, 180, 300])
ax.set_xticklabels(['60$^\circ$E', '180$^\circ$', '60$^\circ$W'])
plt.ylabel('High resolution')

ax = fig.add_subplot(2,3,5)
plt.contourf(newtopolons, topolats, newtopodepths, [0, 5000], linewidths = 1.0, colors = ['grey'], zorder = 4)
plt.contourf(newtopolons, topolats, newtopodepths, [-2000, -1000, 0], linewidths = 1.0, colors = ['lightgrey', 'darkgrey'])
plt.pcolor(hrc07_lon, hrc07_sh_lat, hrc_annual_o3/100, cmap='bwr')
plt.clim(cmin, cmax)
ax.set_yticks([-75,-60, -45, -30, -15])
ax.set_yticklabels(['75$^\circ$S', '60$^\circ$S', '45$^\circ$S', '30$^\circ$S','15$^\circ$S'])
ax.set_xticks([60, 180, 300])
ax.set_xticklabels(['60$^\circ$E', '180$^\circ$', '60$^\circ$W'])
ax.set_yticklabels('')

ax = fig.add_subplot(2,3,6)
plt.contourf(newtopolons, topolats, newtopodepths, [0, 15000], linewidths = 1.0, colors = ['grey'], zorder = 4)
plt.contourf(newtopolons, topolats, newtopodepths, [-2000, -1000, 0], linewidths = 1.0, colors = ['lightgrey', 'darkgrey'])
test = plt.pcolor(hrc07_lon, hrc07_sh_lat, hrc_annual/100,  cmap='bwr')
plt.clim(cmin, cmax)
ax.set_yticks([-75,-60, -45, -30, -15])
ax.set_yticklabels(['75$^\circ$S', '60$^\circ$S', '45$^\circ$S', '30$^\circ$S','15$^\circ$S'])
ax.set_xticks([60, 180, 300])
ax.set_xticklabels(['60$^\circ$E', '180$^\circ$', '60$^\circ$W'])
ax.set_yticklabels('')

cb_ax = fig.add_axes([0.91, 0.1, 0.02, 0.8])
v = np.arange(cmin, cmax+0.1, step)
cbar = fig.colorbar(test, ticks=v, cax=cb_ax)
cbar.ax.set_title('m/$^\circ$C')