#!/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
import xskillscore as xs

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']

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

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)
	tSAMeT = SAMeT_MK.temp
	tSAMeT = tSAMeT.assign_coords({"lon": (tSAMeT.lon -360) })
	tSAMeT = tSAMeT.salem.roi(shape=br).sel(lat=slice(lai, laf), lon=slice(loi, lof))
	lon, lat = np.meshgrid(tSAMeT.lon, tSAMeT.lat)
	lev=np.arange(0.1,1.025,0.025)
	bounds = [0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9,1.0]
	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
		Tmod = Tmod.assign_coords({"lon": (Tmod.lon -360) })
		Tmod = Tmod.salem.roi(shape=br).sel(lat=slice(lai, laf), lon=slice(loi, lof))

		COR = xs.pearson_r(tSAMeT,Tmod, dim='time',skipna=True)

		COR_MEDIO = COR.mean().values
		COR_MEDIO = np.round(COR_MEDIO,2)
		print(modelo,COR_MEDIO,HH,i,j)

		im = ax[i,j].contourf(lon, lat, COR,levels=lev,extend='min', 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, COR_MEDIO ,fontsize = 14, horizontalalignment='center', verticalalignment='center', transform=ax[i,j].transAxes)
#		if HH=="F24" and j==0:
#			ax[i,j].text(-90,-30,'FCT 24 HOURS',rotation='vertical',fontsize = 14)
#			ax[i,j].text(-90,-80,'FCT 48 HOURS',rotation='vertical',fontsize = 14)
#			ax[i,j].text(-90,-130,'FCT 72 HOURS',rotation='vertical',fontsize = 14)
		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)
#plt.show()
plt.savefig('CORR_'+VAR+'.png', format='png', bbox_inches='tight',dpi=130)
sys.exit()