"""Phase 1 steps 6-7: run GRIB->canonical preprocessing on the real sample and
VERIFY temporal alignment against VnGAD observations.

Sample: data/Tx_05-01-2022.grib (mx2t6, control reforecast, init 2022-01-05).
Produces:
  data/processed/s2s_sample_mx2t6_init2022-01-05.nc  (canonical coarse Tx, degC)
  data/processed/s2s_sample_dailymaxTx_init2022-01-05.nc
  metadata/preprocess_sample_report.json
  outputs/PREPROCESS_SAMPLE.html
"""
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")
PROC = os.path.join(DATA, "processed")
GRIB = os.path.join(DATA, "Tx_05-01-2022.grib")
TZ = 7


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


log("[1/5] load + standardize GRIB -> canonical")
pp = S2SReforecastPreprocessor(tz_hours=TZ)
field = pp.load_canonical(GRIB, to_degC=True)
log(f"   var={field.name} units={field.units} shape={field.values.shape} "
    f"leads[{field.lead[0]}..{field.lead[-1]}]h n={len(field.lead)}")
log(f"   init={field.init_time} model_version={field.model_version_date} "
    f"lat_asc={field.lat[0]}->{field.lat[-1]} lon={field.lon[0]}->{field.lon[-1]}")

rep = {
    "grib": os.path.basename(GRIB), "variable": field.name, "units": field.units,
    "n_leads": int(len(field.lead)), "lead_first_h": int(field.lead[0]),
    "lead_last_h": int(field.lead[-1]), "lead_step_h": int(field.lead[1] - field.lead[0]),
    "init_time": str(field.init_time), "model_version_date": str(field.model_version_date),
    "lat_ascending": bool(field.lat[0] < field.lat[-1]),
    "lat_range": [float(field.lat.min()), float(field.lat.max())],
    "lon_range": [float(field.lon.min()), float(field.lon.max())],
    "grid_deg": float(round(field.lat[1] - field.lat[0], 3)),
    "tx_min_degC": round(float(np.nanmin(field.values)), 2),
    "tx_max_degC": round(float(np.nanmax(field.values)), 2),
    "member_type": field.member_type, "level_type": field.level_type,
}

log("[2/5] valid-time checks (valid = init + lead)")
vt = field.valid_time
rep["valid_first"] = str(vt[0]); rep["valid_last"] = str(vt[-1])
# spot checks against the known GRIB truth
exp_first = np.datetime64("2022-01-05T06")
exp_last = np.datetime64("2022-02-20T00")
rep["valid_first_ok"] = bool(vt[0].astype("datetime64[h]") == exp_first)
rep["valid_last_ok"] = bool(vt[-1].astype("datetime64[h]") == exp_last)
log(f"   valid[0]={vt[0]} (exp {exp_first}) ok={rep['valid_first_ok']}")
log(f"   valid[-1]={vt[-1]} (exp {exp_last}) ok={rep['valid_last_ok']}")

log("[3/5] daily max on local (UTC+7) day")
days, txday, cov = daily_max_from_6h(field, tz_hours=TZ)
rep["n_local_days"] = int(len(days))
rep["local_day_first"] = str(days[0]); rep["local_day_last"] = str(days[-1])
rep["full_coverage_days"] = int((cov == 4).sum())
rep["partial_days"] = [int(c) for c in cov if c != 4]
log(f"   {len(days)} local days, {int((cov==4).sum())} full (4 windows), "
    f"partial coverage counts={sorted(set(int(c) for c in cov if c!=4))}")

log("[4/5] align forecast local days to VnGAD observation dates")
obs = xr.open_dataset(os.path.join(DATA, "VnGAD_daily-0.25_1961-2023_tasmax.nc"))
obs_days = obs["time"].values.astype("datetime64[D]")
idx, matched = align_by_date(days, obs_days)
rep["obs_matched_days"] = int(matched.sum())
rep["obs_total_days"] = int(len(days))
rep["all_forecast_days_have_obs"] = bool(matched.all())
log(f"   {int(matched.sum())}/{len(days)} forecast local days present in VnGAD obs; "
    f"all matched={matched.all()}")
obs.close()

log("[5/5] write canonical NetCDF + daily-max + report/HTML")
xr.Dataset(
    {field.name: (("lead", "lat", "lon"), field.values)},
    coords={"lead": field.lead.astype("int32"), "lat": field.lat, "lon": field.lon,
            "valid_time": ("lead", field.valid_time)},
    attrs={"title": "S2S cf reforecast (canonical, coarse 1.5deg)", "variable": field.name,
           "units": field.units, "init_time": str(field.init_time),
           "model_version_date": str(field.model_version_date), "member": "cf",
           "note": "lat ascending; degC; valid_time=init+lead; native 1.5deg (no regrid)"},
).to_netcdf(os.path.join(PROC, "s2s_sample_mx2t6_init2022-01-05.nc"))

xr.Dataset(
    {"tx_daily": (("day", "lat", "lon"), txday.astype("float32")),
     "coverage": (("day",), cov.astype("int8"))},
    coords={"day": days.astype("datetime64[ns]"), "lat": field.lat, "lon": field.lon},
    attrs={"title": "S2S cf daily-max Tx (local UTC+7), coarse 1.5deg", "units": "degC",
           "init_time": str(field.init_time)},
).to_netcdf(os.path.join(PROC, "s2s_sample_dailymaxTx_init2022-01-05.nc"))

with open(os.path.join(ROOT, "metadata", "preprocess_sample_report.json"), "w") as f:
    json.dump(rep, f, indent=2, default=str)

rows = "".join(f"<tr><td>{k}</td><td>{rep[k]}</td></tr>" for k in rep)
html = f"""<!DOCTYPE html><html><head><meta charset="utf-8">
<title>S2S-FM-VN Phase 1 preprocessing sample</title>
<style>body{{font:14px/1.5 -apple-system,Segoe UI,Arial,sans-serif;max-width:860px;margin:24px auto;color:#1a2233}}
h1{{font-size:21px}} table{{border-collapse:collapse}} td{{border:1px solid #d7dde6;padding:5px 10px}}
td:first-child{{font-weight:600;background:#eef2f8}}</style></head><body>
<h1>Phase 1 preprocessing &mdash; real GRIB sample</h1>
<p>GRIB &rarr; cdo NetCDF &rarr; canonical (lat flipped ascending, K&rarr;&deg;C, valid=init+lead),
daily-max on local UTC+7 day, aligned to VnGAD observation dates. All checks below.</p>
<table>{rows}</table>
<p style="color:#5b6472">Sources: data/processed/s2s_sample_*.nc, metadata/preprocess_sample_report.json.</p>
</body></html>"""
with open(os.path.join(ROOT, "outputs", "PREPROCESS_SAMPLE.html"), "w") as f:
    f.write(html)
log("DONE"); log(json.dumps(rep, indent=2, default=str))