""""
Visualize input and output weights of trained classifiers and their response to samples from WLASL dataset
"""
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
import numpy as np
import cv2
import pandas as pd

import matplotlib.pyplot as plt

from modules.dataset_utils import init_dataset, load_json
from modules.model import SignRecognitionModel
from modules.config import CFG


def check_wlasl_input_weights():
    """Visualize input weights of trained models
    """
    # get input weights from each model
    classes = (100, 300, 1000, 2000)
    weights = list()
    for n_classes in classes:
        model = SignRecognitionModel((None, 2*21*2+1*98*2), n_classes)
        model.load_weights(CFG.WLASL_WEIGHTS[n_classes]).expect_partial()
        weights.append(model.layers[1].weights[0].numpy())
        del model
    # reduce last weight dimension to get weights of each input feature
    weights = [w.sum(axis=1) for w in weights]
    
    # image tileset transformation function 
    transform = lambda x: cv2.cvtColor((np.delete(plt.get_cmap('jet').__call__(cv2.normalize(x, None, 0, 1, norm_type=cv2.NORM_MINMAX)), 3, 2) * 255).astype(np.uint8), cv2.COLOR_RGB2BGR)
    # number of pixels per feature
    scale = 10
    # construct visualization tiles of features
    tiles = [np.kron(np.stack((w[0::2], w[1::2]), axis=0), np.ones((scale, scale), dtype=w.dtype)) for w in weights]
    tiles = [transform(t) for t in tiles]
    
    # weights of each keypoint normed to <0, 1>
    data = [cv2.normalize(w[0::2] + w[1::2], None, 0, 1, norm_type=cv2.NORM_MINMAX).reshape(-1) for w in weights]
    # prepare figure
    fig, axes = plt.subplots(len(classes), 1)
    fig.set_size_inches(18.5, 10.5)
    fig.suptitle('Input weights', fontsize=16)
    # init axes
    for ax in axes:
        ax.set_ylim(0, 1)
        ax.set_xlim(0, 2*21+1*98)
        ax.yaxis.set_ticks(np.arange(0, 1.25, 0.25))
        ax.xaxis.set_ticks(np.arange(0, 2*21+1*98, 1))
        labels = [f'H{i+1}' for i in range(21)] + [f'H{i+1}' for i in range(21)] + [f'F{i+1}' for i in range(98)]
        for i in range(len(labels)):
            labels[i] = labels[i] if i % 10 == 0 or i == len(labels) - 1 else ''
        ax.set_xticklabels(labels)
        ax.grid(axis='y')
    # draw data
    for ax, s in zip(axes, data):
        ax.bar(np.arange(0, 2*21+1*98, 1), s)
    ax.legend()
    plt.show()
    # show tileset
    for t, c in zip(tiles, classes):
        cv2.imshow(f'hands_{c}', t[:, :scale*2*21])
        cv2.imshow(f'face_{c}', t[:, scale*2*21:])
    cv2.waitKey(0)
    cv2.destroyAllWindows()


def check_wlasl_output_weights():
    """Visualize output weights of trained models and save them to file
    """
    # get output weights from each model
    classes = (100, 300, 1000, 2000)
    weights = list()
    for n_classes in classes:
        model = SignRecognitionModel((None, 2*21*2+1*98*2), n_classes)
        model.load_weights(CFG.WLASL_WEIGHTS[n_classes]).expect_partial()
        weights.append(model.layers[-1].weights[0].numpy())
        del model
    # reduce first weight dimension to get weights of each output class
    weights = [w.sum(axis=0) for w in weights]

    # image tileset transformation function 
    transform = lambda x: cv2.cvtColor((np.delete(plt.get_cmap('jet').__call__(cv2.normalize(x, None, 0, 1, norm_type=cv2.NORM_MINMAX)), 3, 2) * 255).astype(np.uint8), cv2.COLOR_RGB2BGR)
    # number of pixels per feature
    scale = 10
    # construct visualization tiles of classes
    tiles = [np.kron(np.stack((w[0::2], w[1::2]), axis=0), np.ones((scale, scale), dtype=w.dtype)) for w in weights]
    tiles = [transform(t) for t in tiles]

    # weights of each class normed to <0, 1>
    data = [cv2.normalize(w, None, 0, 1, norm_type=cv2.NORM_MINMAX).reshape(-1) for w in weights]
    # prepare figure
    fig, axes = plt.subplots(len(classes), 1)
    fig.set_size_inches(18.5, 10.5)
    fig.suptitle('Output weights', fontsize=16)
    # init axes
    for ax, cls in zip(axes, classes):
        ax.set_ylim(0, 1)
        ax.set_xlim(0, cls)
        ax.xaxis.set_ticks([])
        ax.grid(axis='y')
    # draw data
    for ax, s, cls in zip(axes, data, classes):
        ax.bar(np.arange(0, cls, 1), s)
    ax.legend()
    plt.show()
    # show tileset
    for t, c in zip(tiles, classes):
        cv2.imshow(f'hands_{c}', t[:, :scale*2*21])
        cv2.imshow(f'face_{c}', t[:, scale*2*21:])
    cv2.waitKey(0)
    cv2.destroyAllWindows()
    
    # load class map and invert to to ID: name
    class_dict = load_json(CFG.WLASL_DATASET_PATHS['class_map'])
    class_dict = {val: key for key, val in class_dict.items()}
    # ordered class names
    glosses = [class_dict[i] for i in range(2000)]
    # prepare table, first column are class names
    columns, headers = [glosses], ['glosses']
    # pad to 2000 vector length function
    pad_to_shape = lambda x: np.pad(x.astype(np.float32), ((0, 2000-x.shape[0])))
    # pad output weights and add to table
    for d, cls in zip(data, classes):
        d = pad_to_shape(d)
        columns.append(d)
        headers.append(f'weights_{cls}')
    # create pandas table and save to csv
    df = pd.DataFrame(data=np.array(columns).T, columns=headers)
    if os.path.isfile(CFG.WLASL_EVAL_WEIGHTS_PATH):
        os.remove(CFG.WLASL_EVAL_WEIGHTS_PATH)
    df.to_csv(CFG.WLASL_EVAL_WEIGHTS_PATH, float_format='%.4f', sep=';')


def visualize_outputs():
    """Visualize models' response to WLASL samples
    """
    # load dataset
    test_dataset = init_dataset(
        CFG.WLASL_TFRECORDS_PATH[100]['test'],
        1,
        100
    )
    # get trained models
    models = list()
    classes = [100, 300, 1000, 2000]
    for n_classes in classes:
        model = SignRecognitionModel((None, 2*21*2+1*98*2), n_classes)
        model.load_weights(CFG.WLASL_WEIGHTS[n_classes]).expect_partial()
        models.append(model)
    # prepare figure, init axes - overall bliting preparation
    fig, axes = plt.subplots(len(classes), 1)
    fig.set_size_inches(18.5, 10.5)
    fig.suptitle('Output class distribution', fontsize=16)
    pred_bars, true_bars = list(), list()
    for ax, cls in zip(axes, classes):
        ax.set_ylim(0, 1)
        ax.set_xlim(0, classes[-1])
        ax.xaxis.set_ticks(classes)
        ax.tick_params(axis='both', which='major', labelsize=12)
        ax.set_ylabel(f'p_{cls}(c)', fontsize=16)
        ax.grid(axis='y')
        bar = ax.bar(np.arange(0, cls, 1), np.zeros((cls, )), width=3)
        pred_bars.append(bar)
        line, = ax.plot([], [], linewidth=0.8, color='g')
        true_bars.append(line)
    axes[-1].set_xlabel('predikované třídy', fontsize=16)
    plt.show(block=False)
    plt.pause(0.1)
    # init figure canvas
    background = fig.canvas.copy_from_bbox(fig.bbox)
    for ax, bar in zip(axes, pred_bars):
        for artist in bar.patches:
            ax.draw_artist(artist)
    fig.canvas.blit(fig.bbox)
    # parse dataset
    for x, y in test_dataset:
        # true class
        y_true = np.argmax(y[0, ...], axis=-1)
        # visualize only samples of classes included in all models
        if y_true >= 100:
            continue
        # get predictions
        y_preds = [model(x)[0, ...] for model in models]
        # visualize predictions on canvas
        fig.canvas.restore_region(background)
        for ax, bar, line, cls, y_pred in zip(axes, pred_bars, true_bars, classes, y_preds):
            line.set_data([y_true, y_true], [0, 1])
            ax.draw_artist(line)
            for artist, val in zip(bar.patches, y_pred):
                artist.set_height(val)
                ax.draw_artist(artist)
        fig.canvas.blit(fig.bbox)
        fig.canvas.flush_events()
        plt.waitforbuttonpress()


if __name__ == '__main__':
    check_wlasl_input_weights()
    # check_wlasl_output_weights()
    # visualize_outputs()
