#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Create stable-climate forcings (1850-1869) for E3SM/ELM runs.

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

Usage: python create_stable_forcings.py [Prec] [Solr] [TPQWL]
       (no arguments = all streams)
"""

import os
import sys
import numpy as np
import xarray as xr

# ------------------------------------------------------------
# User settings
# ------------------------------------------------------------
YEAR_START, YEAR_END = 1850, 1869
OVERWRITE = True            # False: skip output files that already exist

RAW    = "/global/cfs/cdirs/e3sm/inputdata/atm/datm7/WIEMIP/raw_data/stable/05deg"
PREFIX = "stable"
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/stable"
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 -> (directory, [(raw variable, ELM variable, unit scale factor), ...])
# Change a scale factor only if raw units differ from ELM's expectations:
#   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                       # target grid uses 0..360

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 source lat
lon_idx = d_lon.argmin(axis=1)              # (720,) index into source 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
# ------------------------------------------------------------
selected = sys.argv[1:]

for sname, (sdir, varlist) in STREAMS.items():
    if selected and sname not in selected:
        continue
    os.makedirs(f"{OUTDIR}/{sdir}", exist_ok=True)

    for j in range(YEAR_START, YEAR_END + 1):
        # open (lazily) one raw 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}/{PREFIX}.{sname}.{j:04d}-{k:02d}.nc"
            if os.path.exists(outfile) and not OVERWRITE:
                continue

            # template: CRUJRA monthly file (1901 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:
                da = raws[mv][mv]
                # force (time, lat, lon) order regardless of how the file stores it
                da = da.transpose("time", "lat", "lon")
                d = da.isel(time=slice(s, e)).values
                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

                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()