#!/usr/bin/env python3
'''
Generate routing histograms.
'''
from glob import glob
from os import listdir, makedirs, remove
from os.path import splitext

import matplotlib.pyplot as plt
import matplotlib.ticker as tkr
import numpy as np
import seaborn as sns

def get_p_ev(net_path, net_i=0):
    log_paths = glob('nets/%s/%.4i-stats/*.npy' % (net_path, net_i))
    logs = [np.load(p)[()] for p in sorted(log_paths)]
    p_ev = np.zeros((len(logs), 8))
    for i, log in enumerate(logs):
        ℓ = log['root']['sinks'][0]
        for j in range(8):
            p_cor = ℓ['sinks'][0]['stats_ts']['p_cor']
            p_inc = ℓ['sinks'][0]['stats_ts']['p_inc']
            p_ev[i, j] = p_cor + p_inc
            try:
                ℓ = ℓ['sinks'][1]
            except:
                break
    return p_ev[::-1]

sns.set_style('ticks')
sns.set_color_codes()
plt.rcParams['text.usetex'] = True
plt.rcParams['text.latex.unicode'] = True
plt.rcParams['font.family'] = 'serif'
plt.figure(figsize=(5, 3))

plt.subplot(1, 4, 1)
plt.imshow(
    get_p_ev('hybrid-ac', 0),
    vmin=0, vmax=1, cmap='hot', aspect='auto')
plt.title(r'$k_\mathrm{cpt}=0$', fontsize='x-small')
plt.xticks([3, 7], ['4', '8'])
plt.yticks([16, 0], ['40k', '80k'])
plt.ylabel('Epoch Index')

plt.subplot(1, 4, 2)
plt.imshow(
    get_p_ev('hybrid-ac', 1),
    vmin=0, vmax=1, cmap='hot', aspect='auto')
plt.title(r'$k_\mathrm{cpt}=1{\times}10^{-9}$', fontsize='x-small')
plt.xticks([3, 7], ['4', '8'])
plt.yticks([])

plt.subplot(1, 4, 3)
plt.imshow(
    get_p_ev('hybrid-ac', 2),
    vmin=0, vmax=1, cmap='hot', aspect='auto')
plt.title(r'$k_\mathrm{cpt}=2{\times}10^{-9}$', fontsize='x-small')
plt.xticks([3, 7], ['4', '8'])
plt.yticks([])

plt.subplot(1, 4, 4)
plt.imshow(
    get_p_ev('hybrid-ac', 3),
    vmin=0, vmax=1, cmap='hot', aspect='auto')
plt.title(r'$k_\mathrm{cpt}=4{\times}10^{-9}$', fontsize='x-small')
plt.xticks([3, 7], ['4', '8'])
plt.yticks([])

plt.tight_layout(pad=0)
plt.colorbar(cax=plt.axes([0.9, 0.16, 0.025, 0.765]))
plt.text(
    0.5, 0, 'Layer Index',
    transform=plt.gcf().transFigure,
    va='bottom', ha='center', fontsize=11)
plt.subplots_adjust(bottom=0.16, top=0.925, right=0.88)
plt.savefig('figures/routing-hists.pdf')
