# -*- coding: utf-8 -*-

import xarray as xr
import matplotlib.pyplot as plt
import numpy as np
import cartopy.feature as cfeature
import cartopy.crs as ccrs
from cartopy.mpl.gridliner import LATITUDE_FORMATTER, LONGITUDE_FORMATTER
import matplotlib as mpl
import cartopy.io.shapereader as shpreader
import os
import cartopy.io.shapereader as shpreader
import cartopy.feature as cfeature

# Garante que os dados necessarios do Cartopy estao baixados
cfeature.COASTLINE.geometries()
cfeature.BORDERS.geometries()

# --- CONFIGURACOES INICIAIS ---
tempos = 5

label_rodada = "Eta05_M01"
RUNDATE = "2025072900"

path_dados = r"/dados/grpeta/Team/DiegoChagas/Eta_Ensemble_5km/CENAPAD/netcdf/2025072900/M01"
path_shapefile = r"/dados/grpeta/Team/DiegoChagas/WORKETAVIII/RECURSOS/SHAPES/contornos_brasil/estados_2010.shp"
path_figuras_base = r"/dados/grpeta/Team/DiegoChagas/WORKETAVIII/figuras_horarias"

# Define o diretorio principal da rodada
path_rodada = os.path.join(path_figuras_base, f"{label_rodada}_{RUNDATE}")

nome_base_prec = f"{label_rodada}_PREC_{RUNDATE}"
nome_base_tp2m = f"{label_rodada}_TP2M_{RUNDATE}"
nome_base_ocis = f"{label_rodada}_OCIS_{RUNDATE}"
nome_base_pslm = f"{label_rodada}_PSLM_{RUNDATE}"
nome_base_zgeo = f"{label_rodada}_ZGEO_{RUNDATE}"
# Adicionado para vento
nome_base_u10m = f"{label_rodada}_U10M_{RUNDATE}"
nome_base_v10m = f"{label_rodada}_V10M_{RUNDATE}"
# Adicionado para temperatura em altitude
nome_base_temp = f"{label_rodada}_TEMP_{RUNDATE}"
# Adicionado para temperatura maxima e minima
nome_base_mxtp = f"{label_rodada}_MXTP_{RUNDATE}"
nome_base_mntp = f"{label_rodada}_MNTP_{RUNDATE}"
# --- NOVAS VARIAVEIS DE FLUXO DE CALOR ---
nome_base_cssf = f"{label_rodada}_CSSF_{RUNDATE}"
nome_base_clsf = f"{label_rodada}_CLSF_{RUNDATE}"


data_str = f"{RUNDATE[:4]}-{RUNDATE[4:6]}-{RUNDATE[6:8]}T{RUNDATE[8:]}"
rodada_data = np.datetime64(data_str, 'h')

arquivo_nc_prec = os.path.join(path_dados, f"{nome_base_prec}.nc")
arquivo_nc_tp2m = os.path.join(path_dados, f"{nome_base_tp2m}.nc")
arquivo_nc_ocis = os.path.join(path_dados, f"{nome_base_ocis}.nc")
arquivo_nc_pslm = os.path.join(path_dados, f"{nome_base_pslm}.nc")
arquivo_nc_zgeo = os.path.join(path_dados, f"{nome_base_zgeo}.nc")
# Adicionado para vento
arquivo_nc_u10m = os.path.join(path_dados, f"{nome_base_u10m}.nc")
arquivo_nc_v10m = os.path.join(path_dados, f"{nome_base_v10m}.nc")
# Adicionado para temperatura em altitude
arquivo_nc_temp = os.path.join(path_dados, f"{nome_base_temp}.nc")
# Adicionado para temperatura maxima e minima
arquivo_nc_mxtp = os.path.join(path_dados, f"{nome_base_mxtp}.nc")
arquivo_nc_mntp = os.path.join(path_dados, f"{nome_base_mntp}.nc")
# --- ABERTURA DOS NOVOS DATASETS ---
arquivo_nc_cssf = os.path.join(path_dados, f"{nome_base_cssf}.nc")
arquivo_nc_clsf = os.path.join(path_dados, f"{nome_base_clsf}.nc")

# Dicionario para mapear o prefixo da figura para o nome da variavel
variaveis_para_diretorios = {
    'prec': 'prec',
    'tp2m': 'tp2m',
    'mxtp': 'mxtp',
    'mntp': 'mntp',
    'ocis': 'ocis',
    'pslm': 'pslm',
    'zgeo': 'zgeo',
    'temp': 'temp',
    'vento10m': 'vento10m',
    'clsf': 'clsf',
    'cssf': 'cssf'
}

# Cria a estrutura de diretorios
for var_dir in variaveis_para_diretorios.values():
    path_final = os.path.join(path_rodada, var_dir)
    os.makedirs(path_final, exist_ok=True)
    print(f"Diretorio criado: {path_final}")


# --- LEITURA DOS DADOS ---
ds_prec = xr.open_dataset(arquivo_nc_prec)
ds_tp2m = xr.open_dataset(arquivo_nc_tp2m)
ds_ocis = xr.open_dataset(arquivo_nc_ocis)
ds_pslm = xr.open_dataset(arquivo_nc_pslm)
ds_zgeo = xr.open_dataset(arquivo_nc_zgeo)
ds_u10m = xr.open_dataset(arquivo_nc_u10m)
ds_v10m = xr.open_dataset(arquivo_nc_v10m)
ds_temp = xr.open_dataset(arquivo_nc_temp)
ds_mxtp = xr.open_dataset(arquivo_nc_mxtp)
ds_mntp = xr.open_dataset(arquivo_nc_mntp)
ds_cssf = xr.open_dataset(arquivo_nc_cssf)
ds_clsf = xr.open_dataset(arquivo_nc_clsf)


# --- PALETAS E NIVEIS ---

# Precipitacao
cmap_prec = mpl.colors.ListedColormap([
    '#cccccc', '#b3b3b3', '#999999', '#80807e',
    '#00007c', '#0000cb', '#010afd', '#045af3', '#6297fd', '#a1bcff',
    '#054e03', '#006600', '#007501', '#058405', '#009700', '#00b000',
    '#00c800', '#00dd00', '#01ff01', '#fcff4d', '#ffd263', '#ffb463',
    '#fd9846', '#ff6e3b', '#ff511a', '#fd1500', '#dd0a03', '#c60200',
    '#b80000'
])
cmap_prec.set_under('#fefefe')
cmap_prec.set_over('#7c0104')
levels_prec = [
    0.2, 0.4, 0.6, 0.8, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10,
    12, 14, 16, 18, 20, 25, 30, 40, 50, 60, 70, 80, 90, 100, 125, 150
]

# Temperatura 2m
cmap_tp2m = mpl.colors.ListedColormap([
    '#326799', '#3F70AC', '#4D85BE', '#5CA2D3', '#6FC0E3', '#80D3E5',
    '#9CE4E8', '#CCF5F0', '#F7FEE9', '#FCEEA3', '#FBD058', '#F0B33A',
    '#E08F22', '#CC6E18', '#B44F14', '#973511', '#7D200E',
])
cmap_tp2m.set_under('#255C92')
cmap_tp2m.set_over('#5C0E0B')
levels_tp2m = np.arange(4, 40, 2)


# OCIS
cmap_ocis = mpl.colors.ListedColormap([
    '#fffca3', '#fdec68', '#f6d548', '#edc029',
    '#efad00', '#e68816', '#d85d00', '#c34200'
])
cmap_ocis.set_under('#fff8e3')
cmap_ocis.set_over('#a22d01')
levels_ocis = np.arange(400, 1300, 100)

# PSLM
cmap_pslm = mpl.colors.ListedColormap([
    '#cccccc', '#b3b3b3', '#999999', '#80807e',
    '#00007c', '#0000cb', '#010afd', '#045af3', '#6297fd', '#a1bcff',
    '#054e03', '#006600', '#007501', '#058405', '#009700', '#00b000',
    '#00c800', '#00dd00', '#01ff01', '#fcff4d', '#ffd263', '#ffb463',
    '#fd9846', '#ff6e3b', '#ff511a', '#fd1500', '#dd0a03', '#c60200',
    '#b80000'
])
levels_pslm = np.arange(1000, 1026, 2)

# ZGEO
cmap_zgeo = mpl.colors.ListedColormap([
    '#cccccc', '#b3b3b3', '#999999', '#80807e',
    '#00007c', '#0000cb', '#010afd', '#045af3', '#6297fd', '#a1bcff',
    '#054e03', '#006600', '#007501', '#058405', '#009700', '#00b000',
    '#00c800', '#00dd00', '#01ff01', '#fcff4d', '#ffd263', '#ffb463',
    '#fd9846', '#ff6e3b', '#ff511a', '#fd1500', '#dd0a03', '#c60200',
    '#b80000',
])

# Vento 10m
colors_vento10m = [
    '#FFFFFF', '#FFFFB3', '#FFE066', '#FFC200', '#FFA500', '#FF8C00',
    '#FF4500', '#DC143C', '#A02020',
]
levels_vento10m = np.arange(2, 18, 2)
cmap_vento10m = mpl.colors.ListedColormap(colors_vento10m[1:8])
cmap_vento10m.set_under(colors_vento10m[0])
cmap_vento10m.set_over(colors_vento10m[-1])

# Temperatura em altitude
colors_temp_suave = [
    '#0000CD', '#1E90FF', '#6495ED', '#ADD8E6', '#F0F8FF', '#FFFACD',
    '#FFD700', '#FF8C00', '#FF4500', '#DC143C', '#8B0000'
]
cmap_temp = mpl.colors.LinearSegmentedColormap.from_list(
    "temp_custom_suave", colors_temp_suave, N=256
)
cmap_temp.set_under(colors_temp_suave[0])
cmap_temp.set_over(colors_temp_suave[-1])

# Fluxo de calor latente e sensivel
levels_clsf = np.array([
    0, 50, 100, 115, 135, 150, 200, 250, 300, 350, 400, 450, 500, 550, 600, 650, 700
])
levels_cssf = np.arange(-200, 650, 50)
colors_fluxos_distinct = [
    '#b866be', '#cf99d4', '#272ae6', '#6d89df', '#4368D6', '#00BFFF',
    '#00CED1', '#20B2AA', '#32CD32', '#ADFF2F', '#FFD700', '#FFC107',
    '#FF8C00', '#FF4500', '#DC143C', '#B22222'
]
cmap_fluxos_clsf = mpl.colors.ListedColormap(colors_fluxos_distinct)
cmap_fluxos_clsf.set_under('#880093')
cmap_fluxos_clsf.set_over('#8B0000')
cmap_cssf = cmap_fluxos_clsf
cmap_clsf = cmap_fluxos_clsf


# --- Calculo de limites fixos ---
niveis_zgeo = [200, 500, 850]
limites_nivel = {}
for nivel in niveis_zgeo:
    dados_nivel = ds_zgeo['zgeo'].sel(lev=nivel, method='nearest')
    min_val = float(dados_nivel.min())
    max_val = float(dados_nivel.max())
    limites_nivel[nivel] = (min_val, max_val)

niveis_temp = [200, 500, 850]
limites_temp_nivel = {}
for nivel in niveis_temp:
    temp_data_all_times_at_level = ds_temp['temp'].sel(lev=nivel, method='nearest') - 273.15
    min_val = float(temp_data_all_times_at_level.min())
    max_val = float(temp_data_all_times_at_level.max())
    limites_temp_nivel[nivel] = (min_val, max_val)


# --- Funcoes de plotagem ---

def plotar_variavel(ds, var_name, valtime, cmap, levels, extend, label_cbar, unidade_cbar, nome_base, RUNDATE, rodada_data, path_figuras_output, nome_fig_prefixo, corr_escala=0):
    dados = ds[var_name].isel(time=valtime)
    fator_interp = 2
    lat_interp = np.linspace(dados.lat.values.min(), dados.lat.values.max(), dados.lat.size * fator_interp)
    lon_interp = np.linspace(dados.lon.values.min(), dados.lon.values.max(), dados.lon.size * fator_interp)
    dados_interp = dados.interp(lat=lat_interp, lon=lon_interp, method='linear')

    if var_name == 'prec':
        dados_interp = dados_interp * 1000
    elif corr_escala != 0:
        dados_interp = dados_interp - corr_escala
    
    if var_name in ['clsf', 'cssf']:
        dados_interp = dados_interp * -1

    fig = plt.figure(figsize=(12, 12))
    ax = fig.add_subplot(111, projection=ccrs.PlateCarree())

    img = dados_interp.plot(
        cmap=cmap,
        add_colorbar=False,
        ax=ax,
        levels=levels,
        extend=extend
    )

    data_prevista = rodada_data + np.timedelta64(valtime, 'h')
    data_prevista_str = data_prevista.astype(str).replace('-', '').replace('T', '').replace(':', '')[:10]
    fct_horas = f"{valtime:03d}"
    data_rodada_formatada = f"{RUNDATE[6:8]}/{RUNDATE[4:6]}/{RUNDATE[0:4]} 00UTC"

    titulo_mapa = f'{label_rodada}\n'
    if nome_fig_prefixo == 'prec':
        titulo_mapa += 'Precipitacao Acumulada'
    elif nome_fig_prefixo == 'tp2m':
        titulo_mapa += 'Temperatura 2m'
    elif nome_fig_prefixo == 'mxtp':
        titulo_mapa += 'Temperatura Maxima 2m'
    elif nome_fig_prefixo == 'mntp':
        titulo_mapa += 'Temperatura Minima 2m'
    elif nome_fig_prefixo == 'cssf':
        titulo_mapa += 'Fluxo de Calor Sensivel'
    elif nome_fig_prefixo == 'clsf':
        titulo_mapa += 'Fluxo de Calor Latente'
    elif nome_fig_prefixo == 'ocis':
        titulo_mapa += 'OCIS'
    
    ax.set_title(f'{titulo_mapa}\nRodada: {data_rodada_formatada} - Previsao: {fct_horas}h', fontsize=20)


    shapefile_geom = list(shpreader.Reader(path_shapefile).geometries())
    ax.add_geometries(shapefile_geom, ccrs.PlateCarree(), edgecolor='black', facecolor='none', linewidth=0.5)

    cbar = plt.colorbar(
        img,
        ax=ax,
        cmap=cmap,
        pad=0.05,
        fraction=0.25 if var_name == 'prec' else 0.15,
        orientation='horizontal',
        extend=extend,
        spacing='uniform',  
    )

    cbar.set_ticks(levels)
    cbar.ax.set_xticklabels([str(int(l)) if l == int(l) else str(l) for l in levels])
    cbar.set_label(f'{label_cbar} ({unidade_cbar})')

    gl = ax.gridlines(crs=ccrs.PlateCarree(), color='black', alpha=1.0,
                     linestyle='dotted', linewidth=0.50, draw_labels=True)
    gl.top_labels = False
    gl.right_labels = False
    gl.yformatter = LATITUDE_FORMATTER
    gl.xformatter = LONGITUDE_FORMATTER

    ax.add_feature(cfeature.COASTLINE)
    ax.add_feature(cfeature.BORDERS)

    fig_name = f"{nome_fig_prefixo}_{RUNDATE}_{data_prevista_str}.png"
    fig_path = os.path.join(path_figuras_output, fig_name)
    fig.savefig(fig_path, bbox_inches='tight')
    plt.close(fig)

def plotar_isolinha(ds, var_name, valtime, cmap, levels, label_cbar, unidade_cbar,
                    nome_base, RUNDATE, rodada_data, path_figuras_output, nome_fig_prefixo):
    dados = ds[var_name].isel(time=valtime)

    fig = plt.figure(figsize=(12, 12))
    ax = fig.add_subplot(111, projection=ccrs.PlateCarree())

    data_prevista = rodada_data + np.timedelta64(valtime, 'h')
    data_prevista_str = data_prevista.astype(str).replace('-', '').replace('T', '').replace(':', '')[:10]  
    fct_horas = f"{valtime:03d}"
    data_rodada_formatada = f"{RUNDATE[6:8]}/{RUNDATE[4:6]}/{RUNDATE[0:4]} 00UTC"

    ax.set_title(
        f'{label_rodada}\n{label_cbar}\nRodada: {data_rodada_formatada} - Previsao: {fct_horas}h',
        fontsize=20
    )

    shapefile_geom = list(shpreader.Reader(path_shapefile).geometries())
    ax.add_geometries(shapefile_geom, ccrs.PlateCarree(), edgecolor='black', facecolor='none', linewidth=0.5)

    cs = ax.contour(
        dados['lon'], dados['lat'], dados.values,
        levels=levels,
        colors=cmap.colors,
        linewidths=1.2,
        transform=ccrs.PlateCarree()
    )
    ax.clabel(cs, inline=True, fontsize=10, fmt='%i', colors='black')

    gl = ax.gridlines(crs=ccrs.PlateCarree(), color='black', alpha=1.0,
                     linestyle='dotted', linewidth=0.50, draw_labels=True)
    gl.top_labels = False
    gl.right_labels = False
    gl.yformatter = LATITUDE_FORMATTER
    gl.xformatter = LONGITUDE_FORMATTER

    ax.add_feature(cfeature.COASTLINE)
    ax.add_feature(cfeature.BORDERS)

    ax.set_xlim([-53.7, -40.5])
    ax.set_ylim([-29.4, -19.6])

    fig_name = f"{nome_fig_prefixo}_{RUNDATE}_{data_prevista_str}_{fct_horas}.png"
    fig_path = os.path.join(path_figuras_output, fig_name)
    fig.savefig(fig_path, bbox_inches='tight')
    plt.close(fig)

def plotar_zgeo(ds, var_name, valtime, nivel, cmap, nome_base, RUNDATE,
                rodada_data, path_figuras_output, limites_nivel):
    dados = ds[var_name].isel(time=valtime).sel(lev=nivel, method='nearest')

    min_val, max_val = limites_nivel[nivel]
    levels = np.linspace(min_val, max_val, 15)

    fig = plt.figure(figsize=(12, 12))
    ax = fig.add_subplot(111, projection=ccrs.PlateCarree())

    data_prevista = rodada_data + np.timedelta64(valtime, 'h')
    data_prevista_str = data_prevista.astype(str).replace('-', '').replace('T', '').replace(':', '')[:10]  
    fct_horas = f"{valtime:03d}"
    data_rodada_formatada = f"{RUNDATE[6:8]}/{RUNDATE[4:6]}/{RUNDATE[0:4]} 00UTC"

    ax.set_title(
        f'{label_rodada}\nGeopotential Height {nivel} hPa\nRodada: {data_rodada_formatada} - Previsao: {fct_horas}h',
        fontsize=20
    )

    shapefile_geom = list(shpreader.Reader(path_shapefile).geometries())
    ax.add_geometries(shapefile_geom, ccrs.PlateCarree(), edgecolor='black', facecolor='none', linewidth=0.5)

    cs = ax.contour(
        dados['lon'], dados['lat'], dados.values,
        levels=levels,
        colors=cmap.colors,
        linewidths=1.5,
        transform=ccrs.PlateCarree()
    )
    ax.clabel(cs, inline=True, fontsize=10, fmt='%i')

    gl = ax.gridlines(crs=ccrs.PlateCarree(), color='black', alpha=1.0,
                     linestyle='dotted', linewidth=0.50, draw_labels=True)
    gl.top_labels = False
    gl.right_labels = False
    gl.yformatter = LATITUDE_FORMATTER
    gl.xformatter = LONGITUDE_FORMATTER

    ax.add_feature(cfeature.COASTLINE)
    ax.add_feature(cfeature.BORDERS)

    ax.set_xlim([-53.7, -40.5])
    ax.set_ylim([-29.4, -19.6])

    fig_name = f"zgeo{nivel}_{RUNDATE}_{data_prevista_str}_{fct_horas}.png"
    fig_path = os.path.join(path_figuras_output, fig_name)
    fig.savefig(fig_path, bbox_inches='tight')
    plt.close(fig)


def plotar_vento10m(ds_u10m, ds_v10m, valtime, cmap, levels, nome_base, RUNDATE, rodada_data, path_figuras_output, nome_fig_prefixo):
    u10m_data = ds_u10m['u10m'].isel(time=valtime)
    v10m_data = ds_v10m['v10m'].isel(time=valtime)

    fator_interp = 2
    lat_interp = np.linspace(u10m_data.lat.values.min(), u10m_data.lat.values.max(), u10m_data.lat.size * fator_interp)
    lon_interp = np.linspace(u10m_data.lon.values.min(), u10m_data.lon.values.max(), u10m_data.lon.size * fator_interp)
    u10m_data_interp = u10m_data.interp(lat=lat_interp, lon=lon_interp, method='linear')
    v10m_data_interp = v10m_data.interp(lat=lat_interp, lon=lon_interp, method='linear')

    wind_speed = np.sqrt(u10m_data_interp**2 + v10m_data_interp**2)

    fig = plt.figure(figsize=(12, 12))
    ax = fig.add_subplot(111, projection=ccrs.PlateCarree())

    img = ax.contourf(
        wind_speed['lon'], wind_speed['lat'], wind_speed.values,
        cmap=cmap,
        levels=levels,
        extend='both',
        transform=ccrs.PlateCarree()
    )

    skip = 10
    ax.quiver(
        u10m_data['lon'].values[::skip], u10m_data['lat'].values[::skip],
        u10m_data.values[::skip, ::skip], v10m_data.values[::skip, ::skip],
        transform=ccrs.PlateCarree(),
        color='black',
        scale=200,
        width=0.003,
        alpha=0.5
    )

    data_prevista = rodada_data + np.timedelta64(valtime, 'h')
    data_prevista_str = data_prevista.astype(str).replace('-', '').replace('T', '').replace(':', '')[:10]
    fct_horas = f"{valtime:03d}"
    data_rodada_formatada = f"{RUNDATE[6:8]}/{RUNDATE[4:6]}/{RUNDATE[0:4]} 00UTC"

    ax.set_title(
        f'{label_rodada}\nVelocidade do Vento 10m\nRodada: {data_rodada_formatada} - Previsao: {fct_horas}h',
        fontsize=20
    )

    shapefile_geom = list(shpreader.Reader(path_shapefile).geometries())
    ax.add_geometries(shapefile_geom, ccrs.PlateCarree(), edgecolor='black', facecolor='none', linewidth=0.5)

    cbar = plt.colorbar(
        img,
        ax=ax,
        pad=0.05,
        fraction=0.15,
        orientation='horizontal',
        extend='both',
        spacing='proportional',
    )
    cbar.set_ticks(levels)
    cbar.ax.set_xticklabels([str(int(l)) for l in levels])
    cbar.set_label('Velocidade do Vento (m/s)')

    gl = ax.gridlines(crs=ccrs.PlateCarree(), color='black', alpha=1.0,
                     linestyle='dotted', linewidth=0.50, draw_labels=True)
    gl.top_labels = False
    gl.right_labels = False
    gl.yformatter = LATITUDE_FORMATTER
    gl.xformatter = LONGITUDE_FORMATTER

    ax.add_feature(cfeature.COASTLINE)
    ax.add_feature(cfeature.BORDERS)

    ax.set_xlim([-53.7, -40.5])
    ax.set_ylim([-29.4, -19.6])

    fig_name = f"{nome_fig_prefixo}_{RUNDATE}_{data_prevista_str}_{fct_horas}.png"
    fig_path = os.path.join(path_figuras_output, fig_name)
    fig.savefig(fig_path, bbox_inches='tight')
    plt.close(fig)

def plotar_temp(ds, var_name, valtime, nivel, cmap, nome_base, RUNDATE,
                rodada_data, path_figuras_output, limites_temp_nivel):
    dados = ds[var_name].isel(time=valtime).sel(lev=nivel, method='nearest') - 273.15
    
    fator_interp = 2
    lat_interp = np.linspace(dados.lat.values.min(), dados.lat.values.max(), dados.lat.size * fator_interp)
    lon_interp = np.linspace(dados.lon.values.min(), dados.lon.values.max(), dados.lon.size * fator_interp)
    dados_interp = dados.interp(lat=lat_interp, lon=lon_interp, method='linear')

    min_val_global, max_val_global = limites_temp_nivel[nivel]

    min_level_int = int(np.floor(min_val_global))
    max_level_int = int(np.ceil(max_val_global))

    if (max_level_int - min_level_int) <= 15:  
        level_step = 1
    elif (max_level_int - min_level_int) <= 30:
        level_step = 2
    else:
        level_step = 5

    levels = np.arange(min_level_int, max_level_int + level_step, level_step)
    levels = np.unique(levels)


    fig = plt.figure(figsize=(12, 12))
    ax = fig.add_subplot(111, projection=ccrs.PlateCarree())

    img = ax.contourf(
        dados_interp['lon'], dados_interp['lat'], dados_interp.values,
        levels=levels,
        cmap=cmap,
        extend='both',
        transform=ccrs.PlateCarree()
    )

    data_prevista = rodada_data + np.timedelta64(valtime, 'h')
    data_prevista_str = data_prevista.astype(str).replace('-', '').replace('T', '').replace(':', '')[:10]  
    fct_horas = f"{valtime:03d}"
    data_rodada_formatada = f"{RUNDATE[6:8]}/{RUNDATE[4:6]}/{RUNDATE[0:4]} 00UTC"

    ax.set_title(
        f'{label_rodada}\nAbsolute Temperature {nivel} hPa\nRodada: {data_rodada_formatada} - Previsao: {fct_horas}h',
        fontsize=20
    )

    shapefile_geom = list(shpreader.Reader(path_shapefile).geometries())
    ax.add_geometries(shapefile_geom, ccrs.PlateCarree(), edgecolor='black', facecolor='none', linewidth=0.5)

    cbar = plt.colorbar(
        img,
        ax=ax,
        pad=0.05,
        fraction=0.15,
        orientation='horizontal',
        extend='both',
        spacing='uniform',  
    )
    cbar.set_ticks(levels)
    cbar.ax.set_xticklabels([str(int(l)) for l in levels])
    cbar.set_label("Temperatura ('\u00B0C')")

    gl = ax.gridlines(crs=ccrs.PlateCarree(), color='black', alpha=1.0,
                     linestyle='dotted', linewidth=0.50, draw_labels=True)
    gl.top_labels = False
    gl.right_labels = False
    gl.yformatter = LATITUDE_FORMATTER
    gl.xformatter = LONGITUDE_FORMATTER

    ax.add_feature(cfeature.COASTLINE)
    ax.add_feature(cfeature.BORDERS)

    fig_name = f"temp{nivel}_{RUNDATE}_{data_prevista_str}_{fct_horas}.png"
    fig_path = os.path.join(path_figuras_output, fig_name)
    fig.savefig(fig_path, bbox_inches='tight')
    plt.close(fig)


# --- LOOP DE PLOTAGEM ---
for valtime in range(tempos):
    print(f"Gerando figuras para o tempo: {valtime}")

    path_prec = os.path.join(path_rodada, variaveis_para_diretorios['prec'])
    path_tp2m = os.path.join(path_rodada, variaveis_para_diretorios['tp2m'])
    path_mxtp = os.path.join(path_rodada, variaveis_para_diretorios['mxtp'])
    path_mntp = os.path.join(path_rodada, variaveis_para_diretorios['mntp'])
    path_ocis = os.path.join(path_rodada, variaveis_para_diretorios['ocis'])
    path_pslm = os.path.join(path_rodada, variaveis_para_diretorios['pslm'])
    path_zgeo = os.path.join(path_rodada, variaveis_para_diretorios['zgeo'])
    path_temp = os.path.join(path_rodada, variaveis_para_diretorios['temp'])
    path_vento10m = os.path.join(path_rodada, variaveis_para_diretorios['vento10m'])
    path_clsf = os.path.join(path_rodada, variaveis_para_diretorios['clsf'])
    path_cssf = os.path.join(path_rodada, variaveis_para_diretorios['cssf'])


    plotar_variavel(ds=ds_prec, var_name='prec', valtime=valtime, cmap=cmap_prec, levels=levels_prec,
                    extend='both', label_cbar='Precipitacao', unidade_cbar='mm',
                    nome_base=nome_base_prec, RUNDATE=RUNDATE, rodada_data=rodada_data,
                    path_figuras_output=path_prec, nome_fig_prefixo='prec', corr_escala=0)

    plotar_variavel(ds=ds_tp2m, var_name='tp2m', valtime=valtime, cmap=cmap_tp2m, levels=levels_tp2m,
                    extend='both', label_cbar='Temperatura 2m', unidade_cbar='\u00B0C',
                    nome_base=nome_base_tp2m, RUNDATE=RUNDATE, rodada_data=rodada_data,
                    path_figuras_output=path_tp2m, nome_fig_prefixo='tp2m', corr_escala=273.15)
    
    plotar_variavel(ds=ds_mxtp, var_name='mxtp', valtime=valtime, cmap=cmap_tp2m, levels=levels_tp2m,
                    extend='both', label_cbar='Temperatura Maxima 2m', unidade_cbar='\u00B0C',
                    nome_base=nome_base_mxtp, RUNDATE=RUNDATE, rodada_data=rodada_data,
                    path_figuras_output=path_mxtp, nome_fig_prefixo='mxtp', corr_escala=273.15)

    plotar_variavel(ds=ds_mntp, var_name='mntp', valtime=valtime, cmap=cmap_tp2m, levels=levels_tp2m,
                    extend='both', label_cbar='Temperatura Minima 2m', unidade_cbar='\u00B0C',
                    nome_base=nome_base_mntp, RUNDATE=RUNDATE, rodada_data=rodada_data,
                    path_figuras_output=path_mntp, nome_fig_prefixo='mntp', corr_escala=273.15)

    plotar_variavel(ds=ds_ocis, var_name='ocis', valtime=valtime, cmap=cmap_ocis, levels=levels_ocis,
                    extend='both', label_cbar='OCIS', unidade_cbar='W/m\u00B2',
                    nome_base=nome_base_ocis, RUNDATE=RUNDATE, rodada_data=rodada_data,
                    path_figuras_output=path_ocis, nome_fig_prefixo='ocis', corr_escala=0)

    plotar_isolinha(ds=ds_pslm, var_name='pslm', valtime=valtime, cmap=cmap_pslm, levels=levels_pslm,
                    label_cbar='Pressao ao nivel do mar', unidade_cbar='hPa',
                    nome_base=nome_base_pslm, RUNDATE=RUNDATE, rodada_data=rodada_data,
                    path_figuras_output=path_pslm, nome_fig_prefixo='pslm')

    for nivel in niveis_zgeo:
        plotar_zgeo(ds=ds_zgeo, var_name='zgeo', valtime=valtime, nivel=nivel,
                    cmap=cmap_zgeo, nome_base=nome_base_zgeo, RUNDATE=RUNDATE,
                    rodada_data=rodada_data, path_figuras_output=path_zgeo,
                    limites_nivel=limites_nivel)
    
    for nivel in niveis_temp:
        plotar_temp(ds=ds_temp, var_name='temp', valtime=valtime, nivel=nivel,
                    cmap=cmap_temp, nome_base=nome_base_temp, RUNDATE=RUNDATE,
                    rodada_data=rodada_data, path_figuras_output=path_temp,
                    limites_temp_nivel=limites_temp_nivel)

    plotar_vento10m(ds_u10m=ds_u10m, ds_v10m=ds_v10m, valtime=valtime,
                    cmap=cmap_vento10m, levels=levels_vento10m,
                    nome_base=f"{label_rodada}_VENTO10M_{RUNDATE}", RUNDATE=RUNDATE,
                    rodada_data=rodada_data, path_figuras_output=path_vento10m,
                    nome_fig_prefixo='vento10m')

    plotar_variavel(ds=ds_clsf, var_name='clsf', valtime=valtime, cmap=cmap_clsf, levels=levels_clsf,
                    extend='both', label_cbar='Fluxo de Calor Latente', unidade_cbar='W/m\u00B2',
                    nome_base=nome_base_clsf, RUNDATE=RUNDATE, rodada_data=rodada_data,
                    path_figuras_output=path_clsf, nome_fig_prefixo='clsf', corr_escala=0)

    plotar_variavel(ds=ds_cssf, var_name='cssf', valtime=valtime, cmap=cmap_cssf, levels=levels_cssf,
                    extend='both', label_cbar='Fluxo de Calor Sensivel', unidade_cbar='W/m\u00B2',
                    nome_base=nome_base_cssf, RUNDATE=RUNDATE, rodada_data=rodada_data,
                    path_figuras_output=path_cssf, nome_fig_prefixo='cssf', corr_escala=0)


print(f"? Todas as figuras foram geradas e salvas em:\n{path_rodada}")