#!/usr/bin/env python3
"""Shared paths, calendar conventions and data inventory for the CFSv2 EHD project.

Verified data conventions (evidence in developing.txt, section 0):
  * cfs_flux: flxf<INIT10>.<VALID10>.grb2, 6-hourly, one 00Z init per day,
    Gaussian T126 grid 384x190 (~0.9375 deg), global, every init runs >= 90 days.
  * The file whose valid time equals the init time is the ANALYSIS: 96 records
    instead of 98 and its TMAX/TMIN carry ":anl:" (not a 6-hour extreme).
    It is therefore EXCLUDED; lead-day aggregation never needs it.
  * cfs_pgb: pgbf<VALID10>.grb2 under the init-date directory, 1.0 deg lat-lon
    360x181, 37 pressure levels. NO TMAX/TMIN and NO 2 m TMP.
  * cfs_pgb/{2001,2005..2011} are SYMLINKS to /home/tanpv/Data/NAS. Always walk
    the tree with follow-symlinks (os.walk(followlinks=True) / find -L), otherwise
    those eight years are silently reported as empty.
  * Observations (target only, never a feature): VnGAD tasmax 0.25 deg, degC,
    1961-2023, complete over 1999-2011 at all 541 land points.

Day alignment (see developing.txt 0.6) -- flux TMAX is a 6-hour maximum stamped in
UTC while VnGAD is a LOCAL-day maximum (UTC+7). Alignment A is the default:
    Tx(day D) = max over the buckets valid 06Z, 12Z, 18Z of D and 00Z of D+1
              = 07 LT of D through 07 LT of D+1
Alignment B (00Z..18Z of D) is kept as an option and is NOT available at lead 0.
The choice is settled against observations, not assumed.
"""

from __future__ import annotations

import datetime as dt
import os
import re

import numpy as np

# ---------------------------------------------------------------- paths
ROOT = "/home/tanpv/Progs/claude/HW_CFS_raw"
OUT = os.path.join(ROOT, "outputs")
LOGS = os.path.join(ROOT, "logs")
META = os.path.join(OUT, "meta")

CFS = "/home/tanpv/Data/QG.25.196/CFS_Rfcs"
FLUX_ROOT = os.path.join(CFS, "cfs_flux")
PGB_ROOT = os.path.join(CFS, "cfs_pgb")

OBS_DIR = "/home/tanpv/Data/QG.25.196/OBS"
OBS_TX = os.path.join(OBS_DIR, "VnGAD_daily-0.25_1961-2023_tasmax.nc")
MASK_FILE = os.path.join(OBS_DIR, "Mask_VN_S2S_EC_0.25.nc")   # vn, reg; 41x65
MASK_VNGAD = os.path.join(OBS_DIR, "mask7reg_VnGAD_0.25.nc")  # 31x61
SHAPE_7REG = "/home/tanpv/GrADS_SCR/VN_SubRegs1/Vietnam7"

WGRIB2 = "/work/apps/gnu_4.8.5/grads/2.2.0/bin/wgrib2"

FLUX6H = os.path.join(OUT, "Flux6h")
FLUX_GRID = os.path.join(META, "flux_grid.npz")
INIT_LIST = os.path.join(META, "init_both.txt")

# ---------------------------------------------------------------- calendar
NLT = 46                 # lead days 0..45, valid day = init + k
YEARS = list(range(1999, 2012))
# Validation years, FROZEN. Chosen so the validation base rate matches the full
# sample (0.1293 vs 0.1274 init-weighted, |diff| 0.0015 -- the best of every
# three-year split taking one year from each of 1999-2002 / 2003-2007 / 2008-2011
# and holding 15-30 % of the inits). The EHD threshold used to rank the candidate
# splits comes from the 1969-1998 baseline, i.e. from data outside the sample, so
# the split was never chosen on the labels it is judged against.
VALID_YEARS = [2001, 2006, 2009]
TRAIN_YEARS = [y for y in YEARS if y not in VALID_YEARS]

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

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",
}

# ---------------------------------------------------------------- flux variables
# (name, daily aggregation). The wgrib2 -match below selects exactly these 16
# records; their order inside a GRIB file is fixed by the file layout, so the
# extractor verifies it against the inventory instead of trusting it.
FLUX_VARS = [
    ("shtfl", "mean"),   # sensible heat flux, surface
    ("lhtfl", "mean"),   # latent heat flux, surface
    ("soilw1", "mean"),  # soil moisture 0-0.1 m
    ("dlwrf", "mean"),   # downward longwave, surface
    ("dswrf", "mean"),   # downward shortwave, surface
    ("prate", "mean"),   # precipitation rate, surface
    ("u10", "mean"),
    ("v10", "mean"),
    ("t2m", "mean"),
    ("q2m", "mean"),     # specific humidity 2 m
    ("psfc", "mean"),
    ("tmax", "max"),     # 6-hour maximum 2 m temperature -> daily Tx
    ("tmin", "min"),
    ("hpbl", "max"),     # boundary layer height
    ("pwat", "mean"),
    ("tcdc", "mean"),
]
FLUX_NAMES = [v for v, _ in FLUX_VARS]
NFVAR = len(FLUX_VARS)

# Expected (variable, level) order as written in the GRIB files; asserted at run
# time so a silent reordering can never scramble the variable axis.
FLUX_ORDER = [
    ("SHTFL", "surface"), ("LHTFL", "surface"),
    ("SOILW", "0-0.1 m below ground"),
    ("DLWRF", "surface"), ("DSWRF", "surface"), ("PRATE", "surface"),
    ("UGRD", "10 m above ground"), ("VGRD", "10 m above ground"),
    ("TMP", "2 m above ground"), ("SPFH", "2 m above ground"),
    ("PRES", "surface"),
    ("TMAX", "2 m above ground"), ("TMIN", "2 m above ground"),
    ("HPBL", "surface"),
    ("PWAT", "entire atmosphere (considered as a single layer)"),
    ("TCDC", "entire atmosphere (considered as a single layer)"),
]

FLUX_MATCH = (
    ":(TMAX|TMIN|TMP|SPFH):2 m above ground:"
    "|:(PRATE|DSWRF|DLWRF|PRES|LHTFL|SHTFL|HPBL):surface:"
    "|:(TCDC|PWAT):entire atmosphere"
    "|:SOILW:0-0.1 m"
    "|:(UGRD|VGRD):10 m above ground:"
)

NSTEP = NLT * 4          # 184 six-hourly steps: valid = init + 6h .. init + 46 d


# ---------------------------------------------------------------- inventory
_D8 = re.compile(r"^\d{8}$")


def _init_dates(root):
    """Init dates under root/YYYY/YYYYMM/YYYYMMDD, following symlinked years."""
    out = []
    for y in sorted(os.listdir(root)):
        ydir = os.path.join(root, y)
        if not (y.isdigit() and os.path.isdir(ydir)):
            continue
        for m in sorted(os.listdir(ydir)):
            mdir = os.path.join(ydir, m)
            if not os.path.isdir(mdir):
                continue
            for d in sorted(os.listdir(mdir)):
                if _D8.match(d) and os.path.isdir(os.path.join(mdir, d)):
                    out.append(d)
    return sorted(set(out))


def flux_inits():
    return _init_dates(FLUX_ROOT)


def pgb_inits():
    return _init_dates(PGB_ROOT)


def both_inits():
    """Init dates that have BOTH flux and pgb -- the working sample (3088)."""
    if os.path.exists(INIT_LIST):
        with open(INIT_LIST) as fh:
            return [ln.strip() for ln in fh if ln.strip()]
    return sorted(set(flux_inits()) & set(pgb_inits()))


def to_date(s):
    return dt.date(int(s[:4]), int(s[4:6]), int(s[6:8]))


def flux_dir(s):
    return os.path.join(FLUX_ROOT, s[:4], s[:6], s)


def pgb_dir(s):
    return os.path.join(PGB_ROOT, s[:4], s[:6], s)


def flux_files(s):
    """The NSTEP six-hourly flux files of init `s`, valid init+6h .. init+46d.

    The analysis file (valid == init) is deliberately excluded: it holds 96
    records instead of 98 and its TMAX/TMIN are ":anl:", not 6-hour extremes.
    """
    d0 = to_date(s)
    base = flux_dir(s)
    out = []
    for k in range(1, NSTEP + 1):
        v = dt.datetime(d0.year, d0.month, d0.day) + dt.timedelta(hours=6 * k)
        out.append(os.path.join(base, "flxf%s00.%s.grb2" % (s, v.strftime("%Y%m%d%H"))))
    return out


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


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


# ---------------------------------------------------------------- grid
def flux_grid():
    """Full-grid shape plus the extraction window (lon 95-120E, lat 0-30N)."""
    z = np.load(FLUX_GRID)
    return dict(nlon=int(z["nlon"]), nlat=int(z["nlat"]),
                i0=int(z["i0"]), i1=int(z["i1"]),
                j0=int(z["j0"]), j1=int(z["j1"]),
                wlon=z["wlon"], wlat=z["wlat"])


def flux6h_path(init):
    return os.path.join(FLUX6H, init[:4], "flx6h_%s.npy" % init)


def tx_thres(qp):
    """EHD threshold = min(Tx90p, 35 C): the union of hot days and heat-wave days."""
    return np.minimum(qp, HOT_DAY_C)