#!/usr/bin/env python3
"""
nc_boxplot.py
-------------
Create boxplots of spatial distributions for selected times from 3D variables
(e.g., Var(lon, lat, time)) in a NetCDF file.

USAGE (examples)
----------------
# Example 1: Pick variables Var1, Var2; pick times by index notation t3 t4 t5
python nc_boxplot.py data.nc Var1 Var2 --times t3 t4 t5

# Example 2: Pick all 3D variables automatically; pick times by zero-based indices
python nc_boxplot.py data.nc --times 2 5 10

# Example 3: Pick times by datetime (requires the file to have a time coordinate)
python nc_boxplot.py data.nc Var1 --times 2001-06-01 2001-06-10T12:00

# Example 4: Subset a region before computing boxplots
python nc_boxplot.py data.nc Var1 Var2 --times t3 t6 --region 8 24 102 110 --outdir figs

NOTES
-----
- The script will create one PNG per variable, saved as <var>_boxplot.png in --outdir.
- Time selection accepts:
  * 'tK' (1-based, e.g., t3 -> index 2),
  * integers (0-based indices),
  * ISO-like date/time strings (e.g., '2001-06-01' or '2001-06-10T12:00').
- Missing values are dropped automatically.
- Works when the time dimension is named 'time' or 't'. Spatial dims can be 'lat/latitude' and 'lon/longitude'.
- X-axis labels are formatted as dd/mm (day/month) when a datetime coordinate is available.
"""

import argparse
from pathlib import Path
import numpy as np
import xarray as xr
import matplotlib.pyplot as plt

# -------- Utilities --------

def _detect_dims(da):
    """Return (time_dim, lat_dim, lon_dim, space_dims) for a 3D DataArray."""
    dims = list(da.dims)
    dimset = set(dims)

    # Heuristics for time, lat, lon
    time_dim = None
    for cand in ('time', 't'):
        if cand in dimset:
            time_dim = cand
            break

    lat_dim = None
    for cand in ('lat', 'latitude', 'y'):
        if cand in dimset:
            lat_dim = cand
            break

    lon_dim = None
    for cand in ('lon', 'longitude', 'x'):
        if cand in dimset:
            lon_dim = cand
            break

    # Fallbacks if lat/lon not matched by name: assume last two non-time dims
    if time_dim is None and len(dims) == 3:
        # try to infer: if there's a coord with datetime dtype
        for d in dims:
            coord = da.coords.get(d, None)
            if coord is not None and np.issubdtype(coord.dtype, np.datetime64):
                time_dim = d
                break

    if lat_dim is None or lon_dim is None:
        others = [d for d in dims if d != time_dim]
        if len(others) >= 2 and (lat_dim is None or lon_dim is None):
            # Assume the first two of the remaining are spatial (order doesn't matter for stacking)
            if lat_dim is None:
                lat_dim = others[0]
            if lon_dim is None:
                lon_dim = others[1]

    if time_dim is None:
        raise ValueError(f"Could not detect a time dimension among {dims}. Tried 'time' and 't'.")

    space_dims = tuple([d for d in dims if d != time_dim])
    if len(space_dims) != 2:
        raise ValueError(f"Expected 3D variable with exactly 2 spatial dims and 1 time dim, got dims {dims}")

    return time_dim, lat_dim, lon_dim, space_dims


def _parse_time_selection(ds, time_dim, time_tokens):
    """
    Parse a list of time tokens into coordinate positions.
    time_tokens can contain:
      - 'tK' (1-based index, e.g., 't3' -> 2)
      - integer strings (0-based index)
      - ISO-like date/time strings to be matched exactly in the time coordinate
    Returns:
      - indices_or_coords: list suitable for xarray .isel(time=...) or .sel(time=...)
      - use_sel: bool, True if we should use .sel (datetime tokens), False for .isel
    """
    coord_values = ds[time_dim].values

    idx_list = []
    sel_list = []
    use_sel = False

    def is_int_like(s):
        try:
            int(s)
            return True
        except Exception:
            return False

    for tok in time_tokens:
        tok_str = str(tok).strip()

        # 'tK' notation (1-based)
        if (tok_str.lower().startswith('t') and is_int_like(tok_str[1:])):
            k = int(tok_str[1:]) - 1
            if k < 0 or k >= coord_values.shape[0]:
                raise IndexError(f"time index from token '{tok_str}' out of range [0, {coord_values.shape[0]-1}]")
            idx_list.append(k)
            continue

        # zero-based integer index
        if is_int_like(tok_str):
            k = int(tok_str)
            if k < 0 or k >= coord_values.shape[0]:
                raise IndexError(f"time index {k} out of range [0, {coord_values.shape[0]-1}]")
            idx_list.append(k)
            continue

        # Otherwise, treat as datetime-like string: use .sel
        use_sel = True
        sel_list.append(tok_str)

    if use_sel and idx_list:
        raise ValueError("Do not mix index-based tokens (e.g., t3, 5) with datetime tokens in one call.")

    if use_sel:
        return sel_list, True
    else:
        return idx_list, False


def _flatten_valid(arr):
    """Return a 1D numpy array with NaNs and infs removed."""
    a = np.asarray(arr).ravel()
    return a[np.isfinite(a)]


def _format_label_from_coord_val(cv):
    """
    Format a time coordinate value to 'dd/mm' if possible.
    Supports: Python datetime, cftime datetime (has .strftime), numpy.datetime64.
    Fallback: str(cv).
    """
    # cftime / python datetime
    try:
        if hasattr(cv, 'strftime'):
            return cv.strftime('%d/%m')
    except Exception:
        pass

    # numpy.datetime64
    try:
        if isinstance(cv, np.datetime64) or 'datetime64' in str(type(cv)):
            s = np.datetime_as_string(cv, unit='D')  # 'YYYY-MM-DD'
            parts = s.split('-')
            if len(parts) >= 3:
                y, m, d = parts[0], parts[1], parts[2]
                return f"{d}/{m}"
            return s
    except Exception:
        pass

    # fallback
    return str(cv)


def boxplot_for_variable(da, var_name, times, outdir):
    """Create a single PNG boxplot for one variable across selected times."""
    time_dim, lat_dim, lon_dim, space_dims = _detect_dims(da)

    # Parse times (tokens only determine selection method; labels will come from actual coords)
    indices_or_coords, use_sel = _parse_time_selection(da.to_dataset(name=var_name), time_dim, list(times))

    # Collect data vectors and labels
    data = []
    labels = []
    for tok in indices_or_coords:
        if use_sel:
            slice_da = da.sel({time_dim: tok})
        else:
            slice_da = da.isel({time_dim: tok})

        # Label from the actual time coordinate value
        try:
            tval = slice_da[time_dim].values
            if np.ndim(tval) > 0:
                tval = tval.item()
            labels.append(_format_label_from_coord_val(tval))
        except Exception:
            labels.append(str(tok))

        # Stack spatial dims to 1D and drop missing
        stacked = slice_da.stack(space=space_dims).values
        data.append(_flatten_valid(stacked))

    # Plot
    units = da.attrs.get('units', '')
    fig, ax = plt.subplots(figsize=(8, 5), dpi=150)
    ax.boxplot(
        data,
        labels=labels,
        showmeans=True,
        whis=1.5
    )
    ax.set_title(f"{var_name} — spatial distribution at selected times")
    ax.set_xlabel("time (dd/mm)")
    if units:
        ax.set_ylabel(f"{var_name} [{units}]")
    else:
        ax.set_ylabel(var_name)

    fig.tight_layout()
    outpath = Path(outdir) / f"{var_name}_boxplot.png"
    fig.savefig(outpath)
    plt.close(fig)
    return str(outpath)


def main():
    import warnings
    warnings.filterwarnings("ignore", category=FutureWarning)

    p = argparse.ArgumentParser(description="Create spatial boxplots for selected times from 3D variables in a NetCDF file.")
    p.add_argument("ncfile", help="Path to NetCDF file")
    p.add_argument("variables", nargs="*", help="Variable names (3D). If omitted, all 3D variables will be used.")
    p.add_argument("--times", nargs="+", required=True,
                   help="Time selection tokens (e.g., t3 t4 5 10 or ISO datetimes like 2001-06-01). "
                        "Do not mix index-based and datetime tokens in one call.")
    p.add_argument("--region", nargs=4, type=float, metavar=("latmin", "latmax", "lonmin", "lonmax"),
                   help="Optional lat/lon box to subset before plotting.")
    p.add_argument("--outdir", default=".", help="Output directory for PNGs (default: current directory)")
    args = p.parse_args()

    ncpath = Path(args.ncfile)
    if not ncpath.exists():
        raise FileNotFoundError(ncpath)

    outdir = Path(args.outdir)
    outdir.mkdir(parents=True, exist_ok=True)

    # Open file
    ds = xr.open_dataset(ncpath)

    # Optional regional subset if the dataset carries lat/lon coords
    if args.region is not None:
        latmin, latmax, lonmin, lonmax = args.region
        # Try common coord names
        lat_name = None
        lon_name = None
        for cand in ("lat", "latitude", "y"):
            if cand in ds.coords:
                lat_name = cand
                break
        for cand in ("lon", "longitude", "x"):
            if cand in ds.coords:
                lon_name = cand
                break

        if lat_name and lon_name:
            ds = ds.sel({lat_name: slice(latmin, latmax)})
            ds = ds.sel({lon_name: slice(lonmin, lonmax)})
        else:
            print("Warning: lat/lon coordinates not found; skipping regional subset.")

    # Determine variables to plot
    varnames = args.variables
    if not varnames:
        # pick all 3D data variables
        varnames = []
        for v, da in ds.data_vars.items():
            if len(da.dims) == 3:
                varnames.append(v)
        if not varnames:
            raise ValueError("No 3D variables found. Specify variable names explicitly.")

    produced = []
    for v in varnames:
        if v not in ds.data_vars:
            print(f"Skip '{v}': not found in dataset.")
            continue
        da = ds[v]
        if len(da.dims) != 3:
            print(f"Skip '{v}': expected 3D (lon,lat,time)-like, got dims {da.dims}.")
            continue
        try:
            out = boxplot_for_variable(da, v, args.times, outdir)
            produced.append(out)
            print(f"Saved: {out}")
        except Exception as e:
            print(f"Failed for variable '{v}': {e}")

    ds.close()

    if not produced:
        raise SystemExit("No figures produced.")
    else:
        print("Done.")

if __name__ == "__main__":
    main()