"""Process one multi-hindcast reforecast GRIB (20 init dates) into per-year
daily-max Tx on the coarse grid, and verify temporal alignment vs VnGAD for
every hindcast year. Confirms the downloaded 2003-2022 file is usable.

Input : data/raw/ecmwf_s2s_1p5deg/s2s_cf_mx2t6_mv2023-12-28_hy2003-2022.grib
Output: data/processed/s2s_dailymaxTx_mv2023-12-28_hy2003-2022.nc
        metadata/reforecast_file_report.json
"""
import json
import os
import sys
import numpy as np
import xarray as xr

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
from data.grib_preprocess import S2SReforecastPreprocessor, daily_max_from_6h  # noqa: E402
from utils.timealign import align_by_date  # noqa: E402

ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
DATA = os.path.join(ROOT, "data")
GRIB = os.path.join(DATA, "raw", "ecmwf_s2s_1p5deg",
                    "s2s_cf_mx2t6_mv2023-12-28_hy2003-2022.grib")
TZ = 7


def log(*a):
    print(*a, flush=True)


log("[1/4] split by init + load canonical (20 hindcast years)")
pp = S2SReforecastPreprocessor(tz_hours=TZ)
fields = pp.load_canonical_by_init(GRIB, to_degC=True)
log(f"   {len(fields)} hindcast-year fields; inits "
    f"{str(fields[0].init_time)}..{str(fields[-1].init_time)}")

log("[2/4] daily-max Tx per year + obs alignment")
obs = xr.open_dataset(os.path.join(DATA, "VnGAD_daily-0.25_1961-2023_tasmax.nc"))
obs_days = obs["time"].values.astype("datetime64[D]")
lat, lon = fields[0].lat, fields[0].lon
per_year, years, matched_all = [], [], True
report_years = []
for f in fields:
    days, txday, cov = daily_max_from_6h(f, tz_hours=TZ)
    idx, matched = align_by_date(days, obs_days)
    matched_all = matched_all and bool(matched.all())
    yr = int(str(f.init_time)[:4])
    years.append(yr)
    per_year.append(txday.astype("float32"))
    report_years.append({"hyear": yr, "init": str(f.init_time),
                         "n_days": int(len(days)), "obs_matched": int(matched.sum()),
                         "all_matched": bool(matched.all()),
                         "tx_max_degC": round(float(np.nanmax(txday)), 2)})
    log(f"   {yr}: {len(days)} days, obs matched {int(matched.sum())}/{len(days)}")
obs.close()

# all years share the same lead structure -> same number of local days; stack
nday = min(a.shape[0] for a in per_year)
stack = np.stack([a[:nday] for a in per_year], axis=0)   # (year, day, lat, lon)

log("[3/4] write per-year daily-max NetCDF")
xr.Dataset(
    {"tx_daily": (("hyear", "day_index", "lat", "lon"), stack)},
    coords={"hyear": np.array(years, dtype="int16"),
            "day_index": np.arange(nday, dtype="int16"), "lat": lat, "lon": lon},
    attrs={"title": "S2S cf daily-max Tx per hindcast year (mv 2023-12-28), coarse 1.5deg",
           "units": "degC", "tz": "UTC+7 local day", "leads": "6..1104h (46 days)"},
).to_netcdf(os.path.join(DATA, "processed", "s2s_dailymaxTx_mv2023-12-28_hy2003-2022.nc"))

log("[4/4] report")
rep = {"grib": os.path.basename(GRIB), "n_hindcast_years": len(fields),
       "years": years, "all_years_aligned_to_obs": matched_all,
       "grid_deg": float(round(lat[1] - lat[0], 3)),
       "lat_range": [float(lat.min()), float(lat.max())],
       "lon_range": [float(lon.min()), float(lon.max())],
       "days_per_year": nday, "per_year": report_years}
with open(os.path.join(ROOT, "metadata", "reforecast_file_report.json"), "w") as f:
    json.dump(rep, f, indent=2, default=str)
log("DONE  all_years_aligned_to_obs =", matched_all)