#!/usr/bin/env python3
"""Shared paths, calendar conventions and I/O helpers for the EHD S2S project.

Verified data conventions (see developing.txt for the evidence):
  * lt index k = 0..45; the valid calendar day is init_date + k (init day included).
  * The `time` coordinate stored inside EC_Combined / OBS_4EC files is written one
    day late; it is a label only. Pairing is by index.
  * Raw ECMWF files (H500, T850, T_plev) carry 47 steps; steps 0..45 line up with
    EC_Combined lt 0..45, so the LAST step is dropped.
  * Raw ECMWF files store latitude descending (24 -> 8); EC_Combined stores it
    ascending (8 -> 24). Raw fields must be flipped.
  * OBS_4EC follows the real calendar and does contain 29 Feb.
    OBS/Tx_EC_0.25_2002-2023.nc is the same series with 29 Feb removed.
"""

from __future__ import annotations

import datetime as dt
import glob
import os
import re

import numpy as np
from netCDF4 import Dataset

ROOT = "/home/tanpv/Progs/claude/HW"
OUT = os.path.join(ROOT, "outputs")
DATA = "/home/tanpv/Data/QG.25.196"

EC_COMBINED = os.path.join(DATA, "EC_Combined")
OBS_4EC = os.path.join(DATA, "OBS_4EC")
OBS_TX_CONT = os.path.join(DATA, "OBS", "Tx_EC_0.25_2002-2023.nc")
MASK_FILE = os.path.join(DATA, "OBS", "Mask_VN_S2S_EC_0.25.nc")
SHAPE_7REG = "/home/tanpv/GrADS_SCR/VN_SubRegs1/Vietnam7"

YEARS = list(range(2003, 2023))
# Fixed validation years -- one per five-year block so the split spans the whole
# record instead of clustering. Keep this list frozen across all experiments.
VALID_YEARS = [2006, 2011, 2016, 2021]
TRAIN_YEARS = [y for y in YEARS if y not in VALID_YEARS]

NLT = 46
HOT_DAY_C = 35.0
PCTL = 90.0
WIN_HALF = 15  # 31-day climatological window

# Week name -> lt indices (0-based), per section 4.1 of HW_HPC_1.
WEEKS = [
    ("W1", np.arange(0, 11)),
    ("W2", np.arange(11, 18)),
    ("W3", np.arange(18, 25)),
    ("W4", np.arange(25, 32)),
    ("W5", np.arange(32, 39)),
    ("W6", np.arange(39, 46)),
]
WEEK_NAMES = [w for w, _ in WEEKS]

REGION_NAMES = {
    1: "S1 Tay Bac", 2: "S2 Dong Bac", 3: "S3 DB Bac Bo", 4: "S4 Bac Trung Bo",
    5: "S5 Nam Trung Bo", 6: "S6 Tay Nguyen", 7: "S7 Nam Bo",
}

_INIT_RE = re.compile(r"(\d{4})-(\d{2})-(\d{2})")


def tx_thres(qp):
    """Tx_thres from Tx90p, per section 2 of HW_HPC_1 and its figure 1.

    EHD = 1 when Tx >= 35 C (hot day) OR Tx >= Tx90p (heat-wave day), which is
    the union of the two definitions and collapses to the single threshold
    min(Tx90p, 35): the 90th percentile in the cool season, capped at 35 C
    through the hot season where Tx90p already sits above it.
    """
    return np.minimum(qp, HOT_DAY_C)


def is_capable(thr):
    """Hot-season frame: cells/days where the 35 C cap sets the threshold.

    Equivalent to Tx90p >= 35 C, so this selects the same (doy, point) pairs it
    always did -- the ones whose local climatology can actually reach 35 C.
    """
    return thr >= HOT_DAY_C - 1e-6


def doy_noleap(date):
    """Day of year on the 365-day convention; 29 Feb is folded onto 28 Feb (59)."""
    if date.month == 2 and date.day == 29:
        return 59
    doy = date.timetuple().tm_yday
    is_leap = date.year % 4 == 0 and (date.year % 100 != 0 or date.year % 400 == 0)
    if is_leap and doy > 60:
        doy -= 1
    return doy


def is_leap_day(date):
    return date.month == 2 and date.day == 29


def init_dates(year):
    """Sorted init dates of the 104 EC_Combined forecasts of `year`."""
    out = []
    for path in sorted(glob.glob(os.path.join(EC_COMBINED, str(year), "EC_all_*.nc"))):
        y, m, d = _INIT_RE.search(os.path.basename(path)).groups()
        out.append(dt.date(int(y), int(m), int(d)))
    return out


def ec_file(date):
    return os.path.join(EC_COMBINED, str(date.year), "EC_all_%s.nc" % date.isoformat())


def obs_file(date):
    return os.path.join(OBS_4EC, str(date.year), "Obs_%s.nc" % date.isoformat())


def valid_days(init, nlt=NLT):
    """Calendar days the 46 lead times of `init` are valid for."""
    return [init + dt.timedelta(days=k) for k in range(nlt)]


class Grid:
    """Model grid plus the land mask, in EC_Combined orientation (lat ascending)."""

    def __init__(self):
        with Dataset(MASK_FILE) as ds:
            self.lon = np.asarray(ds["lon"][:], dtype="f8")
            self.lat = np.asarray(ds["lat"][:], dtype="f8")
            vn = np.ma.filled(ds["vn"][:], np.nan)
            reg = np.ma.filled(ds["reg"][:], np.nan)
        self.nlat, self.nlon = vn.shape
        self.land2d = vn == 1.0                       # (nlat, nlon) bool
        self.iy, self.ix = np.nonzero(self.land2d)    # land point coordinates
        self.npt = self.iy.size
        self.reg = reg[self.land2d].astype(int)       # (npt,) region code 1..7
        self.reg2d = reg

    def to_points(self, field):
        """(..., nlat, nlon) -> (..., npt), keeping only masked land points."""
        return field[..., self.iy, self.ix]

    def to_grid(self, pts, fill=np.nan):
        """(..., npt) -> (..., nlat, nlon) with `fill` off the mask."""
        shape = pts.shape[:-1] + (self.nlat, self.nlon)
        out = np.full(shape, fill, dtype="f4")
        out[..., self.iy, self.ix] = pts
        return out


def read_ec_points(date, var, grid):
    """One EC_Combined variable for one init, as (46, npt) float32."""
    with Dataset(ec_file(date)) as ds:
        arr = np.ma.filled(ds[var][:], np.nan).astype("f4")
    return grid.to_points(arr)


def read_obs_points(date, var, grid):
    """One OBS_4EC variable for one init, as (46, npt) float32."""
    with Dataset(obs_file(date)) as ds:
        arr = np.ma.filled(ds[var][:], np.nan).astype("f4")
    return grid.to_points(arr)


def read_obs_continuous(grid, years=YEARS):
    """Continuous daily obs Tx (29 Feb removed) as (ndays, npt) plus its doy vector."""
    import netCDF4 as nc
    with Dataset(OBS_TX_CONT) as ds:
        tvar = ds["time"]
        dates = nc.num2date(tvar[:], tvar.units, calendar="365_day")
        keep = np.array([d.year in years for d in dates])
        arr = np.ma.filled(ds["tasmax"][keep, :, :], np.nan).astype("f4")
    dates = [d for d, k in zip(dates, keep) if k]
    doy = np.array([(d.timetuple().tm_yday) for d in dates], dtype="i4")
    return grid.to_points(arr), doy, dates


def window_doys(doy, half=WIN_HALF):
    """The 2*half+1 doys centred on `doy`, wrapping around the 365-day year."""
    return ((np.arange(doy - half, doy + half + 1) - 1) % 365) + 1


def write_nc(path, dims, variables, attrs=None):
    """Minimal netCDF4 writer.

    dims: list of (name, values_or_size)
    variables: list of (name, dims_tuple, data, attrs_dict)
    """
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with Dataset(path, "w", format="NETCDF4") as ds:
        for name, val in dims:
            if np.isscalar(val):
                ds.createDimension(name, int(val))
            else:
                val = np.asarray(val)
                ds.createDimension(name, val.size)
                v = ds.createVariable(name, val.dtype.str[1:], (name,))
                v[:] = val
        for name, vdims, data, vattrs in variables:
            data = np.asarray(data)
            v = ds.createVariable(name, "f4", vdims, zlib=True, complevel=4,
                                  fill_value=np.float32(np.nan))
            v[:] = data.astype("f4")
            for k, val in (vattrs or {}).items():
                setattr(v, k, val)
        for k, val in (attrs or {}).items():
            setattr(ds, k, val)