import xarray as xr
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap, BoundaryNorm
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from cartopy.mpl.ticker import LongitudeFormatter, LatitudeFormatter

# ============================================================
# 1. Abrir o arquivo e selecionar a variável
# ============================================================
ds = xr.open_dataset('imerg_regrid.nc')
print(ds)

prec = ds['precipitation']

# Remove a dimensão de tempo, se existir
if 'time' in prec.dims:
    prec = prec.isel(time=0)

# Detecta automaticamente os nomes das coordenadas
lon_name = 'lon' if 'lon' in prec.coords else 'longitude'
lat_name = 'lat' if 'lat' in prec.coords else 'latitude'

# Garante a ordem (lat, lon) para o contourf
prec = prec.transpose(lat_name, lon_name)

print(prec.attrs)
print("Min:", float(prec.min()), "Max:", float(prec.max()))
print("Dims:", prec.dims)

# ============================================================
# 2. Níveis e paleta estilo MONAN / NCL
# ============================================================
levels = np.array([0, 1, 2, 3, 4, 5, 10, 15, 20, 35, 50, 75, 100, 120])

# Uma cor por intervalo: len(levels)-1 = 13 cores
colors = [
    '#ffffff',  # 0-1     branco (sem chuva)
    '#cce5ff',  # 1-2     azul muito claro
    '#99ccff',  # 2-3     azul claro
    '#66b3ff',  # 3-4     azul
    '#3399ff',  # 4-5     azul médio
    '#00ccff',  # 5-10    ciano
    '#00e68a',  # 10-15   verde-ciano
    '#33cc33',  # 15-20   verde
    '#99ff33',  # 20-35   verde-amarelo
    '#ffff00',  # 35-50   amarelo
    '#ffaa00',  # 50-75   laranja
    '#ff3300',  # 75-100  vermelho-laranja
    '#cc0000',  # 100-120 vermelho
    '#660000',  # >120    vinho (extend)
]

cmap = ListedColormap(colors)
norm = BoundaryNorm(levels, ncolors=cmap.N)

# ============================================================
# 3. Área da América do Sul
# ============================================================
lon_min, lon_max = -70, -30
lat_min, lat_max = -45, -10

# ============================================================
# 4. Figura com Cartopy
# ============================================================
fig = plt.figure(figsize=(10, 8))
ax = plt.axes(projection=ccrs.PlateCarree())
ax.set_extent([lon_min, lon_max, lat_min, lat_max], crs=ccrs.PlateCarree())
ax.set_facecolor('white')

ax.add_feature(cfeature.COASTLINE, linewidth=0.7, edgecolor='black')
ax.add_feature(cfeature.BORDERS, linewidth=0.5, edgecolor='black')
ax.add_feature(cfeature.STATES, linewidth=0.3, edgecolor='gray')

# ============================================================
# 5. Plotar precipitação
# ============================================================
plot = ax.contourf(
    prec[lon_name], prec[lat_name], prec,
    levels=levels,
    cmap=cmap,
    norm=norm,
    transform=ccrs.PlateCarree(),
    extend='max'
)

# ============================================================
# 6. Colorbar — mostrando TODOS os níveis
# ============================================================
cbar = plt.colorbar(
    plot, ax=ax,
    orientation='vertical',
    pad=0.03, shrink=0.85,
    ticks=levels,
    spacing='uniform'
)
cbar.set_label('Precipitação (mm/dia)', fontsize=10)
cbar.ax.tick_params(labelsize=8)

# ============================================================
# 7. Gridlines
# ============================================================
gl = ax.gridlines(draw_labels=True, linewidth=0.3,
                  color='gray', alpha=0.4, linestyle='--')
gl.top_labels = False
gl.right_labels = False
gl.xformatter = LongitudeFormatter()
gl.yformatter = LatitudeFormatter()
gl.xlabel_style = {'size': 8}
gl.ylabel_style = {'size': 8}

# ============================================================
# 8. Título (agora indicando IMERG)
# ============================================================
ax.set_title('GPM IMERG | valid:2026083100 | 24h',
             fontsize=13, loc='left')

plt.tight_layout()
plt.savefig('IMERG_prec.png', dpi=150, bbox_inches='tight')
print("Figura salva: IMERG_prec.png")