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
# ============================================================
ds = xr.open_dataset('MONAN_10km_regrid.nc')
print(ds)

# ============================================================
# 2. Somar rainc + rainnc -> precipitação total
# ============================================================
prec = ds['rainc'] + ds['rainnc']
if 'time' in prec.dims:
    prec = prec.isel(time=0)

prec.name = 'prec'
prec.attrs['units'] = 'mm/dia'
prec.attrs['long_name'] = 'Precipitação total (convectiva + não-convectiva)'

# Limpa ruído numérico negativo
prec = prec.where(prec > 0, 0)

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

# ============================================================
# 3. Converter lon 0-360 -> -180..180
#    ATENÇÃO: a coord se chama 'lon' (não 'longitude')
# ============================================================
prec = prec.assign_coords(
    lon=(((prec.lon + 180) % 360) - 180)
).sortby('lon')

# ============================================================
# 4. Fatiar a América do Sul ANTES do contourf
#    (evita o bug do Cartopy com dado global)
# ============================================================
lon_min, lon_max = -70, -30
lat_min, lat_max = -45, -10

prec_sa = prec.sel(
    lon=slice(lon_min, lon_max),
    lat=slice(lat_min, lat_max)
)

print("Min SA:", float(prec_sa.min()), "Max SA:", float(prec_sa.max()))

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

colors = [
    '#ffffff',  # 0-1
    '#cce5ff',  # 1-2
    '#99ccff',  # 2-3
    '#66b3ff',  # 3-4
    '#3399ff',  # 4-5
    '#00ccff',  # 5-10
    '#00e68a',  # 10-15
    '#33cc33',  # 15-20
    '#99ff33',  # 20-35
    '#ffff00',  # 35-50
    '#ffaa00',  # 50-75
    '#ff3300',  # 75-100
    '#cc0000',  # 100-120
    '#660000',  # >120
]
cmap = ListedColormap(colors)
norm = BoundaryNorm(levels, ncolors=cmap.N)

# ============================================================
# 6. Figura
# ============================================================
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')

# ============================================================
# 7. Plotar (usando as coords 'lon' e 'lat')
# ============================================================
plot = ax.contourf(
    prec_sa.lon, prec_sa.lat, prec_sa,
    levels=levels,
    cmap=cmap,
    norm=norm,
    transform=ccrs.PlateCarree(),
    extend='max'
)

# ============================================================
# 8. Colorbar
# ============================================================
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)

# ============================================================
# 9. 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}

# ============================================================
# 10. Título
# ============================================================
ax.set_title('MONAN 10km | init:2026083000 valid:2026083100 fct:24h',
             fontsize=13, loc='left')

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