"""Regional CSVs + per-glacier LIA maxima of a round 5 ensemble dataset.

Parallel version of round4/preprocess_mira.py (+ the lia_max cache of
round4/ensemble/plot_lia_max.py) for round5/output_ens/<name>/:

1. common subset: 11-01 glaciers with output in every member (HIST and NAT);
2. per member, on that subset: the regional volume (km3) and area_min_h
   (km2) series, and the sum of the per-glacier maxima over the LIA windows
   (area max, volume max, volume at the year of the area max).

    python preprocess_ens.py <name>        # ~10 min with 8 processes

-> filtered_ens/<name>/glacier_statistics_reg11.csv (the subset),
   {volume_km3,area_km2}_{hist,nat}_ens_reg11.csv, lia_max_members.csv
"""
import glob
import os
import re
import sys
from concurrent.futures import ProcessPoolExecutor

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

HERE = os.path.dirname(os.path.abspath(__file__))
NAME = sys.argv[1]
DATA = f'{HERE}/output_ens/{NAME}'
OUT = f'{HERE}/filtered_ens/{NAME}'
WINDOWS = [(1800, 1875), (1800, 1900), (1800, 1975)]
N_JOBS = int(os.environ.get('N_JOBS', 8))
# The statistics of the gdirs the dataset was run with (11-01 only)
STATS = os.environ.get('STATS', f'{HERE}/../lmr_gdirs_v7_202610/exps/'
                       'sw_pf3.5_mf5_lam12/RGI62/b_240/L3/summary/glacier_statistics_11.csv')


def valid_ids(f):
    with xr.open_dataset(f) as ds:
        return ds.rgi_id.values[ds.volume.isel(time=0).notnull().values]


def member(args):
    f, ids = args
    kind, num = re.search(r'lmr_mira_(\w+?)_(ens\d+)_reg11', f).groups()
    with xr.open_dataset(f) as ds:
        ds = ds[['volume', 'area_min_h']].sel(rgi_id=ids).load()
    vol = (ds.volume.sum('rgi_id') * 1e-9).to_series()
    area = (ds.area_min_h.sum('rgi_id') * 1e-6).to_series()
    r = {'kind': kind, 'member': num}
    for y0, y1 in WINDOWS:
        w = ds.sel(time=slice(y0, y1))
        r[f'area_max_{y0}_{y1}'] = float(w.area_min_h.max('time').sum()) * 1e-6
        r[f'vol_max_{y0}_{y1}'] = float(w.volume.max('time').sum()) * 1e-9
        iy = w.area_min_h.argmax('time')
        r[f'vol_at_amax_{y0}_{y1}'] = float(w.volume.isel(time=iy).sum()) * 1e-9
    return kind, num, vol, area, r


if __name__ == '__main__':
    os.makedirs(OUT, exist_ok=True)
    st = pd.read_csv(STATS, index_col=0, low_memory=False)
    st = st.loc[st.rgi_subregion == '11-01']
    files = sorted(glob.glob(f'{DATA}/lmr_mira_*_ens*_reg11.nc'))
    print(f'{len(files)} members in {DATA}')

    with ProcessPoolExecutor(N_JOBS) as ex:
        valids = list(ex.map(valid_ids, files))
    cov = pd.Series({os.path.basename(f)[9:-9]: st.rgi_area_km2.reindex(v).sum() / st.rgi_area_km2.sum()
                     for f, v in zip(files, valids)})
    print(f'per-member area coverage: min {cov.min():.2%} ({cov.idxmin()}), median {cov.median():.2%}')
    ids = valids[0]
    for v in valids[1:]:
        ids = np.intersect1d(ids, v)
    ids = np.intersect1d(ids, st.index.values)
    print(f'common subset: {len(ids)} glaciers, area {st.rgi_area_km2.reindex(ids).sum() / st.rgi_area_km2.sum():.2%}, '
          f'inv. volume {st.inv_volume_km3.reindex(ids).sum() / st.inv_volume_km3.sum():.2%}')
    st.loc[ids].to_csv(f'{OUT}/glacier_statistics_reg11.csv')

    with ProcessPoolExecutor(N_JOBS) as ex:
        res = list(ex.map(member, [(f, ids) for f in files]))
    for kind in ['hist', 'nat']:
        vol = pd.DataFrame({n: v for k, n, v, a, r in res if k == kind})
        area = pd.DataFrame({n: a for k, n, v, a, r in res if k == kind})
        vol.index.name = area.index.name = 'year'
        vol.to_csv(f'{OUT}/volume_km3_{kind}_ens_reg11.csv')
        area.to_csv(f'{OUT}/area_km2_{kind}_ens_reg11.csv')
        print(f'{kind}: {vol.shape[1]} members, years {vol.index.min()}-{vol.index.max()}')
    lm = pd.DataFrame({f'{r["kind"]}_{r["member"]}': r for k, n, v, a, r in res}).T
    lm.to_csv(f'{OUT}/lia_max_members.csv')
    print(f'Written to {OUT}')