import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import math
import pickle

def plot_data(text, data, annotation, elem_image):
    len_text=len(text)
    n_top_char=len(data[0])

    data=data.transpose()
    annotation=annotation.transpose()

    labels = (np.asarray(["{0}\n{1:.2f}".format(annotation,data) for annotation, data in zip(annotation.flatten(), data.flatten())])).reshape(n_top_char,len_text)

    n_images=math.ceil(len_text/elem_image)
    for i in range(0,n_images):
        f, ax = plt.subplots(figsize=(25, 3))
        ax.xaxis.set_tick_params(labeltop='on')
        ax.xaxis.set_tick_params(labelbottom='')

        partial_text=text[i*elem_image:(i*elem_image)+elem_image]
        partial_labels=labels[:,i*elem_image:(i*elem_image)+elem_image]
        partial_data=data[:,i*elem_image:(i*elem_image)+elem_image]

        x_axis_labels = list(partial_text)

        heat_map = sns.heatmap(partial_data, annot=partial_labels, fmt='', xticklabels=x_axis_labels, yticklabels=False,  cmap="YlGnBu")

text,annotation, data=pickle.load(open( "20200518-223934.pickle", "rb" ))
annotation=np.array(annotation)
data=np.array(data)
plot_data(text,data,annotation,30)