"""HIST forcing with one shared precipitation series after 1850.

Temperature: Mira's BE_on_all strands (member's own LMR before 1850, shared
from 1850; symlinked). Precipitation: the member's own LMR before 1850, then
SOURCE's from 1850 on (all HIST members share it, as temperature). Elevation:
oldBE (symlinked). NAT runs keep their own precipitation (user, 2026-10-24).

    python build_shared_precip.py <source member> <member> [<member> ...]
    e.g. python build_shared_precip.py 148 066 008 010 055 041

-> unstacked/lmr_seas_fullsplit_oldBE_BE_on_all_p<source>/
"""
import os
import sys
import numpy as np
import xarray as xr

U = '/home/www/fmaussion/runs/lmr_2025/data/unstacked'
T_ROUND = f'{U}/lmr_seas_fullsplit_oldBE_BE_on_all'
P_ROUND = f'{U}/lmr_seas_fullsplit_oldBE'
YEAR = 1850

src = int(sys.argv[1])
members = [int(m) for m in sys.argv[2:]]
out = f'{U}/lmr_seas_fullsplit_oldBE_BE_on_all_p{src:03d}'
os.makedirs(f'{out}/ensemble', exist_ok=True)
if not os.path.exists(f'{out}/elev.nc'):
    os.symlink(f'{P_ROUND}/elev.nc', f'{out}/elev.nc')

coder = xr.coders.CFDatetimeCoder(use_cftime=True)
with xr.open_dataset(f'{P_ROUND}/ensemble/p_monthly_ens{src:03d}.nc', decode_times=coder) as ds:
    psrc = ds.load()
for m in members:
    t = f'{out}/ensemble/tp_monthly_ens{m:03d}.nc'
    if not os.path.exists(t):
        os.symlink(f'{T_ROUND}/ensemble/tp_monthly_ens{m:03d}.nc', t)
    with xr.open_dataset(f'{P_ROUND}/ensemble/p_monthly_ens{m:03d}.nc', decode_times=coder) as ds:
        pm = ds.load()
    assert (pm.time.values == psrc.time.values).all()
    assert np.allclose(pm.lon, psrc.lon) and np.allclose(pm.lat, psrc.lat)
    late = pm['time.year'] >= YEAR
    new = pm.copy()
    new['pr'] = pm.pr.where(~late, psrc.pr)
    new.attrs['history'] = (f'precip of ens{m:03d} before {YEAR}, of ens{src:03d} '
                            f'from {YEAR} (data/build_shared_precip.py)')
    f = f'{out}/ensemble/p_monthly_ens{m:03d}.nc'
    new.to_netcdf(f + '.tmp')
    os.replace(f + '.tmp', f)
    # checks: own precip before YEAR, the source's from YEAR (NaN-safe)
    early = ~late.values
    assert np.array_equal(new.pr.values[..., early], pm.pr.values[..., early], equal_nan=True) \
        if new.pr.dims[-1] == 'time' else np.array_equal(
            new.pr.isel(time=early).values, pm.pr.isel(time=early).values, equal_nan=True)
    assert np.array_equal(new.pr.isel(time=late.values).values,
                          psrc.pr.isel(time=late.values).values, equal_nan=True)
    print('written', f)