"""Build an OGGM forcing: Berkeley Earth 1750-2024, inflated LMR v2.1 before.

Output (same layout as data/unstacked/lmr_seas_fullsplit_oldBE, readable
by lmr_forcing.process_mira_data with ens_member=None):

    data/unstacked/be_lmr21_v<VERSION>/tp_monthly.nc   (lat, lon, time) tas [K]
    data/unstacked/be_lmr21_v<VERSION>/elev.nc         copy of the oldBE one

Method (all per 1° grid cell, the BE grid = the oldBE temperature grid):

- Absolute = CLIM + anomaly vs 1961-1990, where CLIM is the 1961-1990
  monthly mean of the oldBE HIST forcing (Mira's observed climatology,
  identical for all HIST members), so the new file has the same 1961-1990
  climatology and elevation as the forcing used so far. BE's own absolute
  temperatures are 1-4 K warmer than GSWP3-W5E5 in the Alps cells and
  fit no sensible elevation, hence the anomaly method.
- 1750-2024: BE monthly anomalies vs their 1961-1990 monthly means. The
  missing 1752-09 is interpolated from Aug/Oct. Cells without BE data (sea)
  take the anomaly of the nearest valid cell.
- 1-1749: LMR v2.1 annual anomaly (mean of the 20 MC runs, bilinearly
  interpolated from 2° to the 1° grid) vs its 1961-1990 mean. Variance
  inflation in the spirit of OGGM's process_gcm_data(scale_stddev=True),
  but on annual values: the 31-yr centered running mean is kept and the
  deviations from it are scaled to the std of BE's deviations over
  1750-2000. LMR's own variance shrinks back in time (fewer proxies:
  deviation std 0.30 K in years 1-500 against 0.47 K in 1500-1749,
  regionally), so the factor is time-dependent: target std / LMR's
  101-yr running std of the deviations (~2.4 around 1750, more earlier).
  No blending: LMR is used until 1749, BE from 1750.
- Months before 1750: LMR gives only annual values. Applying the annual
  anomaly to every month (as OGGM's process_lmr_data does) leaves no
  intra-annual variability, and scaling each month separately would make
  the annual variance far too large. Instead, each year gets the
  within-year monthly residuals (monthly anomaly minus its annual mean) of
  a BE analog year from 1850-2024: one of the K_ANALOG years whose
  (glacier-area weighted, high-passed) annual anomaly is closest to the
  LMR year's, drawn at random (fixed seed, same draw for all cells). The
  analog keeps the anti-correlation between the annual anomaly and the
  summer residual (warm years are warm mostly in winter); random years
  made JJAS variability ~20% too high. Annual means are exactly the
  inflated LMR values.

Precipitation: BE has none. The output folder gets a p_monthly.nc symlink
to the oldBE ens183 precipitation (the calibration member).

Optional LMR offset (sensitivity to the level mismatch at the switch):
with OFFSET_PERIOD, e.g. 1750-1779, the whole LMR part (years 1-1749) is
shifted by one constant, the glacier-area weighted mean of BE minus
inflated LMR (annual anomalies) over that period. Output then goes to
be_lmr21_off<y0>-<y1>_v<VERSION>. Without it there is no shift and the
31-yr mean steps down by ~0.35 K at 1750 (LMR's 1750s are warm vs BE's).
LMR and BE disagree by +-0.3-0.4 K on 31-yr means throughout the overlap
(e.g. -0.32 K in 1850-1899), so the offset is uncertain: variants are
  - none: no shift
  - 1750-1779: offset of the first 30 years (BE right, LMR biased; ~-0.38)
  - 1750-2000: offset of the full overlap (~-0.14)

Usage: python build_be_lmr_forcing.py [OFFSET_PERIOD]
"""
import os
import shutil
import sys

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

VERSION = '20261009'
ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
OLDBE = os.path.join(ROOT, 'data/unstacked/lmr_seas_fullsplit_oldBE')
CLIM_FILE = os.path.join(OLDBE, 'jjas_hists_v20261005/ensemble/'
                                'tp_monthly_ens183.nc')
PCP_FILE = os.path.join(OLDBE, 'ensemble/p_monthly_ens183.nc')
OUT = os.path.join(ROOT, f'data/unstacked/be_lmr21_v{VERSION}')


def out_dir(offset_period=None):
    if offset_period is None:
        return OUT
    y0, y1 = offset_period
    return os.path.join(ROOT,
                        f'data/unstacked/be_lmr21_off{y0}-{y1}_v{VERSION}')

CLIMDIR = os.path.expanduser('~/www_oggm/climate')
BE_FILE = os.path.join(CLIMDIR, 'berkeley-earth/Complete_TAVG_LatLong1.nc')
LMR_FILE = os.path.join(CLIMDIR,
                        'lmr/LMRv2.1/air_MCruns_ensemble_mean_LMRv2.1.nc')

REF = (1961, 1990)          # anomaly reference period
SWITCH = 1750               # first BE year
OVERLAP = (1750, 2000)      # period for the inflation factor
WINDOW = 31                 # running mean window (yr), as OGGM for 30-yr ref
STD_WINDOW = 101            # window (yr) for LMR's running std
K_ANALOG = 10               # number of nearest analog years to draw from
RESID_POOL = (1850, 2024)   # BE years to draw monthly residuals from
SEED = 0

TC = xr.coders.CFDatetimeCoder(use_cftime=True)


def fill_nearest(a, lat, lon):
    """Fill all-NaN cells of a (..., lat, lon) array with nearest valid."""
    bad = np.isnan(a).all(axis=tuple(range(a.ndim - 2)))
    if not bad.any():
        return a
    yy, xx = np.meshgrid(lat, lon, indexing='ij')
    coslat = np.cos(np.deg2rad(yy))
    vy, vx = np.nonzero(~bad)
    for i, j in zip(*np.nonzero(bad)):
        d2 = ((xx[vy, vx] - xx[i, j]) * coslat[i, j]) ** 2 + \
             (yy[vy, vx] - yy[i, j]) ** 2
        k = np.argmin(d2)
        a[..., i, j] = a[..., vy[k], vx[k]]
    return a


def running(a, window, how='mean'):
    """Centered running mean/std along axis 0, shrinking at the ends."""
    r = pd.DataFrame(a.reshape(len(a), -1)).rolling(
        window, center=True, min_periods=window // 2 + 1)
    return getattr(r, how)().values.reshape(a.shape)


def running_mean(a, window):
    return running(a, window, 'mean')


def glacier_weights(lat, lon):
    """RGI6.2 11-01 glacier area per cell (nearest cell), normalised."""
    gl = pd.read_csv(os.path.join(os.path.dirname(__file__),
                                  'glaciers_rgi62_1101.csv'), index_col=0)
    i = np.abs(gl.CenLat.values[:, None] - lat[None]).argmin(1)
    j = np.abs(gl.CenLon.values[:, None] - lon[None]).argmin(1)
    w = np.zeros((len(lat), len(lon)))
    np.add.at(w, (i, j), gl.Area.values)
    return w / w.sum()


def main(offset_period=None):
    out = out_dir(offset_period)
    # Target grid and climatology from the oldBE HIST forcing
    with xr.open_dataset(CLIM_FILE, decode_times=TC) as ds:
        lat, lon = ds.lat.values, ds.lon.values
        tas = ds.tas.transpose('time', 'lat', 'lon')
        yrs = ds['time.year'].values
        sel = (yrs >= REF[0]) & (yrs <= REF[1])
        clim = tas.values[sel].reshape(-1, 12, len(lat), len(lon)).mean(0)
    clim = clim.astype(np.float64)  # K, (12, lat, lon)

    # --- BE monthly anomalies vs REF, 1750-2024
    with xr.open_dataset(BE_FILE) as ds:
        ds = ds.sel(latitude=lat, longitude=lon)
        be = ds.temperature.values.astype(np.float64)
        y0 = int(np.floor(ds.time.values[0]))
    nyb = be.shape[0] // 12
    be_years = np.arange(y0, y0 + nyb)
    assert y0 == SWITCH and be_years[-1] == 2024
    be = be.reshape(nyb, 12, len(lat), len(lon))
    # interpolate isolated missing months (1752-09) from the neighbours
    flat = be.reshape(nyb * 12, len(lat), len(lon))
    valid_cell = ~np.isnan(flat).all(axis=0)
    for t in np.nonzero(np.isnan(flat[:, valid_cell]).any(axis=1))[0]:
        flat[t] = 0.5 * (flat[t - 1] + flat[t + 1])
    be = flat.reshape(nyb, 12, len(lat), len(lon))
    sel = (be_years >= REF[0]) & (be_years <= REF[1])
    be = be - be[sel].mean(axis=0)
    be = fill_nearest(be, lat, lon)
    assert not np.isnan(be).any()

    # --- LMR annual anomalies vs REF, years 0-2000, on the 1° grid
    with xr.open_dataset(LMR_FILE, decode_times=TC) as ds:
        ds = ds.sel(lat=slice(lat.min() - 4, lat.max() + 4),
                    lon=slice(lon.min() - 4, lon.max() + 4))
        lmr = ds.air.mean('MCrun').load()
        lmr_years = np.array([t.year for t in ds.time.values])
    lmr = lmr.interp(lat=lat, lon=lon).values.astype(np.float64)
    sel = (lmr_years >= REF[0]) & (lmr_years <= REF[1])
    lmr = lmr - lmr[sel].mean(axis=0)

    # --- Variance inflation of the LMR deviations from the running mean
    lmr_lo = running_mean(lmr, WINDOW)
    lmr_hi = lmr - lmr_lo
    be_ann = be.mean(axis=1)
    be_hi = be_ann - running_mean(be_ann, WINDOW)
    ol = (OVERLAP[0], OVERLAP[1])
    ol_l = (lmr_years >= ol[0]) & (lmr_years <= ol[1])
    ol_b = (be_years >= ol[0]) & (be_years <= ol[1])
    target = be_hi[ol_b].std(axis=0, ddof=1)
    fac = target / running(lmr_hi, STD_WINDOW, 'std')
    lmr_infl = lmr_lo + lmr_hi * fac

    # --- Monthly values before SWITCH: annual LMR + BE within-year residuals
    pre_years = np.arange(1, SWITCH)
    ann = lmr_infl[np.isin(lmr_years, pre_years)]       # (ny, lat, lon)
    resid = be - be_ann[:, None]                         # (nyb, 12, lat, lon)
    pool = np.nonzero((be_years >= RESID_POOL[0]) &
                      (be_years <= RESID_POOL[1]))[0]
    w = glacier_weights(lat, lon)
    x_lmr = ((lmr_infl - lmr_lo) * w).sum(axis=(1, 2))
    x_lmr = x_lmr[np.isin(lmr_years, pre_years)]
    x_be = (be_hi * w).sum(axis=(1, 2))[pool]
    rng = np.random.default_rng(SEED)
    near = np.argsort(np.abs(x_lmr[:, None] - x_be[None]), axis=1)
    pick = rng.integers(0, K_ANALOG, size=len(pre_years))
    draw = pool[near[np.arange(len(pre_years)), pick]]
    pre = ann[:, None] + resid[draw]

    # --- Optional constant offset of the LMR part
    offset = 0.
    if offset_period is not None:
        sl = (lmr_years >= offset_period[0]) & (lmr_years <= offset_period[1])
        sb = (be_years >= offset_period[0]) & (be_years <= offset_period[1])
        offset = float(((be_ann[sb].mean(0) - lmr_infl[sl].mean(0)) * w).sum())
        pre = pre + offset
    print(f'LMR offset: {offset:+.3f} K')

    anom = np.concatenate([pre, be], axis=0)             # years 1-2024
    years = np.arange(1, be_years[-1] + 1)
    assert len(anom) == len(years)
    absT = (anom + clim[None]).reshape(-1, len(lat), len(lon))

    time = [cftime.DatetimeNoLeap(y, m, 1) for y in years
            for m in range(1, 13)]
    da = xr.DataArray(absT.transpose(1, 2, 0).astype(np.float32),
                      dims=('lat', 'lon', 'time'),
                      coords={'lat': lat.astype(np.float32),
                              'lon': lon.astype(np.float32), 'time': time})
    da.attrs = {
        'units': 'K',
        'description': 'Monthly absolute temperatures: oldBE HIST 1961-1990 '
                       'monthly climatology + anomalies vs 1961-1990. '
                       f'Berkeley Earth from {SWITCH}, before that LMR v2.1 '
                       'annual anomalies (variance-inflated, see '
                       'climate_analysis/build_be_lmr_forcing.py) with BE '
                       'within-year monthly residuals of analog years.'}
    ds = xr.Dataset({'tas': da})
    ds.lat.attrs = {'units': 'degrees_north',
                    'description': 'Latitude of Alps grid'}
    ds.lon.attrs = {'units': 'degrees_east',
                    'description': 'Longitude of Alps grid'}
    ds.attrs = {'title': 'Berkeley Earth (1750-2024) + inflated LMR v2.1 '
                         '(1-1749) temperature for the Alps',
                'version': VERSION,
                'inflation_target': 'std of BE deviations from 31-yr '
                                    'running mean, 1750-2000',
                'inflation_lmr_std_window': STD_WINDOW,
                'analog_k': K_ANALOG,
                'inflation_overlap': f'{OVERLAP[0]}-{OVERLAP[1]}',
                'running_mean_window': WINDOW,
                'residual_pool': f'{RESID_POOL[0]}-{RESID_POOL[1]}',
                'seed': SEED,
                'lmr_offset_K': offset,
                'lmr_offset_period': ('none' if offset_period is None else
                                      '%d-%d' % offset_period)}
    os.makedirs(out, exist_ok=True)
    enc = {'time': {'units': 'days since 0001-01-01 00:00:00',
                    'calendar': 'noleap', 'dtype': 'int64'}}
    ds.to_netcdf(os.path.join(out, 'tp_monthly.nc'), encoding=enc)
    shutil.copy(os.path.join(OLDBE, 'elev.nc'), os.path.join(out, 'elev.nc'))
    pcp = os.path.join(out, 'p_monthly.nc')
    if not os.path.lexists(pcp):
        os.symlink(os.path.relpath(PCP_FILE, out), pcp)

    # Diagnostics
    np.set_printoptions(precision=2, suppress=True, linewidth=150)
    iy = np.isin(lmr_years, [1, 500, 1000, 1500, 1749])
    print('inflation factor, glacier-weighted, years 1/500/1000/1500/1749:',
          (fac[iy] * w).sum(axis=(1, 2)))
    pd.DataFrame({'year': years.repeat(1)[:len(pre_years)],
                  'be_year_drawn': be_years[draw]}).to_csv(
        os.path.join(out, 'residual_draws.csv'), index=False)
    print('written:', out)


if __name__ == '__main__':
    period = None
    if len(sys.argv) > 1 and sys.argv[1] != 'none':
        period = tuple(int(y) for y in sys.argv[1].split('-'))
    main(period)