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

- All chosen 3D variables are plotted on the SAME chart (grouped by time).
- X-axis labels are formatted as dd/mm (no year) when a datetime coordinate is available.
- Output filename keeps only the DATE-like tag found in the input filename.

USAGE (examples)
----------------
# Example 1: Plot tp5day and tpx5d together at t3 t4 t5 t10 t20
python nc_boxplot.py Prec_5days_2025-08-16.nc tp5day tpx5d --times t3 t4 t5 t10 t20

# Example 2: Plot ALL 3D variables together; select times by zero-based indices
python nc_boxplot.py data.nc --times 2 5 10

# Example 3: Select by datetime strings (must match the time coordinate)
python nc_boxplot.py data.nc Var1 --times 2001-06-01 2001-06-10T12:00

# Optional: subset region before plotting
python nc_boxplot.py data.nc Var1 Var2 --times t3 t6 --region 8 24 102 110 --outdir figs
"""

import argparse
import re
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 _extract_date_tag_from_filename(ncfile: Path) -> str:
    """
    Return a date-like tag from filename, e.g. '2025-08-16' from 'Prec_5days_2025-08-16.nc'.
    Falls back to stem if no date found.
    """
    base = ncfile.stem
    # Try patterns: YYYY[-_]MM[-_]DD, YYYYMMDD, YYYY[-_]MM
    for pat in (r'(\d{4}[-_]\d{2}[-_]\d{2})', r'(\d{8})', r'(\d{4}[-_]\d{2})'):
        m = re.search(pat, base)
        if m:
            tag = m.group(1)
            if '_' in tag:
                tag = tag.replace('_', '-')
            if len(tag) == 8 and tag.isdigit():
                tag = f"{tag[0:4]}-{tag[4:6]}-{tag[6:8]}"
            return tag
    return base


def grouped_boxplot(ds, varnames, time_tokens, outdir: Path, ncfile: Path):
    """
    Make a single grouped boxplot figure for multiple variables over selected times.
    """
    # Prepare data arrays per variable and time, and X tick labels from the actual time coords
    labels = None            # dd/mm per time in order
    data_by_var = []         # list of list: [var_idx][time_idx] -> 1D values
    units_candidates = []

    # We do parsing of times using first variable as reference for time coord
    ref_da = ds[varnames[0]]
    time_dim, lat_dim, lon_dim, space_dims = _detect_dims(ref_da)
    indices_or_coords, use_sel = _parse_time_selection(ref_da.to_dataset(name=varnames[0]), time_dim, list(time_tokens))

    # Build labels (dd/mm) using the reference variable's time coord
    labels = []
    for tok in indices_or_coords:
        if use_sel:
            slice_da = ref_da.sel({time_dim: tok})
        else:
            slice_da = ref_da.isel({time_dim: tok})
        tval = slice_da[time_dim].values
        if np.ndim(tval) > 0:
            tval = tval.item()
        labels.append(_format_label_from_coord_val(tval))

    # Collect values for each variable
    for v in varnames:
        da = ds[v]
        # Check dims
        _td, _la, _lo, _sp = _detect_dims(da)  # raise early if incompatible
        units_candidates.append(da.attrs.get('units', ''))
        series_per_time = []
        for tok in indices_or_coords:
            if use_sel:
                sda = da.sel({time_dim: tok})
            else:
                sda = da.isel({time_dim: tok})
            vals = sda.stack(space=_sp).values
            vals = _flatten_valid(vals)
            series_per_time.append(vals)
        data_by_var.append(series_per_time)

    # Compute positions for grouped boxplots
    T = len(labels)
    V = len(varnames)
    base = np.arange(1, T+1)
    group_width = min(0.8, 0.15 * V)  # total width per time group
    step = group_width / max(V, 1)
    offsets = np.linspace(-group_width/2 + step/2, group_width/2 - step/2, V) if V>1 else np.array([0.0])

    positions = []
    data_flat = []
    for j in range(T):
        for k in range(V):
            positions.append(base[j] + offsets[k])
            data_flat.append(data_by_var[k][j])

    # Draw figure
    fig, ax = plt.subplots(figsize=(10, 5), dpi=150)
    bp = ax.boxplot(data_flat, positions=positions, widths=step*0.9, showmeans=True, whis=1.5)

    # X ticks at the center of each group with dd/mm labels
    ax.set_xticks(base)
    ax.set_xticklabels(labels)
    # Title and labels
    date_tag = _extract_date_tag_from_filename(ncfile)
    ax.set_title(f"Spatial boxplots ({', '.join(varnames)}) — {date_tag}")
    # y-label: prefer first non-empty units
    y_units = next((u for u in units_candidates if u), '')
    ax.set_ylabel(f"Value [{y_units}]" if y_units else "Value")
    ax.set_xlabel("time (dd/mm)")

    # Overlay mean markers per variable with distinct marker shapes (no explicit colors)
    marker_cycle = ['o', 's', '^', 'D', 'P', 'X', '*', 'v', '<', '>']
    handles = []
    for k, v in enumerate(varnames):
        mk = marker_cycle[k % len(marker_cycle)]
        pos_k = [base[j] + offsets[k] for j in range(T)]
        means_k = [float(np.nanmean(data_by_var[k][j])) if len(data_by_var[k][j])>0 else np.nan for j in range(T)]
        h, = ax.plot(pos_k, means_k, linestyle='None', marker=mk, label=v)
        handles.append(h)

    if V > 1:
        ax.legend(handles=handles, title="Variable", loc='best')

    fig.tight_layout()

    # Output path uses only the date-like part of input filename
    outname = f"{date_tag}_boxplot.png"
    outpath = outdir / outname
    fig.savefig(outpath)
    plt.close(fig)
    return str(outpath)


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

    p = argparse.ArgumentParser(description="Create a single grouped boxplot for multiple 3D variables over selected times.")
    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 PNG (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.")

    try:
        out = grouped_boxplot(ds, varnames, args.times, outdir, ncpath)
        print(f"Saved: {out}")
    except Exception as e:
        raise SystemExit(f"Failed to create figure: {e}")
    finally:
        ds.close()

    print("Done.")

if __name__ == "__main__":
    main()