"""
Pull ERA5 variables from NSF NCAR GDEX (dataset d633000) via the GDEX
THREDDS Data Server (TDS) using OPeNDAP -- no download of full files,
no login required for plain reads.

Requires:
    pip install siphon xarray "pydap>=3.4" pandas netCDF4

Uses the "pydap" backend instead of netCDF4/libcurl for the OPeNDAP
transport. pydap is a pure-Python DAP client, so it installs cleanly
everywhere (no compiled libnetcdf/libcurl-with-DAP-support needed) --
handy on Windows or in minimal environments. Functionally it reads the
exact same data as the netCDF4 engine. (netCDF4 is still listed below
because xarray uses it to *write* the final local .nc file -- swap
to h5netcdf, or ds_out.to_zarr(...), if you'd rather drop it too.)

How it works
------------
GDEX publishes a THREDDS catalog for every monthly file group, e.g.:
    https://tds.gdex.ucar.edu/thredds/catalog/files/g/d633000/e5.oper.an.sfc/202006/catalog.xml
Rather than hard-coding each file's exact name (the "analysis" and
"forecast mean-flux" products use different date-range suffixes),
this script asks the catalog for the file that matches each variable's
ECMWF table/parameter code, then opens it lazily via OPeNDAP.

Your variables split into two GDEX products:
  - msl, u10, v10, t2, d2            -> e5.oper.an.sfc        (analysis, table 128)
  - metss, mntss, mtpr, mslhf,
    msnlwrf, msdwlwrf, mer,
    msnswrf, msshf                    -> e5.oper.fc.sfc.meanflux (forecast, table 235)
"""

import re
import warnings

# silence library warnings (pydap / xarray / siphon) so only the
# script's own progress messages are printed
warnings.filterwarnings("ignore")

import xarray as xr
import pandas as pd
from siphon.catalog import TDSCatalog

CATALOG_BASE = "https://tds.gdex.ucar.edu/thredds/catalog/files/g/d633000"

# variable -> (GDEX product directory, ECMWF table, parameter code, netCDF var name)
VARIABLES = {
    "msl":      ("e5.oper.an.sfc",          128, 151, "MSL"),
    "u10":      ("e5.oper.an.sfc",          128, 165, "VAR_10U"),
    "v10":      ("e5.oper.an.sfc",          128, 166, "VAR_10V"),
    "t2":       ("e5.oper.an.sfc",          128, 167, "VAR_2T"),
    "d2":       ("e5.oper.an.sfc",          128, 168, "VAR_2D"),
    "msshf":    ("e5.oper.fc.sfc.meanflux", 235,  33, "MSSHF"),
    "mslhf":    ("e5.oper.fc.sfc.meanflux", 235,  34, "MSLHF"),
    "msdwlwrf": ("e5.oper.fc.sfc.meanflux", 235,  36, "MSDWLWRF"),
    "msnswrf":  ("e5.oper.fc.sfc.meanflux", 235,  37, "MSNSWRF"),
    "msnlwrf":  ("e5.oper.fc.sfc.meanflux", 235,  38, "MSNLWRF"),
    "metss":    ("e5.oper.fc.sfc.meanflux", 235,  41, "METSS"),
    "mntss":    ("e5.oper.fc.sfc.meanflux", 235,  42, "MNTSS"),
    "mer":      ("e5.oper.fc.sfc.meanflux", 235,  43, "MER"),
    "mtpr":     ("e5.oper.fc.sfc.meanflux", 235,  55, "MTPR"),
}


def opendap_url_for(varname: str, year: int, month: int) -> str:
    """Look up the exact OPeNDAP URL for one variable's monthly file."""
    product, table, code, _ = VARIABLES[varname]
    yyyymm = f"{year:04d}{month:02d}"
    cat_url = f"{CATALOG_BASE}/{product}/{yyyymm}/catalog.xml"

    cat = TDSCatalog(cat_url)
    pattern = re.compile(rf"^{re.escape(product)}\.{table}_{code:03d}_")
    matches = [name for name in cat.datasets if pattern.match(name)]

    if not matches:
        raise FileNotFoundError(
            f"No file found for '{varname}' ({table}_{code:03d}) in {yyyymm}. "
            f"Checked: {cat_url}"
        )
    return cat.datasets[matches[0]].access_urls["OPENDAP"]


def _normalize_lon_bounds(lon_min: float, lon_max: float) -> tuple:
    """
    GDEX's ERA5 files use 0-360 longitude. Accept bbox input in either
    -180/180 or 0-360 and convert to 0-360 so .sel() lines up with the
    file's actual coordinate values (mirrors what Wilkin's MATLAB code
    does automatically for negative-longitude bboxes).
    """
    lon_min = lon_min % 360
    lon_max = lon_max % 360
    return lon_min, lon_max


def load_variable(
    varname: str,
    start: str,
    end: str,
    bbox: tuple | None = None,
) -> xr.DataArray:
    """
    Load one variable for the EXACT time span [start, end] (inclusive)
    and EXACT spatial bbox via OPeNDAP.

    start, end : str
        Timestamps, e.g. "2025-03-01T00:00" and "2025-03-31T23:00".
        Only the calendar month(s) spanned are ever opened/read.
    bbox : (lon_min, lon_max, lat_min, lat_max) in degrees, or None for
        the full global field. Longitude may be given in -180/180 or
        0-360 -- it's normalized to GDEX's 0-360 convention internally.
    """
    _, _, _, ncvar = VARIABLES[varname]
    t0, t1 = pd.Timestamp(start), pd.Timestamp(end)
    months = pd.period_range(t0, t1, freq="M")

    pieces = []
    for p in months:
        url = opendap_url_for(varname, p.year, p.month)
        print(f"  {varname}: opening {url}")
        ds = xr.open_dataset(url, engine="pydap")  # lazy -- no data pulled yet
        da = ds[ncvar]

        if bbox is not None:
            lon_min, lon_max, lat_min, lat_max = bbox
            lon_min, lon_max = _normalize_lon_bounds(lon_min, lon_max)

            # latitude may be ascending or descending depending on file;
            # sel() needs the slice given in the same order as the coord
            lat_vals = da["latitude"].values
            lat_slice = (
                slice(lat_max, lat_min) if lat_vals[0] > lat_vals[-1]
                else slice(lat_min, lat_max)
            )

            if lon_min <= lon_max:
                da = da.sel(longitude=slice(lon_min, lon_max), latitude=lat_slice)
            else:
                # bbox straddles the 0/360 seam -- wrap around
                west = da.sel(longitude=slice(lon_min, 360), latitude=lat_slice)
                east = da.sel(longitude=slice(0, lon_max), latitude=lat_slice)
                da = xr.concat([west, east], dim="longitude")

        pieces.append(da)

    full = xr.concat(pieces, dim="time") if len(pieces) > 1 else pieces[0]
    
    try:
        return full.sel(time=slice(t0, t1))  # trim to the exact requested time span
    except KeyError:
        # find if the dataset uses 'forecast_initial_time' instead of 'time' for the time dimension
        return full.sel(forecast_initial_time=slice(t0, t1))


def main():
    variables = [
        "msl", "u10", "v10", "t2", "d2",
        "metss", "mntss", "mtpr", "mslhf",
        "msnlwrf", "msdwlwrf", "mer", "msnswrf", "msshf",
    ]

    # ---- EXACT time span and spatial domain -- edit these -----------
    start, end = "2025-03-01T00:00", "2025-03-31T23:00"   # hourly, full month of March 2025
    bbox = (-80.0, -67.25, 33.25, 42.0)                    # lon_min, lon_max, lat_min, lat_max
                                                            # (either -180/180 or 0-360 is fine)
    out_file = "era5_202503_subset.nc"
    # ------------------------------------------------------------------

    data = {}
    for v in variables:
        print(f"Fetching {v} ...")
        data[v] = load_variable(v, start, end, bbox=bbox)

    ds_out = xr.Dataset(data)
    print(ds_out)

    ds_out.to_netcdf(out_file)
    print(f"Saved subset to {out_file}")


if __name__ == "__main__":
    main()