"""LAPrec-based precipitation: observations from 1900, resampled years otherwise.

Decided 2026-10-31 (see climate_analysis/README.md, LAPrec sections):
- HIST precipitation = LAPrec1871 v1.3 from 1900 to 2023 (before 1900 its
  seasonal cycle is suspicious), ERA5 for 2024-2025, rescaled to the LAPrec
  climatology.
- Every random series (HIST before 1900, all NAT) is made of whole LAPrec
  calendar years (1900-2023, the "pool"), all months and all cells together,
  as ratios to the 1900-2023 climatology. That keeps LAPrec's monthly
  variability, seasonal shape and spatial covariance.
- Coupling to temperature: the pool year is drawn among the third of the pool
  in the same tercile of interannual JJAS temperature as the target year
  (wet summers are cool: observed r -0.40 for JJA; the terciles keep about
  -0.31). Interannual = the glacier-weighted JJAS temperature minus its
  centred 31-yr running mean, computed separately before and after 1850, so
  that the classes before 1850 only depend on the (shared) pre-1850 years.
- One draw per (member, year), seeded by both: HIST and NAT of the same
  member get the same precipitation wherever their temperature classes are
  the same, in particular before 1850.

Everything is on the 0.25 deg grid of Mira's precipitation files. Cells that
LAPrec does not cover (none of them has a glacier) take the values of the
nearest covered cell.

The pool file (POOL_FILE) is built once with `python laprec_generator.py`.
"""
import os
import sys

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

ROOT = '/home/www/fmaussion/runs/lmr_2025'
sys.path.insert(0, f'{ROOT}/climate_analysis')
import laprec  # noqa: E402

U = f'{ROOT}/data/unstacked'
OUT = f'{U}/lmr_BE_on_all_laprec_v20261031'
POOL_FILE = f'{OUT}/laprec_pool.nc'
GRID_FILE = f'{U}/lmr_seas_fullsplit_oldBE/ensemble/p_monthly_ens010.nc'    # the 0.25 deg grid
HIST_T = f'{U}/lmr_seas_fullsplit_oldBE_BE_on_all/ensemble/tp_monthly_ens{{:03d}}.nc'
NAT_T = f'{U}/lmr_seas_fullsplit_oldBE_newnat/ensemble/tp_nat_monthly_ens{{:03d}}.nc'
STATS = f'{ROOT}/round5/filtered_ens/ens_pf3.5_mf5_lam12/glacier_statistics_reg11.csv'
ERA5 = ('/home/users/fmaussion/www_oggm/climate/era5/monthly/v1.2/flattened/'
        'era5_tp_global_monthly_1940_2025_flat_glaciers_v2026.01.14.nc')
POOL_Y0, POOL_Y1 = 1900, 2023
OBS_Y0 = 1900                  # HIST is observed from this year
SPLIT = 1850                   # members share their temperature before this year
ERA5_YEARS = (2024, 2025)
ERA5_REF = (1991, 2020)        # climatology period for rescaling ERA5 to LAPrec
N_CLASS = 3
WINDOW = 31
SEED = 20261031
SHARED_ID = 0                  # "member" id of the shared HIST series 1850-1899
TC = xr.coders.CFDatetimeCoder(use_cftime=True)


# ------------------------------------------------------------------ temperature

def jjas_glacier_weighted(f):
    """Glacier-area weighted JJAS mean temperature per year, from a T file."""
    st = pd.read_csv(STATS, index_col=0)
    with xr.open_dataset(f, decode_times=TC) as ds:
        lat, lon = ds.lat.values, ds.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)
        w /= w.sum()
        t = ds.tas.transpose('lat', 'lon', 'time').values.astype(np.float64)
        yrs = ds['time.year'].values
        mon = ds['time.month'].values
    reg = np.einsum('ij,ijt->t', w, t)
    s = pd.Series(reg, index=pd.MultiIndex.from_arrays([yrs, mon]))
    jjas = s[s.index.get_level_values(1).isin([6, 7, 8, 9])]
    return jjas.groupby(level=0).mean()


def interannual(t, split=SPLIT):
    """T minus its centred 31-yr running mean, separately before / from split."""
    out = []
    for seg in [t.loc[:split - 1], t.loc[split:]]:
        if len(seg):
            out.append(seg - seg.rolling(WINDOW, center=True, min_periods=WINDOW // 2 + 1).mean())
    return pd.concat(out)


def classes(anom, ref=None):
    """Tercile class (0 cold, 1, 2 warm) of each value, thresholds from ref."""
    ref = anom if ref is None else ref
    q = np.quantile(ref.values, np.linspace(0, 1, N_CLASS + 1)[1:-1])
    return pd.Series(np.digitize(anom.values, q), index=anom.index)


def target_classes(t):
    """Classes of a member's years: thresholds per segment (before / from SPLIT)."""
    a = interannual(t)
    return pd.concat([classes(a.loc[:SPLIT - 1]), classes(a.loc[SPLIT:])]
                     if a.index.max() >= SPLIT else [classes(a)])


# ------------------------------------------------------------------ the pool

def _fill_uncovered(a, covered):
    """Copy the nearest covered cell (grid index distance) into uncovered ones."""
    ci, cj = np.where(covered)
    for i, j in zip(*np.where(~covered)):
        k = np.argmin((ci - i) ** 2 + (cj - j) ** 2)
        a[..., i, j] = a[..., ci[k], cj[k]]
    return a


def build_pool():
    with xr.open_dataset(GRID_FILE, decode_times=TC) as ds:
        glat, glon = ds.lat.values, ds.lon.values
    data, lplon, lplat, years = laprec.load()
    valid = np.isfinite(data).all(axis=(0, 1))
    ny = len(years)
    box = np.full((ny, 12, len(glat), len(glon)), np.nan, dtype=np.float64)
    for i, la in enumerate(glat):
        for j, lo in enumerate(glon):
            m = valid & (np.abs(lplat - la) <= 0.125) & (np.abs(lplon - lo) <= 0.125)
            if m.any():
                box[:, :, i, j] = data[:, :, m].mean(axis=-1)
    covered = np.isfinite(box).all(axis=(0, 1))
    box = _fill_uncovered(box, covered)
    sel = (years >= POOL_Y0) & (years <= POOL_Y1)
    obs = box[sel]                                   # (pool year, month, lat, lon)
    clim = obs.mean(axis=0)                          # (month, lat, lon)
    assert (clim > 0).all()
    ratio = obs / clim

    # Classes of the pool years: the shared HIST temperature (== BE from 1900)
    th = jjas_glacier_weighted(HIST_T.format(10))
    a = interannual(th)
    pool_years = years[sel]
    pcls = classes(a.loc[pool_years])

    # ERA5 2024-2025, rescaled: LAPrec clim (ERA5_REF) x ERA5 / its own clim
    e5 = xr.open_dataset(ERA5)
    plon = np.where(e5.longitude.values > 180, e5.longitude.values - 360, e5.longitude.values)
    plat = e5.latitude.values
    e5_ratio = np.full((len(ERA5_YEARS), 12, len(glat), len(glon)), np.nan)
    ry0, ry1 = ERA5_REF
    for i, la in enumerate(glat):
        for j, lo in enumerate(glon):
            k = np.where((np.abs(plat - la) < 1e-3) & (np.abs(plon - lo) < 1e-3))[0]
            if len(k) == 0:
                continue
            s = e5.tp.isel(points=k[0]).to_series()
            mm = (s * 1000 * s.index.days_in_month).values.reshape(-1, 12)
            e5y = np.arange(1940, 1940 + len(mm))
            c = mm[(e5y >= ry0) & (e5y <= ry1)].mean(axis=0)
            for n, y in enumerate(ERA5_YEARS):
                e5_ratio[n, :, i, j] = mm[e5y == y][0] / c
    has_e5 = np.isfinite(e5_ratio).all(axis=(0, 1))
    e5_ratio = _fill_uncovered(e5_ratio, has_e5)
    lp_ref = box[(years >= ry0) & (years <= ry1)].mean(axis=0)
    e5_mapped = lp_ref[None] * e5_ratio

    ds = xr.Dataset(
        {'clim': (('month', 'lat', 'lon'), clim),
         'ratio': (('pool_year', 'month', 'lat', 'lon'), ratio),
         'pool_class': (('pool_year',), pcls.values.astype(np.int8)),
         'pool_tanom': (('pool_year',), a.loc[pool_years].values),
         'obs': (('pool_year', 'month', 'lat', 'lon'), obs),
         'era5_mapped': (('era5_year', 'month', 'lat', 'lon'), e5_mapped),
         'laprec_covered': (('lat', 'lon'), covered),
         'era5_covered': (('lat', 'lon'), has_e5)},
        coords={'pool_year': pool_years, 'month': np.arange(1, 13), 'lat': glat, 'lon': glon,
                'era5_year': list(ERA5_YEARS)})
    ds.attrs['description'] = (
        f'LAPrec1871 v1.3, mean of the 5 km cells in each 0.25 deg cell (uncovered cells: '
        f'nearest covered cell), {POOL_Y0}-{POOL_Y1}: obs (mm/month), clim = their mean, '
        f'ratio = obs / clim. pool_class: tercile of the interannual glacier-weighted JJAS '
        f'temperature (HIST BE_on_all strand minus its 31-yr centred mean). era5_mapped: '
        f'ERA5 {ERA5_YEARS} as ratio to its {ry0}-{ry1} clim, times the LAPrec {ry0}-{ry1} '
        f'clim (mm/month). Built by data/laprec_generator.py')
    os.makedirs(OUT, exist_ok=True)
    ds.to_netcdf(POOL_FILE + '.tmp')
    os.replace(POOL_FILE + '.tmp', POOL_FILE)
    return ds


def load_pool():
    return xr.open_dataset(POOL_FILE).load()


# ------------------------------------------------------------------ drawing

def draw_years(member, cls, pool):
    """Pool year drawn for each year of `cls` (Series year -> class)."""
    py = pool.pool_year.values
    pc = pool.pool_class.values
    out = {}
    for y, c in cls.items():
        rng = np.random.default_rng([SEED, int(member), int(y)])
        out[y] = int(rng.choice(py[pc == c]))
    return pd.Series(out)


def generated(draws, pool):
    """mm/month (year, month, lat, lon) for a Series year -> pool year."""
    idx = pd.Index(pool.pool_year.values).get_indexer(draws.values)
    return pool.clim.values[None] * pool.ratio.values[idx]


def hist_precip(member, pool):
    """HIST precipitation of a member, years 1..2025, mm/month, plus the draw table.

    Before SPLIT: the member's own draws (temperature classes from its own
    pre-SPLIT years, == NAT). SPLIT..OBS_Y0-1: the shared series (SHARED_ID,
    classes from the shared HIST temperature). OBS_Y0..2023: LAPrec. 2024-2025:
    ERA5 mapped to LAPrec.
    """
    th = jjas_glacier_weighted(HIST_T.format(member))
    cls = target_classes(th)
    d_own = draw_years(member, cls.loc[:SPLIT - 1], pool)
    d_sh = draw_years(SHARED_ID, cls.loc[SPLIT:OBS_Y0 - 1], pool)
    obs = pool.obs.values
    e5 = pool.era5_mapped.values
    p = np.concatenate([generated(d_own, pool), generated(d_sh, pool), obs, e5])
    years = np.arange(1, ERA5_YEARS[-1] + 1)
    assert len(p) == len(years)
    table = pd.DataFrame({'year': years, 'source': 'own'})
    table['pool_year'] = np.nan
    table['t_class'] = cls.reindex(years).values
    table.loc[table.year < SPLIT, 'pool_year'] = d_own.values
    m = (table.year >= SPLIT) & (table.year < OBS_Y0)
    table.loc[m, 'source'] = 'shared'
    table.loc[m, 'pool_year'] = d_sh.values
    table.loc[(table.year >= OBS_Y0) & (table.year <= POOL_Y1), 'source'] = 'laprec'
    table.loc[table.year > POOL_Y1, 'source'] = 'era5'
    return p, years, table


def nat_precip(member, pool):
    """NAT precipitation of a member (all years drawn), mm/month, + draws."""
    tn = jjas_glacier_weighted(NAT_T.format(member))
    cls = target_classes(tn)
    d = draw_years(member, cls, pool)
    return generated(d, pool), d.index.values, d


def to_dataset(p_mm, years, pool, history):
    """mm/month (year, month, lat, lon) -> Mira-style dataset, pr in kg m-2 s-1."""
    import cftime
    ny = len(years)
    time = [cftime.DatetimeNoLeap(int(y), m, 1, has_year_zero=True)
            for y in years for m in range(1, 13)]
    days = np.tile([31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31], ny)
    pr = p_mm.reshape(ny * 12, *p_mm.shape[2:]) / (days[:, None, None] * 86400.)
    ds = xr.Dataset({'pr': (('lat', 'lon', 'time'), np.moveaxis(pr, 0, -1))},
                    coords={'lat': pool.lat.values, 'lon': pool.lon.values, 'time': time})
    ds.pr.attrs = {'units': 'kg m-2 s-1', 'description': history}
    ds.time.encoding = {'units': 'days since 0001-01-01 00:00:00.000000', 'calendar': 'noleap'}
    return ds


if __name__ == '__main__':
    ds = build_pool()
    print(ds)
    print('class sizes:', np.bincount(ds.pool_class.values))
    print('cells covered by LAPrec:', int(ds.laprec_covered.sum()), 'of', ds.laprec_covered.size,
          '; with ERA5 points:', int(ds.era5_covered.sum()))