# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.19.3
#   kernelspec:
#     display_name: Python 3 (ipykernel)
#     language: python
#     name: python3
# ---

# %%
import os

import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from oggm import utils

FILTERED_DIR = 'filtered'

# %%
experiments = sorted(
    d for d in os.listdir(FILTERED_DIR)
    if os.path.isdir(os.path.join(FILTERED_DIR, d)) and not d.startswith('.')
)
print(experiments)

# %%
data = {}
for exp in experiments:
    exp_dir = os.path.join(FILTERED_DIR, exp)
    d = {
        'vol': pd.read_csv(f'{exp_dir}/volume_km3_hist_reg11.csv', index_col=0)['hist'],
        'area': pd.read_csv(f'{exp_dir}/area_km2_hist_reg11.csv', index_col=0)['hist'],
    }
    vol_ens_f = f'{exp_dir}/volume_km3_nat_ens_reg11.csv'
    if os.path.exists(vol_ens_f):
        d['vol_ens'] = pd.read_csv(vol_ens_f, index_col=0)
        d['area_ens'] = pd.read_csv(f'{exp_dir}/area_km2_nat_ens_reg11.csv', index_col=0)
    dfok = pd.read_csv(f'{exp_dir}/glacier_statistics_reg11.csv', index_col=0)
    d['inv_volume_2003'] = dfok.inv_volume_km3.sum()
    data[exp] = d

# %%
dfo = pd.read_csv("../round1/filtered_update/Area_volume_change_LIA_Alps.csv",
                   encoding="ISO-8859-1", index_col=0, skiprows=[1]).loc[17]

dfg = utils.get_geodetic_mb_dataframe(regional=True)
dfg = dfg.loc[dfg.period == '2000-01-01_2020-01-01']
dfg = dfg.loc[dfg.reg == 11]
dvol = dfg.iloc[0]['dmdt'] * 20 / 900 * 1000  # Gt -> km3
dvol_err = dfg.iloc[0]['err_dmdt'] * 20 / 900 * 1000
dvol, dvol_err


# %%
def add_volume_refs(ax, hist_ser, inv_volume_2003):
    ax.plot(1850, dfo['Volume LIA'], 'o', c='grey', label='Volume Ref. LIA')
    ax.plot(2003, inv_volume_2003, 'x', c='k', label='Volume OGGM 2003')
    v_2000 = hist_ser.loc[2000]
    v_2019 = v_2000 + dvol
    ax.plot([2000, 2019], [v_2000, v_2019], '-', c='k', label='Obs. dV (Hugonnet)')
    ax.fill_between([2000, 2019], [v_2000, v_2019 - dvol_err], [v_2000, v_2019 + dvol_err],
                     alpha=0.2, color='k', label='Obs. dV uncertainty')
    ax.plot(2015, dfo['Volume 2015'], 'o', c='k', label='Volume Ref. 2015')


def add_area_refs(ax):
    ax.plot(1850, dfo['Area LIA'], 'o', c='grey', label='Area Reinthaler LIA')
    ax.plot(2003, dfo['Area 2003'], 'x', c='k', label='Area Reinthaler 2003')
    ax.plot(2015, dfo['Area 2015'], 'o', c='k', label='Area Reinthaler 2015')


# %% [markdown]
# ## Per-experiment plots (hist + ensemble where available)

# %%
for exp in experiments:
    d = data[exp]
    with sns.axes_style("whitegrid"):
        plt.figure(figsize=(12, 5.5))
        if 'vol_ens' in d:
            for i, col in enumerate(d['vol_ens'].columns):
                plt.plot(d['vol_ens'].index, d['vol_ens'][col].values,
                         color='C0', alpha=0.15, linewidth=0.8,
                         label='Mira Ensemble' if i == 0 else '_nolegend_')
            d['vol_ens'].mean(axis=1).plot(label='Mira Nat Ensemble Mean', color='C0', linewidth=2)
        d['vol'].plot(label='Mira Hist', color='C3', linewidth=2)
        add_volume_refs(plt.gca(), d['vol'], d['inv_volume_2003'])
        plt.xlabel('Year'); plt.ylabel('Total volume (km$^3$)')
        plt.xlim([1200, d['vol'].index.max()])
        plt.legend(loc='lower left'); plt.title(f'OGGM simulations - {exp} - Volume')
        plt.savefig(fname=f'oggm_round2_{exp}_volume.png', bbox_inches='tight', dpi=150)

# %%
for exp in experiments:
    d = data[exp]
    with sns.axes_style("whitegrid"):
        plt.figure(figsize=(12, 5))
        if 'area_ens' in d:
            for i, col in enumerate(d['area_ens'].columns):
                plt.plot(d['area_ens'].index, d['area_ens'][col].values,
                         color='C0', alpha=0.2, linewidth=1,
                         label='Mira Ensemble' if i == 0 else '_nolegend_')
            d['area_ens'].mean(axis=1).plot(label='Mira Nat Ensemble Mean', color='C0', linewidth=2)
        d['area'].plot(label='Mira Hist', color='C3', linewidth=2)
        add_area_refs(plt.gca())
        plt.xlabel('Year'); plt.ylabel('Total area (km$^2$)')
        plt.xlim([1200, d['area'].index.max()])
        plt.legend(loc='lower left'); plt.title(f'OGGM simulations - {exp} - Area')
        plt.savefig(fname=f'oggm_round2_{exp}_area.png', bbox_inches='tight', dpi=150)

# %% [markdown]
# ## Hist-only comparison across all experiments

# %%
with sns.axes_style("whitegrid"):
    plt.figure(figsize=(12, 5.5))
    for exp in experiments:
        data[exp]['vol'].plot(label=exp, linewidth=2)
    plt.plot(1850, dfo['Volume LIA'], 'o', c='grey', label='Volume Ref. LIA')
    plt.plot(2015, dfo['Volume 2015'], 'o', c='k', label='Volume Ref. 2015')
    plt.xlabel('Year'); plt.ylabel('Total volume (km$^3$)')
    plt.xlim([1200, max(data[exp]['vol'].index.max() for exp in experiments)])
    plt.legend(loc='lower left'); plt.title('OGGM simulations - Round 2 - Hist comparison - Volume')
    plt.savefig(fname='oggm_round2_hist_comparison_volume.png', bbox_inches='tight', dpi=150)

# %%
with sns.axes_style("whitegrid"):
    plt.figure(figsize=(12, 5))
    for exp in experiments:
        data[exp]['area'].plot(label=exp, linewidth=2)
    plt.plot(1850, dfo['Area LIA'], 'o', c='grey', label='Area Reinthaler LIA')
    plt.plot(2003, dfo['Area 2003'], 'x', c='k', label='Area Reinthaler 2003')
    plt.plot(2015, dfo['Area 2015'], 'o', c='k', label='Area Reinthaler 2015')
    plt.xlabel('Year'); plt.ylabel('Total area (km$^2$)')
    plt.xlim([1200, max(data[exp]['area'].index.max() for exp in experiments)])
    plt.legend(loc='lower left'); plt.title('OGGM simulations - Round 2 - Hist comparison - Area')
    plt.savefig(fname='oggm_round2_hist_comparison_area.png', bbox_inches='tight', dpi=150)

# %%