#!/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 + '/mergetime_vertsum.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 + '/mergetime_vertsum.nc4')

# Make a plot of SZA and KPP integration steps
print("Making total steps to sza plot")

fig, axs = plt.subplots(nrows=1, ncols=1,figsize = (8,8))

ymax = np.max(nc_control["KppTotSteps"].values)
ymin = 0

axs.scatter(np.ndarray.flatten(nc_control["Met_SUNCOSmid"].values[0:,:,:]),
                   np.ndarray.flatten(nc_control["KppTotSteps"].values[0:,:,:]),
                   s = 0.04, label = "Control")
axs.scatter(np.ndarray.flatten(nc_relax_tol["Met_SUNCOSmid"].values[0:,:,:]),
                   np.ndarray.flatten(nc_relax_tol["KppTotSteps"].values[0:,:,:]),
                   s = 0.04, color = "red",label = "Relaxed Tolerances")
axs.vlines(np.cos(98*np.pi/180),ymin,ymax,colors='purple', linestyles='--',label='cos(98˚)')
axs.vlines(np.cos(82*np.pi/180),ymin,ymax,colors='orange', linestyles='--',label='cos(82˚)')
axs.set_ylabel("KPP integration counts",fontsize=25)
axs.set_xlabel("Cosine of solar zenith angle",fontsize=25)
axs.legend(markerscale=30,fontsize=12)
plt.savefig('figs/SZA_KppTotSteps_comparison.png',dpi=300)