"""HIST forcing with LAPrec precipitation (see laprec_generator.py).

Temperature: Mira's BE_on_all strand of the member (symlinked). Elevation:
oldBE (symlinked). Precipitation: generated before 1900 (the member's own
draws before 1850, the shared series 1850-1899), LAPrec 1900-2023, ERA5
mapped to LAPrec 2024-2025.

    python build_laprec_precip.py <member> [<member> ...]     e.g. 10

-> unstacked/lmr_BE_on_all_laprec_v20261031/ensemble/p_monthly_ens###.nc
   + draws_hist_ens###.csv (which pool year each year comes from)
Needs the pool file (python laprec_generator.py). Prints checks.
"""
import os
import sys

import numpy as np
import pandas as pd
import xarray as xr

import laprec_generator as g

P_OLD = f'{g.U}/lmr_seas_fullsplit_oldBE'
T_ROUND = f'{g.U}/lmr_seas_fullsplit_oldBE_BE_on_all'


def glacier_cells(pool):
    """(i, j) of the 0.25 deg cells nearest to all RGI 11-01 glaciers + weights."""
    st = pd.read_csv(g.STATS, index_col=0)
    lat, lon = pool.lat.values, pool.lon.values
    w = np.zeros((len(lat), len(lon)))
    il = np.abs(st.cenlat.values[:, None] - lat).argmin(axis=1)
    jl = np.abs(st.cenlon.values[:, None] - lon).argmin(axis=1)
    np.add.at(w, (il, jl), st.rgi_area_km2.values)
    allg = pd.read_csv(f'{g.ROOT}/climate_analysis/glaciers_rgi62_1101.csv')
    ia = np.abs(allg.CenLat.values[:, None] - lat).argmin(axis=1)
    ja = np.abs(allg.CenLon.values[:, None] - lon).argmin(axis=1)
    return w / w.sum(), (ia, ja)


def stats(reg):
    """reg: DataFrame year x month (mm), glacier weighted."""
    ann = reg.sum(axis=1)
    d = {'annual mean': ann.mean(), 'annual CV %': 100 * ann.std() / ann.mean(),
         'DJF CV %': 100 * (lambda x: x.std() / x.mean())(reg[12].shift(1) + reg[1] + reg[2]),
         'JJA CV %': 100 * (lambda x: x.std() / x.mean())(reg[[6, 7, 8]].sum(axis=1)),
         'month CV % (median)': 100 * np.median(reg.std() / reg.mean()),
         'JJAS share %': 100 * reg[[6, 7, 8, 9]].sum(axis=1).sum() / ann.sum(),
         'lag-1 annual': ann.autocorr(1)}
    return d


if __name__ == '__main__':
    members = [int(m) for m in sys.argv[1:]]
    pool = g.load_pool()
    os.makedirs(f'{g.OUT}/ensemble', exist_ok=True)
    if not os.path.exists(f'{g.OUT}/elev.nc'):
        os.symlink(f'{P_OLD}/elev.nc', f'{g.OUT}/elev.nc')

    w, (ia, ja) = glacier_cells(pool)
    assert pool.laprec_covered.values[ia, ja].all(), 'a glacier cell is not covered by LAPrec'
    assert pool.era5_covered.values[ia, ja].all(), 'a glacier cell has no ERA5 point'
    print('all', len(ia), 'RGI 11-01 glaciers sit in LAPrec- and ERA5-covered cells')

    for m in members:
        t = f'{g.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)
        p, years, table = g.hist_precip(m, pool)
        hist = (f'HIST precipitation of ens{m:03d}: whole LAPrec years (1900-2023) drawn '
                f'within JJAS temperature terciles before {g.OBS_Y0} (own draws before '
                f'{g.SPLIT}, the shared series {g.SPLIT}-{g.OBS_Y0 - 1}), LAPrec1871 v1.3 '
                f'{g.OBS_Y0}-{g.POOL_Y1}, ERA5 mapped to LAPrec {g.ERA5_YEARS[0]}-'
                f'{g.ERA5_YEARS[-1]}. data/build_laprec_precip.py')
        ds = g.to_dataset(p, years, pool, hist)
        f = f'{g.OUT}/ensemble/p_monthly_ens{m:03d}.nc'
        ds.to_netcdf(f + '.tmp')
        os.replace(f + '.tmp', f)
        table.to_csv(f'{g.OUT}/ensemble/draws_hist_ens{m:03d}.csv', index=False)
        print('written', f)

        # --- checks
        with xr.open_dataset(f, decode_times=g.TC) as chk, \
                xr.open_dataset(f'{P_OLD}/ensemble/p_monthly_ens{m:03d}.nc', decode_times=g.TC) as old:
            assert (chk.time.values == old.time.values).all()
            assert np.allclose(chk.lat, old.lat) and np.allclose(chk.lon, old.lon)
            assert chk.pr.dims == old.pr.dims
            assert np.isfinite(chk.pr.values).all() and (chk.pr.values >= 0).all()
        reg = pd.DataFrame(np.einsum('ij,ymij->ym', w, p), index=years, columns=range(1, 13))
        rows = {'generated 1-1849': stats(reg.loc[1:1849]),
                'generated 1850-1899': stats(reg.loc[1850:1899]),
                'LAPrec 1900-2023': stats(reg.loc[1900:2023])}
        print(pd.DataFrame(rows).round(2).to_string())
        print('2024 / 2025 annual (mm):', reg.loc[2024].sum().round(0), reg.loc[2025].sum().round(0),
              '; LAPrec 1991-2020 mean', reg.loc[1991:2020].sum(axis=1).mean().round(0))
        # obs part == pool obs exactly
        assert np.array_equal(p[(years >= 1900) & (years <= 2023)], pool.obs.values)
        # coupling kept: JJA precip vs interannual JJAS temperature, before 1850
        th = g.jjas_glacier_weighted(g.HIST_T.format(m))
        a = g.interannual(th)
        y = np.arange(1, 1850)
        r = np.corrcoef(reg.loc[y, [6, 7, 8]].sum(axis=1), a.loc[y])[0, 1]
        yo = np.arange(1900, 2024)
        ro = np.corrcoef(reg.loc[yo, [6, 7, 8]].sum(axis=1), a.loc[yo])[0, 1]
        print(f'r(JJA P, interannual JJAS T): generated 1-1849 {r:.2f}, LAPrec 1900-2023 {ro:.2f}')
        use = table.loc[table.year < 1850, 'pool_year'].value_counts()
        print(f'pool use before 1850: {use.size} of 124 years drawn, '
              f'min/max {use.min()}/{use.max()} times (mean {use.mean():.1f})')
        # HIST == NAT before 1850 (if the member has a NAT temperature file)
        if os.path.exists(g.NAT_T.format(m)):
            pn, yn, _ = g.nat_precip(m, pool)
            k = yn < g.SPLIT
            same = np.array_equal(pn[k], p[(years >= yn[0]) & (years < g.SPLIT)])
            print(f'NAT ens{m:03d} precip == HIST before {g.SPLIT}: {same}')
            assert same