%load_ext autoreload
%autoreload 2

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import glob
import tabulate
import pprint
import click
import numpy as np
import pandas as pd
from ray.tune.commands import *
from nupic.research.frameworks.dynamic_sparse.common.browser import *

import matplotlib
import matplotlib.pyplot as plt
from matplotlib import rcParams

%config InlineBackend.figure_format = 'retina'

import seaborn as sns
sns.set(style="whitegrid")
sns.set_palette("colorblind")

"""
Explanation: Experiment:
Evaluate pruning by magnitude weighted by coactivations (more thorough evaluation), compare it to baseline (SET), in GSC. Applied only to linear layers
Motivation.
Check if results are consistently above baseline.
Conclusion
End of explanation
"""


exps = ['improved_mag_gsc_eval3', 'improved_mag_gsc_eval4']
paths = [os.path.expanduser("~/nta/results/{}".format(e)) for e in exps]
df = load_many(paths)

df.head(5)

# replace hebbian prine
df['hebbian_prune_perc'] = df['hebbian_prune_perc'].replace(np.nan, 0.0, regex=True)
df['weight_prune_perc'] = df['weight_prune_perc'].replace(np.nan, 0.0, regex=True)

df.columns

df.shape

df.iloc[1]

df.groupby('model')['model'].count()

"""
Explanation: Load and check data
End of explanation
"""


# Did any  trials failed?
df[df["epochs"]<30]["epochs"].count()

# Removing failed or incomplete trials
df_origin = df.copy()
df = df_origin[df_origin["epochs"]>=30]
df.shape

# which ones failed?
# failed, or still ongoing?
df_origin['failed'] = df_origin["epochs"]<30
df_origin[df_origin['failed']]['epochs']

# helper functions
def mean_and_std(s):
    return "{:.3f} ± {:.3f}".format(s.mean(), s.std())

def round_mean(s):
    return "{:.0f}".format(round(s.mean()))

stats = ['min', 'max', 'mean', 'std']

def agg(columns, filter=None, round=3):
    if filter is None:
        return (df.groupby(columns)
             .agg({'val_acc_max_epoch': round_mean,
                   'val_acc_max': stats,                
                   'model': ['count']})).round(round)
    else:
        return (df[filter].groupby(columns)
             .agg({'val_acc_max_epoch': round_mean,
                   'val_acc_max': stats,                
                   'model': ['count']})).round(round)


def agg2(columns, filter=None, round=3):
    if filter is None:
        return (df2.groupby(columns)
             .agg({'val_acc_max_epoch': round_mean,
                   'val_acc_max': stats,                
                   'model': ['count']})).round(round)
    else:
        return (df2[filter].groupby(columns)
             .agg({'val_acc_max_epoch': round_mean,
                   'val_acc_max': stats,                
                   'model': ['count']})).round(round)

base = os.path.join('gsc-smalldense-2019-10-11-exp1')
exps = [
    os.path.join(base, exp) for exp in [
        'gsc-smalldense'
    ]
]
paths = [os.path.expanduser("~/nta/results/{}".format(e)) for e in exps]
df2 = load_many(paths)

df_agg = agg2(['equivalent_on_perc'])
df_agg

equivalent_on_percs = df_agg.index.values

val_means = df_agg['val_acc_max']['mean']
val_means = np.array(val_means)
val_stds = df_agg['val_acc_max']['std']
val_stds = np.array(val_stds)

"""
Explanation: ## Analysis
Experiment Details
End of explanation
"""


agg(['model'])

agg(['on_perc'])

agg(['on_perc', 'model'], df['on_perc'] == 0.04)

# translate model names
rcParams['figure.figsize'] = 16, 8
d = {
    'DSNNWeightedMag': 'Dynamic Sparse Neural Network (DSNN)',
    'DSNNMixedHeb': 'Sparse Evolutionary Training (SET)',
    'SparseModel': 'Static Sparse',        
}
df_plot = df.copy()
df_plot['model'] = df_plot['model'].apply(lambda x: d[x])

# sns.scatterplot(data=df_plot, x='on_perc', y='val_acc_max', hue='model')
sns.lineplot(data=df_plot, x='on_perc', y='val_acc_max', hue='model')

fltr = df_plot['model'] != 'Sparse'
agg(['on_perc', 'model'], filter=fltr)

!source ~/.bash_profile
!echo $GOOGLE_DRIVE_PATH
!echo hi

def plot_for_paper():
    rcParams['figure.figsize'] = 10,6
    matplotlib.rc('xtick', labelsize=14)
    matplotlib.rc('ytick', labelsize=14)
    matplotlib.rc('ytick', labelsize=14)
    plt.rcParams.update({'text.usetex': True})
    plt.rcParams.update({'axes.labelsize': 14})
    plt.rcParams.update({'legend.fontsize': 14, 'legend.loc': 'lower right'})
    plt.rcParams.update({'font.size': 14})
    plt.rcParams.update({"axes.grid": True, "grid.linewidth": 0.5})
    plt.rcParams.update({'font.family': 'serif'})
    plt.rcParams.update({'text.usetex': True})    
    
    
plot_for_paper()

fig, ax = plt.subplots()
sns.lineplot(data=df_plot, x='on_perc', y='val_acc_max', hue='model')
plt.errorbar(
    x=equivalent_on_percs,
    y=val_means,
    yerr=val_stds,
    color='k',
    marker='.',
    lw=0,
    elinewidth=2,
    capsize=4,
    markersize=5,
    label="Small-Dense Equivalents"
)

plt.xlabel("% of active weights ($\sigma$)")
plt.ylabel("test accuracy")
plt.ylim((0.4, 1.0))
handles, labels = ax.get_legend_handles_labels()
ax.legend(handles=handles[1:], labels=labels[1:])
fig_path = r"~/Google Drive/papers/dynamic_sparse/figures/neurips_2019/active_weights.png"
fig_path = os.path.expanduser(fig_path)

plt.savefig(fig_path, dpi=1000, bbox_inches='tight')

    # plt.rcParams['figure.figsize'] = (8, 3)
    font_size=12
#     plt.rcParams.update({'font.size': font_size})
#     plt.rcParams.update({'font.family': 'Times New Roman'})
#     plt.rcParams.update({'axes.labelsize': font_size})
#     plt.rcParams.update({'axes.titlesize': 1.5*font_size})
#     plt.rcParams.update({'legend.fontsize': font_size})
#     plt.rcParams.update({'xtick.labelsize': font_size})
#     plt.rcParams.update({'ytick.labelsize': font_size})
#     plt.rcParams.update({'savefig.dpi': 2*plt.rcParams['savefig.dpi']})
#     plt.rcParams.update({'xtick.major.size': 3})
#     plt.rcParams.update({'xtick.minor.size': 3})
#     plt.rcParams.update({'xtick.major.width': 1})
#     plt.rcParams.update({'xtick.minor.width': 1})
#     plt.rcParams.update({'ytick.major.size': 3})
#     plt.rcParams.update({'ytick.minor.size': 3})
#     plt.rcParams.update({'ytick.major.width': 1})
#     plt.rcParams.update({'ytick.minor.width': 1 })
    # plt.rcParams.update({})

# sns.scatterplot(data=df_plot, x='on_perc', y='val_acc_max', hue='model')
sns.lineplot(data=df_plot, x='on_perc', y='val_acc_max', hue='model')
plt.ylim(0.7,0.98)
plt.errorbar(
    x=equivalent_on_percs,
    y=val_means,
    yerr=val_stds,
    color='k',
    marker='*',
    lw=0,
    elinewidth=2,
    capsize=2,
    markersize=10,
    label="Small-Dense Equivalents"
)

rcParams['figure.figsize'] = 16, 8
filter = df_plot['model'] != 'Static'
sns.lineplot(data=df_plot[filter], x='on_perc', y='val_acc_max_epoch', hue='model')

sns.lineplot(data=df_plot, x='on_perc', y='val_acc_last', hue='model')

"""
Explanation: Does improved weight pruning outperforms regular SET
End of explanation
"""


