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('BAM_OPER_regrid.nc')
print(ds)

prec = ds['prec']
if 'time' in prec.dims:
    prec = prec.isel(time=0)

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

# ============================================================
# 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-2     branco (sem chuva)
    '#cce5ff',  # 2-4     azul muito claro
    '#99ccff',  # 4-6     azul claro
    '#66b3ff',  # 6-8     azul
    '#3399ff',  # 8-10    azul médio
    '#00ccff',  # 10-20   ciano
    '#00e68a',  # 20-30   verde-ciano
    '#33cc33',  # 30-40   verde
    '#99ff33',  # 40-50   verde-amarelo
    '#ffff00',  # 50-75   amarelo
    '#ffaa00',  # 75-100  laranja
    '#ff3300',  # 100-125 vermelho-laranja
    '#cc0000',  # 125-150 vermelho
    '#660000',  # >150    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, prec.lat, 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,        # todos os níveis aparecem na barra
    spacing='uniform'    # mantém o espaçamento visual uniforme entre cores
)
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 (mantido como estava)
# ============================================================
ax.set_title('BAM NCEP | init:2026083000 valid:2026083100 fct:24h',
             fontsize=13, loc='left')

plt.tight_layout()
plt.savefig('BAM_OPER_prec.png', dpi=150, bbox_inches='tight')