"""Regional series of a round 6 permutation dataset.

On the common subset of the round5 ensemble (round5/filtered_ens/<ENS>/
glacier_statistics_reg11.csv, 3784 glaciers), as round5/preprocess_ens.py:
regional volume (km3), area_min_h (km2) and length (sum of the glacier
lengths, km) of every permutation run and of the HIST strands, plus each
state's own history 1800-1849 (from the round5 ensemble NAT run of the
member, identical to HIST before 1850). Glaciers of the subset without output in a run (NaN) are
counted and reported.

    python preprocess_perm.py <name> [<round5 ensemble name>]

-> filtered/<name>/{volume_km3,area_km2,length_km}_perm_reg11.csv (one column
   per run: state###_nat_ens###, state###_hist), pre1850_reg11.csv (per state
   and variable, 1800-1849), n_missing_reg11.csv
"""
import glob
import os
import re
import sys
from concurrent.futures import ProcessPoolExecutor

import pandas as pd
import xarray as xr

HERE = os.path.dirname(os.path.abspath(__file__))
NAME = sys.argv[1]
ENS = sys.argv[2] if len(sys.argv) > 2 else 'ens_pf3.5_mf5_lam12'
DATA = f'{HERE}/output/{NAME}'
OUT = f'{HERE}/filtered/{NAME}'
IDS = pd.read_csv(f'{HERE}/../round5/filtered_ens/{ENS}/glacier_statistics_reg11.csv',
                  index_col=0).index.values


def regional(f):
    col = re.search(r'perm_(state\d{3}_\w+?)_reg11', f).group(1)
    with xr.open_dataset(f) as ds:
        ds = ds[['volume', 'area_min_h', 'length']].sel(rgi_id=IDS).load()
    n_missing = int(ds.volume.isnull().any('time').sum())
    vol = (ds.volume.sum('rgi_id') * 1e-9).to_series()
    area = (ds.area_min_h.sum('rgi_id') * 1e-6).to_series()
    length = (ds.length.sum('rgi_id') * 1e-3).to_series()
    # NAT runs end with their climate: keep NaN after the last year
    last = ds.volume.notnull().any('rgi_id').values
    return col, vol.where(last), area.where(last), length.where(last), n_missing


def pre1850(state):
    f = f'{HERE}/../round5/output_ens/{ENS}/lmr_mira_nat_ens{state}_reg11.nc'
    with xr.open_dataset(f) as ds:
        ds = ds[['volume', 'area_min_h', 'length']].sel(rgi_id=IDS, time=slice(1800, 1849)).load()
    return pd.DataFrame({f'volume_km3_state{state}': ds.volume.sum('rgi_id').to_series() * 1e-9,
                         f'area_km2_state{state}': ds.area_min_h.sum('rgi_id').to_series() * 1e-6,
                         f'length_km_state{state}': ds.length.sum('rgi_id').to_series() * 1e-3})


if __name__ == '__main__':
    os.makedirs(OUT, exist_ok=True)
    files = sorted(glob.glob(f'{DATA}/perm_state*_reg11.nc'))
    print(f'{len(files)} files in {DATA}, subset {len(IDS)} glaciers')
    with ProcessPoolExecutor(int(os.environ.get('N_JOBS', 8))) as ex:
        res = list(ex.map(regional, files))
    states = sorted({c[5:8] for c, *_ in res})
    with ProcessPoolExecutor(len(states)) as ex:
        pre = pd.concat(list(ex.map(pre1850, states)), axis=1)
    pre.index = pre.index.astype(int)
    pre.index.name = 'year'
    pre.to_csv(f'{OUT}/pre1850_reg11.csv')
    for i, name in [(1, 'volume_km3'), (2, 'area_km2'), (3, 'length_km')]:
        df = pd.DataFrame({r[0]: r[i] for r in res})
        df.index = df.index.astype(int)
        df.index.name = 'year'
        df.to_csv(f'{OUT}/{name}_perm_reg11.csv')
    miss = pd.Series({r[0]: r[4] for r in res}, name='n_missing')
    miss.to_csv(f'{OUT}/n_missing_reg11.csv')
    print(f'glaciers of the subset with NaN somewhere: max {miss.max()} '
          f'({miss.idxmax()}), {int((miss > 0).sum())} of {len(miss)} runs affected')
    print(f'written to {OUT}')