#!/usr/bin/env python
# coding: utf-8

import os
import sys
import glob
import numpy as np
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import cartopy.feature as cfeature
import cartopy as cart
import matplotlib.colors as mcolors
from netCDF4 import Dataset,num2date
print('All modules imported..')



#just gets the lat and lon from a file       
def get_lonlat_2d(filepath):
    #should be a                                                                                                                                  
    try:
        nc = Dataset(filepath)
    except:
        print('no file '+filepath)
    return nc.variables['lon'],nc.variables['lat']


#Retrieve variable from nc file
def get_var(filepath,varname):
    try:
        nc = Dataset(filepath)
    except:
        print('no file '+filepath)
    return nc.variables[varname]


#----------------Program start-------------------
#
#Plots multiple panels of maps
#
#Based on time merged monthly mean data files for variables TREFHT and PBLH:
#cdo -selname,TREFHT,PBLH cat BASE_fsst/*h0* BASE_fsst_merged.h0.monthlymean.nc
#For variables Ri and N, values are first monthly averaged then also averaged in time:
#cdo monmean RiN_fsst/BASE_fsst/* RiN_BASE_fsst_monthlymean.nc
#cdo timavg RiN_BASE_fsst_monthlymea.nc RiN_BASE_fsst_monthlymea_avg.nc
#


#Leave at 0 for absolute diff as in Fig.1
reldiff=0

dooceanmask=0

dir_data = ''#fill in path of post processed data


variables=["TREFHT","N","Ri","PBLH"]
units=['[K]','[1/s]','[ ]','[m]']
nvars=len(variables)


#exptype can be set to "fsst" or "cpl", FIG1 is cpl
exptype="cpl"
if exptype=="cpl": cases=["BASE_cpl","BCx10_cpl","SO4x5_cpl","CO2x2_cpl"]
if exptype=="fsst": cases=["BASE_fsst","BCx10_fsst","SO4x5_fsst","CO2x2_fsst"]
casenames=["BCx10","SO4x5","CO2x2"]
allcases=cases
ncases=len(cases)

figname='FIG_diffMaps_'+exptype+'.pdf'
if reldiff==1: figname='FIG_diffMaps_'+exptype+'_rel.png'


if exptype=="fsst":
    LFfile='landfrac_cam6.nc'
if exptype=="cpl":
    LFfile='landfrac_cam6_cpl.nc'
lon_values,lat_values = get_lonlat_2d(LFfile)
nlon=len(lon_values)
nlat=len(lat_values)
landfrac = np.squeeze(get_var(LFfile,'landfrac'))
lon_valuesLF = np.squeeze(get_var(LFfile,'lon'))
lat_valuesLF = np.squeeze(get_var(LFfile,'lat'))
oceanfrac = 1.-landfrac

# First we need to convert the latitudes to radians
latr = np.deg2rad(lat_values)
# Use the cosine of the converted latitudes as weights for the average
weights = np.cos(latr)

#Declare variable to store all data in:                   
allvalues=np.zeros((ncases,nvars,nlat,nlon))
#------------Read in the data--------------
for c in range (0,ncases): #Loop through the cases
    print('Going through '+cases[c]+'...')
    #
    for filename in glob.glob(dir_data+cases[c]+'*merge.h0.*monthlymean.nc'):
        print('\t Using file: '+filename)
        for v in range (0,nvars):
            print('      Going through '+variables[v]+'...')
            if variables[v]=="PRECT":
                precc = get_var(filename,"PRECC")
                precl = get_var(filename,"PRECL")
#                invar=np.add(precc,precl)*86400000.
                invar=np.multiply(precc,86400000.)
            if ((variables[v]=="Ri") or (variables[v]=="N")):
                filenameRN=glob.glob(dir_data+"RiN_"+cases[c]+'*monthlymean_avg.nc')
                filenameRN=filenameRN[0]
                #CHOOSE: pick out certain level or vert. average?
                rn = get_var(filenameRN,variables[v])
                invar=np.squeeze(rn[:,1,:,:]) #936 hPa
            if variables[v]=="LTS":
                t = get_var(filename,"T")
                lev_values = get_var(filename,'lev')
                index850=min(range(len(lev_values)), key=lambda i: abs(lev_values[i]-850))
                index1000=min(range(len(lev_values)), key=lambda i: abs(lev_values[i]-1000))
                s=np.shape(invar)
                ntime=s[0]
                invar=np.zeros((ntime,nlat,nlon))
                invar=np.squeeze(t[:,index1000,:,:])-np.squeeze(t[:,index850,:,:])
            if (variables[v]!="LTS") and (variables[v]!="PRECT") and (variables[v]!="Ri") and (variables[v]!="N"):
                invar = get_var(filename,variables[v])
                s=np.shape(invar)
                if len(s)==4:
                    invar=np.squeeze(invar[:,31,:,:])
            allvalues[c,v,:,:]=np.squeeze(invar)
       
        
#   Recall: [allvalues_diffs] = [ncases,nvars,nlat,nlon]                                                                                          
#   Now make a variable with differences between given cases and the BASE run                                                                     
allvalues_diffs=np.zeros((ncases-1,nvars,nlat,nlon))
if reldiff==0:
    allvalues_diffs[0,:,:,:]=allvalues[1,:,:,:]-allvalues[0,:,:,:]
    allvalues_diffs[1,:,:,:]=allvalues[2,:,:,:]-allvalues[0,:,:,:]
    allvalues_diffs[2,:,:,:]=allvalues[3,:,:,:]-allvalues[0,:,:,:]
if reldiff==1:
    allvalues_diffs[0,:,:,:]=((allvalues[1,:,:,:]-allvalues[0,:,:,:])*100.)/allvalues[0,:,:,:]
    allvalues_diffs[1,:,:,:]=((allvalues[2,:,:,:]-allvalues[0,:,:,:])*100.)/allvalues[0,:,:,:]
    allvalues_diffs[2,:,:,:]=((allvalues[3,:,:,:]-allvalues[0,:,:,:])*100.)/allvalues[0,:,:,:]


if variables[0]=='PBLH':
    np.save('/div/qbo/great/analysis/CAM6-paper/BCx10_PBLH_diffs',np.squeeze(allvalues_diffs[0,0,:,:]))
    np.save('/div/qbo/great/analysis/CAM6-paper/CO2x2_PBLH_diffs',np.squeeze(allvalues_diffs[2,0,:,:]))

  
cases=allcases[1:len(allcases)] #As I want to plot differences, skip the BASE case here                                                       
ncases=len(cases)
print(cases)

allvalues_pval=np.zeros((ncases,nvars,nlat,nlon))
allvalues_pval=np.load('signif_4vars.npy') 
#allvalues_pval[allvalues_pval>0.05]=0  #Not significant
#allvalues_pval[allvalues_pval<=0.05]=1 #significant

#Specify what color scheme I want to use:                                                                                                     
#cmap =plt.get_cmap('RdBu')  #         

##Combining two color schemes, see https://matplotlib.org/tutorials/colors/colormaps.html
colors1 = plt.cm.pink(np.linspace(0, 1, 128)) #works when the lower one, as it fades to white
##colors2 = plt.cm.BuGn(np.linspace(0., 1, 128)) #works when the upper one, as it fades FROM white
colors2 = np.flip(plt.cm.bone(np.linspace(0., 1, 128)),axis=0) #must be flipped when the upper one, as it fades TO white
## combine them and build a new colormap
colors = np.vstack((colors1, colors2))
cmap = mcolors.LinearSegmentedColormap.from_list('my_colormap', colors)
cmap_orig=cmap
cmap_reversed=cmap.reversed()

#xlon,xlat=np.meshgrid(lon_values[:],lat_values[:])
xlon=lon_values
xlat=lat_values

#Define colorbar limits and extensions
variables=["Near-surface temp.","Brunt-Väisälä freq.","Richardson number","PBL height"]
minlimits = [-10., -2.5e-3,-0.08, -150.]
maxlimits = [10.,2.5e-3,0.08, 150.]


if reldiff:
    minlimits = [-20., -30., -50.,-60., -1.,-120.]
    maxlimits = [20., 30., 50.,60.,1.,120.]
    units=['[%]','[%]','[%]','[%]','[%]','[%]']    

print('Create the figure..')

fig, axes = plt.subplots(nrows=nvars, ncols=ncases, figsize=(10, 8), subplot_kw=dict(projection=ccrs.Robinson()))

print('Getting ready to plot the plot..')
mainfontsize=10.#18.
density=5
#-------------Loop over cases------------
for r in range(nvars): #the rows

    #---------Loop over variables--------
    for c in range(ncases): #the columns
        print('       ..plotting row '+str(r)+', col '+str(c)+'..')

        #Define colormap
        cmap=cmap_orig
        #if r==0: cmap=cmap_reversed
        #if r>0: cmap=cmap_orig
        cmap=cmap_reversed
        if r==3: cmap=cmap_orig
        #if r==0: cmap=cmap_orig
        
        #Axes and coastlines and ocean mask
        ax = axes[r, c]
        if dooceanmask: ax.add_feature(cart.feature.OCEAN,zorder=100, edgecolor='k',facecolor='white',linewidth=0.2) #black(k) coastlines and white oceans, set high zorder to put masking on top of all other data
        #If I don't want the ocean mask, I have to plot the coastlines separately:
        #ax.add_feature(cfeature.COASTLINE, linewidth=0.2)
        if dooceanmask!=1: ax.add_feature(cfeature.COASTLINE, linewidth=0.8,edgecolor='k')
        ax.axis('off') #In this case a thin line around the globe plot remains, as it's part of the ocean mask thing above, I think
        ax.set_global()
       
        #Plot
        val = np.squeeze(allvalues_diffs[c,r,:,:])

        #Calculate global land-only average value:        
        if dooceanmask: var=np.multiply(val,landfrac)
        if dooceanmask!=1: var=val

        var_zonavg=np.squeeze(var.mean(axis=1))
        glbval = np.average(var_zonavg, weights=weights)
        print(np.nanmean(val))
        print(np.nanmean(var))
        print(glbval)
        ax.set_title(casenames[c]+', '+variables[r],fontsize=mainfontsize)

        h1 = ax.pcolormesh(xlon, xlat, val, transform=ccrs.PlateCarree(), cmap=cmap, vmin=minlimits[r], vmax=maxlimits[r], rasterized=True)

        signif=np.squeeze(allvalues_pval[c,r,:,:])
        sigwrap=np.zeros((nlat,nlon+1))
        sigwrap[:,0:nlon]=signif
        sigwrap[:,nlon]=signif[:,0]
        lonwrap=np.zeros((nlon+1))
        lonwrap[0:nlon]=lon_values[:]
        lonwrap[nlon]=360.

        ax.contourf(lonwrap, lat_valuesLF, sigwrap > 0.1,transform=ccrs.PlateCarree(),colors='none',levels=[.5,1.5],hatches=[density*'.',density*'.'],)


        if c == ncases-1 :
            print('plotting colorbars')
            posplot = ax.get_position() #Get the position of the current panel plot
            poscb   = [posplot.x1+0.01, posplot.y0,  posplot.width/20., posplot.height] #Wanted position of colorbar
            cb_ax = fig.add_axes(poscb) #Add an extra axes only for the colorbar, to avoid it shrinking only the panels with colorbars
            cbar = fig.colorbar(h1, cax=cb_ax,orientation="vertical",extend='both')
            cbar.ax.tick_params(size=0,labelsize=mainfontsize*0.9) #size=0 to remove the ticks on the righthand side of the colorbar
            cbar.set_label(units[r], rotation=90,fontsize=mainfontsize*0.8)


#Set distance between subplots
plt.subplots_adjust(hspace=0.05, wspace=0.04)


fig.savefig(figname, bbox_inches='tight')