"""Read the single-member (control) S2S reforecast and derive daily Tx per initialization.

Forecast file convention (verified, see docs/DATA_INVENTORY.md):
  * variable `mx2t6` = max 2 m T in the last 6 h, packed short -> Kelvin (xarray unpacks).
  * 184 time steps at 6 h; first stamp = init + 06:00, last = init + 45.75 d.
  * daily Tx for forecast day d = max over the 4 six-hourly values with indices [4d:4d+4].
  * valid calendar date of forecast day d = init_date + d days (day 0 = init date).
  * axes are named longitude/latitude; we rename to lon/lat and sort latitude ascending
    so forecast and obs share an identical, comparable grid.
"""
from __future__ import annotations
import os, glob, re
import numpy as np
import pandas as pd
import xarray as xr

_DATE_RE = re.compile(r"(\d{4}-\d{2}-\d{2})")
STEPS_PER_DAY = 4
FORECAST_DAYS = 46


def init_date_from_path(path: str) -> pd.Timestamp:
    m = _DATE_RE.search(os.path.basename(path))
    if not m:
        raise ValueError(f"no init date in {path}")
    return pd.Timestamp(m.group(1))


def list_inits(forecast_tx_dir: str, year: int) -> list[str]:
    return sorted(glob.glob(os.path.join(forecast_tx_dir, str(year), "Tx_tm_6hours_cf_*.nc")))


def load_forecast_daily_tx(path: str, var: str = "mx2t6", units_out: str = "C") -> xr.DataArray:
    """Return daily Tx (lead_day 0..45, lat, lon) in degC for one init file."""
    ds = xr.open_dataset(path)
    da = ds[var].astype("float64")
    rename = {}
    if "longitude" in da.dims:
        rename["longitude"] = "lon"
    if "latitude" in da.dims:
        rename["latitude"] = "lat"
    da = da.rename(rename)
    n = da.sizes["time"]
    ndays = n // STEPS_PER_DAY
    if ndays != FORECAST_DAYS:
        raise ValueError(f"{path}: expected {FORECAST_DAYS} days, got {ndays}")
    da = da.isel(time=slice(0, ndays * STEPS_PER_DAY))
    # reshape time -> (lead_day, step) and take the daily max
    arr = da.values.reshape(ndays, STEPS_PER_DAY, da.sizes["lat"], da.sizes["lon"])
    daily = np.nanmax(arr, axis=1)
    if units_out == "C" and str(ds[var].attrs.get("units", "K")).strip() in ("K", "kelvin"):
        daily = daily - 273.15
    init = init_date_from_path(path)
    valid = pd.to_datetime([init + pd.Timedelta(days=int(d)) for d in range(ndays)])
    out = xr.DataArray(
        daily,
        coords={"lead_day": np.arange(ndays), "lat": da["lat"].values, "lon": da["lon"].values,
                "valid_date": ("lead_day", valid)},
        dims=("lead_day", "lat", "lon"), name="tx",
        attrs={"units": "C", "init_date": str(init.date()), "source": os.path.basename(path)},
    )
    # sort latitude ascending for a canonical, obs-comparable grid
    out = out.sortby("lat")
    return out