"""NAT forcing with LAPrec-resampled precipitation (see laprec_generator.py).

The companion of build_laprec_precip.py (HIST). Every year of a NAT member is
a whole LAPrec year (1900-2023, as ratio to its climatology) drawn within the
tercile of the member's own interannual NAT JJAS temperature, seeded by
(member, year). Before 1850 the NAT and HIST temperatures of a member are the
same, so are the draws: NAT == HIST precipitation before 1850, bit for bit
(checked here).

Written to their own "server" folder, so that lmr_forcing.process_mira_data
reads them unchanged (server = .../nat, temp_server = newnat, temp_var =
tp_nat_monthly):

    python build_laprec_nat.py <member> [<member> ...]

-> unstacked/lmr_BE_on_all_laprec_v20261031/nat/ensemble/p_monthly_ens###.nc
   + draws_nat_ens###.csv, and nat/elev.nc (symlink to oldBE).
"""
import os
import sys

import numpy as np
import xarray as xr

import laprec_generator as g

NAT_OUT = f'{g.OUT}/nat'
P_OLD = f'{g.U}/lmr_seas_fullsplit_oldBE'

if __name__ == '__main__':
    members = [int(m) for m in sys.argv[1:]]
    pool = g.load_pool()
    os.makedirs(f'{NAT_OUT}/ensemble', exist_ok=True)
    if not os.path.lexists(f'{NAT_OUT}/elev.nc'):
        os.symlink(f'{P_OLD}/elev.nc', f'{NAT_OUT}/elev.nc')
    for m in members:
        p, years, draws = g.nat_precip(m, pool)
        hist = (f'NAT precipitation of ens{m:03d}: whole LAPrec years (1900-2023) drawn '
                f'within the terciles of its interannual NAT JJAS temperature, seeded by '
                f'(member, year); == HIST before {g.SPLIT}. data/build_laprec_nat.py')
        ds = g.to_dataset(p, years, pool, hist)
        f = f'{NAT_OUT}/ensemble/p_monthly_ens{m:03d}.nc'
        ds.to_netcdf(f + '.tmp')
        os.replace(f + '.tmp', f)
        draws.rename('pool_year').rename_axis('year').to_csv(
            f'{NAT_OUT}/ensemble/draws_nat_ens{m:03d}.csv')
        # NAT == HIST before 1850, where a HIST file exists
        fh = f'{g.OUT}/ensemble/p_monthly_ens{m:03d}.nc'
        if os.path.exists(fh):
            with xr.open_dataset(fh, use_cftime=True) as h, xr.open_dataset(f, use_cftime=True) as n:
                sel = slice(None, f'{g.SPLIT - 1}-12-31')
                same = np.array_equal(h.pr.sel(time=sel).values, n.pr.sel(time=sel).values)
            print(f'ens{m:03d}: NAT == HIST before {g.SPLIT}: {same}')
            assert same
        print('written', f, f'years {years[0]}-{years[-1]}')