"""Glacier-area weighted regional temperature series from the flat_*.nc files.

Writes to ``series/``:

- ``tas_monthly.csv``: monthly regional mean (°C), one column per dataset
  (LMR is annual only and not in this file)
- ``tas_annual.csv``: calendar-year means (°C; LMR: anomaly as delivered,
  mean over the 20 Monte-Carlo runs); a year needs all 12 months
- ``tas_annual_anom_1961-1990.csv``: the same as anomalies to 1961-1990
- ``tas_jjas.csv`` / ``tas_jjas_anom_1961-1990.csv``: Jun-Sep means (no LMR
  v2.1; for LMR Seasonal the column is its JJA season, Jun-Aug)
- LMR Seasonal (4 seasons per year): its annual value is the mean of DJF,
  MAM, JJA and SON of that year, i.e. Dec of the year before to Nov
- ``lmr_v2.1_annual_mcruns.csv``: LMR per Monte-Carlo run (anomaly to
  1961-1990)

Usage: python regional.py
"""
import os

import pandas as pd
import xarray as xr

HERE = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.join(HERE, 'series')
DATASETS = ['lmr_v2.1', 'lmr_seas', 'berkeley-earth', 'histalp', 'gswp3-w5e5',
            'cru']
REF = (1961, 1990)


def regional(name):
    """Area-weighted mean over points: DataArray (time[, mcrun])."""
    ds = xr.open_dataset(os.path.join(HERE, f'flat_{name}.nc'))
    w = ds.weight / ds.weight.sum()
    return (ds.tas * w).sum('points', skipna=False)


def anom(df):
    return df - df.loc[REF[0]:REF[1]].mean()


def seasonal_mean(s, months):
    s = s[s.index.month.isin(months)]
    g = s.groupby(s.index.year)
    out = g.mean()
    return out[(g.count() == len(months))]


if __name__ == '__main__':
    os.makedirs(OUT, exist_ok=True)
    monthly, annual, jjas = {}, {}, {}
    for name in DATASETS:
        da = regional(name)
        if name == 'lmr_v2.1':
            df = da.to_pandas()
            df.columns = [f'mcrun{i:02d}' for i in df.columns]
            df.index.name = 'year'
            anom(df).round(4).to_csv(
                os.path.join(OUT, 'lmr_v2.1_annual_mcruns.csv'))
            annual[name] = df.mean(axis=1)
            continue
        if name == 'lmr_seas':
            s = da.to_pandas()
            s.index = pd.MultiIndex.from_arrays([da.year.values,
                                                 da.season.values])
            s = s.unstack()
            annual[name] = s.mean(axis=1)
            jjas[name] = s['JJA']
            continue
        s = da.to_pandas().dropna()
        monthly[name] = s
        annual[name] = seasonal_mean(s, range(1, 13))
        jjas[name] = seasonal_mean(s, [6, 7, 8, 9])

    monthly = pd.DataFrame(monthly)
    monthly.index.name = 'time'
    monthly.round(4).to_csv(os.path.join(OUT, 'tas_monthly.csv'))
    for lab, d in [('annual', annual), ('jjas', jjas)]:
        df = pd.DataFrame(d)
        df.index.name = 'year'
        df.round(4).to_csv(os.path.join(OUT, f'tas_{lab}.csv'))
        anom(df).round(4).to_csv(
            os.path.join(OUT, f'tas_{lab}_anom_{REF[0]}-{REF[1]}.csv'))
    print(pd.DataFrame(annual).describe().round(2))