"""Build elevation (DEM) features on the project's 0.25 deg model grid.

Source: ETOPO1 ice-surface, grid-registered (1 arc-minute global relief).
For Vietnam (no permanent ice) ice-surface == bedrock == land surface.

We area-average ETOPO1 into each 0.25 deg model cell, then pull the 541
land-point vector THROUGH the project Grid so the ordering matches AI_Inp
memmap exactly (see STATUS.md bug #2 on axis/order).

Outputs outputs/DEM/dem_features.nc with, per land point:
  elev_mean  mean ETOPO1 elevation in the 0.25 deg cell            [m]
  elev_std   sub-grid std of ETOPO1 (terrain ruggedness)           [m]
  elev_min / elev_max  cell extremes                               [m]
Also stores the 2-D (nlat,nlon) fields for plotting/QC.
"""
import os
import sys

import numpy as np
from netCDF4 import Dataset

sys.path.insert(0, os.path.join(os.path.dirname(__file__)))
from ehd_common import Grid, OUT  # noqa: E402

ETOPO = os.path.join(OUT, "DEM", "ETOPO1_Ice_g_gmt4.grd")
OUT_NC = os.path.join(OUT, "DEM", "dem_features.nc")
HALF = 0.125  # half of 0.25 deg cell


def load_etopo1(path):
    with Dataset(path) as ds:
        names = set(ds.variables)
        # GMT4 grd: coords x (lon), y (lat); data z (y,x)
        xn = "x" if "x" in names else ("lon" if "lon" in names else "longitude")
        yn = "y" if "y" in names else ("lat" if "lat" in names else "latitude")
        zn = "z" if "z" in names else [v for v in names if v not in (xn, yn)][0]
        x = np.asarray(ds[xn][:], dtype="f8")
        y = np.asarray(ds[yn][:], dtype="f8")
        z = np.asarray(ds[zn][:])
        if z.ndim == 1:  # some grd store z flat as (ny*nx)
            z = z.reshape(y.size, x.size)
    return x, y, z


def main():
    grid = Grid()
    lon = grid.lon
    lat = grid.lat
    nlat, nlon = grid.nlat, grid.nlon

    x, y, z = load_etopo1(ETOPO)
    print("ETOPO1 x:", x.min(), x.max(), x.size, " y:", y.min(), y.max(), y.size)
    # subset to VN bbox with margin for speed
    xm = (x >= lon.min() - 0.5) & (x <= lon.max() + 0.5)
    ym = (y >= lat.min() - 0.5) & (y <= lat.max() + 0.5)
    xs, ys, zs = x[xm], y[ym], z[np.ix_(ym, xm)]
    print("subset z shape:", zs.shape)

    emean = np.full((nlat, nlon), np.nan)
    estd = np.full((nlat, nlon), np.nan)
    emin = np.full((nlat, nlon), np.nan)
    emax = np.full((nlat, nlon), np.nan)
    for i, la in enumerate(lat):
        jy = (ys >= la - HALF) & (ys < la + HALF)
        if not jy.any():
            continue
        for j, lo in enumerate(lon):
            jx = (xs >= lo - HALF) & (xs < lo + HALF)
            if not jx.any():
                continue
            blk = zs[np.ix_(jy, jx)].astype("f8")
            land = blk[blk >= 0.0]          # land-only: drop ocean bathymetry
            use = land if land.size else blk  # fallback for all-ocean cells
            emean[i, j] = use.mean()
            estd[i, j] = use.std()
            emin[i, j] = use.min()
            emax[i, j] = use.max()

    # pull land-point vectors through the Grid (identical order to AI_Inp)
    p_mean = grid.to_points(emean)
    p_std = grid.to_points(estd)
    p_min = grid.to_points(emin)
    p_max = grid.to_points(emax)
    npt = p_mean.size
    print("n land points:", npt)
    print("elev_mean  min/med/max: %.1f / %.1f / %.1f m" % (
        np.nanmin(p_mean), np.nanmedian(p_mean), np.nanmax(p_mean)))
    print("elev_std   min/med/max: %.1f / %.1f / %.1f m" % (
        np.nanmin(p_std), np.nanmedian(p_std), np.nanmax(p_std)))
    nan_pts = int(np.isnan(p_mean).sum())
    print("NaN land points:", nan_pts)

    os.makedirs(os.path.dirname(OUT_NC), exist_ok=True)
    with Dataset(OUT_NC, "w") as ds:
        ds.createDimension("lat", nlat)
        ds.createDimension("lon", nlon)
        ds.createDimension("pt", npt)
        vlat = ds.createVariable("lat", "f8", ("lat",)); vlat[:] = lat
        vlon = ds.createVariable("lon", "f8", ("lon",)); vlon[:] = lon
        for name, arr2d in (("elev_mean", emean), ("elev_std", estd),
                            ("elev_min", emin), ("elev_max", emax)):
            v = ds.createVariable(name, "f4", ("lat", "lon"))
            v[:] = arr2d
        for name, arr1d in (("p_elev_mean", p_mean), ("p_elev_std", p_std),
                            ("p_elev_min", p_min), ("p_elev_max", p_max)):
            v = ds.createVariable(name, "f4", ("pt",))
            v[:] = arr1d
        ds.source = "ETOPO1 Ice surface, grid-registered, area-averaged to 0.25deg"
    print("wrote", OUT_NC)


if __name__ == "__main__":
    main()