#!/usr/bin/env python3
"""Align gridded Tx observations to the EC reforecast sampling.

For every (hindcast year, init date, lead 0..45) the matching observed Tx is
extracted, giving an array shaped exactly like the forecasts:

    tx_obs(year, init, lead, lat, lon)      20 x 104 x 46 x 65 x 41

`sev`/`EHD` are deliberately NOT stored here: their threshold must be rebuilt from
the training years of each fold, so they are derived downstream from
Qp_tx_obs_<fold>.nc. What is stored is the doy of every (init, lead) so the right
threshold slice can be picked without recomputing calendars.

The observation file uses a 365_day calendar, so day-of-year arithmetic is exact.
"""
import os
import numpy as np
import xarray as xr

OBS = "/home/tanpv/Data/QG.25.196/OBS/Tx_EC_0.25_2002-2023.nc"
HDATES = "/home/tanpv/Data/QG.25.196/hdate_2023.txt"
OUT = "/home/tanpv/Progs/claude/HW_EC_10/outputs/OBS_4EC15/tx_obs_aligned.nc"

YEARS = list(range(2003, 2023))
NLEAD = 46                       # lead 0..45
OBS_YEAR0 = 2002                 # first year in the observation file
MLEN = [31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31]   # 365-day calendar
CUM = np.concatenate([[0], np.cumsum(MLEN)])


def mmdd_to_doy(mmdd):
    mm, dd = int(mmdd[:2]), int(mmdd[3:])
    return CUM[mm - 1] + dd          # 1-based day of year


def main():
    with open(HDATES) as fh:
        inits = sorted({t for t in fh.read().split() if len(t) == 5 and t[2] == "-"})
    print(f"{len(inits)} init dates, {len(YEARS)} years, {NLEAD} leads")

    ds = xr.open_dataset(OBS, decode_times=xr.coders.CFDatetimeCoder(use_cftime=True))
    tx = ds.tasmax.values.astype("float32")       # (ndays, lat, lon), 365-day calendar
    nday, nlat, nlon = tx.shape
    print(f"obs {tx.shape}, {nday/365:.0f} years from {OBS_YEAR0}")

    init_doy = np.array([mmdd_to_doy(s) for s in inits])          # (ninit,)
    # absolute day index into the obs array for every (year, init, lead)
    base = (np.array(YEARS) - OBS_YEAR0)[:, None, None] * 365 \
        + (init_doy - 1)[None, :, None] \
        + np.arange(NLEAD)[None, None, :]                          # (nyear, ninit, nlead)

    if base.min() < 0 or base.max() >= nday:
        raise RuntimeError(f"obs record too short: need [{base.min()},{base.max()}], "
                           f"have [0,{nday-1}]")

    out_tx = tx[base.ravel()].reshape(len(YEARS), len(inits), NLEAD, nlat, nlon)
    doy = (base % 365) + 1                                         # 1..365

    res = xr.Dataset(
        {
            "tx_obs": (("year", "init", "lead", "lat", "lon"), out_tx,
                       {"units": "C", "long_name": "observed daily maximum temperature"}),
            "doy": (("year", "init", "lead"), doy.astype("int16"),
                    {"long_name": "day of year of the valid date (365-day calendar)"}),
            "init_doy": (("init",), init_doy.astype("int16")),
        },
        coords={"year": YEARS, "init": [s for s in inits],
                "lead": np.arange(NLEAD), "lat": ds.lat.values, "lon": ds.lon.values},
    )
    res.attrs["source"] = OBS
    res.attrs["note"] = ("lead 0 is the init date itself; lead 45 is +45 days. "
                         "sev/EHD are derived downstream with fold-specific thresholds.")

    os.makedirs(os.path.dirname(OUT), exist_ok=True)
    res.to_netcdf(OUT, encoding={"tx_obs": {"zlib": True, "complevel": 4,
                                            "dtype": "float32"}})
    print(f"wrote {OUT} ({os.path.getsize(OUT)/1e9:.2f} GB)")

    finite = np.isfinite(out_tx)
    print(f"  samples per point per lead: {len(YEARS)*len(inits)}")
    print(f"  finite fraction: {finite.mean():.4f}")
    print(f"  tx range: {np.nanmin(out_tx):.1f} .. {np.nanmax(out_tx):.1f} C")


if __name__ == "__main__":
    main()