#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Apr  6 09:55:56 2023

@author: shenzheqi
"""
case_dir = "./"
from netCDF4 import Dataset
TMPdata = Dataset(case_dir+'sensitive2008-2012.nc')
z_t = TMPdata['depth'][:]
temp_std = TMPdata['temp_std'][:]
lon = TMPdata['lon'][:]
lat = TMPdata['lat'][:]
SLTdata = Dataset(case_dir+'sensitive2008-2012salt.nc')
salt_std = SLTdata['salt_std'][:]

#%%
import numpy as np
Mstds_temp = np.zeros([6,180,360])
Mstds_salt = np.zeros([6,180,360])
zlevs = np.array([0,5,9,16,20,22])
for j in range(6):
    Mstds_temp[j] = np.nanmean(temp_std[36::,zlevs[j]],axis=0)
    Mstds_salt[j] = np.nanmean(salt_std[36::,zlevs[j]],axis=0)
    
#%%
LONg,LATg=np.meshgrid(lon,lat)
import matplotlib.pyplot as plt
import numpy as np
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from cartopy.mpl.ticker import LongitudeFormatter, LatitudeFormatter
import pandas as pd
from matplotlib.colors import ListedColormap 
rgb1 = pd.read_csv('WhiteBlueGreenYellowRedgood.rgb',sep='\s+',skiprows=2,names=['r','g','b']).values/255
#将rgb信息映射为colormap
colormap1 = ListedColormap(rgb1)
rgb2 = pd.read_csv('cmocean_deep.rgb',sep='\s+',skiprows=2,names=['r','g','b']).values/255
#将rgb信息映射为colormap
colormap2 = ListedColormap(rgb2)

scale = '110m'
land = cfeature.NaturalEarthFeature('physical', 'land', scale, edgecolor='face',
facecolor=cfeature.COLORS['land'])
fig=plt.figure(figsize=(12,11))
for j in range(5):
    ax1 = plt.subplot(5,2,2*j+1,projection=ccrs.PlateCarree(central_longitude=180))
    ax1.set_extent([0,420,-70,70], crs=ccrs.PlateCarree())
    ax1.add_feature(land, facecolor='0.75')
    ax1.coastlines(scale)
    ax1.set_xticks([0, 60, 120, 180, 240, 300, 360],[], crs=ccrs.PlateCarree())
    ax1.set_yticks([-60, -30, 0, 30, 60],[], crs=ccrs.PlateCarree())
    lon_formatter = LongitudeFormatter(zero_direction_label=False)
    lat_formatter = LatitudeFormatter()
    ax1.xaxis.set_major_formatter(lon_formatter)
    ax1.yaxis.set_major_formatter(lat_formatter)
    ax1.coastlines(resolution='110m')
    cor = plt.contourf(LONg,LATg,Mstds_temp[j],levels=np.arange(0,2.1,0.2), transform=ccrs.PlateCarree(),cmap =colormap1,extend='both')
    plt.grid(linestyle='dashed')
    if j==4:
        cbar = plt.colorbar(cor,cax = fig.add_axes([0.485, 0.25, 0.012, 0.5]),orientation='vertical',pad=0.05,label='degC')
    
    ax2 = plt.subplot(5,2,2*j+2,projection=ccrs.PlateCarree(central_longitude=180))
    ax2.set_extent([0,420,-70,70], crs=ccrs.PlateCarree())
    ax2.add_feature(land, facecolor='0.75')
    ax2.coastlines(scale)
    ax2.set_xticks([0, 60, 120, 180, 240, 300, 360],[], crs=ccrs.PlateCarree())
    ax2.set_yticks([-60, -30, 0, 30, 60],[], crs=ccrs.PlateCarree())
    ax2.xaxis.set_major_formatter(lon_formatter)
    ax2.yaxis.set_major_formatter(lat_formatter)
    ax2.coastlines(resolution='110m')
    cor = plt.contourf(LONg,LATg,Mstds_salt[j],levels=np.arange(0,0.55,0.05), transform=ccrs.PlateCarree(),cmap =colormap2,extend='both')
    plt.grid(linestyle='dashed')
    if j==4:
        cbar = plt.colorbar(cor,cax = fig.add_axes([0.915, 0.25, 0.012, 0.5]),orientation='vertical',pad=0.05,label='psu')
# plt.savefig("test.eps")
#%%   

import numpy as np
TTg, DDg = np.meshgrid(range(1,61),z_t)
STDeq = np.zeros([60,31])
STDglobal = np.zeros([60,31])
STDeqS = np.zeros([60,31])
STDglobalS = np.zeros([60,31])
for i in range(60):
    for j in range(31):
        STDeq[i,j] = np.nanmean(temp_std[i,j,85:95,:])
        STDglobal[i,j] = np.nanmean(temp_std[i,j,30:150,:])
        STDeqS[i,j] = np.nanmean(salt_std[i,j,85:95,:])
        STDglobalS[i,j] = np.nanmean(salt_std[i,j,30:150,:])

import matplotlib.pyplot as plt
plt.figure(figsize=(12,5))

plt.subplot(2,2,1)
plt.contourf(TTg,DDg,STDglobal.T,levels=np.arange(0,0.7,0.05),extend='both',cmap=colormap1)
plt.yscale('log')
plt.yticks([10,100,1000],[10,100,1000],fontsize=14)
plt.xticks(np.arange(12,62,12),[],fontsize=14)
plt.title('Temperature spread',fontsize=14);plt.ylabel('Depth',fontsize=14)
plt.grid(linestyle='dashed')
ax = plt.gca()
ax.invert_yaxis()
plt.colorbar(label='degC')
# plt.text(50,4,'Global Mean',fontsize=12)
plt.subplot(2,2,3)
plt.contourf(TTg,DDg,STDeq.T,levels=np.arange(0,1.4,0.1),extend='both',cmap=colormap1)
plt.yscale('log')
plt.yticks([10,100,1000],[10,100,1000],fontsize=14)
plt.xticks(np.arange(12,62,12),fontsize=14);plt.ylabel('Depth',fontsize=14)
plt.grid(linestyle='dashed')
ax = plt.gca()
ax.invert_yaxis()
plt.colorbar(label='degC')
plt.xlabel('Months',fontsize=14)
# plt.text(50,4,'Equator Mean',fontsize=12)
plt.subplot(2,2,2)
plt.contourf(TTg,DDg,STDglobalS.T,levels=np.arange(0,0.14,0.01),extend='both',cmap=colormap2)
plt.yscale('log')
plt.yticks([10,100,1000],[],fontsize=14)
plt.title('Salinity spread',fontsize=14);
ax = plt.gca()
ax.invert_yaxis()
plt.colorbar(label='psu')
plt.xticks(np.arange(12,62,12),[])
plt.grid(linestyle='dashed')
plt.subplot(2,2,4)
plt.contourf(TTg,DDg,STDeqS.T,levels=np.arange(0,0.28,0.02),extend='both',cmap=colormap2)
plt.yscale('log')
plt.yticks([10,100,1000],[],fontsize=14)
plt.xticks(np.arange(12,62,12),fontsize=14)
plt.grid(linestyle='dashed')
ax = plt.gca()
ax.invert_yaxis()
plt.colorbar(label='psu')
plt.xlabel('Months',fontsize=14);
# plt.savefig("fig1.eps")