#! /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/2026/04/04/00/pos/dataout/GPOSNMC20260404002026040800P.fct.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='Absolute temperature',level=3)[0]
grb = grib.select(name='Geopotential height',level=500)[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)

clim_ranges = {
    "Zonal_wind_(u)":                      {"min": -60,     "max": 60},       # m/s
    "Meridional_wind_(v)":                 {"min": -60,     "max": 60},       # m/s
    "Omega":                             {"min": -2,      "max": 2},        # Pa/s
    "Stream_function":                   {"min": -3e7,    "max": 3e7},      # m^2/s
    "Velocity_potential":                {"min": -2e7,    "max": 2e7},      # m^2/s
    "Geopotential_height":               {"min": 0,       "max": 16000},    # m (depende do nível!)
    "Absolute_temperature":              {"min": 180,     "max": 320},      # K
    "Specific_humidity":                 {"min": 0,       "max": 0.025},    # kg/kg
    "Vertical_dist_total_cloud_cover":   {"min": 0,       "max": 100},      # %
    "Time_mean_surface_relative_humidity":{"min": 0,      "max": 100}       # %
}

geopotential_ranges = {
    1000: (0, 2000),
    925:  (500, 2500),
    850:  (1000, 3500),
    500:  (4800, 6000),
    250:  (9000, 11000),
    100:  (14000, 17000)
}

temperature_ranges = {
    1000: (250, 320),
    850:  (240, 310),
    500:  (220, 270),
    250:  (200, 240)
}

#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)

var_name = str(grb.name).replace(' ', '_')

if var_name == "Geopotential_height":
    vmin, vmax = geopotential_ranges.get(grb.level, (0, 16000))
elif var_name == "Absolute_temperature":
    vmin, vmax = temperature_ranges.get(grb.level, (180, 320))
else:
    vmin = clim_ranges[var_name]["min"]
    vmax = clim_ranges[var_name]["max"]

levels = np.linspace(vmin, vmax, 50)

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('/lustre_xc50/carlos_bastarz/SMNAMonitoringApp/cron_scripts/anls_img/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='Absolute temperature', orientation='vertical', pad=0.01, fraction=0.05)

plt.title('BAM: Absolute temperature @ 3 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(2026040900), str(grb.name).replace(' ','_').replace('(','').replace(')',''), str(3))

os.makedirs(fig_dir, exist_ok=True)

#fname_fig = 'BAM_' + str(grb.name).replace(' ','_').replace('(','').replace(')','') + '_level_' + str(3) + '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(3) + '/' + str(0)
#fname_fig = os.path.join('/home/carlos/GitHub/SMNAMonitoringApp/cron_scripts/anls/imgs', 'BAM', str(grb.name).replace(' ','_').replace('(','').replace(')',''), str(3), str(0))
#fname_fig = os.path.join(fig_dir, str(0))
fname_fig = os.path.join('/lustre_xc50/carlos_bastarz/SMNAMonitoringApp/cron_scripts/anls_img', str(96))

#plt.savefig(fname_fig + '.png', dpi=150, bbox_inches='tight')
#plt.savefig(fname_fig + '.png', bbox_inches='tight')
plt.savefig(
    fname_fig + '.jpg',
    bbox_inches='tight',
    dpi=100,
    format='jpg',
    pil_kwargs={
        'quality': 70,
        'optimize': True,
        'progressive': True,
        'subsampling': 2
    }
)

grib.close()