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

import argparse
import sys
import numpy as np
import datetime as dt
from calendar import monthrange
import os
from netCDF4 import Dataset,num2date
import pandas as pd
import glob
import matplotlib as mpl
# mpl.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.lines as mlines
import matplotlib.dates as mdates
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
from cartopy.mpl.ticker import LongitudeFormatter, LatitudeFormatter

from pandas.plotting import register_matplotlib_converters
import xarray as xr
register_matplotlib_converters()

sys.path.insert(1,'/discover/nobackup/cakelle2/GEOS_CF/Apps/cftools/external/cmcrameri')
from cmcrameri import cm
print("loading data")
path = '/discover/nobackup/projects/gmao/geos_cf_dev/psturm/CFv2_c90/holding/gcc_kpp'
# nc = xr.open_mfdataset( path + '/CFv2_ctrl_c90.gcc_kpp.20181001*.nc4')
# nc = xr.open_mfdataset( path + '/CFv2_ctrl_c90*.nc4')
nc = xr.open_dataset( path + '/mergetime_vertsum.nc4')

# Make a plot of SZA and KPP integration steps
print("Making total steps to sza plot")
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
fig = plt.figure(figsize=(10,3.5))
gs  = GridSpec(1,2)
lons = nc.variables['lon'][:]
lats = nc.variables['lat'][:]

timestep = 0
max_count = 100*np.ceil(np.max(nc["KppTotSteps"].values[0,:,:]) / 100 )
ax1 = _make_subplot(fig,gs[0],lons,lats,
                    nc["Met_SUNCOSmid"].values[timestep,:,:],
                    levels=np.arange(-1.0,1.01,0.01),
                    cmap = cm.vik,
                    cbar_label = "Cosine of Solar Zenith Angle",
                    cbar_ticks =[-1.0, -0.5, 0, 0.5, 1.0] )
ax2 = _make_subplot(fig,gs[1],lons,lats,
                    nc["KppTotSteps"].values[0,:,:],
                    levels=np.arange(0,np.max(nc["KppTotSteps"].values[timestep,:,:]),1),
                    cbar_label = "Number of KPP integration steps",
                    cbar_ticks =[0, max_count/4, max_count/2, max_count*0.75, max_count ] )

fig.suptitle("GEOS with Full Chemistry")
plt.savefig('figs/TotSteps_to_SZA.png',dpi=300)

# Make a plot of KppCPUSteps
print("Making KppCPUSteps plot")
max_count = 100*np.ceil(np.max(nc["KppCPUSteps"].values[0,:,:]) / 100 )
fig = plt.figure(figsize=(5,3.5))
gs = GridSpec(1,1)
ax1 = _make_subplot(fig,gs[0],lons,lats,
                    nc["KppCPUSteps"].values[timestep,:,:],
                    levels=np.arange(0,np.max(nc["KppCPUSteps"].values[timestep,:,:]),100),
                    cbar_label = "Total KPP CPU Steps",
                    cbar_ticks =[0, max_count/4, max_count/2, max_count*0.75, max_count ] )
plt.savefig('figs/KppCPUSteps.pdf',dpi=300)