visualize.py

"""
visualize.py
============

Funciones de visualización rápida para inspeccionar los datos descargados
(HYCOM y viento ERA5/GFS) antes y después de convertirlos al formato FVCOM.

Funciones principales
---------------------
- :func:`plot_hycom_surface`     : SST + SSS en superficie de un NetCDF HYCOM.
- :func:`plot_hycom_section`     : sección vertical de T (o S) en una longitud dada.
- :func:`plot_wind`              : magnitud + flechas de viento (ERA5 o GFS).
- :func:`plot_wind_timeseries`   : serie temporal de viento en un punto.

Uso típico:

    from visualize import plot_hycom_surface, plot_wind

    plot_hycom_surface('datos/hycom/hycom_GoC_forecast.nc')
    plot_wind('datos/gfs/gfs_GoC_forecast.nc', source='gfs')
    plot_wind('datos/era5/era5_GoC.nc', source='era5')

Autor: Liesvy
"""

from __future__ import annotations

from typing import Optional, Union
from pathlib import Path

import numpy as np
import pandas as pd
import xarray as xr
import matplotlib.pyplot as plt
import matplotlib as mpl

try:
    import cartopy.crs as ccrs
    import cartopy.feature as cfeature
    _HAS_CARTOPY = True
except ImportError:
    _HAS_CARTOPY = False

try:
    import cmocean.cm as cmo
    CMAP_T = cmo.thermal
    CMAP_S = cmo.haline
    CMAP_W = cmo.speed
except ImportError:
    CMAP_T = 'inferno'
    CMAP_S = 'viridis'
    CMAP_W = 'viridis'


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _apply_style():
    mpl.rcParams.update({
        'font.family':    'serif',
        'font.size':      12,
        'axes.labelsize': 12,
        'axes.titlesize': 13,
        'figure.titlesize': 15,
    })


def _add_map_features(ax, extent=None):
    if not _HAS_CARTOPY:
        return
    ax.add_feature(cfeature.LAND,      facecolor='lightgray', zorder=10)
    ax.add_feature(cfeature.COASTLINE, linewidth=0.5,         zorder=11)
    ax.add_feature(cfeature.BORDERS,   linewidth=0.3,         zorder=11)
    if extent is not None:
        ax.set_extent(extent, crs=ccrs.PlateCarree())
    gl = ax.gridlines(draw_labels=True, linewidth=0.3, alpha=0.5)
    gl.top_labels = False
    gl.right_labels = False


def _projection():
    return ccrs.PlateCarree() if _HAS_CARTOPY else None


def _subplot_kw():
    return {'projection': ccrs.PlateCarree()} if _HAS_CARTOPY else {}


def _transform_kw():
    return {'transform': ccrs.PlateCarree()} if _HAS_CARTOPY else {}


def _lon_to_180(lon):
    lon = np.asarray(lon)
    return ((lon + 180) % 360) - 180


# ---------------------------------------------------------------------------
# HYCOM
# ---------------------------------------------------------------------------

def plot_hycom_surface(
    nc_file: Union[str, Path],
    time_index: int = 0,
    depth_index: int = 0,
    temp_var: str = 'water_temp',
    sal_var: str = 'salinity',
    lon_var: str = 'lon',
    lat_var: str = 'lat',
    figsize: tuple = (14, 6),
    title: Optional[str] = None,
    show: bool = True,
) -> plt.Figure:
    """Pinta SST y SSS de un NetCDF HYCOM en un instante y un nivel."""
    _apply_style()
    ds = xr.open_dataset(nc_file)

    t  = ds[temp_var].isel(time=time_index, depth=depth_index)
    s  = ds[sal_var].isel(time=time_index, depth=depth_index)
    lon = ds[lon_var].values
    lat = ds[lat_var].values
    date = pd.to_datetime(ds.time.values[time_index])
    z = float(ds.depth.values[depth_index])

    fig, axes = plt.subplots(1, 2, figsize=figsize, subplot_kw=_subplot_kw())

    p1 = axes[0].pcolormesh(lon, lat, t, cmap=CMAP_T, shading='auto',
                            **_transform_kw())
    axes[0].set_title(f'T (°C) — z={z:.0f} m — {date:%Y-%m-%d %H UTC}')
    plt.colorbar(p1, ax=axes[0], orientation='horizontal', pad=0.05, label='°C')

    p2 = axes[1].pcolormesh(lon, lat, s, cmap=CMAP_S, shading='auto',
                            **_transform_kw())
    axes[1].set_title(f'S (PSU) — z={z:.0f} m — {date:%Y-%m-%d %H UTC}')
    plt.colorbar(p2, ax=axes[1], orientation='horizontal', pad=0.05, label='PSU')

    extent = [float(lon.min()), float(lon.max()),
              float(lat.min()), float(lat.max())]
    for ax in axes:
        _add_map_features(ax, extent=extent)

    fig.suptitle(title or f'HYCOM — {Path(nc_file).name}', fontsize=15)
    plt.tight_layout()
    ds.close()
    if show:
        plt.show()
    return fig


def plot_hycom_section(
    nc_file: Union[str, Path],
    lon_target: float,
    time_index: int = 0,
    var: str = 'water_temp',
    lon_var: str = 'lon',
    lat_var: str = 'lat',
    figsize: tuple = (9, 5),
    cmap=None,
    max_depth: Optional[float] = None,
    title: Optional[str] = None,
    show: bool = True,
) -> plt.Figure:
    """Sección vertical de una variable HYCOM en la longitud `lon_target`."""
    _apply_style()
    ds = xr.open_dataset(nc_file)

    da = ds[var].isel(time=time_index)
    sec = da.sel({lon_var: lon_target}, method='nearest')
    lat = ds[lat_var].values
    depth = ds.depth.values

    fig, ax = plt.subplots(figsize=figsize)
    p = ax.pcolormesh(lat, depth, sec.values,
                      cmap=cmap or (CMAP_T if 'temp' in var.lower() else CMAP_S),
                      shading='auto')
    ax.invert_yaxis()
    if max_depth is not None:
        ax.set_ylim(max_depth, 0)
    ax.set_xlabel('Latitud (°N)')
    ax.set_ylabel('Profundidad (m)')
    plt.colorbar(p, ax=ax, label=ds[var].attrs.get('units', ''))

    date = pd.to_datetime(ds.time.values[time_index])
    ax.set_title(title or f'{var} — lon={lon_target}°  —  {date:%Y-%m-%d %H UTC}')
    plt.tight_layout()
    ds.close()
    if show:
        plt.show()
    return fig


# ---------------------------------------------------------------------------
# Viento (ERA5 / GFS)
# ---------------------------------------------------------------------------

def _wind_var_names(source: str) -> dict:
    if source == 'era5':
        return dict(u_var='u10', v_var='v10',
                    lon_var='longitude', lat_var='latitude',
                    time_var='valid_time')
    if source == 'gfs':
        return dict(u_var='u10', v_var='v10',
                    lon_var='longitude', lat_var='latitude',
                    time_var='time')
    if source == 'fvcom':
        # Archivo ya convertido (XLAT/XLONG, U10/V10)
        return dict(u_var='U10', v_var='V10',
                    lon_var='XLONG', lat_var='XLAT',
                    time_var='time')
    raise ValueError(f"source debe ser 'era5', 'gfs' o 'fvcom', no {source!r}")


def plot_wind(
    nc_file: Union[str, Path],
    source: str = 'gfs',
    time_index: int = 0,
    skip: int = 2,
    vmax: float = 15.0,
    figsize: tuple = (10, 10),
    title: Optional[str] = None,
    show: bool = True,
) -> plt.Figure:
    """Pinta magnitud y flechas de viento a 10 m de un NetCDF.

    Parameters
    ----------
    source : 'era5', 'gfs' o 'fvcom' (post-conversión).
    skip   : submuestreo de las flechas (1 = todas).
    vmax   : tope de la barra de color (m/s).
    """
    _apply_style()
    names = _wind_var_names(source)
    ds = xr.open_dataset(nc_file)

    u_da = ds[names['u_var']].isel({names['time_var']: time_index})
    v_da = ds[names['v_var']].isel({names['time_var']: time_index})

    if source == 'fvcom':
        # XLAT, XLONG son 2D
        lon2d = ds[names['lon_var']].values
        lat2d = ds[names['lat_var']].values
        u = u_da.values
        v = v_da.values
    else:
        lon = ds[names['lon_var']].values
        lat = ds[names['lat_var']].values
        lon = _lon_to_180(lon)
        lon2d, lat2d = np.meshgrid(lon, lat)
        u = u_da.values
        v = v_da.values

    mag = np.sqrt(u**2 + v**2)
    date = pd.to_datetime(ds[names['time_var']].values[time_index])

    fig, ax = plt.subplots(figsize=figsize, subplot_kw=_subplot_kw())
    p = ax.pcolormesh(lon2d, lat2d, mag,
                      cmap=CMAP_W, shading='auto', vmin=0, vmax=vmax,
                      **_transform_kw())
    plt.colorbar(p, ax=ax, orientation='vertical', pad=0.05,
                 label='|viento| (m/s)', shrink=0.7)

    ax.quiver(lon2d[::skip, ::skip] if lon2d.ndim == 2 else lon2d[::skip],
              lat2d[::skip, ::skip] if lat2d.ndim == 2 else lat2d[::skip],
              u[::skip, ::skip], v[::skip, ::skip],
              scale=200, color='white', alpha=0.8, width=0.003,
              **_transform_kw())

    extent = [float(lon2d.min()), float(lon2d.max()),
              float(lat2d.min()), float(lat2d.max())]
    _add_map_features(ax, extent=extent)

    ax.set_title(title or f'{source.upper()} viento 10 m — {date:%Y-%m-%d %H UTC}')
    plt.tight_layout()
    ds.close()
    if show:
        plt.show()
    return fig


def plot_wind_timeseries(
    nc_file: Union[str, Path],
    lon_target: float,
    lat_target: float,
    source: str = 'gfs',
    figsize: tuple = (10, 4),
    title: Optional[str] = None,
    show: bool = True,
) -> plt.Figure:
    """Serie temporal de u10/v10 y magnitud en el punto más cercano."""
    _apply_style()
    names = _wind_var_names(source)
    ds = xr.open_dataset(nc_file)

    if source == 'fvcom':
        # Buscar el índice (i, j) del nodo más cercano en la malla 2D
        lon2d = ds[names['lon_var']].values
        lat2d = ds[names['lat_var']].values
        dist2 = (lon2d - lon_target) ** 2 + (lat2d - lat_target) ** 2
        i, j = np.unravel_index(np.argmin(dist2), dist2.shape)
        u = ds[names['u_var']][:, i, j].values
        v = ds[names['v_var']][:, i, j].values
    else:
        lon_vals = _lon_to_180(ds[names['lon_var']].values)
        # Selección por coordenada más cercana
        ds_pt = ds.assign_coords({names['lon_var']: lon_vals}).sortby(names['lon_var'])
        ds_pt = ds_pt.sel({names['lon_var']: lon_target,
                           names['lat_var']: lat_target}, method='nearest')
        u = ds_pt[names['u_var']].values
        v = ds_pt[names['v_var']].values

    times = pd.to_datetime(ds[names['time_var']].values)
    mag = np.sqrt(u**2 + v**2)

    fig, ax = plt.subplots(figsize=figsize)
    ax.plot(times, u,   label='u10', linewidth=1.2)
    ax.plot(times, v,   label='v10', linewidth=1.2)
    ax.plot(times, mag, label='|w|', color='k', linewidth=1.5)
    ax.axhline(0, color='gray', linewidth=0.5)
    ax.set_ylabel('m s⁻¹')
    ax.legend(loc='best')
    ax.grid(alpha=0.3)
    fig.autofmt_xdate()
    ax.set_title(title or f'{source.upper()} — viento 10 m en '
                          f'({lon_target}°, {lat_target}°)')
    plt.tight_layout()
    ds.close()
    if show:
        plt.show()
    return fig


if __name__ == '__main__':
    import argparse
    p = argparse.ArgumentParser(description=__doc__)
    sub = p.add_subparsers(dest='cmd', required=True)

    h = sub.add_parser('hycom', help='Surface HYCOM plot')
    h.add_argument('file')
    h.add_argument('--time',  type=int, default=0)
    h.add_argument('--depth', type=int, default=0)

    w = sub.add_parser('wind', help='Wind plot')
    w.add_argument('file')
    w.add_argument('--source', choices=['era5', 'gfs', 'fvcom'], default='gfs')
    w.add_argument('--time', type=int, default=0)

    args = p.parse_args()
    if args.cmd == 'hycom':
        plot_hycom_surface(args.file, time_index=args.time,
                           depth_index=args.depth)
    elif args.cmd == 'wind':
        plot_wind(args.file, source=args.source, time_index=args.time)