"""This script plots the results generated by benchmark.py
"""
import numpy as np
import matplotlib.pyplot as plt
from benchmark_v2 import ks, nr_refs, nr_queries
from matplotlib.colors import LogNorm
from matplotlib.ticker import LogFormatterExponent # <-- one new import here
import matplotlib.ticker as ticker

plt.set_cmap('jet')
timing_results = np.load("benchmark_results.npz")["timing_results"]

nr_ks = ks.size
nr_algs = timing_results.shape[0]
grid = np.stack(np.meshgrid(nr_refs, nr_queries, indexing='ij'), axis=-1)

valid_timing_results = timing_results.ravel()[~np.isnan(timing_results.ravel())]
max_timing = np.minimum(np.max(valid_timing_results), 1.)
cont_lines = np.logspace(np.log10(np.min(valid_timing_results)), np.log10(max_timing), num=15)
#cont_lines = np.log(cont_lines) / np.log(10)
#Create figure matrix
fig, axes = plt.subplots(nrows=nr_algs, ncols=nr_ks, sharex='col', sharey='row') #, gridspec_kw={'hspace': 0, 'wspace': 0})

#https://stackoverflow.com/questions/25983218/scientific-notation-colorbar-in-matplotlib
def fmt(x, pos):
    a, b = '{:.1e}'.format(x).split('e')
    b = int(b)
    return r'${} \cdot 10^{{{}}}$'.format(a, b)

for alg_i in range(nr_algs):
    for k_i, k in enumerate(ks):
        ax = axes[alg_i, k_i]
        #ax.contourf(origin='lower', )
        cont = ax.contourf(grid[..., 1], grid[..., 0], timing_results[alg_i, ..., k_i].T, 
                            cont_lines, norm=LogNorm(vmin=np.min(valid_timing_results), 
                                                     vmax=max_timing))
        ax.set_xscale("log")
        ax.set_yscale("log")

        if alg_i == 0:
            ax.set_title("k = {:d}".format(k))

        if k_i == 0:
            ax.set_ylabel("#Refs")

        if alg_i == nr_algs - 1:
            ax.set_xlabel("#Queries")


fig.subplots_adjust(right=0.8)
cbar_ax = fig.add_axes([0.85, 0.15, 0.05, 0.7])
cbar = fig.colorbar(cont, cax=cbar_ax, format=ticker.FuncFormatter(fmt))
cbar.set_label('Time [$\\log_{10}$(s)]', rotation=270, labelpad=12)

#https://stackoverflow.com/questions/24814490/how-to-label-rows-cols-of-a-matrix-of-plots
#Add some annotations
for ax, row in zip(axes[:,0], ['Numpy', 'Scipy KD-Tree', 'TF KD-Tree']):
    ax.annotate(row, (0, 0.5), xytext=(-45, 0), ha='right', va='center',
                size=20, rotation=90, xycoords='axes fraction',
                textcoords='offset points')
fig.suptitle("Benchmark Timings")
#fig.text(0.5, 0.04, 'K Variation', ha='center')
#fig.text(0.04, 0.5, 'Algorithms', va='center', rotation='vertical')
fig.set_size_inches((10, 8))
fig.savefig("benchmark.png")