"""Build the forcing package for the E3SMv3 LR historical SamudrACE-style emulators.

One package serves all three flavors; each checkpoint reads only the names it needs.
Per year YYYY it writes, into OUT_DIR:

  atmosphere-forcing-YYYY.nc   6-hourly (1460 steps, noleap), from the E3SM eam.h0 files:
                               SOLIN (time, lat, lon); global_mean_co2 (time) = h0 co2vmr;
                               PHIS, LANDFRAC (lat, lon), static in time;
                               ak_0..8, bk_0..8, P0 (vertical coordinate, scalars).
  aerosol/<window>/aerosol-YYYY.nc  (window = 60mo, 12mo, 01mo for methods clim5yr, run12mo, year)
                               6-hourly, same time axis: aerindexall, colccn.3 replaced by a
                               smoothed version (three methods, below). Only the C1A1
                               flavor reads one of them.
  ocean-forcing-YYYY.nc        5-daily time axis (73 steps, noleap), static data:
                               sea_surface_fraction (time, lat, lon; repeated, the loader
                               needs one such field), atmosphere_flux_fraction,
                               momentum_flux_fraction, LANDFRAC (lat, lon);
                               mask_0..18, mask_2d (lat, lon); idepth_0..19 (interface depths, m).

Aerosol methods (--aero-method, default: all three). Per cell and per 6-hourly step, linear
interpolation in time between values assigned to the mid-point of each month, plus a mean
diurnal composite (4 phases, deviation from the monthly mean); colccn.3 is clipped at 0:
  clim5yr  mid-month value = 5-year centred calendar-month mean (Y-2..Y+2, shifted to stay inside
           the full years 1940-2064); diurnal composite = that 5-year window's mean. Identical to
           the smooth part S of cft-diag/cpl_e3x/skill_shuf/scripts/build_overlays.py. Smears
           volcanic years into their neighbours (1985 is 2-20% above the raw 1985 monthly means
           because the window includes El Chichon, 1983).
  year     mid-month value = that month's own mean; diurnal composite = that month's own.
           Keeps each month's own anomaly (the run's weather at monthly scale).
  run12mo  mid-month value = centred 12-month running mean of the monthly means (2x12 moving
           average; the slow and volcanic part) + the 5-year centred mean seasonal anomaly of
           that calendar month (monthly mean minus the year's mean, averaged over Y-2..Y+2);
           diurnal composite = the 5-year window's mean. Windows are shifted to stay inside
           1940-01..2064-12.
Per-month statistics (monthly mean and diurnal composite, float64 npz) are cached in
--stats-cache; directories in --stats-dirs (e.g. skill_shuf's _stats/, same format) are
searched first.

Usage:
  python build_forcing.py --years 1985 2015                 # selected years
  python build_forcing.py --start 1950 --end 2064 --nproc 8 # full span, all kinds/methods
  options: --out DIR, --kinds atmosphere aerosol ocean, --aero-method clim5yr year run12mo,
           --stats-cache DIR, --stats-dirs DIR [DIR ...], --overwrite
"""

import argparse
import glob
import os
import sys
import time
from multiprocessing import Pool

import netCDF4 as nc
import numpy as np

RUN = "/pscratch/sd/m/mahf708/v3.LR.historical_0101.aigo/run"
CASE = "v3.LR.historical_0101.aigo"
INPUTS = "/pscratch/sd/m/mahf708/e3sm-hist-inputs"
FULL_YEAR_RANGE = (1940, 2064)  # years with all 12 h0 months
AER = ("aerindexall", "colccn.3")
NDAYS = np.array([31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31])
CUM = np.concatenate([[0], np.cumsum(NDAYS)[:-1]])
COMP = dict(zlib=True, complevel=4, shuffle=True)
PROVENANCE = (
    "Built by build_forcing.py from the E3SMv3 v3.LR.historical_0101 simulation "
    "(1 degree Gaussian grid remap of EAM/MPAS output)."
)


def h0(y, m):
    return f"{RUN}/{CASE}.eam.h0.{y}-{m:02d}.nc"


# The run starts at 1940-01-01 00:00 and writes no h0 record for that instant, so January 1940 has 123
# 6-hourly records instead of 124. read_month() adds the missing first record: SOLIN from 1941-01-01 00:00
# (same calendar position) scaled by the 1940/1941 ratio of global-mean SOLIN at 06:00; co2vmr from 06:00;
# other fields NaN (monthly statistics skip it).
FIRST_RECORD_FILL = ("1940-01-01 00:00 is not in the run's output (the run starts there); this record was filled: "
                     "SOLIN = SOLIN(1941-01-01 00:00) x [global-mean SOLIN 1940-01-01 06:00 / 1941-01-01 06:00], "
                     "co2vmr = its 1940-01-01 06:00 value.")


def read_month(y, m, names):
    """Times, bounds and the named h0 fields for one month, always 4 records per day starting at 00:00."""
    with nc.Dataset(h0(y, m)) as d:
        d.set_auto_mask(False)
        t, b = d["time"][:], d["time_bnds"][:]
        f = {v: d[v][:] for v in names}
    if len(t) == NDAYS[m - 1] * 4 - 1 and (y, m) == FULL_YEAR_RANGE[:1] + (1,):
        t = np.concatenate([[t[0] - 0.25], t])
        b = np.concatenate([[[b[0, 0] - 0.25, b[0, 0]]], b])
        for v, a in f.items():
            if v == "SOLIN":
                with nc.Dataset(h0(y + 1, 1)) as d1:
                    d1.set_auto_mask(False)
                    s0, s6 = d1["SOLIN"][0], d1["SOLIN"][1]
                    w = np.cos(np.deg2rad(d1["lat"][:]))[:, None]
                ratio = float((a[0] * w).sum() / (s6 * w).sum())
                first = s0 * ratio
            elif v == "co2vmr":
                first = a[0]
            else:
                first = np.full_like(a[0], np.nan, dtype=np.float64)
            f[v] = np.concatenate([np.asarray(first, dtype=np.float64)[None], a.astype(np.float64)])
    return t, b, f


def _copy_attrs(src, dst, skip=("_FillValue",)):
    for a in src.ncattrs():
        if a not in skip:
            dst.setncattr(a, src.getncattr(a))


def _new_file(path, ref, lat_lon_src, time_vals, time_attrs, time_bnds=None):
    out = nc.Dataset(path + ".tmp", "w", format="NETCDF4")
    out.createDimension("time", len(time_vals))
    out.createDimension("lat", len(lat_lon_src["lat"]))
    out.createDimension("lon", len(lat_lon_src["lon"]))
    for c in ("lat", "lon"):
        v = out.createVariable(c, "f8", (c,), fill_value=False)
        _copy_attrs(lat_lon_src[c], v)
        v[:] = lat_lon_src[c][:]
    t = out.createVariable("time", "f8", ("time",), fill_value=False)
    for k, val in time_attrs.items():
        t.setncattr(k, val)
    t[:] = time_vals
    if time_bnds is not None:
        out.createDimension("nbnd", 2)
        tb = out.createVariable("time_bnds", "f8", ("time", "nbnd"), fill_value=False)
        tb.setncattr("long_name", "time interval endpoints")
        tb[:] = time_bnds
    for a in ("case", "source", "source_id", "institution_id", "institution", "contact", "Conventions"):
        if a in ref.ncattrs():
            out.setncattr(a, ref.getncattr(a))
    return out


def _finish(out, path):
    out.close()
    os.replace(path + ".tmp", path)


# ---------------------------------------------------------------- atmosphere
def build_atmosphere(y, out_dir):
    path = f"{out_dir}/atmosphere-forcing-{y}.nc"
    times, bnds, solin, co2 = [], [], [], []
    for m in range(1, 13):
        t, b, f = read_month(y, m, ("SOLIN", "co2vmr"))
        times.append(t)
        bnds.append(b)
        solin.append(f["SOLIN"])
        co2.append(f["co2vmr"])
    times, bnds = np.concatenate(times), np.concatenate(bnds)
    assert len(times) == 1460 and np.allclose(np.diff(times), 0.25), (y, len(times))
    with nc.Dataset(h0(y, 1)) as d:
        d.set_auto_mask(False)
        tattrs = {a: d["time"].getncattr(a) for a in d["time"].ncattrs()}
        out = _new_file(path, d, d.variables, times, tattrs, bnds)
        v = out.createVariable("SOLIN", "f4", ("time", "lat", "lon"), chunksizes=(1, 180, 360), fill_value=False, **COMP)
        _copy_attrs(d["SOLIN"], v)
        v[:] = np.concatenate(solin)
        v = out.createVariable("global_mean_co2", "f8", ("time",), fill_value=False)
        _copy_attrs(d["co2vmr"], v)
        v.setncattr("note", "E3SM co2vmr (volume mixing ratio), renamed to the name the emulator uses")
        v[:] = np.concatenate(co2)
        for name in ("PHIS", "LANDFRAC"):
            v = out.createVariable(name, "f4", ("lat", "lon"), fill_value=False, **COMP)
            _copy_attrs(d[name], v)
            v.setncattr("note", "static in time in this simulation; first January record")
            v[:] = d[name][0]
        for name in [f"{p}_{k}" for k in range(9) for p in ("ak", "bk")] + ["P0"]:
            v = out.createVariable(name, "f8", (), fill_value=False)
            _copy_attrs(d[name], v)
            v.assignValue(d[name][...])
        out.title = f"Atmosphere forcing for {y} (6-hourly)"
        out.history = PROVENANCE + " Variables copied unchanged except co2vmr->global_mean_co2 and PHIS/LANDFRAC stored once."
        if y == FULL_YEAR_RANGE[0]:
            out.first_record_note = FIRST_RECORD_FILL
    _finish(out, path)
    return path


# ---------------------------------------------------------------- aerosol climatology
class Stats:
    def __init__(self, cache, extra_dirs):
        self.cache, self.dirs, self.mem = cache, [cache] + list(extra_dirs), {}

    def _compute(self, y, m):
        out = f"{self.cache}/{y}-{m:02d}.npz"
        _, _, f = read_month(y, m, AER)
        r = np.stack([f[v] for v in AER]).astype(np.float64)
        assert r.shape[1] == NDAYS[m - 1] * 4, (y, m, r.shape)
        mm = np.nanmean(r, 1)
        dc = np.stack([np.nanmean(r[:, h::4], 1) for h in range(4)], 1) - mm[:, None]
        np.savez(out + ".tmp.npz", mm=mm, dc=dc)
        os.replace(out + ".tmp.npz", out)
        return out

    def path(self, y, m, compute=True):
        for dd in self.dirs:
            p = f"{dd}/{y}-{m:02d}.npz"
            if os.path.exists(p):
                return p
        return self._compute(y, m) if compute else None

    def get(self, y, m):
        if (y, m) not in self.mem:
            z = np.load(self.path(y, m))
            self.mem[(y, m)] = (z["mm"], z["dc"])
        return self.mem[(y, m)]


def window(y):
    lo, hi = FULL_YEAR_RANGE
    start = min(max(y - 2, lo), hi - 4)
    return range(start, start + 5)


def clim5(st, y, m):
    yrs = window(y)
    return (np.mean([st.get(yy, m)[0] for yy in yrs], 0), np.mean([st.get(yy, m)[1] for yy in yrs], 0))


# full months with h0 data: 1940-01 .. 2064-12 (2065-01 holds a single record)
FIRST_MONTH, LAST_MONTH = (1940, 1), (2064, 12)


def _mi(y, m):
    return y * 12 + m - 1


def _ym(i):
    return i // 12, i % 12 + 1


def _clamp_month(y, m):
    i = min(max(_mi(y, m), _mi(*FIRST_MONTH)), _mi(*LAST_MONTH))
    return _ym(i)


def run12_months(y, m):
    """Weights of the centred 2x12 running mean around (y, m): months -6 and +6 get 1/24,
    months -5..+5 get 1/12. The 13-month window is shifted to stay inside FIRST..LAST_MONTH."""
    c = _mi(y, m)
    lo, hi = _mi(*FIRST_MONTH), _mi(*LAST_MONTH)
    c = min(max(c, lo + 6), hi - 6)
    return [(_ym(i), (0.5 if abs(i - c) == 6 else 1.0) / 12.0) for i in range(c - 6, c + 7)]


def seasonal_anomaly5(st, y, m):
    """5-year centred mean seasonal anomaly for calendar month m: monthly mean minus that
    year's annual mean (mean of its 12 monthly means), averaged over window(y)."""
    out = []
    for yy in window(y):
        annual = np.mean([st.get(yy, k)[0] for k in range(1, 13)], 0)
        out.append(st.get(yy, m)[0] - annual)
    return np.mean(out, 0)


def base_mean(method, st, y, m):
    """Smooth value assigned to the mid-point of month (y, m)."""
    if method == "clim5yr":
        return clim5(st, y, m)[0]
    if method == "year":
        return st.get(*_clamp_month(y, m))[0]
    if method == "run12mo":
        slow = sum(w * st.get(yy, mm)[0] for (yy, mm), w in run12_months(y, m))
        return slow + seasonal_anomaly5(st, y, m)
    raise ValueError(method)


def diurnal(method, st, y, m):
    """Mean diurnal composite (4 x 6-hourly phases, deviation from the monthly mean) for (y, m)."""
    if method in ("clim5yr", "run12mo"):
        return clim5(st, y, m)[1]
    return st.get(*_clamp_month(y, m))[1]  # year: that month's own composite


def center(y, m):
    return (y - 1850) * 365 + CUM[m - 1] + NDAYS[m - 1] / 2.0


def prev_next(y, m):
    return ((y, m - 1) if m > 1 else (y - 1, 12)), ((y, m + 1) if m < 12 else (y + 1, 1))


def smooth(st, y, m, times, method="clim5yr"):
    c0, d0 = base_mean(method, st, y, m), diurnal(method, st, y, m)
    (yp, mp), (yn, mn) = prev_next(y, m)
    cp, cn = base_mean(method, st, yp, mp), base_mean(method, st, yn, mn)
    tc, tp, tn = center(y, m), center(yp, mp), center(yn, mn)
    o = np.empty((2, len(times)) + c0.shape[1:])
    for i, t in enumerate(times):
        if t < tc:
            w = (t - tp) / (tc - tp)
            base = (1 - w) * cp + w * c0
        else:
            w = (t - tc) / (tn - tc)
            base = (1 - w) * c0 + w * cn
        o[:, i] = base + d0[:, i % 4]
    return o


class _Recorder:
    """Stand-in for Stats that records which (year, month) statistics a method needs."""

    def __init__(self):
        self.keys = set()

    def get(self, y, m):
        self.keys.add((y, m))
        return np.zeros((2, 1)), np.zeros((2, 4, 1))


def needed_stat_months(y, method="clim5yr"):
    rec = _Recorder()
    for m in range(1, 13):
        smooth(rec, y, m, [center(y, m)], method)
    return sorted(rec.keys)


AERO_NOTES = {
    "clim5yr": "linear interpolation between mid-month points of the 5-year centred calendar-month mean "
    "(years Y-2..Y+2, shifted to stay within 1940-2064) plus that 5-year window's mean diurnal composite "
    "for the month",
    "year": "linear interpolation between mid-month points of this year's own monthly means (the adjacent "
    "months of the neighbouring years at the year edges; Dec 2064 is held flat to the end) plus the "
    "month's own mean diurnal composite",
    "run12mo": "linear interpolation between mid-month points of [centred 12-month running mean of the "
    "monthly means (2x12 moving average: months -6 and +6 weighted 1/24, months -5..+5 weighted 1/12; "
    "window shifted to stay within 1940-01..2064-12)] + [5-year centred mean seasonal anomaly for the "
    "calendar month (monthly mean minus that year's mean of 12 monthly means, averaged over Y-2..Y+2, "
    "shifted to stay within 1940-2064)], plus the 5-year window's mean diurnal composite for the month",
}
AERO_METHODS = tuple(AERO_NOTES)


# output folder per method, named by the averaging window that sets the year-to-year values
AERO_DIRS = {"clim5yr": "60mo", "run12mo": "12mo", "year": "01mo"}


def aerosol_path(out_dir, method, y):
    return f"{out_dir}/aerosol/{AERO_DIRS[method]}/aerosol-{y}.nc"


def build_aerosol(y, out_dir, st, method="clim5yr"):
    path = aerosol_path(out_dir, method, y)
    os.makedirs(os.path.dirname(path), exist_ok=True)
    times, parts = [], []
    for m in range(1, 13):
        t, _, _ = read_month(y, m, ())
        times.append(t)
        parts.append(smooth(st, y, m, t, method))
    times = np.concatenate(times)
    x = np.concatenate(parts, axis=1)
    nclip = int((x[1] < 0).sum())
    x[1] = np.maximum(x[1], 0.0)
    with nc.Dataset(h0(y, 1)) as d:
        d.set_auto_mask(False)
        tattrs = {a: d["time"].getncattr(a) for a in d["time"].ncattrs() if a != "bounds"}
        out = _new_file(path, d, d.variables, times, tattrs)
        for k, name in enumerate(AER):
            v = out.createVariable(name, "f4", ("time", "lat", "lon"), chunksizes=(1, 180, 360), fill_value=False, **COMP)
            _copy_attrs(d[name], v)
            v.setncattr(
                "note",
                f"Smoothed version ({method}) of the simulation's field: " + AERO_NOTES[method]
                + ("; clipped at 0" if name == "colccn.3" else ""),
            )
            v[:] = x[k].astype(np.float32)
        out.title = f"Smoothed aerosol diagnostics ({method}) for {y} (6-hourly)"
        out.aerosol_method = method
        out.history = PROVENANCE + f" Aerosol method {method}; colccn.3 values clipped at 0: {nclip}."
    _finish(out, path)
    return path, nclip


# ---------------------------------------------------------------- ocean
OCN = f"{RUN}/{CASE}.mpaso.hist.am"


def build_ocean(y, out_dir):
    path = f"{out_dir}/ocean-forcing-{y}.nc"
    times = (y - 1940) * 365 + 5.0 * np.arange(73)
    src_dc = f"{OCN}.fmeDepthCoarsening5D.1985-01.remapped.nc"
    src_df = f"{OCN}.fmeDerivedFields5D.1985-01.remapped.nc"
    with nc.Dataset(src_dc) as d, nc.Dataset(src_df) as f, nc.Dataset(
        f"{INPUTS}/landfrac5d/landfrac5d.{y}.nc"
    ) as lf, nc.Dataset(f"{INPUTS}/momfrac5d/momfrac5d.{y}.nc") as mf:
        for x in (d, f, lf, mf):
            x.set_auto_mask(False)
        tattrs = {"long_name": "time", "units": "days since 1940-01-01 00:00:00", "calendar": "noleap", "axis": "T"}
        out = _new_file(path, d, d.variables, times, tattrs)
        for src, name in [(lf, "sea_surface_fraction"), (lf, "atmosphere_flux_fraction"), (mf, "momentum_flux_fraction"), (lf, "LANDFRAC")]:
            a = src[name][:]
            assert np.array_equal(a, np.broadcast_to(a[0], a.shape)), (y, name, "not static")
            if name == "sea_surface_fraction":
                # fme's loader takes the horizontal dims from the first (time, lat, lon)
                # variable, so one field keeps its (repeated) time dimension.
                v = out.createVariable(name, "f4", ("time", "lat", "lon"), chunksizes=(1, 180, 360), fill_value=False, **COMP)
                v[:] = np.broadcast_to(a[0], (len(times),) + a[0].shape)
            else:
                v = out.createVariable(name, "f4", ("lat", "lon"), fill_value=False, **COMP)
                v[:] = a[0]
            _copy_attrs(src[name], v)
            v.setncattr("note", "static in time in this simulation")
        for k in range(19):
            v = out.createVariable(f"mask_{k}", "f4", ("lat", "lon"), fill_value=False, **COMP)
            _copy_attrs(d[f"mask_{k}"], v)
            v[:] = d[f"mask_{k}"][:]
        v = out.createVariable("mask_2d", "f4", ("lat", "lon"), fill_value=False, **COMP)
        _copy_attrs(f["mask_2d"], v)
        v[:] = f["mask_2d"][:]
        for k in range(20):  # 20 interfaces for the 19 layers
            v = out.createVariable(f"idepth_{k}", "f8", (), fill_value=False)
            _copy_attrs(d[f"idepth_{k}"], v)
            v.assignValue(d[f"idepth_{k}"][...])
        out.title = f"Ocean forcing (static fields on a 5-daily time axis) for {y}"
        out.history = PROVENANCE + " Fields are static; the time axis only tells the loader which dates exist."
    _finish(out, path)
    return path


# ---------------------------------------------------------------- driver
def _out_path(kind, y, out_dir, method):
    if kind == "aerosol":
        return aerosol_path(out_dir, method, y)
    return f"{out_dir}/{ {'atmosphere': 'atmosphere-forcing', 'ocean': 'ocean-forcing'}[kind] }-{y}.nc"


def _job(args):
    kind, method, y, out_dir, cache, extra = args
    t0 = time.time()
    if os.path.exists(_out_path(kind, y, out_dir, method)):
        return f"{kind}{'/' + method if method else ''} {y} exists, skipped"
    if kind == "atmosphere":
        p = build_atmosphere(y, out_dir)
        info = ""
    elif kind == "ocean":
        p = build_ocean(y, out_dir)
        info = ""
    else:
        p, nclip = build_aerosol(y, out_dir, Stats(cache, extra), method)
        info = f" colccn.3 clipped={nclip}"
    return f"{kind}{'/' + method if method else ''} {y} {os.path.getsize(p) / 1e6:.1f} MB {time.time() - t0:.0f}s{info}"


def _stat_job(args):
    y, m, cache, extra = args
    st = Stats(cache, extra)
    return st.path(y, m)


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--years", type=int, nargs="*")
    ap.add_argument("--start", type=int)
    ap.add_argument("--end", type=int)
    ap.add_argument("--out", default=os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "forcing_data"))
    ap.add_argument("--kinds", nargs="*", default=["atmosphere", "aerosol", "ocean"])
    ap.add_argument("--aero-method", nargs="+", default=list(AERO_METHODS), choices=AERO_METHODS)
    ap.add_argument("--stats-cache", default=os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "_build", "aero_stats"))
    ap.add_argument("--stats-dirs", nargs="*", default=[])
    ap.add_argument("--nproc", type=int, default=4)
    ap.add_argument("--overwrite", action="store_true")
    a = ap.parse_args()
    years = a.years or list(range(a.start, a.end + 1))
    lo, hi = FULL_YEAR_RANGE
    if min(years) < lo or max(years) > hi:
        sys.exit(f"years must be within {lo}..{hi} (aerosol windows are shifted to stay inside these years)")
    out = os.path.abspath(a.out)
    cache = os.path.abspath(a.stats_cache)
    os.makedirs(out, exist_ok=True)
    os.makedirs(cache, exist_ok=True)
    jobs = []
    for y in years:
        for k in a.kinds:
            for method in (a.aero_method if k == "aerosol" else [None]):
                jobs.append((k, method, y, out, cache, a.stats_dirs))
    if a.overwrite:
        for j in jobs:
            f = _out_path(j[0], j[2], out, j[1])
            if os.path.exists(f):
                os.remove(f)
    t0 = time.time()
    with Pool(a.nproc) as p:
        if "aerosol" in a.kinds:
            need = sorted({ym for y in years for meth in a.aero_method for ym in needed_stat_months(y, meth)})
            st = Stats(cache, a.stats_dirs)
            todo = [(y, m, cache, a.stats_dirs) for (y, m) in need if st.path(y, m, compute=False) is None]
            print(f"aerosol stats: {len(need)} months needed, {len(todo)} to compute", flush=True)
            for i, _ in enumerate(p.imap_unordered(_stat_job, todo)):
                if i % 60 == 0:
                    print(f"  stats {i + 1}/{len(todo)} {time.time() - t0:.0f}s", flush=True)
        for msg in p.imap_unordered(_job, jobs):
            print(msg, f"(elapsed {time.time() - t0:.0f}s)", flush=True)
    print(f"done {len(jobs)} files in {time.time() - t0:.0f}s", flush=True)


if __name__ == "__main__":
    main()