import sys, datetime, copy, os, statistics
from dateutil.relativedelta import relativedelta
import numpy as np; np.set_printoptions(threshold=sys.maxsize)
import matplotlib; matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.colors as clr
from mpl_toolkits.basemap import Basemap

import seaborn as sns
from matplotlib.colors import ListedColormap
from netCDF4 import Dataset

nc   = Dataset('/home/tanpv/Progs/Flash_Dr/NC/FD_Char_1961-2020_1yrs_Daily.nc','r')
lon  = nc.variables['lon'][:]   ; nx = len(lon)
lat  = nc.variables['lat'][:]   ; ny = len(lat)
cha  = nc.variables['cha'][:]
year = nc.variables['year'][:]  ; nt = len(year)
fd   = nc.variables['fd'][:].filled(np.nan)

for k in range(len(cha)):
    #

    var_2d = np.empty((nt,ny))

    for t in range(nt):
        for j in range(ny):

            rtmp = (fd[t,k,j,:][~np.isnan(fd[t,k,j,:])])
            if len(rtmp) > 0:
                var_2d[t,j] = statistics.mean(rtmp)
            else:
                var_2d[t,j] = np.nan

    # plot
    fig, ax = plt.subplots()
    my_color = ['#ffffff', '#d6e2ff', '#b5c9ff', '#8eb2ff', '#7f96ff', '#6370f7', '#0063ff', '#009696', '#00c633', '#63ff00', '#96ff00', '#c6ff33', '#ffff00', '#ffc600', '#ffa000', '#ff7c00', '#ff1900']
    cmap = clr.LinearSegmentedColormap.from_list('custom cu', my_color, N=256)

    #map = ax.pcolormesh(var_2d.T,cmap=cmap)
    map = ax.imshow(var_2d.T[::-1],interpolation='bilinear',cmap=cmap)
    ax.contour(var_2d.T[::-1])
    ax.set_aspect(0.35)
    ax.set_ylabel('lat')
    ax.set_xlabel('year')
    ax.set_yticks(np.arange(0, ny, 10))
    ax.set_yticklabels([str(lat[i]) for i in np.arange(0, ny, 10)])

    cbar = fig.colorbar(map, cax=fig.add_axes([0.915, 0.29, 0.025, 0.41]))
    #cbar.ax.set_title('(mm/6hr)',fontsize=9)
    #cbar.ax.tick_params(labelsize=9)
    fig.savefig('demo_hov_cha'+str(k+1)+'.pdf', bbox_inches="tight", pad_inches=.05, dpi=600)

    plt.cla()
    plt.clf()
    plt.close()

    sys.exit()