#!/usr/bin/env python3
"""Build daily forecast Tx from 6-hourly mx2t6, aligned to (year, init, lead).

mx2t6 at valid time t covers the interval (t-6h, t]. Forecast day `lead` therefore
spans hours (24*lead, 24*lead+24], i.e. the four 6-hourly maxima at indices
4*lead .. 4*lead+3 of the 184-step record. Daily Tx is their maximum.

The same layout is produced by the old 0.25 deg archive and by the new 1.5 deg
`txtm` group, so one code path serves both. `--to-ec15` averages the result onto the true
1.5 deg MARS grid (100..113.5E, 6..24N), producing the proxy used to develop the
downscaling method before the real 1.5 deg data has finished downloading. Averaging
onto the real grid rather than 6x6 blocks matters: 6x6 blocks would silently drop the
southern and eastern edge of the domain, and the code would then need rewriting.

Usage:
  build_ec_tx.py --root <dir> --pattern <glob> --out <file.nc> [--to-ec15] [--jobs 10]
"""
import argparse
import glob
import os
import re
import sys
from concurrent.futures import ProcessPoolExecutor

import numpy as np
import xarray as xr

YEARS = list(range(2003, 2023))
NLEAD = 46
NSTEP_PER_DAY = 4

# the grid MARS actually returns for area [25,100,6,114] at 1.5 deg
EC15_LON = 100.0 + 1.5 * np.arange(10)          # 100.0 .. 113.5
EC15_LAT = 24.0 - 1.5 * np.arange(13)           # 24.0 .. 6.0 (north to south)


def ec15_weights(lat, lon):
    """Map fine (lat,lon) points onto their enclosing 1.5 deg cell."""
    jl = np.abs(lat[:, None] - EC15_LAT[None, :]).argmin(axis=1)
    il = np.abs(lon[:, None] - EC15_LON[None, :]).argmin(axis=1)
    return jl, il


def to_ec15(daily, jl, il):
    """Average a (lead, lat, lon) field onto the 1.5 deg grid."""
    nl = daily.shape[0]
    out = np.full((nl, len(EC15_LAT), len(EC15_LON)), np.nan, dtype="float32")
    flat = jl[:, None] * len(EC15_LON) + il[None, :]
    fl = flat.ravel()
    n = np.bincount(fl, minlength=len(EC15_LAT) * len(EC15_LON))
    for k in range(nl):
        s = np.bincount(fl, weights=daily[k].ravel().astype("float64"),
                        minlength=len(EC15_LAT) * len(EC15_LON))
        with np.errstate(invalid="ignore", divide="ignore"):
            out[k] = (s / np.maximum(n, 1)).reshape(len(EC15_LAT), len(EC15_LON))
        out[k][n.reshape(out[k].shape) == 0] = np.nan
    return out


def load_year(a):
    root, pattern, year, inits, to15 = a
    out = None
    for k, mmdd in enumerate(inits):
        matches = glob.glob(os.path.join(root, str(year), pattern.format(year=year, mmdd=mmdd)))
        if not matches:
            raise FileNotFoundError(f"{year}-{mmdd}: no file for pattern {pattern}")
        with xr.open_dataset(matches[0]) as d:
            v = d["mx2t6"]
            arr = v.values.astype("float32")            # (184, lat, lon), Kelvin
            if arr.shape[0] < NLEAD * NSTEP_PER_DAY:
                raise ValueError(f"{matches[0]}: only {arr.shape[0]} steps")
            arr = arr[:NLEAD * NSTEP_PER_DAY]
            daily = arr.reshape(NLEAD, NSTEP_PER_DAY, arr.shape[1], arr.shape[2]).max(axis=1)
            daily = daily - 273.15                      # -> deg C
            if to15:
                la = d.latitude.values if "latitude" in d.coords else d.lat.values
                lo = d.longitude.values if "longitude" in d.coords else d.lon.values
                daily = to_ec15(daily, *ec15_weights(la, lo))
            if out is None:
                out = np.empty((len(inits),) + daily.shape, dtype="float32")
            out[k] = daily
    return year, out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", required=True)
    ap.add_argument("--pattern", required=True,
                    help="glob inside <root>/<year>/, may use {year} and {mmdd}")
    ap.add_argument("--out", required=True)
    ap.add_argument("--to-ec15", action="store_true",
                    help="average onto the true 1.5 deg MARS grid")
    ap.add_argument("--jobs", type=int, default=10)
    ap.add_argument("--hdates", default="/home/tanpv/Data/QG.25.196/hdate_2023.txt")
    args = ap.parse_args()

    with open(args.hdates) as fh:
        inits = sorted({t for t in fh.read().split() if len(t) == 5 and t[2] == "-"})

    # coordinates from the first file
    first = glob.glob(os.path.join(args.root, str(YEARS[0]),
                                   args.pattern.format(year=YEARS[0], mmdd=inits[0])))
    if not first:
        print(f"no files under {args.root}", file=sys.stderr)
        return 1
    with xr.open_dataset(first[0]) as d:
        lat0 = d.latitude.values if "latitude" in d.coords else d.lat.values
        lon0 = d.longitude.values if "longitude" in d.coords else d.lon.values
    if args.to_ec15:
        lat0, lon0 = EC15_LAT, EC15_LON
    print(f"{len(inits)} inits, {len(YEARS)} years, grid {len(lat0)}x{len(lon0)}, "
          f"to_ec15={args.to_ec15}", flush=True)

    data = np.empty((len(YEARS), len(inits), NLEAD, len(lat0), len(lon0)), dtype="float32")
    tasks = [(args.root, args.pattern, y, inits, args.to_ec15) for y in YEARS]
    with ProcessPoolExecutor(max_workers=args.jobs) as pool:
        for year, arr in pool.map(load_year, tasks):
            data[YEARS.index(year)] = arr
            print(f"  {year} done", flush=True)

    ds = xr.Dataset(
        {"tx_ec": (("year", "init", "lead", "lat", "lon"), data,
                   {"units": "C", "long_name": "forecast daily maximum 2m temperature",
                    "note": "daily max of the four 6-hourly mx2t6 values of each lead day"})},
        coords={"year": YEARS, "init": inits, "lead": np.arange(NLEAD),
                "lat": lat0, "lon": lon0},
    )
    ds.attrs["source_root"] = args.root
    ds.attrs["grid"] = "ec15" if args.to_ec15 else "native"
    os.makedirs(os.path.dirname(args.out), exist_ok=True)
    ds.to_netcdf(args.out, encoding={"tx_ec": {"zlib": True, "complevel": 4}})
    print(f"wrote {args.out} ({os.path.getsize(args.out)/1e9:.2f} GB)")
    print(f"  tx range {np.nanmin(data):.1f} .. {np.nanmax(data):.1f} C")
    return 0


if __name__ == "__main__":
    sys.exit(main())