#!/home/ioper/bin/miniconda3/envs/geada/bin/python
# coding: utf-8

import geopandas as gpd
import cartopy.crs as ccrs
import cartopy.feature as cfeature
import matplotlib.pyplot as plt
import xarray as xr
import sys
import salem
import numpy as np
import rasterio
import math

VAR = sys.argv[1]
#HH = sys.argv[2]

VAR = str(VAR)
#HH = str(HH)

plt.rcParams.update({'font.size': 14})

lai=-35
laf=6.5
loi=-75.0
lof=-33.5
shp = salem.read_shapefile('a__031_001_americaDoSul.shp')
br = shp.loc[shp['ADMIN'] == 'Brazil']

fig, ax = plt.subplots(nrows=3,ncols=5,subplot_kw={'projection': ccrs.PlateCarree()},figsize=(16,8))
i=0
k=0
hh_list = ("F24","F48","F72")
for HH in hh_list:
	print(HH)
	SAMeT = xr.open_dataset('/mnt/beegfs/ioper/users/rozante/ENS_CNN/MODEL_ATMOS/'+VAR+'/SAMET_'+VAR+'_'+HH+'.nc',decode_times=False)
	SAMeT_MK = SAMeT.where(SAMeT['temp'] < 100.0)  #troca undeff por NAN
	tSAMeT = SAMeT_MK.temp
	lon, lat = np.meshgrid(tSAMeT.lon, tSAMeT.lat)
	lon=lon-360
	lev=np.arange(-5.0,5.1,1)
	lev=np.arange(0,5.1,0.1)
	bounds = [-5,-4,-3,-2,-1,0, 1,2, 3, 4 ,5]
	bounds = [0, 1, 2, 3, 4, 5 ]
	letters=['a','b','c','d','e','f','g','h','i','j','l','m','n','o','p']
	j=0
	mod_list = ("ETA","WRF","BRAMS","SMEC","CNN")
	for modelo in mod_list:
		if modelo == 'CNN':
			mod = xr.open_dataset('/mnt/beegfs/ioper/users/rozante/ENS_CNN/MODEL_OUTPUT/'+modelo+'_'+VAR+'_'+HH+'.nc',decode_times=False)
			modelo = 'MAPNet'
		else:
			mod = xr.open_dataset('/mnt/beegfs/ioper/users/rozante/ENS_CNN/MODEL_ATMOS/'+VAR+'/'+modelo+'_'+VAR+'_'+HH+'.nc',decode_times=False)
		mod_MK = mod.where(mod['temp'] < 100.0)
		Tmod = mod_MK.temp
		tmp = ((Tmod-tSAMeT)*(Tmod-tSAMeT)).mean("time",skipna=True)**(1/2)
		tmp = tmp.assign_coords({"lon": (tmp.lon -360) })
		rmse = tmp.salem.roi(shape=br).sel(lat=slice(lai, laf), lon=slice(loi, lof))
		rmse_medio = rmse.mean().values
		rmse_medio = np.round(rmse_medio,2)

		print(modelo,rmse_medio,HH,i,j)
		im = ax[i,j].contourf(lon, lat, rmse,levels=lev,extend='max', cmap='RdBu_r', transform=ccrs.PlateCarree())
		g1 = ax[i,j].gridlines(crs=ccrs.PlateCarree(), linestyle='--', color='gray', linewidth=0.01, draw_labels=True)
		g1.top_labels = False
		g1.right_labels = False
		g1.left_labels = False
		g1.bottom_labels  = False
		if j == 0:
			g1.left_labels = True
		if i == 2:
			g1.bottom_labels = True
		
		ax[i,j].add_feature(cfeature.BORDERS)
		ax[i,j].add_feature(cfeature.COASTLINE)
		if i == 0 :
				ax[i,j].set_title(modelo,fontsize = 20)
# define o tamanho do label X
		g1.xlabel_style = {'size': 10}
# define o tamanho do label Y
		g1.ylabel_style = {'size': 10}
		if i==0 and j==0:
			ax[i,j].text(-0.3, 0.55,'FCT 24 HOURS', horizontalalignment='center', verticalalignment='center',rotation = 90, transform=ax[i,j].transAxes)

		if i==1 and j==0:
			ax[i,j].text(-0.3, 0.55,'FCT 48 HOURS', horizontalalignment='center', verticalalignment='center',rotation = 90, transform=ax[i,j].transAxes)

		if i==2 and j==0:
			ax[i,j].text(-0.3, 0.55,'FCT 72 HOURS', horizontalalignment='center', verticalalignment='center',rotation = 90, transform=ax[i,j].transAxes)

		ax[i,j].text(0.9, 0.9, '('+letters[k]+')',fontsize = 14, horizontalalignment='center', verticalalignment='center', transform=ax[i,j].transAxes)
		ax[i,j].text(0.83, 0.09, rmse_medio ,fontsize = 14, horizontalalignment='center', verticalalignment='center', transform=ax[i,j].transAxes)
		j=j+1
		k=k+1
	i=i+1
cbar_ax = fig.add_axes([0.13, 0.065, 0.75, 0.02])
fig.colorbar(im, cax=cbar_ax, orientation="horizontal",ticks=bounds,label="[\N{DEGREE SIGN}C]")
plt.savefig('RMSE_'+VAR+'.png', format='png', bbox_inches='tight',dpi=130)
sys.exit()