#!/usr/bin/env python3

import glob, os, sys
import numpy as np
from netCDF4 import Dataset
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap, cm, maskoceans
import time
import calendar
from datetime import datetime, timedelta
import matplotlib.cm as cm
from matplotlib import gridspec

# ============= READ FOO ============
def readff(fname):
    nc = Dataset(fname, mode='r')
    foo = nc.variables['footprint'][:]    # seconds
    lon = nc.variables['longitude'][:]
    lat = nc.variables['latitude'][:]
    return lon,lat,foo

foot = []
first=True
for AN in np.arange(2019,2023,1):
    lons = readff('FOOTPRINT_MONAVE_'+str(AN)+'.nc')[0]
    lats = readff('FOOTPRINT_MONAVE_'+str(AN)+'.nc')[1]
    first=False
    foot.append(readff('FOOTPRINT_MONAVE_'+str(AN)+'.nc')[2])


avef = np.mean(foot,axis=0); print('shape of foot',avef.shape)
########

aveSON = np.mean(avef[8:11,:,:],axis=0)
aveDJF = np.mean(np.concatenate((avef[0:2,:,:],avef[None,-1,:,:]),axis=0), axis=0)
aveMAM = np.mean(avef[2:5,:,:],axis=0)
aveJJA = np.mean(avef[5:8,:,:],axis=0)

def myformat(x):
    return ('%.2f' % x).rstrip('0').rstrip('.')

from flexpart.flex_plot import custom_cmap
# ========== PLOT =========
def plot(fig,tit_str,xx,yy,var,ax,colflag=False):
    ax.set_title(tit_str ,fontsize=16)
    m = Basemap(projection='ortho',lon_0=50,lat_0=90,resolution='l')
    lon, lat = np.meshgrid(xx, yy)
    xi, yi = m(lon, lat)
    levs = np.arange(-2,6.01,0.01)
    clevs = 2**levs
    boundlevs = [str(('%.2f' % v).rstrip('0').rstrip('.')) if v <=1 else str(int(round(v))) for v in clevs]; print(boundlevs[::100]) 
    custom_cmap.custom_cmap(); cmap = plt.get_cmap('FlexpartCmap')
    norm = cm.colors.BoundaryNorm(boundaries=clevs,ncolors=50)
    m.drawcoastlines()
    m.drawparallels(np.arange(-90.,91.,45.) )
    m.drawmeridians(np.arange(-180.,181.,60.))
    m.drawmapboundary()
    cs = m.contourf(xi,yi,var,clevs,cmap=cmap,extend='both', norm=norm)
    
    #cbar = m.colorbar(cs,ticks=clevs[::100], location='bottom', pad="10%")
    #cbar.set_label('BC emissions (mg m'+r'$\mathrm{\mathsf{^{-2}}}$ h'+r'$\mathrm{\mathsf{^{-1}}}$)',fontsize=15)
    #cbar.ax.set_xticklabels(boundlevs[::100])
#cbar.ax.set_xticklabels(['0','','0.012','','0.05','','0.2','0.4','0.8','1.6','3.2','6.4'])
    if colflag == True:
        cbaxes = fig.add_axes([0.64, 0.3, 0.03, 0.4])
        cbar = fig.colorbar(cs,ticks=clevs[::100], cax = cbaxes,orientation='vertical')
        cbar.ax.set_yticklabels(boundlevs[::100])
        cbar.set_label('Emission sensitivity (sec)',fontsize=14)
        cbar.ax.tick_params(labelsize=13)
#===========================



fig = plt.figure(figsize=(8,5))
#plt.suptitle('GLOBAL ANNUAL EMISSIONS OF $NH_3$',fontsize=18)
#gs = gridspec.GridSpec(5, 1,height_ratios=[1,0.1,1,0.1,1,0.1,1], width_ratios=[1])
#gs00 = gridspec.GridSpecFromSubplotSpec(1, 3, subplot_spec=gs[0], width_ratios=[3,0.01,3])
#gs01 = gridspec.GridSpecFromSubplotSpec(1, 3, subplot_spec=gs[2], width_ratios=[3,0.01,3])
#gs02 = gridspec.GridSpecFromSubplotSpec(1, 3, subplot_spec=gs[4], width_ratios=[3,0.01,3])

F1 = fig.add_subplot(231)
plot(fig,'(a) DJF',lons,lats,aveDJF,F1)
F2 = fig.add_subplot(232)
plot(fig,'(b) MAM',lons,lats,aveMAM,F2)
F3 = fig.add_subplot(234)
plot(fig,'(c) JJA',lons,lats,aveJJA,F3)
F4 = fig.add_subplot(235)
plot(fig,'(d) SON',lons,lats,aveSON,F4,True)
#=====================
#plt.tight_layout()

plt.savefig('seas_foot.png', dpi=450)
plt.show()