import sys, os, datetime
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('got_coastalradar_surfcurrent_20230101_20230706.nc','r')
lons         = nc.variables['lon'][:]
lats         = nc.variables['lat'][:]
times        = nc.variables['time'][:]
time_units   = "hours since 2023-1-1 00:00:00"
S_DATE       = time_units.replace('hours','').replace('since','')
yyyy         = int(S_DATE.split('-')[0])
mm           = int(S_DATE.split('-')[1])
dd           = int(S_DATE.split('-')[2].split(' ')[0])
hh           = int(S_DATE.split('-')[2].split(' ')[1].split(':')[0])
xtimes       = range(len(times))
vmin         = 0
vmax         = 66
list_datimes = [(datetime.datetime(yyyy,mm,dd,hh) + relativedelta(hours=int(i))) for i in xtimes]

for i in xtimes:
    fout = 'final/RADAR_'+list_datimes[i].strftime("%Y%m%d%H").upper()+'.png'
    if os.path.exists(fout):

        u = nc.variables['u'][:].filled(np.nan)[i,:,:]
        v = nc.variables['v'][:].filled(np.nan)[i,:,:]
        speed = np.sqrt(u*u + v*v)
        speed[speed>vmax] = vmax

        fig, ax = plt.subplots()
        my_color = ['#640096', '#2730d7', '#598dfc', '#8befd9', '#60cf91', '#1a9850', '#91cf60', '#d9ef8b', '#fee08b', '#fc8d59', '#d73027', '#960064']
        cmap = clr.LinearSegmentedColormap.from_list('custom cu', my_color, N=256)

        m = Basemap(projection='merc',llcrnrlat=17.5,urcrnrlat=20,llcrnrlon=106,urcrnrlon=108.5,ax=ax)
        m.drawparallels(np.arange(-80., 81., 0.5), labels=[1,0,0,0], fontsize=9, linewidth=0.75, color='gray', textcolor='k')
        m.drawmeridians(np.arange(-180., 181., 0.5), labels=[0,0,0,1], fontsize=9, linewidth=0.75, color='gray', textcolor='k')
        m.readshapefile('/home/longtt/apps/shp/ne_50m_admin_0_countries_lakes', 'ne_50m_admin_0_countries_lakes', color='k', linewidth=1)
        m.readshapefile('/home/longtt/apps/shp/haidao', 'haidao', color='k', linewidth=1)

        x, y = np.meshgrid(lons, lats)
        xi, yi = m(x,y)
        yy = np.arange(0, yi.shape[0], 22)
        xx = np.arange(0, xi.shape[1], 22)
        points = np.meshgrid(yy, xx)

        map = m.contourf(xi, yi, speed, vmin=vmin,vmax=vmax,levels=range(vmin,vmax+2,2), cmap=cmap)
        #m.quiver(xi[points], yi[points], u[points], v[points], speed[points], scale=2, scale_units='inches')
        m.quiver(xi, yi, u, v, speed, scale=220, scale_units='inches',cmap=ListedColormap(sns.color_palette(['#000']).as_hex()))
        ax.set_title('RADAR '+list_datimes[i].strftime("%b-%Y").upper(),fontsize=9)
        cbar = fig.colorbar(map, cax=fig.add_axes([0.8, 0.11, 0.025, 0.77]))
        cbar.ax.set_title('(cm/s)',fontsize=9)
        cbar.ax.tick_params(labelsize=9)
        fig.savefig(fout.replace('final/',''), bbox_inches="tight", pad_inches=.05, dpi=150)
        plt.cla()
        plt.clf()
        plt.close()
        #sys.exit()