#!/usr/local/other/python/GEOSpyD/2019.03_py3.7/2019-04-22/bin/python
import os
import sys

import numpy as np
import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.cm import get_cmap
from matplotlib.gridspec import GridSpec
from matplotlib.colors import LightSource
from matplotlib.cm import get_cmap
import cartopy.crs as ccrs

import cartopy.crs as ccrs
import cartopy.feature as cf

import xarray as xr
sys.path.insert(1,'/discover/nobackup/cakelle2/GEOS_CF/Apps/cftools/external/cmcrameri')
from cmcrameri import cm
nc_blacklist = xr.open_dataset('/discover/nobackup/psturm/GEOS-CF/Transport_Project/GCv13_c48blacklist/holding_notransport/geosgcm_gcc/merged_days.201807-01-05_0300z.nc4')
nc_baseline = xr.open_dataset('/discover/nobackup/psturm/GEOS-CF/Terminator_Project/GCv13_c48baseline/holding5day/geosgcm_gcc/merged_days.201807-01-05_0300z.nc4')
fig = plt.figure(figsize=(7,7))
fig.tight_layout()
gs  = GridSpec(2,2)

cmap = cm.batlow
scale = 1e9 # convert to ppb

levs = nc_blacklist.variables['lev'][71]
lons = nc_blacklist.variables['lon'][:]
lats = nc_blacklist.variables['lat'][:]

blacklist_O3 = nc_blacklist.variables['SpeciesConc_O3'][4,71,:,:]*scale
baseline_O3 = nc_baseline.variables['SpeciesConc_O3'][4,71,:,:]*scale
bias_O3 = blacklist_O3-baseline_O3
rel_bias_O3 = 100*(blacklist_O3-baseline_O3)/baseline_O3 # np.mean(baseline_O3)
flev    = np.arange(0.0,110.0,1.0)


def _make_subplot(fig,igs,longitude,latitude,concentration,
        title='',cmap=cm.batlow,levels=np.arange(-100.0,100.0,1.0),
        cbar_label= 'Surface Ozone Concentration [ppb]',cbar_ticks =[0, 25, 50, 75, 100]):
    proj = ccrs.PlateCarree()
    ax = fig.add_subplot(igs,projection=proj)
    ax.coastlines()
    cp = ax.contourf(longitude,latitude,concentration,cmap=cmap,levels=levels)
    plt.colorbar(cp,ax=ax,orientation='horizontal',label=cbar_label, ticks=cbar_ticks)
    ax.set_title(title)
    return ax

ax1 = _make_subplot(fig,gs[0,0],lons,lats,baseline_O3,title='Baseline Simulation',levels=flev)
ax2 = _make_subplot(fig,gs[0,1],lons,lats,blacklist_O3,title='No Advection',levels=flev)
ax3 = _make_subplot(fig,gs[1,0],lons,lats,bias_O3,title='Bias',cmap = cm.vik, levels = np.linspace(-np.max(np.abs(bias_O3)),np.max(np.abs(bias_O3)),100),cbar_label = 'Bias [ppb]', cbar_ticks = [-50, -25, 0, 25, 50])
ax4 = _make_subplot(fig,gs[1,1],lons,lats,rel_bias_O3,title='Relative Bias',cmap=cm.vik,levels = np.linspace(-np.max(np.abs(rel_bias_O3)),np.max(np.abs(rel_bias_O3)),100),cbar_label = 'Relative Bias [%]', cbar_ticks = [-300, -150, 0, 150, 300])

print("Maximum bias:    ", np.max(np.abs(bias_O3)))
print("Maximum relative bias:    ",np.max(np.abs(rel_bias_O3)))

plt.savefig('png/time-averaged-sfc-ozone-notransport.png')
plt.close()