"""Extract glacier-relevant grid points from the climate datasets.

For each dataset, every RGI6.2 11-01 glacier (Alps; 11-02 excluded) is
assigned to its nearest grid point (nearest *valid* point, i.e. non-NaN).
Only the points that hold at least one glacier are kept, in a flat
``flat_<name>.nc`` file with dims (time, points) and:

- ``tas``: temperature at the points (°C; for LMR: anomaly as delivered,
  LMR's own reference period, with an extra ``mcrun`` dimension)
- ``weight``: total RGI area (km²) of the glaciers assigned to the point
- ``n_glaciers``: number of glaciers assigned to the point

``glaciers_rgi62_1101.csv`` lists the glaciers with, for each dataset, the
index of the point they are assigned to (``pt_<name>``), so that subsets of
glaciers can be re-weighted later without touching the big files.

Usage::

    python extract.py [name ...]   # default: all datasets
"""
import bz2
import gzip
import os
import shutil
import sys
import tempfile

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

HERE = os.path.dirname(os.path.abspath(__file__))
CLIM = os.path.expanduser('~/www_oggm/climate')
RGI_SHP = os.path.expanduser(
    '~/run_data/rgi/RGIV62/11_rgi62_CentralEurope/11_rgi62_CentralEurope.shp')
GLACIER_CSV = os.path.join(HERE, 'glaciers_rgi62_1101.csv')

# Bounding box for subsetting the global grids before loading (with margin)
BOX = dict(lon=(3, 19), lat=(42, 50))

SOURCES = {
    'lmr_v2.1': os.path.join(
        CLIM, 'lmr/LMRv2.1/air_MCruns_ensemble_mean_LMRv2.1.nc'),
    # LMR Seasonal (Meng et al. 2025, doi:10.1175/JCLI-D-25-0048.1), the
    # reconstruction behind the lmr_seas forcing of rounds 3 and 4
    'lmr_seas': os.path.join(CLIM, 'lmr/LMR_Seasonal/tas_mean.nc'),
    'berkeley-earth': os.path.join(
        CLIM, 'berkeley-earth/Complete_TAVG_LatLong1.nc'),
    'histalp': os.path.join(
        CLIM, 'histalp/hr/HISTALP_temperature_1780-2014.nc.bz2'),
    'gswp3-w5e5': os.path.join(
        CLIM, 'gswp3-w5e5/flattened/2025.11.25/monthly/'
        'gswp3-w5e5_obsclim_tas_global_monthly_1901_2019_flat_glaciers_'
        'v2025.11.25.nc'),
    # CRU TS 4.04 (1901-2019) from the OGGM download cache; an uncompressed
    # copy in ~/tmp is used if present (saves a 3 GB gunzip)
    'cru': '/home/data/download/crudata.uea.ac.uk/cru/data/hrg/cru_ts_4.04/'
           'cruts.2004151855.v4.04/tmp/cru_ts4.04.1901.2019.tmp.dat.nc.gz',
}
CRU_UNZIPPED = os.path.expanduser('~/tmp/cru_ts4.04.1901.2019.tmp.dat.nc')


def get_glaciers():
    if os.path.exists(GLACIER_CSV):
        return pd.read_csv(GLACIER_CSV, index_col=0)
    df = gpd.read_file(RGI_SHP, ignore_geometry=True)
    df = df.loc[df.O2Region.astype(str) == '1']
    df = df[['RGIId', 'CenLon', 'CenLat', 'Area']].set_index('RGIId')
    df.to_csv(GLACIER_CSV)
    return df


def monthly_time(start_year, n):
    return pd.date_range(f'{start_year}-01-01', periods=n, freq='MS')


def assign_points(glon, glat, plon, plat, valid):
    """Index of the nearest valid point for each glacier.

    plon, plat, valid: 1D arrays over candidate points. Distance in degrees
    with a cos(lat) factor on longitude, which is fine at this scale.
    """
    plon, plat = plon[valid], plat[valid]
    idx_valid = np.nonzero(valid)[0]
    coslat = np.cos(np.deg2rad(glat))[:, None]
    d2 = ((glon[:, None] - plon[None, :]) * coslat) ** 2 + \
         (glat[:, None] - plat[None, :]) ** 2
    return idx_valid[np.argmin(d2, axis=1)]


def to_flat(name, tas, plon, plat, gl, extra_attrs=None):
    """tas: DataArray (time, [mcrun,] point) over candidate points."""
    valid = np.isfinite(tas.isel(time=slice(-240, None)).mean('time'))
    if 'mcrun' in valid.dims:
        valid = valid.all('mcrun')
    valid = valid.values
    gidx = assign_points(gl.CenLon.values, gl.CenLat.values,
                         plon, plat, valid)
    used, inv = np.unique(gidx, return_inverse=True)
    weight = np.bincount(inv, weights=gl.Area.values)
    count = np.bincount(inv)
    out = tas.isel(point=used).rename(point='points')
    ds = xr.Dataset({'tas': out})
    ds['weight'] = ('points', weight)
    ds['weight'].attrs = {'units': 'km2',
                          'long_name': 'RGI6.2 11-01 glacier area at point'}
    ds['n_glaciers'] = ('points', count)
    ds['lon'] = ('points', plon[used])
    ds['lat'] = ('points', plat[used])
    ds = ds.set_coords(['lon', 'lat'])
    ds.attrs = {'source_file': SOURCES[name],
                'description': 'Grid points nearest to RGI6.2 11-01 '
                               'glaciers, produced by climate_analysis/'
                               'extract.py'}
    ds.attrs.update(extra_attrs or {})
    ds.to_netcdf(os.path.join(HERE, f'flat_{name}.nc'))

    gl[f'pt_{name}'] = inv
    print(f'{name}: {len(used)} points, {ds.sizes["time"]} time steps '
          f'({str(ds.time.values[0])[:7]} - {str(ds.time.values[-1])[:7]})')
    return gl


def stack_grid(da, lon='lon', lat='lat'):
    """(time, lat, lon) -> (time, point) with 1D lon/lat arrays."""
    da = da.stack(point=(lat, lon))
    plon = da[lon].values.astype(float)
    plat = da[lat].values.astype(float)
    da = da.drop_vars(['point', lon, lat])
    return da, plon, plat


def box(ds, lon='lon', lat='lat'):
    lo, la = ds[lon], ds[lat]
    sel = {lon: (lo >= BOX['lon'][0]) & (lo <= BOX['lon'][1]),
           lat: (la >= BOX['lat'][0]) & (la <= BOX['lat'][1])}
    return ds.isel({k: np.nonzero(v.values)[0] for k, v in sel.items()})


def do_lmr(gl):
    ds = xr.open_dataset(SOURCES['lmr_v2.1'],
                         decode_times=xr.coders.CFDatetimeCoder(use_cftime=True))
    ds = box(ds)
    da = ds.air.load()
    years = np.array([t.year for t in ds.time.values])
    da = da.assign_coords(time=years, MCrun=np.arange(da.sizes['MCrun']))
    da = da.rename(MCrun='mcrun').transpose('time', 'mcrun', 'lat', 'lon')
    da, plon, plat = stack_grid(da)
    da.attrs = {'units': 'K', 'long_name': 'annual mean air temperature '
                'anomaly (LMR reference), per Monte-Carlo run'}
    return to_flat('lmr_v2.1', da, plon, plat, gl,
                   {'time_resolution': 'annual (time = year)'})


def do_lmr_seas(gl):
    """Seasonal anomalies, 800-2000, four per year.

    The time stamps are 1 Jan / Apr / Jul / Oct. Checked against Berkeley
    Earth (detrended, central Europe, 1850-2000): they are the middle months
    of DJF / MAM / JJA / SON (DJF of year Y = Dec Y-1 to Feb Y). The flat
    file has time = year + (stamp month - 1) / 12, and `year` / `season`
    coordinates.
    """
    ds = xr.open_dataset(SOURCES['lmr_seas'],
                         decode_times=xr.coders.CFDatetimeCoder(use_cftime=True))
    ds = box(ds)
    da = ds.tas.load()
    years = da['time.year'].values
    months = da['time.month'].values
    da = da.assign_coords(time=years + (months - 1) / 12)
    da, plon, plat = stack_grid(da)
    da = da.assign_coords(
        year=('time', years),
        season=('time', np.array(['DJF', 'MAM', 'JJA', 'SON'])[months // 3]))
    da.attrs = {'units': 'K', 'long_name': 'seasonal mean air temperature '
                'anomaly (LMR Seasonal reference)'}
    return to_flat('lmr_seas', da, plon, plat, gl,
                   {'time_resolution': 'seasonal (DJF, MAM, JJA, SON)'})


def do_be(gl):
    ds = xr.open_dataset(SOURCES['berkeley-earth'])
    ds = box(ds, lon='longitude', lat='latitude')
    ds = ds.rename(longitude='lon', latitude='lat')
    anom = ds.temperature.load()
    clim = ds.climatology.load()
    n = anom.sizes['time']
    y0 = int(np.floor(ds.time.values[0]))
    # absolute = anomaly + climatology (Jan 1951 - Dec 1980)
    mon = np.arange(n) % 12
    tas = anom.values + clim.values[mon]
    da = xr.DataArray(tas, dims=('time', 'lat', 'lon'),
                      coords={'time': monthly_time(y0, n),
                              'lat': ds.lat, 'lon': ds.lon})
    da, plon, plat = stack_grid(da)
    da.attrs = {'units': 'degC', 'long_name': 'monthly mean 2m temperature '
                '(BE anomaly + BE 1951-1980 climatology)'}
    return to_flat('berkeley-earth', da, plon, plat, gl)


def do_histalp(gl):
    with tempfile.TemporaryDirectory() as tmp:
        f = os.path.join(tmp, 'histalp.nc')
        with bz2.open(SOURCES['histalp']) as fi, open(f, 'wb') as fo:
            shutil.copyfileobj(fi, fo)
        ds = xr.open_dataset(f, decode_times=False)
        da = ds.T_2M.load()
    da = da.assign_coords(time=monthly_time(1780, da.sizes['time']))
    da, plon, plat = stack_grid(da)
    da.attrs = {'units': 'degC', 'long_name': 'monthly mean 2m temperature'}
    return to_flat('histalp', da, plon, plat, gl)


def do_gswp3(gl):
    ds = xr.open_dataset(SOURCES['gswp3-w5e5'])
    lon = ds.longitude.values
    lon = np.where(lon > 180, lon - 360, lon)
    sel = (lon >= BOX['lon'][0]) & (lon <= BOX['lon'][1]) & \
          (ds.latitude.values >= BOX['lat'][0]) & \
          (ds.latitude.values <= BOX['lat'][1])
    da = ds.tas.isel(points=np.nonzero(sel)[0]).load() - 273.15
    da = da.drop_vars(['latitude', 'longitude']).rename(points='point')
    da = da.assign_coords(time=monthly_time(1901, da.sizes['time']))
    da.attrs = {'units': 'degC', 'long_name': 'monthly mean 2m temperature'}
    return to_flat('gswp3-w5e5', da, lon[sel],
                   ds.latitude.values[sel].astype(float), gl)


def _cru(gl, f):
    ds = xr.open_dataset(f)
    ds = box(ds)
    da = ds.tmp.load()
    da = da.assign_coords(time=monthly_time(1901, da.sizes['time']))
    da, plon, plat = stack_grid(da)
    da.attrs = {'units': 'degC', 'long_name': 'monthly mean 2m temperature'}
    return to_flat('cru', da, plon, plat, gl,
                   {'version': 'CRU TS 4.04'})


def do_cru(gl):
    if os.path.exists(CRU_UNZIPPED):
        return _cru(gl, CRU_UNZIPPED)
    with tempfile.TemporaryDirectory() as tmp:
        f = os.path.join(tmp, 'cru.nc')
        with gzip.open(SOURCES['cru']) as fi, open(f, 'wb') as fo:
            shutil.copyfileobj(fi, fo)
        return _cru(gl, f)


FUNCS = {'lmr_v2.1': do_lmr, 'lmr_seas': do_lmr_seas, 'berkeley-earth': do_be, 'histalp': do_histalp,
         'gswp3-w5e5': do_gswp3, 'cru': do_cru}

if __name__ == '__main__':
    names = sys.argv[1:] or list(FUNCS)
    gl = get_glaciers()
    for name in names:
        gl = FUNCS[name](gl)
        gl.to_csv(GLACIER_CSV)