#!/usr/bin/env python
# coding: utf-8

# In[31]:


import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from oggm import utils
import xarray as xr
import glob, os


# # Reference data

# In[17]:


dfz = pd.read_csv(utils.get_demo_file('zemp_ref_2006_2016.csv'), index_col=0)
dfh = pd.read_csv(utils.get_demo_file('table_hugonnet_regions_10yr_20yr_ar6period.csv'))
dfh = dfh.loc[dfh.period == '2000-01-01_2020-01-01'].set_index('reg')


# In[18]:


dfz['SMB_ZEMP'] = dfz['SMB'] * 1000
dfz['SMB_ZEMP_err'] = dfz['SMB_err'] * 1000
dfz['SMB_HUGUONNET'] = dfh['dmdtda']
dfz['SMB_HUGUONNET_err'] = dfh['err_dmdtda']
dfz.index = ['{:02d}'.format(rgi_reg) for rgi_reg in dfz.index]


# In[19]:


f, ax = plt.subplots()
dfz.plot(ax=ax, y='SMB_ZEMP', yerr='SMB_ZEMP_err', marker='o', linestyle='none');
dfz.plot(ax=ax, y='SMB_HUGUONNET', yerr='SMB_HUGUONNET_err', marker='o', linestyle='none');


# ## Fixed geometry SMB after calibration 

# In[35]:


idir = '/home/www/oggm/gdirs/oggm_v1.4/L3-L5_files/'
exps = glob.glob(idir + '*/*/qc*/pcp*/match_geod')
exps


# In[36]:


for exp in exps:

    dfz['SMB_OGGM'] = np.NaN
    dfz['SMB_OGGM_calv'] = np.NaN

    odir = os.path.join(exp, 'RGI62', 'b_080', 'L5', 'summary')
    
    title = exp.replace(idir, '').replace('/match_geod', '').replace('/elev_bands/', '/').replace('/', ' ')

    for rgi_reg in np.arange(1, 20):
        rgi_reg = '{:02d}'.format(rgi_reg)

        try:
            df = pd.read_csv(odir + '/fixed_geometry_mass_balance_{}.csv'.format(rgi_reg), index_col=0, low_memory=False)
            dfs = pd.read_csv(odir + '/glacier_statistics_{}.csv'.format(rgi_reg), index_col=0, low_memory=False)
        except FileNotFoundError: 
            print(odir, rgi_reg, 'MISSING')
            continue

        df = df.dropna(axis=0, how='all')
        df = df.dropna(axis=1, how='all')

        odf = pd.DataFrame(df.loc[2000:].mean(), columns=['SMB'])
        odf['AREA'] = dfs.rgi_area_km2
        dfz.loc[rgi_reg, 'SMB_OGGM'] = np.average(odf['SMB'], weights=odf['AREA'])

    f, ax = plt.subplots()
    dfz.plot(ax=ax, y='SMB_OGGM', marker='o', linestyle='none', markersize=9);
    dfz.plot(ax=ax, y='SMB_HUGUONNET', yerr='SMB_HUGUONNET_err', marker='o', linestyle='none');
    plt.title(title);


# ## Analysis of the regional bias corrections 

# In[56]:


for exp in exps:

    odir = os.path.join(exp, 'RGI62', 'b_080', 'L5', 'summary')
    title = exp.replace(idir, '').replace('/match_geod', '').replace('/elev_bands/', '/').replace('/', ' ')

    odfs = pd.DataFrame()
    for rgi_reg in np.arange(1, 20):
        rgi_reg = '{:02d}'.format(rgi_reg)
        try:
            dfs = pd.read_csv(odir + '/glacier_statistics_{}.csv'.format(rgi_reg), index_col=0, low_memory=False)
        except FileNotFoundError: 
            print(odir, rgi_reg, 'MISSING')
            continue

        dfs = dfs[['mb_bias_before_geodetic_corr', 'mb_bias', 'rgi_area_km2']].dropna()

        odfs.loc[rgi_reg, 'avg_bias_before_corr'] = np.average(dfs['mb_bias_before_geodetic_corr'], weights=dfs['rgi_area_km2'])
        odfs.loc[rgi_reg, 'avg_bias_after_corr'] = np.average(dfs['mb_bias'], weights=dfs['rgi_area_km2'])
        odfs.loc[rgi_reg, 'bias_corr'] = (dfs.mb_bias.iloc[0] - dfs.mb_bias_before_geodetic_corr.iloc[0])
        odfs.loc[rgi_reg, 'area'] = dfs.rgi_area_km2.sum()

    f, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
    odfs.plot(ax=ax1, y='bias_corr', marker='o', linestyle='none', markersize=9);
    ax1.set_title(title + ' (avg corr: {:.0f})'.format(np.average(odfs.bias_corr, weights=odfs['area'])));

    odfs.plot(ax=ax2, y='avg_bias_before_corr', marker='o', linestyle='none', markersize=9);
    odfs.plot(ax=ax2, y='avg_bias_after_corr', marker='o', linestyle='none', markersize=9);
    ax2.set_title('avg abs bias before: {:.0f} after: {:.0f}'.format(np.average(odfs.avg_bias_before_corr.abs(), weights=odfs.area), 
                                                                     np.average(odfs.avg_bias_after_corr.abs(), weights=odfs.area)));


# # Now compare the SMB of the dynamics run, which is not perfect anymore

# In[59]:


for exp in exps:

    odir = os.path.join(exp, 'RGI62', 'b_080', 'L5', 'summary')
    title = exp.replace(idir, '').replace('/match_geod', '').replace('/elev_bands/', '/').replace('/', ' ')

    for rgi_reg in np.arange(1, 20):
        rgi_reg = '{:02d}'.format(rgi_reg)
        f = odir + '/historical_run_output_extended_{}.nc'.format(rgi_reg)
        try:
            with xr.open_dataset(f) as ds:
                ds = ds.sel(time=slice(2000, None))
                area = ds.area.isel(time=0)
                area = area.where(area > 0)
                dmdtda = (ds.volume.isel(time=-1) - ds.volume.isel(time=0)) / area / len(ds.time) * 900  # ice density
        except FileNotFoundError: 
            pass
        dfz.loc[rgi_reg, 'SMB_OGGM_run'] = np.average(dmdtda.dropna(dim='rgi_id'), weights=area.dropna(dim='rgi_id')) 

    f, ax = plt.subplots()
    dfz.plot(ax=ax, y='SMB_OGGM_run', marker='o', linestyle='none', markersize=9);
    dfz.plot(ax=ax, y='SMB_HUGUONNET', yerr='SMB_HUGUONNET_err', marker='o', linestyle='none');
    plt.xticks(range(0, len(dfz.index)), labels=dfz.index);
    plt.title(title);


# In[ ]: