#!/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)

    # 1. Si no hay más argumentos que el archivo (o si se pidió el header), mostramos información útil.
    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)
        
        # Mostramos las variables y sus dimensiones para orientar al usuario
        for var in ds.data_vars:
            dims = ds[var].dims
            print(f"  - {var} {list(dims)}")
            
        print("\nEjemplo de uso:")
        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)

    time_dim = ds[time_coord].dims[0]
    times = ds[time_coord].values
    total = len(times)
    
    print(f"Usando coordenada temporal: '{time_coord}' (Eje dimensional: '{time_dim}')")
    print(f"Total de frames en el Zarr: {total}\n")

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

    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})...")
                
                ds_plot = ds[var_name].isel({time_dim: -1})
                plot_kwargs = {}
                
                # Búsqueda oportunista de georreferencia
                lat_name, lon_name = None, None
                for lat in ['lat', 'latitude', 'XLAT', 'xlat']:
                    if lat in ds.variables:
                        lat_name = lat
                        break
                for lon in ['lon', 'longitude', 'XLONG', 'xlon']:
                    if lon in ds.variables:
                        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].isel({time_dim: -1}) if time_dim in ds[lat_name].dims else ds[lat_name]
                    lon_arr = ds[lon_name].isel({time_dim: -1}) if time_dim in ds[lon_name].dims else ds[lon_name]
                    
                    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
                
                # Achicar la barra de colores a la mitad
                plot_kwargs['cbar_kwargs'] = {'shrink': 0.5}
                
                ds_plot.plot(**plot_kwargs)
                
                # Mantener proporción georeferenciada 1:1
                plt.gca().set_aspect('equal')
                
                plt.title(f"{var_name} | Frame: {ultimo_time_str}")
                nombre_imagen = 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}:")
                    # Imprimir valores con salto de línea para facilitar el uso
                    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}")
                print(f"-> Extrayendo los últimos {args.npoints} registros...")
                
                ds_plot = ds[var_name].sel({dim_name: casted_val}).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.")
                print(f"-> Extrayendo los últimos {args.npoints} registros...")
                
                ds_plot = ds[var_name].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()
