#! /usr/bin/env python3

import os
import pygrib
import matplotlib.pyplot as plt
import cartopy, cartopy.crs as ccrs
import cartopy.feature as cfeature
import cartopy.io.shapereader as shpreader
import numpy as np
from cartopy.util import add_cyclic_point

grib = pygrib.open('/lustre_xc50/ioper/data/BAM/TQ0666L064/2025/10/01/00/pos/dataout/GPOSNMC20251001002025100100P.icn.TQ0666L064.grb')

nx = 2000
lon0 = 0.0
dx = 0.18
lons = lon0 + dx * np.arange(nx)

ny = 1000
lat0 = -89.910
dy = 0.18
lats = lat0 + dy * np.arange(ny)

#grb = grib.select(name='Omega',level=1000)[0]
grb = grib.select(name='Omega',level=1000)[0]

init  = str(grb.analDate)      # Init date / time
run   = str(grb.hour).zfill(2) # Run
valid = str(grb.validDate)     # Valid date / time

tmtmp = grb.values 

tmtmp_plot = np.flipud(tmtmp)

# remove linha branca vertical no meio da figura
data, lons = add_cyclic_point(tmtmp_plot, coord=lons)

plt.figure(figsize=(30,10))
ax = plt.axes(projection=ccrs.PlateCarree())

lon2d, lat2d = np.meshgrid(lons, lats)

#data_min = 600
#data_max = 1000
#interval = 50
#levels = np.arange(data_min, data_max+interval, interval)

vmin = np.nanmin(data)
vmax = np.nanmax(data)
interval = 50
levels = np.linspace(vmin, vmax, interval)

img1 = ax.contourf(lon2d, lat2d, data, cmap='jet', levels=levels, extend='both')

#img2 = ax.contour(lon2d, lat2d, data, colors='white', linewidths=0.3, levels=levels)
#ax.clabel(img2, inline=1, inline_spacing=15, fontsize=10, fmt='%1.0f', colors='white')

shapefile = list(shpreader.Reader('./br_unidades_da_federacao/BR_UF_2019.shp').geometries())
ax.add_geometries(shapefile, ccrs.PlateCarree(), edgecolor='black', facecolor='none', linewidth=0.3)

ax.coastlines(resolution='10m', color='black', linewidth=0.8)
ax.add_feature(cfeature.BORDERS, edgecolor='black', linewidth=0.5)
gl = ax.gridlines(crs=ccrs.PlateCarree(), color='gray', alpha=1.0, linestyle='--', linewidth=0.25,
xlocs=np.arange(-180, 180, 5), ylocs=np.arange(-90, 90, 5), draw_labels=True)
gl.top_labels = False
gl.right_labels = False

plt.colorbar(img1, label='Omega', orientation='vertical', pad=0.01, fraction=0.05)

plt.title('BAM: Omega @ 1000 hPa, FCT 0 h', fontweight='bold', fontsize=12, loc='left')
plt.title('Valid: ' + valid, fontsize=12, loc='right')

fig_dir = os.path.join('/lustre_xc50/carlos_bastarz/SMNAMonitoringApp/cron_scripts/anls_img/imgs', 'BAM', str(2025100100), str(grb.name).replace(' ','_').replace('(','').replace(')',''), str(1000))

os.makedirs(fig_dir, exist_ok=True)

#fname_fig = 'BAM_' + str(grb.name).replace(' ','_').replace('(','').replace(')','') + '_level_' + str(1000) + 'hPa_valid_for_' + str(valid).replace(' ','_').replace('-','_').replace(':00:00','Z') + '_fct_' + str(0) + 'h'
#fname_fig = 'BAM/' + str(grb.name).replace(' ','_').replace('(','').replace(')','') + '/' + str(1000) + '/' + str(0)
#fname_fig = os.path.join('/home/carlos/GitHub/SMNAMonitoringApp/cron_scripts/anls/imgs', 'BAM', str(grb.name).replace(' ','_').replace('(','').replace(')',''), str(1000), str(0))
fname_fig = os.path.join(fig_dir, str(0))

#plt.savefig(fname_fig + '.png', dpi=150, bbox_inches='tight')
plt.savefig(fname_fig + '.png', bbox_inches='tight')

grib.close()