# -*- 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
from datetime import datetime, timedelta

label_rodada = "Eta05_M01"
RUNDATE = "2025072900" # Ex: 2023021900

# --- CONFIGURAÇÕES INICIAIS ---
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 = r"/dados/grpeta/Team/DiegoChagas/WORKETAVIII"

# Subdiretório para acumulados de precipitação de 24 horas
path_acum_prec_24h = os.path.join(path_figuras, "prec_acum24h")
os.makedirs(path_acum_prec_24h, exist_ok=True)


nome_base_prec = f"{label_rodada}_PREC_{RUNDATE}"
arquivo_nc_prec = os.path.join(path_dados, f"{nome_base_prec}.nc")

# Parse da data e hora da rodada
rodada_data_dt = datetime.strptime(RUNDATE, '%Y%m%d%H')

# --- LEITURA DOS DADOS ---
ds_prec = xr.open_dataset(arquivo_nc_prec)

# --- PALETA E NÍVEIS PARA PRECIPITAÇÃO ---
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
]

# --- FUNÇÃO DE PLOTAGEM PARA ACUMULADOS DE PRECIPITAÇÃO ---
def plotar_acum_prec_24h(prec_acumulada, hp, tii, tfi,
                         cmap, levels, rodada_data_dt, path_figuras_output, RUNDATE_str):
    
    fig = plt.figure(figsize=(12, 12))
    ax = fig.add_subplot(111, projection=ccrs.PlateCarree())

    # Plotar o DataArray interpolado.
    # Usamos o keyword 'robust=True' para evitar outliers que possam distorcer a escala de cores,
    # embora não seja diretamente para suavização.
    img = prec_acumulada.plot(
        cmap=cmap,
        add_colorbar=False,
        ax=ax,
        levels=levels,
        extend='both'
    )

    # Lógica para o título (já corrigida na iteração anterior)
    fct_parte_data_hora_rodada_str = rodada_data_dt.strftime("%d/%m/%Y %H")

    if (hp % 24) == 0:
        horas_from = (hp - 24)
        horas_to = hp
    elif (hp % 12) == 0:
        horas_from = (hp - 24)
        horas_to = hp
    else:
        horas_from = tii
        horas_to = tfi

    data_inicio_acum_titulo_dt = rodada_data_dt + timedelta(hours=horas_from)
    inicio_acum_formatado = data_inicio_acum_titulo_dt.strftime("%d/%m/%Y , %HUTC")

    data_fim_acum_titulo_dt = rodada_data_dt + timedelta(hours=horas_to)
    fim_acum_formatado = data_fim_acum_titulo_dt.strftime("%d/%m/%Y, %HUTC")

    titulo_mapa = (
        f'Model {label_rodada} - Daily Precipitation (mm)\n'
        f'Fct {fct_parte_data_hora_rodada_str}+{hp:03d}h,\nfrom {inicio_acum_formatado} to {fim_acum_formatado}'
    )
    ax.set_title(titulo_mapa, fontsize=20)

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

    # Ajustes na barra de cores: pad e tamanho dos labels
    cbar = plt.colorbar(
        img,
        ax=ax,
        cmap=cmap,
        pad=0.05,  # Reduzindo o padding para diminuir o espaço
        fraction=0.25,
        orientation='horizontal',
        extend='both',
        spacing='uniform',
    )
    cbar.set_label('Precipitacao Acumulada em 24h (mm)', fontsize=14) # Aumentando o fontsize do label da barra
    cbar.set_ticks(levels)
    cbar.ax.set_xticklabels([str(int(l)) if l == int(l) else str(l) for l in levels], fontsize=12) # Aumentando o fontsize dos ticks

    # Ajustes nos labels de latitude e longitude
    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
    gl.xlabel_style = {'size': 12} # Aumentando o tamanho da fonte dos labels de longitude
    gl.ylabel_style = {'size': 12} # Aumentando o tamanho da fonte dos labels de latitude


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

    # Definir limites do mapa
    ax.set_xlim([-53.7, -40.5])
    ax.set_ylim([-29.4, -19.6])

    # Nome do arquivo: precip_24acum_YYYYMMDDHH_HP.png
    fig_name = f"precip_24acum_{RUNDATE_str}_{hp:03d}.png"
    fig_path = os.path.join(path_figuras_output, fig_name)
    fig.savefig(fig_path, bbox_inches='tight', dpi=600) # Aumentei o DPI para máxima qualidade
    plt.close(fig)

# --- REGRAS DE ACUMULADOS (HP, TII, TFI) ---
# TII e TFI são os ÍNDICES do Xarray, que representam as horas desde o início da rodada.
acum_rules = [
    {'HP': 24, 'TII': 2, 'TFI': 25},
    {'HP': 36, 'TII': 14, 'TFI': 37},
    {'HP': 48, 'TII': 26, 'TFI': 49},
    {'HP': 60, 'TII': 38, 'TFI': 61},
    {'HP': 72, 'TII': 50, 'TFI': 73},
    {'HP': 84, 'TII': 62, 'TFI': 85},
    {'HP': 96, 'TII': 74, 'TFI': 97},
    {'HP': 108, 'TII': 86, 'TFI': 109},
    {'HP': 120, 'TII': 98, 'TFI': 121},
    {'HP': 132, 'TII': 110, 'TFI': 133},
    {'HP': 144, 'TII': 122, 'TFI': 145},
    {'HP': 156, 'TII': 134, 'TFI': 157},
    {'HP': 168, 'TII': 146, 'TFI': 169},
    {'HP': 180, 'TII': 158, 'TFI': 181},
    {'HP': 192, 'TII': 170, 'TFI': 193},
    {'HP': 204, 'TII': 182, 'TFI': 205},
    {'HP': 216, 'TII': 194, 'TFI': 217},
    {'HP': 228, 'TII': 206, 'TFI': 229},
    {'HP': 240, 'TII': 218, 'TFI': 241},
    {'HP': 252, 'TII': 230, 'TFI': 252},
    {'HP': 264, 'TII': 242, 'TFI': 265},
]

# Obter o número total de tempos no dataset para verificação de limites
total_times_in_ds = len(ds_prec['time'])

# --- LOOP PARA GERAR ACUMULADOS DE PRECIPITAÇÃO DE 24 HORAS SEGUNDO AS REGRAS ---
for rule in acum_rules:
    hp = rule['HP']
    tii = rule['TII']
    tfi = rule['TFI']

    # Verificar se os índices TII e TFI estão dentro dos limites do dataset
    if tfi >= total_times_in_ds:
        print(f"Atencao: Nao ha dados suficientes no arquivo .nc para o acumulado HP={hp}h (requer TFI={tfi}). "
              f"Total de tempos no dataset: {total_times_in_ds}. Parando.")
        break
    
    if tii >= total_times_in_ds:
        print(f"Atencao: TII={tii} excede o numero de tempos no dataset ({total_times_in_ds}). Pulando acumulado HP={hp}h.")
        continue

    # Selecionar os dados de precipitação no intervalo e SOMAR
    acumulado_24h_raw = ds_prec['prec'].isel(time=slice(tii, tfi + 1)).sum(dim='time') * 1000

    # --- INTERPOLAÇÃO PARA SUAVIZAR OS DADOS ---
    # Pegar as coordenadas de lat/lon do DataArray original
    lats = acumulado_24h_raw['lat'].values
    lons = acumulado_24h_raw['lon'].values

    # Criar novas coordenadas mais densas (o dobro da resolução, por exemplo)
    # Ajuste o fator '2' para maior ou menor suavização
    new_lats = np.linspace(lats.min(), lats.max(), len(lats) * 2)
    new_lons = np.linspace(lons.min(), lons.max(), len(lons) * 2)

    # Realizar a interpolação. 'method='linear'' é uma boa opção padrão.
    # Outros métodos incluem 'cubic', 'nearest'.
    acumulado_24h_interpolado = acumulado_24h_raw.interp(lat=new_lats, lon=new_lons, method='linear')
    print(f"Dados interpolados para o acumulado HP={hp}h.")

    plotar_acum_prec_24h(
        prec_acumulada=acumulado_24h_interpolado, # Passa os dados interpolados
        hp=hp,
        tii=tii,
        tfi=tfi,
        cmap=cmap_prec,
        levels=levels_prec,
        rodada_data_dt=rodada_data_dt,
        path_figuras_output=path_acum_prec_24h,
        RUNDATE_str=RUNDATE
    )
    print(f"? Acumulado de 24h (HP={hp}h) de TII={tii} a TFI={tfi} (soma e interpolado) gerado e salvo.")

print(f"\nFinalizado! Todas as figuras de acumulado de 24h foram salvas em:\n{path_acum_prec_24h}")