#!/usr/bin/env python3
"""Static predictors for downscaling EC 1.5 deg -> 0.25 deg over Vietnam.

The physical core of temperature downscaling is elevation: a 1.5 deg cell sees one
smoothed height, while the 0.25 deg target grid resolves real terrain. The difference
  dz = orog_fine - orog_coarse
is therefore the single most informative static predictor (it carries the lapse-rate
correction the coarse model cannot represent).

Fields written on the 0.25 deg target grid (41 x 65):
  orog        mean GTOPO30 elevation in the 0.25 deg cell
  orog_std    sub-grid roughness (std of GTOPO30 inside the cell)
  orog_coarse mean elevation of the 1.5 deg cell containing the point
  dz          orog - orog_coarse
  landfrac    fraction of GTOPO30 pixels above sea level in the cell
  lat2d/lon2d coordinates
  vn/reg      land mask and 7-region mask (from the project mask file)
"""
import os
import numpy as np
import xarray as xr

DEM = "/home/tanpv/Progs/Topo_map/GTOPO_DEM_30s.nc"
MASK = "/home/tanpv/Data/QG.25.196/OBS/Mask_VN_S2S_EC_0.25.nc"
OUT = "/home/tanpv/Progs/claude/HW_EC_10/outputs/Cli/static_fields_0.25.nc"

# target grid (matches Tx_EC_0.25_2002-2023.nc / grid_VN_S2S_0.25.dat)
LON0, DLON, NLON = 101.0, 0.25, 41
LAT0, DLAT, NLAT = 8.0, 0.25, 65
# coarse EC grid actually delivered by MARS for area [25,100,6,114] at 1.5 deg
CLON0, CDLON = 100.0, 1.5
CLAT0, CDLAT = 6.0, 1.5


def cell_stats(dem, dem_lat, dem_lon, edges_lat, edges_lon):
    """Mean / std / land-fraction of DEM pixels inside each target cell."""
    ilat = np.searchsorted(edges_lat, dem_lat, side="right") - 1
    ilon = np.searchsorted(edges_lon, dem_lon, side="right") - 1
    nlat, nlon = len(edges_lat) - 1, len(edges_lon) - 1
    okla = (ilat >= 0) & (ilat < nlat)
    oklo = (ilon >= 0) & (ilon < nlon)

    z = dem[np.ix_(okla, oklo)].astype("float64")
    ila = ilat[okla]
    ilo = ilon[oklo]
    flat = ila[:, None] * nlon + ilo[None, :]

    n = np.bincount(flat.ravel(), minlength=nlat * nlon)
    s = np.bincount(flat.ravel(), weights=z.ravel(), minlength=nlat * nlon)
    s2 = np.bincount(flat.ravel(), weights=(z ** 2).ravel(), minlength=nlat * nlon)
    land = np.bincount(flat.ravel(), weights=(z.ravel() > 0).astype("float64"),
                       minlength=nlat * nlon)

    n = np.maximum(n, 1)
    mean = s / n
    var = np.maximum((s2 / n) - mean ** 2, 0.0)
    return (mean.reshape(nlat, nlon),
            np.sqrt(var).reshape(nlat, nlon),
            (land / n).reshape(nlat, nlon))


def main():
    lon = LON0 + DLON * np.arange(NLON)
    lat = LAT0 + DLAT * np.arange(NLAT)
    elon = np.append(lon - DLON / 2, lon[-1] + DLON / 2)
    elat = np.append(lat - DLAT / 2, lat[-1] + DLAT / 2)

    ds = xr.open_dataset(DEM)
    sub = ds.sel(lat=slice(elat[0] - 1, elat[-1] + 1),
                 lon=slice(elon[0] - 1, elon[-1] + 1))
    dem = sub.z.values
    dlat, dlon = sub.lat.values, sub.lon.values
    print(f"DEM subset {dem.shape}, lat {dlat[0]:.3f}..{dlat[-1]:.3f}, "
          f"lon {dlon[0]:.3f}..{dlon[-1]:.3f}")

    orog, orog_std, landfrac = cell_stats(dem, dlat, dlon, elat, elon)
    orog = np.maximum(orog, 0.0)      # GTOPO30 sets ocean to <=0

    # coarse (1.5 deg) elevation, then mapped onto every fine point it contains
    nclon = int(np.ceil((elon[-1] - CLON0) / CDLON)) + 1
    nclat = int(np.ceil((elat[-1] - CLAT0) / CDLAT)) + 1
    clon = CLON0 + CDLON * np.arange(nclon)
    clat = CLAT0 + CDLAT * np.arange(nclat)
    eclon = np.append(clon - CDLON / 2, clon[-1] + CDLON / 2)
    eclat = np.append(clat - CDLAT / 2, clat[-1] + CDLAT / 2)
    corog, _, _ = cell_stats(dem, dlat, dlon, eclat, eclon)
    corog = np.maximum(corog, 0.0)

    jc = np.clip(np.searchsorted(eclat, lat, side="right") - 1, 0, len(clat) - 1)
    ic = np.clip(np.searchsorted(eclon, lon, side="right") - 1, 0, len(clon) - 1)
    orog_coarse = corog[np.ix_(jc, ic)]
    dz = orog - orog_coarse

    mask = xr.open_dataset(MASK)
    lon2d, lat2d = np.meshgrid(lon, lat)

    out = xr.Dataset(
        {
            "orog": (("lat", "lon"), orog.astype("float32"),
                     {"units": "m", "long_name": "mean elevation in 0.25 deg cell"}),
            "orog_std": (("lat", "lon"), orog_std.astype("float32"),
                         {"units": "m", "long_name": "sub-grid terrain roughness"}),
            "orog_coarse": (("lat", "lon"), orog_coarse.astype("float32"),
                            {"units": "m",
                             "long_name": "elevation of the enclosing 1.5 deg EC cell"}),
            "dz": (("lat", "lon"), dz.astype("float32"),
                   {"units": "m",
                    "long_name": "orog - orog_coarse: elevation the coarse model misses"}),
            "landfrac": (("lat", "lon"), landfrac.astype("float32"),
                         {"long_name": "fraction of DEM pixels above sea level"}),
            "lat2d": (("lat", "lon"), lat2d.astype("float32")),
            "lon2d": (("lat", "lon"), lon2d.astype("float32")),
            "vn": (("lat", "lon"), mask.vn.values),
            "reg": (("lat", "lon"), mask.reg.values),
        },
        coords={"lat": lat, "lon": lon},
    )
    out.attrs["dem_source"] = DEM
    os.makedirs(os.path.dirname(OUT), exist_ok=True)
    out.to_netcdf(OUT)
    print(f"wrote {OUT}")

    land = np.isfinite(mask.vn.values) & (mask.vn.values > 0)
    for v in ["orog", "orog_std", "orog_coarse", "dz"]:
        a = out[v].values[land]
        print(f"  {v:12s} over {land.sum()} land pts: "
              f"min {a.min():7.1f}  mean {a.mean():7.1f}  max {a.max():7.1f} m")


if __name__ == "__main__":
    main()