%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")

matplotlib.rc('xtick', labelsize=14)
matplotlib.rc('ytick', labelsize=14)
plt.rcParams.update({'font.size': 14})
plt.rcParams.update({"axes.grid": True, "grid.linewidth": 0.5})

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


exps = ['improved_mag_gsc_searchperc2']
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)


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


agg(['model'])

agg(['weight_prune_perc'])

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


agg(['on_perc', 'model'])

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

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({'axes.labelsize': 14})
    plt.rcParams.update({'legend.fontsize': 14, 'legend.loc': 'lower left'})
    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()


# sns.scatterplot(data=df_plot, x='weight_prune_perc', y='val_acc_max', hue='model')
fig, ax = plt.subplots()
sns.lineplot(data=df_plot, x='weight_prune_perc', y='val_acc_max', hue='model')
plt.xlabel("% of weights pruned and grown at each epoch ($\\beta$)")
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/pruning_percentage.png"
fig_path = os.path.expanduser(fig_path)
plt.savefig(fig_path, dpi=1000, bbox_inches='tight')

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

"""
Explanation: Results not in line with what the previous experiment shows - what changed?
- All hyperparameters are set to the same values - expected similar results
- Any relevant changes to the code?
Edit: yes, momentum is set to 0.9, while in the original experiments is set to 0
End of explanation
"""


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

"""
Explanation: Acc where it reaches the max validation accuracy Consistently decreases in the Weight Magnitude - even in cases where the final acc is higher, as in between [0,0.2] of weight_prune_perc
End of explanation
"""


