"""Shifted NAT temperature: a cooler (or warmer) natural counterfactual from 1850 on.

Copies every newnat NAT member (lmr_seas_fullsplit_oldBE_newnat/ensemble/
tp_nat_monthly_ens###.nc) with
    tas(t) - delta * ramp(t),  ramp = 0 before 1850, linear 0 -> 1 over 1850-1900
    (the provider's own LMR -> GCM blend period), 1 from 1900 on.
Before 1850 the files are bit-identical to newnat, so the 1850 glacier states (and
HIST == NAT before 1850) are untouched. Only the NAT temperature changes: precip is read
from the LAPrec NAT folder as usual.

    python build_nat_shift.py 0.5      -> unstacked/lmr_seas_fullsplit_oldBE_newnat_shift-0.5/
    python build_nat_shift.py 0.2 [--ramp0 1850 --ramp1 1900]

delta > 0 = cooler NAT. ~2.1 GB per variant, ~1 min with 8 processes.
"""
import argparse
import glob
import os
from concurrent.futures import ProcessPoolExecutor

import numpy as np
import xarray as xr

HERE = os.path.dirname(os.path.abspath(__file__))
SRC = f'{HERE}/unstacked/lmr_seas_fullsplit_oldBE_newnat'


def ramp(years, y0, y1):
    return np.clip((years - y0) / (y1 - y0), 0, 1)


def one(args):
    f, out, delta, y0, y1 = args
    with xr.open_dataset(f, decode_times=xr.coders.CFDatetimeCoder(use_cftime=True)) as ds:
        ds = ds.load()
    enc = {k: v for k, v in ds.tas.encoding.items()
           if k in ('dtype', 'zlib', 'complevel', '_FillValue', 'contiguous')}
    yrs = np.array([t.year for t in ds.time.values])
    r = ramp(yrs, y0, y1).astype(np.float32)
    old = ds.tas.values.copy()
    ds['tas'] = ds.tas - np.float32(delta) * xr.DataArray(r, dims='time', coords={'time': ds.time})
    ds['tas'] = ds.tas.astype(np.float32)
    ds.tas.attrs['history'] = (f'newnat NAT minus {delta} K x linear ramp '
                               f'({y0}: 0 -> {y1}: 1), build_nat_shift.py')
    pre = yrs < y0
    assert np.array_equal(ds.tas.values[..., pre], old[..., pre], equal_nan=True)
    ds.to_netcdf(out, encoding={'tas': enc})
    return os.path.basename(out)


if __name__ == '__main__':
    p = argparse.ArgumentParser()
    p.add_argument('delta', type=float, help='K, > 0 = cooler')
    p.add_argument('--ramp0', type=int, default=1850)
    p.add_argument('--ramp1', type=int, default=1900)
    a = p.parse_args()
    tag = f'shift-{a.delta:g}' + ('' if (a.ramp0, a.ramp1) == (1850, 1900)
                                   else f'_r{a.ramp0}-{a.ramp1}')
    dst = f'{HERE}/unstacked/lmr_seas_fullsplit_oldBE_newnat_{tag}/ensemble'
    os.makedirs(dst, exist_ok=True)
    files = sorted(glob.glob(f'{SRC}/ensemble/tp_nat_monthly_ens*.nc'))
    jobs = [(f, f'{dst}/{os.path.basename(f)}', a.delta, a.ramp0, a.ramp1) for f in files]
    with ProcessPoolExecutor(8) as ex:
        done = list(ex.map(one, jobs))
    print(f'{len(done)} files -> {dst}')