#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Create UKESM forcings for 1pct CO2 runs.

For each ELM/DATM stream (Prec, Solr, TPQWL) and each month of 1850-2000:
  - read a CRUJRA monthly file as a template (grid, time axis, metadata)
  - replace every variable of that stream with UKESM data (nearest neighbour)
  - write ONE output file containing all variables of the stream
"""

import os
import numpy as np
import xarray as xr

# ------------------------------------------------------------
# User settings (change paths if running on Perlmutter instead of Compy)
# ------------------------------------------------------------
YEAR_START, YEAR_END = 1850, 2000
OVERWRITE = True            # False: skip output files that already exist

RAW    = "/global/cfs/cdirs/e3sm/inputdata/atm/datm7/WIEMIP/raw_data/UKESM1-0-LL/05deg"
PREFIX = "UKESM1-0-LL_clim3_50perc_1pctCO2"
TMPL   = "/global/cfs/cdirs/e3sm/inputdata/atm/datm7/atm_forcing.datm7.cruncep_qianFill.0.5d.V4.c130305"
OUTDIR = "/global/cfs/cdirs/e3sm/inputdata/atm/datm7/WIEMIP/UKESM"
REF_FILE = ("/global/cfs/cdirs/e3sm/inputdata/atm/datm7/atm_forcing.datm7.CRUJRA.0.5d.v1.c190604/"
            "clmforc.CRUJRAV1.1.c2019.0.5x0.5.Prec.1901-01.nc")

FILL = np.float32(1.0e30)

# stream name -> (directory, [(raw UKESM variable, ELM forcing variable, unit scale factor), ...])
# Set a scale factor other than 1.0 only if the raw units differ from what ELM expects:
#   PRECTmms: mm/s, FSDS/FLDS: W/m2, PSRF: Pa, TBOT: K, WIND: m/s, QBOT: kg/kg
STREAMS = {
    "Prec":  ("Precip6Hrly", [("pre",   "PRECTmms", 1.0)]),
    "Solr":  ("Solar6Hrly",  [("dswrf", "FSDS",     1.0)]),
    "TPQWL": ("TPHWL6Hrly",  [("pres",  "PSRF",     1.0),
                              ("tmp",   "TBOT",     1.0),
                              ("wind",  "WIND",     1.0),
                              ("spfh",  "QBOT",     1.0),
                              ("dlwrf", "FLDS",     1.0)]),
}

# ------------------------------------------------------------
# Month slicing for 6-hourly noleap data (1460 steps per year)
# ------------------------------------------------------------
DAYS_IN_MONTH = np.array([31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31])
STEPS = DAYS_IN_MONTH * 4
END_IDX = np.cumsum(STEPS)               # python-exclusive end
START_IDX = END_IDX - STEPS              # 0-based start
assert END_IDX[-1] == 1460

# ------------------------------------------------------------
# Target (ELM) grid and nearest-neighbour mapping
# ------------------------------------------------------------
with xr.open_dataset(REF_FILE) as ds_ref:
    lat_met = ds_ref["LATIXY"].values      # (360, 720)
    lon_met = ds_ref["LONGXY"].values      # (360, 720)

with xr.open_dataset(f"{RAW}/{PREFIX}.pre.{YEAR_START:04d}_6hr.noleap.nc") as ds_geo:
    lat = ds_geo["lat"].values
    lon = ds_geo["lon"].values.copy()
lon[lon < 0] += 360.0                       # same convention as the target grid (0..360)

# The grids are regular, so nearest neighbour in 2D is the same as
# nearest neighbour independently in lat and lon.
lat_t = lat_met[:, 0]
lon_t = lon_met[0, :]
assert np.allclose(lat_met, lat_t[:, None]) and np.allclose(lon_met, lon_t[None, :]), \
    "Reference LATIXY/LONGXY are not regular; use a 2D search instead."

d_lat = np.abs(lat_t[:, None] - lat[None, :])
d_lon = np.abs(lon_t[:, None] - lon[None, :])
lat_idx = d_lat.argmin(axis=1)              # (360,) index into UKESM lat
lon_idx = d_lon.argmin(axis=1)              # (720,) index into UKESM lon
print(f"max nearest-neighbour distance: lat {d_lat.min(axis=1).max():.3f}, "
      f"lon {d_lon.min(axis=1).max():.3f} degrees")
assert d_lat.min(axis=1).max() < 0.5 and d_lon.min(axis=1).max() < 0.5, \
    "Grid mismatch (check lon convention 0..360 vs -180..180)."

# ------------------------------------------------------------
# Main loops
# ------------------------------------------------------------
for sname, (sdir, varlist) in STREAMS.items():
    os.makedirs(f"{OUTDIR}/{sdir}", exist_ok=True)

    for j in range(YEAR_START, YEAR_END + 1):
        # open (lazily) one raw UKESM file per variable for this year
        raws = {mv: xr.open_dataset(f"{RAW}/{PREFIX}.{mv}.{j:04d}_6hr.noleap.nc")
                for mv, _, _ in varlist}

        for k in range(1, 13):
            outfile = f"{OUTDIR}/{sdir}/UKESM.{sname}.{j:04d}-{k:02d}.nc"
            if os.path.exists(outfile) and not OVERWRITE:
                continue

            # template: CRUJRA monthly file (1901 is reused for years <= 1901)
            ty = 1901 if j <= 1901 else j
            template = (f"{TMPL}/{sdir}/"
                        f"clmforc.cruncep.V4.c2011.0.5d.{sname}.{ty:04d}-{k:02d}.nc")
            ds_out = xr.load_dataset(template, decode_cf=False)

            s, e = int(START_IDX[k - 1]), int(END_IDX[k - 1])
            nt = e - s
            assert ds_out.sizes["time"] == nt, \
                f"{template}: template has {ds_out.sizes['time']} time steps, expected {nt}"

            enc = {}
            for mv, ev, scale in varlist:
                d = raws[mv][mv].isel(time=slice(s, e)).values       # (nt, lat, lon)
                assert d.shape[0] == nt, f"{mv} {j}: got {d.shape[0]} steps, expected {nt}"

                # nearest-neighbour remap onto the ELM 360x720 grid
                arr = d[:, lat_idx[:, None], lon_idx[None, :]].astype(np.float32) * np.float32(scale)

                bad = ~np.isfinite(arr)
                if bad.any():
                    print(f"  WARNING {j}-{k:02d} {ev}: {int(bad.sum())} non-finite values -> fill")
                    arr[bad] = FILL

                # keep the template's attributes except those that would conflict
                attrs = {a: b for a, b in ds_out[ev].attrs.items()
                         if a not in ("missing_value", "_FillValue", "scale_factor", "add_offset")}
                ds_out[ev] = (ds_out[ev].dims, arr, attrs)
                enc[ev] = {"_FillValue": FILL, "dtype": "float32"}

            ds_out["time"].attrs["units"] = f"days since {j:04d}-{k:02d}-01 00:00:00"
            ds_out.to_netcdf(outfile, mode="w", encoding=enc)
            print(j, k, sname, [ev for _, ev, _ in varlist])

        for r in raws.values():
            r.close()