"""Area discretisation of the model glaciers (handover to-do 1).

The OGGM area changes one grid point at a time (width x dx, where the ice is
thicker than min_ice_thick_for_area = 2 m). Per glacier, from the RGI-date
model flowlines of the v7 sw_pf3.5_mf5_lam12 L3 gdirs (read from the tars,
plain numpy):

- dA_med5 / dA_max5: median / max area of the last 5 ice points of the main
  (last) flowline, km2;
- n_ice: number of ice grid points on all flowlines;
- ratios to the RGI area, and the area sigma of the handover proposal,
  sigma^2 = (0.05 A)^2 + k dA^2, for k = 1/12 and 1/4.

    python area_quantisation.py     # ~1 min
-> area_quantisation.csv (+ summary printed)
"""
import glob
import io
import json
import os
import tarfile

import numpy as np
import pandas as pd

HERE = os.path.dirname(os.path.abspath(__file__))
L3 = (f'{HERE}/../lmr_gdirs_v7_202610/exps/sw_pf3.5_mf5_lam12/RGI62/b_240/L3')
MIN_H = 2.
N_LAST = 5


def flowline_stats(npz_bytes):
    z = np.load(io.BytesIO(npz_bytes), allow_pickle=False)
    meta = json.loads(str(z['__meta__']))
    fls = meta['root']['items']
    n_ice, last = 0, None
    for i, fl in enumerate(fls):
        f = fl['fields']
        dx_m = f['dx']['v'] * f['map_dx']['v']
        thick = z[f'model_flowlines/{i}/surface_h'] - z[f'model_flowlines/{i}/bed_h']
        w = z[f'model_flowlines/{i}/widths_m']
        ice = thick > MIN_H
        n_ice += int(ice.sum())
        last = (w[ice] * dx_m * 1e-6, dx_m)
    da = last[0][-N_LAST:]
    return dict(dA_med5=float(np.median(da)), dA_max5=float(da.max()),
                dA_last=float(da[-1]), dx_m=last[1], n_ice=n_ice, n_fl=len(fls))


rows = {}
for tf in sorted(glob.glob(f'{L3}/RGI60-11/RGI60-11.*.tar')):
    with tarfile.open(tf) as outer:
        for m in outer.getmembers():
            if not m.name.endswith('.tar.gz'):
                continue
            rid = os.path.basename(m.name)[:-7]
            with tarfile.open(fileobj=outer.extractfile(m), mode='r:gz') as inner:
                try:
                    b = inner.extractfile(f'{rid}/model_flowlines.npz').read()
                except KeyError:
                    continue                      # failed glacier, no flowlines
            rows[rid] = flowline_stats(b)

st = pd.read_csv(f'{L3}/summary/glacier_statistics_11.csv', index_col=0, low_memory=False)
df = pd.DataFrame(rows).T.join(st[['rgi_area_km2', 'rgi_subregion']])
df = df[df.rgi_subregion == '11-01'].drop(columns='rgi_subregion')
a = df.rgi_area_km2
for c in ['dA_med5', 'dA_max5']:
    df[f'{c}_frac'] = df[c] / a
for k, lab in [(1 / 12, 'k12'), (1 / 4, 'k4')]:
    df[f'sigma_area_{lab}_frac'] = np.sqrt(0.05 ** 2 + k * df.dA_max5_frac ** 2)
df.index.name = 'rgi_id'
df.to_csv(f'{HERE}/area_quantisation.csv')

print(f'{len(df)} glaciers with flowlines (11-01)')
bins = [0, 0.1, 0.5, 1, 5, 20, np.inf]
cls = pd.cut(a, bins, labels=['<0.1', '0.1-0.5', '0.5-1', '1-5', '5-20', '>20'])
g = df.groupby(cls, observed=True)
print(pd.DataFrame({'n': g.size(), 'area km2': g.rgi_area_km2.sum().round(0),
                    'n_ice median': g.n_ice.median(),
                    'dA_max5/A median': g.dA_max5_frac.median().round(3),
                    'dA_max5/A P90': g.dA_max5_frac.quantile(0.9).round(3),
                    'sigma k=1/12 median': g.sigma_area_k12_frac.median().round(3),
                    'sigma k=1/4 median': g.sigma_area_k4_frac.median().round(3)}).to_string())
for thr in [0.05, 0.1, 0.2]:
    m = df.dA_max5_frac > thr
    print(f'dA_max5 > {thr:.0%} of A: {m.sum()} glaciers, {a[m].sum() / a.sum():.1%} of the area')
print(df.loc[['RGI60-11.01450', 'RGI60-11.02072', 'RGI60-11.02773', 'RGI60-11.00897'],
             ['rgi_area_km2', 'dx_m', 'n_ice', 'dA_med5', 'dA_max5', 'dA_max5_frac',
              'sigma_area_k4_frac']].round(3).to_string())