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

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.feature as cf

import xarray as xr
sys.path.insert(1,'/discover/nobackup/cakelle2/GEOS_CF/Apps/cftools/external/cmcrameri')
from cmcrameri import cm

path_relax_tol = '/discover/nobackup/projects/gmao/geos_cf_dev/psturm/CFv2_c90/holding_relax_tol/gcc_kpp'
# nc_relax_tol = xr.open_mfdataset(path_relax_tol + '/*.nc4')
nc_relax_tol = xr.open_dataset(path_relax_tol + '/CFv2_ctrl_c90.gcc_kpp.20181005_2100z.nc4')
path_control = '/discover/nobackup/projects/gmao/geos_cf_dev/psturm/CFv2_c90/holding_control/gcc_kpp'
# nc_control = xr.open_mfdataset(path_control + '/*.nc4')
nc_control = xr.open_dataset(path_control + '/CFv2_ctrl_c90.gcc_kpp.20181005_2100z.nc4')

def _make_subplot(fig,igs,longitude,latitude,concentration,
        title='',cmap=cm.batlow,levels=np.arange(0,101.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='vertical',label=cbar_label, ticks=cbar_ticks)
    ax.set_title(title)
    return ax

# Plot surface ozone concentration for control, relax_tol, and difference
print('Plotting surface ozone concentration for control, relax_tol, and difference')
fig = plt.figure(figsize=(8,12))
gs = GridSpec(3,1,figure=fig)

lons = nc_control.variables['lon'][:]
lats = nc_control.variables['lat'][:]

ax1 = _make_subplot(fig,gs[0],lons,lats,nc_control['SpeciesConc_O3'][0,-1,:,:]*1e9,
        title='Control',cmap=cm.batlow,levels=np.arange(0.0,100.0,1.0),
        cbar_label= 'Surface Ozone Concentration [ppb]',cbar_ticks =[0, 25, 50, 75, 100])
ax2 = _make_subplot(fig,gs[1],lons,lats,nc_relax_tol['SpeciesConc_O3'][0,-1,:,:]*1e9,
        title='Relax_tol',cmap=cm.batlow,levels=np.arange(0.0,100.0,1.0),
        cbar_label= 'Surface Ozone Concentration [ppb]',cbar_ticks =[0, 25, 50, 75, 100])
ax3 = _make_subplot(fig,gs[2],lons,lats,(nc_relax_tol['SpeciesConc_O3'][0,-1,:,:] - nc_control['SpeciesConc_O3'][0,-1,:,:])*1e9,
        title='Bias',cmap=cm.vik,levels=np.arange(-.1,.1,0.01),
        cbar_label= 'Bias [ppb]',cbar_ticks =[-.1, -.05, 0, 0.05, 0.1])

plt.savefig('figs/surface_ozone_concentration.png',dpi=300,bbox_inches='tight')

# Repeat for level 56
print('Plotting ozone concentration at L56 for control, relax_tol, and difference')
fig = plt.figure(figsize=(8,12))
gs = GridSpec(3,1,figure=fig)

lons = nc_control.variables['lon'][:]
lats = nc_control.variables['lat'][:]

ax1 = _make_subplot(fig,gs[0],lons,lats,nc_control['SpeciesConc_O3'][0,-56,:,:]*1e9,
        title='Control at L56 (1.9 hPa)',cmap=cm.batlow,levels=np.arange(0.0,6000,1.0),
        cbar_label= 'Ozone Concentration [ppb]',cbar_ticks =[0, 2000, 4000, 6000])
ax2 = _make_subplot(fig,gs[1],lons,lats,nc_relax_tol['SpeciesConc_O3'][0,-56,:,:]*1e9,
        title='Relax_tol at L56 (1.9 hPa)',cmap=cm.batlow,levels=np.arange(0.0,6000,1.0),
        cbar_label= 'Ozone Concentration [ppb]',cbar_ticks =[0, 2000, 4000, 6000])
ax3 = _make_subplot(fig,gs[2],lons,lats,(nc_relax_tol['SpeciesConc_O3'][0,-56,:,:] - nc_control['SpeciesConc_O3'][0,-56,:,:])*1e9,
        title='Bias',cmap=cm.vik,levels=np.arange(-4,4,0.1),
        cbar_label= 'Bias [ppb]',cbar_ticks =[-4, -2, 0, 2, 4])

plt.savefig('figs/ozone_concentration_L56.png',dpi=300,bbox_inches='tight')

# Repeat for level 48


print('Plotting ozone concentration at L48 for control, relax_tol, and difference')
fig = plt.figure(figsize=(8,12))
gs = GridSpec(3,1,figure=fig)

lons = nc_control.variables['lon'][:]
lats = nc_control.variables['lat'][:]

# find the largest absolute bias to set as the symmetric colorbar range
bias_range = np.max(np.abs(nc_relax_tol['SpeciesConc_O3'][0,-48,:,:] - 
                           nc_control['SpeciesConc_O3'][0,-48,:,:])*1e9)
bias_levels = np.arange(-bias_range,bias_range+0.01*bias_range,bias_range/100)
bias_cbar_ticks = np.arange(-bias_range,bias_range+0.1*bias_range,bias_range/5)

ax1 = _make_subplot(fig,gs[0],lons,lats,nc_control['SpeciesConc_O3'][0,-48,:,:]*1e9,
        title='Control at L48 (10 hPa)',cmap=cm.batlow,
        cbar_label= 'Ozone Concentration [ppb]')
ax2 = _make_subplot(fig,gs[1],lons,lats,nc_relax_tol['SpeciesConc_O3'][0,-48,:,:]*1e9,
        title='Relax_tol at L48 (10 hPa)',cmap=cm.batlow,
        cbar_label= 'Ozone Concentration [ppb]')

ax3 = _make_subplot(fig,gs[2],lons,lats,(nc_relax_tol['SpeciesConc_O3'][0,-48,:,:] - nc_control['SpeciesConc_O3'][0,-48,:,:])*1e9,
        title='Bias',cmap=cm.vik, levels=bias_levels, cbar_ticks=bias_cbar_ticks,
        cbar_label= 'Bias [ppb]')
plt.savefig('figs/ozone_concentration_L48.png',dpi=300,bbox_inches='tight')