#!/usr/bin/env python3
'''
Generate accuracy/efficiency plots.
'''
from glob import glob
from os import listdir, makedirs, remove
from os.path import basename, splitext

import matplotlib as mpl
import matplotlib.patches as pch
import matplotlib.pyplot as plt
import numpy as np

mpl.style.use('classic')
import seaborn as sns

################################################################################
# Load experiment results.
################################################################################

logs = {}

for t in range(32):
    e = 2500 * (t + 1)
    logs['hybrid-ac-%.4i' % t] = [
        np.load(p)[()] for p in sorted(glob(
            'nets/hybrid-ac/*-stats/%.8i.npy' % e))]
    logs['hybrid-cr-%.4i' % t] = [
        np.load(p)[()] for p in sorted(glob(
            'nets/hybrid-cr/*-stats/%.8i.npy' % e))]
    logs['hybrid-cr-opt-%.4i' % t] = [
        np.load(p)[()] for p in sorted(glob(
            'nets/hybrid-cr-opt/*-stats/%.8i.npy' % e))]
    logs['hybrid-ac-tree-%.4i' % t] = [
        np.load(p)[()] for p in sorted(glob(
            'nets/hybrid-ac-tree/*-stats/%.8i.npy' % e))]

errs = {name: [1 - net['stats_ts']['acc'] for net in log]
        for name, log in logs.items()}
mocs = {name: [net['stats_ts']['moc'] for net in log]
        for name, log in logs.items()}

################################################################################
# Ensure that the output directory exists.
################################################################################

makedirs('figures', exist_ok=True)

################################################################################
# Generate node-link diagrams.
################################################################################

cls_labels = [
    '0', '1', '2', '3', '4',
    'Airplane', 'Automobile', 'Deer', 'Frog', 'Horse']

def nld_bounds(ℓ, x, y, r):
    x0 = x; x1 = x
    y0 = y; y1 = y
    sinks = ℓ['sinks'][1:]
    for i, s in enumerate(sinks):
        a = ((i + 1) / (len(sinks) + 1) - 0.5) * np.pi
        x_s = x + np.cos(a)
        y_s = y + r * np.sin(a)
        x2, x3, y2, y3 = nld_bounds(s, x_s, y_s, r / len(sinks))
        x0 = min(x0, x2); x1 = max(x1, x3)
        y0 = min(y0, y2); y1 = max(y1, y3)
    return x0, x1, y0, y1

def draw_acc_nld(ℓ, x, y, r):
    p_cor = ℓ['sinks'][0]['stats_ts']['p_cor']
    p_inc = ℓ['sinks'][0]['stats_ts']['p_inc']
    p_tot = p_cor + p_inc
    sinks = ℓ['sinks'][1:]
    if p_tot > 0.0025:
        plt.pie([p_cor / p_tot, p_inc / p_tot],
                center=(x, y), startangle=90,
                radius=(np.sqrt(p_tot) / 2),
                colors=[sns.color_palette()[0],
                        sns.color_palette()[2]])
    else:
        plt.pie([1], center=(x, y),
                radius=(np.sqrt(0.0025) / 2),
                colors=['k'])
    for i, s in enumerate(sinks):
        a = ((i + 1) / (len(sinks) + 1) - 0.5) * np.pi
        x_s = x + np.cos(a)
        y_s = y + r * np.sin(a)
        plt.plot([x, x_s], [y, y_s], 'k', linewidth=2, zorder=-1)
        draw_acc_nld(s, x_s, y_s, r / len(sinks))

def draw_cls_nld(ℓ, x, y, r):
    p_cor = ℓ['sinks'][0]['stats_ts']['p_cor_by_cls']
    p_inc = ℓ['sinks'][0]['stats_ts']['p_inc_by_cls']
    p_cls = list(map(sum, zip(p_cor, p_inc)))
    p_tot = sum(p_cls)
    sinks = ℓ['sinks'][1:]
    if p_tot > 0.0025:
        plt.pie([p / p_tot for p in p_cls],
                center=(x, y), startangle=90,
                radius=(np.sqrt(p_tot) / 2),
                colors=sns.color_palette('hls', len(p_cls)))
    else:
        plt.pie([1], center=(x, y),
                radius=(np.sqrt(0.0025) / 2),
                colors=['k'])
    for i, s in enumerate(sinks):
        a = ((i + 1) / (len(sinks) + 1) - 0.5) * np.pi
        x_s = x + np.cos(a)
        y_s = y + r * np.sin(a)
        plt.plot([x, x_s], [y, y_s], 'k', linewidth=2, zorder=-1)
        draw_cls_nld(s, x_s, y_s, r / len(sinks))

def make_chain_acc_nld(dst, log_name):
    plt.figure()
    n_nets = len(logs[log_name])
    max_depth = nld_bounds(logs[log_name][0]['root']['sinks'][0], 0, 0, 1)[1]
    for i, net in enumerate(logs[log_name]):
        draw_acc_nld(net['root']['sinks'][0], 0, i, 1)
    plt.arrow(0.4 * max_depth, -0.6, 0.2 * max_depth, 0, fc='#909090',
              ec='#909090', linewidth=3, head_width=0.075)
    plt.arrow(-0.6, 0.2 * (n_nets - 1), 0, 0.6 * (n_nets - 1), fc='#909090',
              ec='#909090', linewidth=3, head_width=0.075)
    plt.ylabel('Cost of Computation', fontsize='large', color='#606060')
    plt.xlabel('Layer Index', fontsize='large', color='#606060')
    plt.xlim(-0.7, max_depth + 0.7)
    plt.ylim(-0.7, n_nets - 0.7)
    plt.figlegend([
        pch.Patch(color=sns.color_palette()[0]),
        pch.Patch(color=sns.color_palette()[2])],
        ['Correct Labels', 'Incorrect Labels'],
        'lower right', fontsize='medium')
    plt.gca().set_aspect('equal')
    plt.gcf().set_size_inches(max_depth + 2, n_nets + 0.7)
    plt.tight_layout()
    plt.savefig('figures/' + dst, bbox_inches='tight')
    plt.close()

def make_chain_cls_nld(dst, log_name):
    plt.figure()
    n_nets = len(logs[log_name])
    max_depth = nld_bounds(logs[log_name][0]['root']['sinks'][0], 0, 0, 1)[1]
    for i, net in enumerate(logs[log_name]):
        draw_cls_nld(net['root']['sinks'][0], 0, i, 1)
    plt.arrow(0.4 * max_depth, -0.6, 0.2 * max_depth, 0, fc='#909090',
              ec='#909090', linewidth=3, head_width=0.075)
    plt.arrow(-0.6, 0.2 * (n_nets - 1), 0, 0.6 * (n_nets - 1), fc='#909090',
              ec='#909090', linewidth=3, head_width=0.075)
    plt.ylabel('Cost of Computation', fontsize='large', color='#606060')
    plt.xlabel('Layer Index', fontsize='large', color='#606060')
    plt.xlim(-0.7, max_depth + 0.7)
    plt.ylim(-0.7, n_nets - 0.7)
    plt.figlegend(
        [pch.Patch(color=c) for c in sns.color_palette('hls', 10)],
        cls_labels, 'lower center', ncol=5, columnspacing=1,
        fontsize='medium')
    plt.gca().set_aspect('equal')
    plt.gcf().set_size_inches(max_depth + 2, n_nets + 2)
    plt.tight_layout()
    plt.savefig('figures/' + dst, bbox_inches='tight')
    plt.close()

def make_tree_acc_nld(dst, log_name, i):
    plt.figure()
    net = logs[log_name][i]
    draw_acc_nld(net['root']['sinks'][0], 0, 0, 2)
    x0, x1, y0, y1 = nld_bounds(net['root']['sinks'][0], 0, 0, 2)
    plt.arrow(0.35 * x1, y0 - 0.33, 0.3 * x1, 0, fc='#909090',
              ec='#909090', linewidth=3, head_width=0.075)
    plt.xlabel('Dataflow', fontsize='large', color='#606060')
    plt.xlim(x0, x1)
    plt.ylim(y0 - 0.4, y1 + 0.05)
    plt.figlegend([
        pch.Patch(color=sns.color_palette()[0]),
        pch.Patch(color=sns.color_palette()[2])],
        ['Correct Labels', 'Incorrect Labels'],
        'lower right', fontsize='small')
    plt.gca().set_aspect('equal')
    plt.gcf().set_size_inches(x1 - x0 + 2, y1 - y0 + 2)
    plt.tight_layout()
    plt.savefig('figures/' + dst, bbox_inches='tight')
    plt.close()

def make_tree_cls_nld(dst, log_name, i, t=0):
    plt.figure()
    plt.title('Epoch %.6i' % t, x=0)
    net = logs[log_name][i]
    draw_cls_nld(net['root']['sinks'][0], 0, 0, 2)
    x0, x1, y0, y1 = nld_bounds(net['root']['sinks'][0], 0, 0, 2)
    plt.arrow(0.35 * x1, y0 - 0.33, 0.3 * x1, 0, fc='#909090',
              ec='#909090', linewidth=3, head_width=0.075)
    plt.xlabel('Dataflow', fontsize='x-large', color='#606060')
    plt.xlim(x0, x1)
    plt.ylim(y0 - 0.4, y1 + 0.05)
    plt.figlegend(
        [pch.Patch(color=c) for c in sns.color_palette('hls', 10)],
        cls_labels, 'lower left', fontsize='medium')
    plt.gca().set_aspect('equal')
    plt.gcf().set_size_inches(x1 - x0 + 2, y1 - y0 + 2)
    plt.tight_layout()
    plt.savefig('figures/' + dst, dpi=200)
    plt.close()

def make_chain_acc_and_cls_nld(dst, log_name, t=0):
    plt.figure()
    plt.title('Epoch %.6i' % t, x=0.95, fontsize='large')
    n_nets = len(logs[log_name])
    max_depth = nld_bounds(logs[log_name][0]['root']['sinks'][0], 0, 0, 1)[1]
    for i, net in enumerate(logs[log_name]):
        draw_acc_nld(net['root']['sinks'][0], 0, i, 1)
    for i, net in enumerate(logs[log_name]):
        draw_cls_nld(net['root']['sinks'][0], max_depth + 1, i, 1)
    plt.xlim(-0.9, 2 * max_depth + 1.9)
    plt.ylim(-1.7, n_nets - 0.7)
    plt.arrow(-0.7, 0.2 * (n_nets - 1), 0, 0.6 * (n_nets - 1), fc='#909090',
              ec='#909090', linewidth=3, head_width=0.075)
    plt.text(-0.9, 0.5 * n_nets - 0.5, 'Cost of Computation', va='center',
             ha='center', rotation='vertical', fontsize='large',
             color='#606060')
    plt.arrow(0.4 * max_depth, -0.5, 0.2 * max_depth, 0, fc='#909090',
              ec='#909090', linewidth=3, head_width=0.075)
    plt.text(0.5 * max_depth, -0.725, 'Layer Index', va='center', ha='center',
             fontsize='large', color='#606060')
    plt.arrow(1.4 * max_depth + 1, -0.5, 0.2 * max_depth, 0, fc='#909090',
              ec='#909090', linewidth=3, head_width=0.075)
    plt.text(1.5 * max_depth + 1, -0.725, 'Layer Index', va='center',
             ha='center', fontsize='large', color='#606060')
    plt.gca().add_artist(plt.legend(
        handles=[
            pch.Patch(color=sns.color_palette()[0], label='Correct Labels'),
            pch.Patch(color=sns.color_palette()[2], label='Incorrect Labels')],
        bbox_to_anchor=(0, -1.8, max_depth, 0),
        bbox_transform=plt.gca().transData,
        loc='lower center', fontsize='large'))
    plt.gca().add_artist(plt.legend(
        handles=[
            pch.Patch(color=c, label=cls_labels[i])
            for i, c in enumerate(sns.color_palette('hls', 10))],
        bbox_to_anchor=(max_depth + 1, -1.8, max_depth, 0),
        bbox_transform=plt.gca().transData,
        loc='lower center', fontsize='large',
        ncol=5, columnspacing=1))
    plt.gca().set_aspect('equal')
    plt.gcf().set_size_inches(2 * max_depth + 2, n_nets + 2)
    plt.tight_layout(pad=0)
    plt.savefig('figures/' + dst, bbox_inches='tight', dpi=200)
    plt.close()

makedirs('figures/hybrid-ac', exist_ok=True)
makedirs('figures/hybrid-cr', exist_ok=True)
makedirs('figures/hybrid-cr-opt', exist_ok=True)
makedirs('figures/hybrid-ac-tree-0', exist_ok=True)
makedirs('figures/hybrid-ac-tree-1', exist_ok=True)

plt.rcParams['text.usetex'] = True
plt.rcParams['font.family'] = 'serif'

for i in range(32):
    make_chain_acc_and_cls_nld(
        'hybrid-ac/nld-%.4i.png' % i,
        'hybrid-ac-%.4i' % i, 2500 * (i + 1))
for i in range(32):
    make_chain_acc_and_cls_nld(
        'hybrid-cr/nld-%.4i.png' % i,
        'hybrid-cr-%.4i' % i, 2500 * (i + 1))
for i in range(32):
    make_chain_acc_and_cls_nld(
        'hybrid-cr-opt/nld-%.4i.png' % i,
        'hybrid-cr-opt-%.4i' % i, 2500 * (i + 1))
for i in range(32):
    make_tree_cls_nld(
        'hybrid-ac-tree-0/nld-%.4i.png' % i,
        'hybrid-ac-tree-%.4i' % i, 0, 2500 * (i + 1))
for i in range(32):
    make_tree_cls_nld(
        'hybrid-ac-tree-1/nld-%.4i.png' % i,
        'hybrid-ac-tree-%.4i' % i, 1, 2500 * (i + 1))
