zarrviewer/zarrviewer

#!/usr/bin/env python3
"""
zarrviewer: Visualiza archivos Zarr, muestra su estructura y grafica mapas o series de tiempo.
Uso básico: zarrviewer ruta/al/archivo.zarr/
"""

import argparse
import sys
import os
import re
import xarray as xr
import pandas as pd

def main():
    parser = argparse.ArgumentParser(
        description="Explorador y graficador rápido para archivos Zarr.",
        add_help=False 
    )
    
    parser.add_argument("--help", action="help", default=argparse.SUPPRESS,
                        help="Muestra este mensaje de ayuda y termina")
    parser.add_argument("zarr_path", help="Ruta al directorio del archivo Zarr")
    parser.add_argument("--header", "-h", action="store_true",
                        help="Muestra la estructura/metadata del archivo Zarr (ncdump -h)")
    parser.add_argument("--first", "-f", type=int, default=3,
                        help="Número de tiempos iniciales a mostrar (default: 3)")
    parser.add_argument("--last", "-l", type=int, default=3,
                        help="Número de tiempos finales a mostrar (default: 3)")
    parser.add_argument("--all", "-a", action="store_true",
                        help="Mostrar todos los tiempos")
    parser.add_argument("--coord", "-c", type=str, default=None,
                        help="Nombre de la coord. temporal (auto-detectada por defecto)")
    parser.add_argument("--variable", "-v", type=str, default=None,
                        help="Variable a graficar. Sintaxis: var, var:dim, var.dim o 'var(dim)'")
    parser.add_argument("-n", "--npoints", type=int, default=50,
                        help="Número de registros a graficar en series de tiempo (default: 50)")

    args = parser.parse_args()

    try:
        ds = xr.open_zarr(args.zarr_path)
    except Exception as e:
        print(f"Error al abrir el Zarr: {e}", file=sys.stderr)
        sys.exit(1)

    if len(sys.argv) == 2 or args.header:
        print("="*60)
        print(f" ESTRUCTURA DEL ARCHIVO ZARR: {args.zarr_path}")
        print("="*60)
        ds.info()
        
        print("\n" + "="*60)
        print(" VARIABLES DISPONIBLES PARA GRAFICAR (-v):")
        print("="*60)
        
        for var in ds.data_vars:
            dims = ds[var].dims
            print(f"  - {var} {list(dims)}")
            
        print("\nEjemplo de uso:")
        if len(ds.data_vars) > 0:
            print(f"  zarrviewer {args.zarr_path} -v {list(ds.data_vars)[0]}")
        sys.exit(0)

    # =====================================================================
    # SECCIÓN DE TIEMPOS
    # =====================================================================
    posibles_nombres = ["time", "valid_time", "datetime", "t", "Time", "Valid_time"]
    time_coord = None

    if args.coord:
        if args.coord in ds.coords:
            time_coord = args.coord
        else:
            print(f"La coordenada '{args.coord}' no existe. Disponibles: {list(ds.coords.keys())}", file=sys.stderr)
            sys.exit(1)
    else:
        for nombre in posibles_nombres:
            if nombre in ds.coords:
                time_coord = nombre
                break

    if time_coord is None:
        print("No se encontró una coordenada temporal.", file=sys.stderr)
        sys.exit(1)

    # === CORRECCIÓN VITAL: Manejo de snapshot estático (0 dimensiones temporales) ===
    dims_of_time = ds[time_coord].dims
    if len(dims_of_time) > 0:
        time_dim = dims_of_time[0]
        times = ds[time_coord].values
        total = len(times)
        print(f"Usando coordenada temporal: '{time_coord}' (Eje dimensional: '{time_dim}')")
    else:
        # Es un único frame estático
        time_dim = None
        times = [ds[time_coord].values[()]] 
        total = 1
        print(f"Usando coordenada temporal: '{time_coord}' (Snapshot estático, sin eje)")

    print(f"Total de frames en el Zarr: {total}\n")

    def format_time(t):
        try:
            return pd.to_datetime(t).strftime("%Y-%m-%d %H:%M")
        except:
            return str(t)

    if args.all:
        for t in times:
            print(format_time(t))
    else:
        first_n = min(args.first, total)
        last_n = min(args.last, total)
        for t in times[:first_n]:
            print(format_time(t))
        if total > first_n + last_n:
            print("\n        ...        \n")
        for t in times[-last_n:]:
            print(format_time(t))

    # =====================================================================
    # SECCIÓN DE GRAFICADO INTELIGENTE
    # =====================================================================
    if args.variable:
        print("\n" + "="*60)
        
        match = re.match(r"^([a-zA-Z0-9_]+)(?:[\(\:\.](.+?)\)?)?$", args.variable)
        if not match:
            print(f"[ERROR] Sintaxis inválida para -v. Usa 'var', 'var:dim', 'var.dim' o 'var(dim)'", file=sys.stderr)
            sys.exit(1)
            
        var_name = match.group(1)
        dim_val_request = match.group(2)

        if var_name not in ds.data_vars:
            print(f"[ERROR] La variable '{var_name}' no existe.", file=sys.stderr)
            sys.exit(1)

        print(f" PROCESANDO VARIABLE: '{var_name}'")
        print("="*60)

        var_dims = list(ds[var_name].dims)
        other_dims = [d for d in var_dims if d != time_dim]

        try:
            import matplotlib.pyplot as plt

            # CASO A: Mapa (2D o más)
            if len(other_dims) >= 2:
                ultimo_time_str = format_time(times[-1])
                print(f"-> Detectado formato de mapa (Dimensiones: {other_dims}).")
                print(f"-> Extrayendo último frame ({ultimo_time_str})...")
                
                # Cortamos en tiempo SOLO si la variable depende del tiempo
                ds_plot = ds[var_name]
                if time_dim and time_dim in ds_plot.dims:
                    ds_plot = ds_plot.isel({time_dim: -1})
                    
                plot_kwargs = {}
                
                # Buscamos georreferencia (agregué 'y' y 'x' para que use los de UTM por defecto)
                lat_name, lon_name = None, None
                for lat in ['lat', 'latitude', 'XLAT', 'xlat', 'y']:
                    if lat in ds.variables or lat in ds.coords:
                        lat_name = lat
                        break
                for lon in ['lon', 'longitude', 'XLONG', 'xlon', 'x']:
                    if lon in ds.variables or lon in ds.coords:
                        lon_name = lon
                        break
                
                if lat_name and lon_name:
                    print(f"-> Georreferencia detectada: '{lat_name}', '{lon_name}'. Aplicando ejes...")
                    lat_arr = ds[lat_name]
                    if time_dim and time_dim in lat_arr.dims:
                        lat_arr = lat_arr.isel({time_dim: -1})
                        
                    lon_arr = ds[lon_name]
                    if time_dim and time_dim in lon_arr.dims:
                        lon_arr = lon_arr.isel({time_dim: -1})
                    
                    ds_plot = ds_plot.assign_coords({lat_name: lat_arr, lon_name: lon_arr})
                    plot_kwargs['x'] = lon_name
                    plot_kwargs['y'] = lat_name

                plt.figure(figsize=(8, 10))
                cmap_choice = 'Blues' if var_name in ['tp', 'sd', 'snowc', 'RAINNC', 'RAINC'] else 'viridis'
                plot_kwargs['cmap'] = cmap_choice
                plot_kwargs['cbar_kwargs'] = {'shrink': 0.5}
                
                ds_plot.plot(**plot_kwargs)
                
                plt.gca().set_aspect('equal')
                plt.title(f"{var_name} | Frame: {ultimo_time_str}")
                nombre_imagen = f"mapa_{var_name}.png" if time_dim is None else f"mapa_{var_name}_{pd.to_datetime(times[-1]).strftime('%Y%m%d_%H%M')}.png"

            # CASO B: Puntos / Estaciones (1D adicional)
            elif len(other_dims) == 1:
                dim_name = other_dims[0]
                dim_values = ds[dim_name].values
                
                if dim_val_request is None:
                    print(f"Se detectó que la variable '{var_name}' está asociada a la dimensión: '{dim_name}'")
                    print(f"Valores posibles para {dim_name}:")
                    print(", ".join(map(str, dim_values)) + "\n")
                    print("Ejemplos de uso (cópialos y pégalos):")
                    for ex in dim_values[:3]:
                        print(f"  zarrviewer {args.zarr_path} -v {var_name}:{ex}")
                    sys.exit(0)
                
                try:
                    target_type = type(dim_values[0])
                    casted_val = target_type(dim_val_request)
                except ValueError:
                    casted_val = dim_val_request
                
                if casted_val not in dim_values:
                    print(f"[ERROR] El valor '{casted_val}' no existe en la dimensión '{dim_name}'.", file=sys.stderr)
                    sys.exit(1)

                print(f"-> Detectada serie de tiempo para {dim_name} = {casted_val}")
                ds_plot = ds[var_name].sel({dim_name: casted_val})
                
                if time_dim and time_dim in ds_plot.dims:
                    print(f"-> Extrayendo los últimos {args.npoints} registros...")
                    ds_plot = ds_plot.isel({time_dim: slice(-args.npoints, None)})
                
                plt.figure(figsize=(10, 5))
                ds_plot.plot.line(marker='o', linestyle='-', color='#1f77b4')
                plt.title(f"Serie de tiempo: {var_name} ({dim_name}: {casted_val})")
                plt.grid(True, linestyle='--', alpha=0.6)
                nombre_imagen = f"serie_{var_name}_{casted_val}.png"

            # CASO C: Serie de tiempo global (0 dimensiones)
            else:
                print(f"-> Detectada serie de tiempo unidimensional global.")
                ds_plot = ds[var_name]
                
                if time_dim and time_dim in ds_plot.dims:
                    print(f"-> Extrayendo los últimos {args.npoints} registros...")
                    ds_plot = ds_plot.isel({time_dim: slice(-args.npoints, None)})
                
                plt.figure(figsize=(10, 5))
                ds_plot.plot.line(marker='o')
                plt.title(f"Serie global: {var_name}")
                plt.grid(True, linestyle='--', alpha=0.6)
                nombre_imagen = f"serie_global_{var_name}.png"

            plt.tight_layout()
            plt.savefig(nombre_imagen, dpi=150)
            plt.close()
            
            print(f"-> [ÉXITO] Gráfico generado correctamente.")
            print(f"-> Archivo guardado en: {os.path.abspath(nombre_imagen)}")
            
        except ImportError:
            print("[ERROR] Faltan librerías. Asegúrate de instalar 'matplotlib'", file=sys.stderr)
            sys.exit(1)
        except Exception as e:
            print(f"[ERROR] No se pudo generar el gráfico debido a: {e}", file=sys.stderr)
            sys.exit(1)

if __name__ == "__main__":
    main()