"""D3 + D7: verify mask7reg_VnGAD alignment with the VnGAD target grid, and
handle the valid Tx cells that carry no R1-R7 region code.

D3: programmatically confirm lat/lon exact alignment, orientation, dimensions,
    missing-value conventions and region codes. -> metadata/mask_alignment_report.json
D7: keep the unassigned valid cells in the NATIONAL domain (region code 0),
    exclude them from R1-R7 aggregation, and emit a diagnostic table + map.
    -> data/processed/region_codes_vngad.nc  (0=UNASSIGNED national, 1-7=R1-R7, -1=outside)
       metadata/unassigned_region_cells.csv
       outputs/figures/unassigned_region_cells.png
Reads originals in place; never modifies them.
"""
import json
import os
import sys
import numpy as np
import xarray as xr
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

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"
REF = (2003, 2022)


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


log("[1/4] open VnGAD target grid + region mask")
tx = xr.open_dataset(os.path.join(DATA, "VnGAD_daily-0.25_1961-2023_tasmax.nc"))
rr = xr.open_dataset(os.path.join(SRC, "mask7reg_VnGAD_0.25.nc"))
reg = rr["vn_7regions"].values.astype("float64")  # (lat, lon)

# valid Tx mask (time-stable) over reference period
dates = tx["time"].values
yr = dates.astype("datetime64[Y]").astype(int) + 1970
sel = (yr >= REF[0]) & (yr <= REF[1])
valid = np.isfinite(tx["tasmax"].values[sel]).any(0)  # (lat, lon)

log("[2/4] D3 alignment verification")
def close(a, b):
    return bool(a.size == b.size and np.allclose(np.round(a, 4), np.round(b, 4)))

lat_t, lon_t = tx["lat"].values, tx["lon"].values
lat_r, lon_r = rr["lat"].values, rr["lon"].values
reg_codes = np.unique(reg[np.isfinite(reg)])
align = {
    "lat_exact_match": close(lat_t, lat_r),
    "lon_exact_match": close(lon_t, lon_r),
    "dims_match": bool(reg.shape == valid.shape),
    "lat_orientation_target": "ascending" if lat_t[0] < lat_t[-1] else "descending",
    "lat_orientation_mask": "ascending" if lat_r[0] < lat_r[-1] else "descending",
    "orientation_consistent": bool((lat_t[0] < lat_t[-1]) == (lat_r[0] < lat_r[-1])),
    "region_codes": [float(c) for c in reg_codes],
    "region_code_count": int(reg_codes.size),
    "target_fill_value": (None if tx["tasmax"].encoding.get("_FillValue") is None
                          else float(tx["tasmax"].encoding.get("_FillValue"))),
    "mask_missing_is_nonfinite": bool(np.isnan(reg).any()),
    "n_valid_tx_cells": int(valid.sum()),
    "n_cells_with_region": int(((reg >= 1) & valid).sum()),
}
align["n_unassigned_valid_cells"] = align["n_valid_tx_cells"] - align["n_cells_with_region"]
align["PASS"] = bool(align["lat_exact_match"] and align["lon_exact_match"]
                     and align["dims_match"] and align["orientation_consistent"]
                     and align["region_code_count"] == 7)
with open(os.path.join(ROOT, "metadata", "mask_alignment_report.json"), "w") as f:
    json.dump(align, f, indent=2)
log("   alignment PASS =", align["PASS"], "| unassigned valid cells =", align["n_unassigned_valid_cells"])

log("[3/4] D7 build national region field (0=UNASSIGNED) + diagnostic table")
# region field on the VnGAD grid: -1 outside national domain (invalid Tx),
# 0 = valid Tx but no R1-R7 code (UNASSIGNED, kept in national analysis),
# 1..7 = R1..R7.
region_field = np.full(valid.shape, -1, dtype="int16")
has_reg = np.isfinite(reg) & (reg >= 1)
region_field[valid] = 0                       # national, unassigned by default
region_field[valid & has_reg] = reg[valid & has_reg].astype("int16")
unassigned = valid & ~has_reg                 # the 45 cells

# coast/border heuristic: a valid cell with >=1 invalid 4-neighbour
inv = ~valid
pad = np.pad(inv, 1, constant_values=True)
neighbour_invalid = (pad[:-2, 1:-1] | pad[2:, 1:-1] | pad[1:-1, :-2] | pad[1:-1, 2:])
coast_border = neighbour_invalid & valid

lat2d, lon2d = np.meshgrid(lat_t, lon_t, indexing="ij")
rows = ["lat,lon,vngad_valid,existing_region_code,coast_or_border"]
ii, jj = np.where(unassigned)
for i, j in zip(ii, jj):
    rc = reg[i, j]
    rc = "NaN" if not np.isfinite(rc) else str(int(rc))
    rows.append(f"{lat2d[i,j]:.3f},{lon2d[i,j]:.3f},True,{rc},{bool(coast_border[i,j])}")
with open(os.path.join(ROOT, "metadata", "unassigned_region_cells.csv"), "w") as f:
    f.write("\n".join(rows) + "\n")
log("   wrote", int(unassigned.sum()), "unassigned cells;",
    int((coast_border & unassigned).sum()), "flagged coast/border")

# save region field NetCDF
xr.Dataset({"region": (("lat", "lon"), region_field)},
           coords={"lat": lat_t, "lon": lon_t},
           attrs={"title": "VnGAD region codes for S2S-FM-VN",
                  "codes": "-1=outside national domain; 0=national UNASSIGNED (no R1-R7); 1-7=R1-R7",
                  "source_mask": "mask7reg_VnGAD_0.25.nc",
                  "note": "unassigned cells kept in national analysis, excluded from R1-R7"}
           ).to_netcdf(os.path.join(DATA, "processed", "region_codes_vngad.nc"))

log("[4/4] diagnostic map")
plot = np.where(region_field == -1, np.nan, region_field).astype("float64")
fig, ax = plt.subplots(figsize=(6, 8))
im = ax.pcolormesh(lon_t, lat_t, plot, cmap="tab10", vmin=-0.5, vmax=9.5, shading="nearest")
uy, ux = lat2d[unassigned], lon2d[unassigned]
ax.scatter(ux, uy, s=22, facecolors="none", edgecolors="black", linewidths=1.1,
           label=f"unassigned (region=0), n={int(unassigned.sum())}")
ax.set_title("VnGAD region codes (R1-R7) with unassigned valid cells")
ax.set_xlabel("lon"); ax.set_ylabel("lat"); ax.set_aspect("equal")
fig.colorbar(im, ax=ax, label="region code (0=UNASSIGNED, 1-7=R1-R7)", shrink=0.8)
ax.legend(loc="upper right", fontsize=8)
fig.tight_layout()
fig.savefig(os.path.join(ROOT, "outputs", "figures", "unassigned_region_cells.png"), dpi=130)
log("DONE")