import matplotlib.pyplot as plt
import numpy as np
import geopandas as gpd
import streamlit as st

def plot_precip_com_ferrovia(da, gdf, vermelhos_gdf, nome=""):
    fig, ax = plt.subplots(figsize=(10, 8))
    da.plot(ax=ax, cmap="Blues", cbar_kwargs={"label": "Precipitação (mm)"})
    ferrovia_buffer = gpd.read_file("/home/gabriela-vitelli/shapefiles/BUFFER_EFC_25KM_reconstruido.shp")
    ferrovia_buffer.boundary.plot(ax=ax, color="lightgray", markersize=4, label="Estacas")
    if not vermelhos_gdf.empty:
        vermelhos_gdf.plot(ax=ax, color="red", markersize=10, label="Estacas a cada 50 km")
        for _, row in vermelhos_gdf.iterrows():
            if np.isclose(row["km"] % 50, 0, atol=0.1):
                ax.annotate(f"{int(row['km'])}", xy=(row.geometry.x, row.geometry.y),
                            xytext=(3, 3), textcoords="offset points", fontsize=10, color="darkred")
    ax.set_title(f"Acumulado de Precipitação ({nome})")
    ax.set_xlabel("Longitude")
    ax.set_ylabel("Latitude")
    ax.grid(True, linestyle="--", alpha=0.4)
    ax.legend()
    plt.tight_layout()
    st.pyplot(fig)
import os
import xarray as xr
import matplotlib.pyplot as plt
import imageio.v2 as imageio
import cartopy.crs as ccrs
import geopandas as gpd
from glob import glob

def gerar_gif_passo_a_passo(dir_grib, regiao, shp_ferrovia=None, output_dir="/tmp", output_name="gif_24h.gif"):
    import os
    from glob import glob
    import xarray as xr
    import matplotlib.pyplot as plt
    import imageio.v2 as imageio
    import cartopy.crs as ccrs
    import geopandas as gpd

    arquivos = sorted(glob(os.path.join(dir_grib, "*_tp.grib2")))
    arquivos = [arq for arq in arquivos if any(f"{h:03d}_" in os.path.basename(arq) for h in range(0, 25, 3))]

    if not arquivos:
        raise FileNotFoundError("Nenhum arquivo GRIB2 encontrado para as primeiras 24h.")

    imagens = []
    ferrovia = gpd.read_file(shp_ferrovia) if regiao == "EFC" and shp_ferrovia else None

    for arq in arquivos:
        try:
            print(f"🔍 Lendo {arq}")
            ds = xr.open_dataset(arq, engine="cfgrib")
            lon = np.linspace(0, 360 - 0.25, 1440)
            lat = np.linspace(90, -90, 721)
            ds = ds.assign_coords(longitude=("longitude", lon), latitude=("latitude", lat))
            ds.coords["longitude"] = ("longitude", lon)
            ds.coords["latitude"] = ("latitude", lat)
            ds = ds.rename({"longitude": "lon", "latitude": "lat"}) 
            if "tp" not in ds:
                print(f"⚠️ 'tp' não encontrado em {arq}. Variáveis disponíveis: {list(ds.data_vars)}")
                continue

            tp = ds["tp"].squeeze()

            fig, ax = plt.subplots(figsize=(8, 6), subplot_kw={'projection': ccrs.PlateCarree()})
            tp.plot.pcolormesh(
                ax=ax, cmap="Blues", transform=ccrs.PlateCarree(),
                cbar_kwargs={"label": "Precipitação (m)"}
            )
            ax.coastlines()
            ax.set_title(f"Step {os.path.basename(arq).split('_')[0]}")

            if ferrovia is not None:
                ferrovia.boundary.plot(ax=ax, edgecolor="red", linewidth=1)

            nome_img = f"frame_{os.path.basename(arq).split('_')[0]}.png"
            path_img = os.path.join(output_dir, nome_img)
            fig.savefig(path_img, dpi=100, bbox_inches="tight")
            plt.close(fig)
            imagens.append(path_img)

        except Exception as e:
            print(f"❌ Erro ao processar {arq}: {e}")
            continue

    if not imagens:
        raise RuntimeError("❌ Nenhuma imagem foi gerada. Verifique os arquivos GRIB.")

    gif_path = os.path.join(output_dir, output_name)
    with imageio.get_writer(gif_path, mode="I", duration=0.5) as writer:
        for img in imagens:
            writer.append_data(imageio.imread(img))

    print(f"✅ GIF salvo em: {gif_path}")
    return gif_path