"""Phase 0 data audit + fixed P90/EHD-threshold construction + target prototype.

Runs everything that needs xarray in ONE process (imports are slow on this
NFS home). Writes:
  data/processed/P90_2003_2022.nc
  data/processed/EHD_THRESHOLD_2003_2022.nc
  outputs/metrics/target_event_frequency.csv
  metadata/phase0_data_audit.json
Reads originals in place; never modifies them.
"""
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 targets.climatology import p90_climatology, ehd_threshold, threshold_for_dates, day_of_year_noleap  # noqa: E402
from targets.events import daily_ehd, exceedance, analyze_forecast_sequence  # noqa: E402

ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
DATA = os.path.join(ROOT, "data")
SRC = "/work/projects/G09/QG.25.196/OBS"
PROC = os.path.join(DATA, "processed")
os.makedirs(PROC, exist_ok=True)

REF = (2003, 2022)
audit = {}


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


log("[1/6] open gridded Tx (VnGAD)")
tx = xr.open_dataset(os.path.join(DATA, "VnGAD_daily-0.25_1961-2023_tasmax.nc"))
da = tx["tasmax"]
dates = tx["time"].values
audit["gridded_tx"] = {
    "file": "data/VnGAD_daily-0.25_1961-2023_tasmax.nc",
    "var": "tasmax", "units": da.attrs.get("units"),
    "dims": dict(da.sizes), "grid_deg": 0.25,
    "lon": [float(tx.lon.min()), float(tx.lon.max()), int(tx.lon.size)],
    "lat": [float(tx.lat.min()), float(tx.lat.max()), int(tx.lat.size)],
    "time": [str(dates[0])[:10], str(dates[-1])[:10], int(dates.size)],
    "fill_value": (None if da.encoding.get("_FillValue") is None
                   else float(da.encoding.get("_FillValue"))),
}

log("[2/6] leap / valid-mask / reference subset")
yr = dates.astype("datetime64[Y]").astype(int) + 1970
doy = day_of_year_noleap(dates)
mm = dates.astype("datetime64[M]"); month = (mm.astype(int) % 12) + 1
day = (dates.astype("datetime64[D]") - mm).astype("timedelta64[D]").astype(int) + 1
n_feb29 = int(((month == 2) & (day == 29)).sum())
ref_sel = (yr >= REF[0]) & (yr <= REF[1])
tx_ref = da.values[ref_sel]
valid_any = np.isfinite(tx_ref).any(0)
valid_all = np.isfinite(tx_ref).all(0)
audit["leap_and_mask"] = {
    "feb29_days_present": n_feb29,
    "ref_period": list(REF), "ref_n_days": int(ref_sel.sum()),
    "valid_cells_any": int(valid_any.sum()),
    "valid_cells_all_time": int(valid_all.sum()),
    "n_cells_total": int(valid_any.size),
    "mask_time_stable": bool((valid_any == valid_all).all()),
    "tx_ref_min": float(np.nanmin(tx_ref)), "tx_ref_max": float(np.nanmax(tx_ref)),
}

log("[3/6] region + land masks")
rr = xr.open_dataset(os.path.join(SRC, "mask7reg_VnGAD_0.25.nc"))
reg = rr["vn_7regions"].values
u = np.unique(reg[np.isfinite(reg)])
same_grid = (rr.lon.size == tx.lon.size and rr.lat.size == tx.lat.size
             and np.allclose(np.round(rr.lon.values, 3), np.round(tx.lon.values, 3))
             and np.allclose(np.round(rr.lat.values, 3), np.round(tx.lat.values, 3)))
ecm = xr.open_dataset(os.path.join(DATA, "Mask_VN_S2S_EC_0.25.nc"))
audit["region_mask"] = {
    "file": "OBS/mask7reg_VnGAD_0.25.nc (NOT symlinked into data/)",
    "var": "vn_7regions",
    "region_codes": [float(x) for x in u],
    "region_counts": {str(float(x)): int((reg == x).sum()) for x in u},
    "on_vngad_grid": bool(same_grid),
}
audit["ec_mask"] = {
    "file": "data/Mask_VN_S2S_EC_0.25.nc (symlinked)",
    "lon": [float(ecm.lon.min()), float(ecm.lon.max()), int(ecm.lon.size)],
    "lat": [float(ecm.lat.min()), float(ecm.lat.max()), int(ecm.lat.size)],
    "matches_vngad_grid": bool(ecm.lon.size == tx.lon.size and ecm.lat.size == tx.lat.size),
}

log("[4/6] fixed P90 climatology 2003-2022 (cyclic +/-15d)")
p90 = p90_climatology(da.values, dates, ref_years=REF, window=15, q=90)
thr = ehd_threshold(p90, cap=35.0)
capped_frac = float(np.nanmean((p90 > 35.0)[:, valid_any]))
audit["p90"] = {
    "shape": list(p90.shape), "ref_years": list(REF), "window_days": 15,
    "percentile": 90, "cap_degC": 35.0,
    "p90_min": float(np.nanmin(p90[:, valid_any])),
    "p90_max": float(np.nanmax(p90[:, valid_any])),
    "frac_daycells_capped_by_35": capped_frac,
}

log("[5/6] write P90 + EHD threshold NetCDF")
coords = {"dayofyear": np.arange(1, 366), "lat": tx.lat.values, "lon": tx.lon.values}
p90_ds = xr.Dataset(
    {"P90": (("dayofyear", "lat", "lon"), p90.astype("float32"))}, coords=coords,
    attrs={"title": "Fixed climatological P90 of observed Tx (VnGAD)",
           "reference_period": f"{REF[0]}-{REF[1]}", "window": "cyclic +/-15 days",
           "leap_handling": "29 Feb pooled with 28 Feb (calendar-day index 59); 365 indices",
           "source": "VnGAD_daily-0.25_1961-2023_tasmax.nc", "units": "degC"})
p90_ds["P90"].attrs["units"] = "degC"
p90_ds.to_netcdf(os.path.join(PROC, "P90_2003_2022.nc"))
thr_ds = xr.Dataset(
    {"EHD_THRESHOLD": (("dayofyear", "lat", "lon"), thr.astype("float32"))}, coords=coords,
    attrs={"title": "EHD threshold = min(P90, 35 degC)", "units": "degC",
           "reference_period": f"{REF[0]}-{REF[1]}"})
thr_ds["EHD_THRESHOLD"].attrs["units"] = "degC"
thr_ds.to_netcdf(os.path.join(PROC, "EHD_THRESHOLD_2003_2022.nc"))

log("[6/6] target prototype + event-frequency over reference period")
thr_by_date = threshold_for_dates(thr, dates)          # (T, lat, lon)
ehd_all = daily_ehd(da.values, thr_by_date)            # (T, lat, lon) int8
exc_all = exceedance(da.values, thr_by_date)
# climatological daily EHD frequency over reference period (sanity: ~0.10)
ehd_ref = ehd_all[ref_sel].astype("float64")
ehd_ref[:, ~valid_any] = np.nan
mean_ehd_freq = float(np.nanmean(ehd_ref))
# demonstrate cross-week episode accounting on one observed 52-day window
# (LT0..45 + a 6-day post buffer), starting at a hot period (May 2003).
i0 = int(np.argmax(dates >= np.datetime64("2003-05-01")))
lead_days = np.arange(-6, 52)                          # -6..-1 buffer, 0..45, 46..51 buffer
seg = slice(i0 - 6, i0 - 6 + len(lead_days))
cell = np.unravel_index(np.nanargmax(np.nan_to_num(exc_all[ref_sel].sum(0))), valid_any.shape)
proto_e, proto_w = analyze_forecast_sequence(
    lead_days, ehd_all[seg, cell[0], cell[1]], exc_all[seg, cell[0], cell[1]])
audit["target_prototype"] = {
    "mean_daily_ehd_freq_ref": mean_ehd_freq,
    "demo_cell_latlon": [float(tx.lat.values[cell[0]]), float(tx.lon.values[cell[1]])],
    "demo_n_episodes": len(proto_e),
    "demo_first_episode": (proto_e[0] if proto_e else None),
}

# event-frequency CSV per region (mean EHD frequency, ref period)
lines = ["region,n_cells,mean_ehd_freq"]
for code in u:
    m = (reg == code) & valid_any
    if m.any():
        v = float(np.nanmean(ehd_ref[:, m]))
        lines.append(f"{int(code)},{int(m.sum())},{v:.4f}")
allm = valid_any
lines.append(f"ALL,{int(allm.sum())},{float(np.nanmean(ehd_ref[:, allm])):.4f}")
with open(os.path.join(ROOT, "outputs/metrics/target_event_frequency.csv"), "w") as f:
    f.write("\n".join(lines) + "\n")

with open(os.path.join(ROOT, "metadata/phase0_data_audit.json"), "w") as f:
    json.dump(audit, f, indent=2)
log("DONE")
log(json.dumps(audit, indent=2))