"""
Evaluate trained classification models
"""
import os
import tensorflow as tf
import numpy as np
import pandas as pd

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


def eval_wlasl_classes(n_classes, split):
    """Evaluate WLASL model

    Args:
        n_classes (int): number of classes
        split (str): dataset subset to evaluate on

    Returns:
        dict: dictionary of evaluation results
    """
    # init dataset, model and load model weights
    test_data = init_dataset(
        CFG.WLASL_TFRECORDS_PATH[n_classes][split],
        1,
        n_classes
    )
    model = SignRecognitionModel((None, 2*21*2+1*98*2), n_classes)
    model.load_weights(CFG.WLASL_WEIGHTS[n_classes]).expect_partial()
    # evaluation metrics
    metrics = [
        tf.keras.metrics.TopKCategoricalAccuracy(k=1, name='t1_acc'),
        tf.keras.metrics.TopKCategoricalAccuracy(k=5, name='t5_acc'),
        tf.keras.metrics.TopKCategoricalAccuracy(k=10, name='t10_acc'),
        ClassSpecificMetrics(n_classes, name='classes'),
    ]
    model.compile(metrics=metrics)
    # eval and return
    results = model.evaluate(test_data, return_dict=True)
    return results


def eval_wlasl(to_file=True, only_test=True):
    """WLASL evaluation

    Args:
        to_file (bool, optional): flag wether to save detailed evaluation results. Defaults to True.
        only_test (bool, optional): flag wether to perform evaluation only on testing subset or on training and validation subsets. Defaults to True.
    """
    # evaluate all models
    splits = ['test'] if only_test else ['train', 'val', 'test']
    results = {n_classes: {split: eval_wlasl_classes(n_classes, split) for split in splits} for n_classes in (100, 300, 1000, 2000)}
    # if not logging, return
    if not to_file:
        return
    # load class map and inverse it 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 list of sign names
    glosses = [class_dict[i] for i in range(2000)]
    # prepare table, first column is name of words
    columns, headers = [glosses], ['glosses']
    # zero-pad a vector to 2000 length function
    pad_to_shape = lambda x: np.pad(x.astype(np.float32), ((0, 2000-x.shape[0])))
    # for every model
    for n_classes in (100, 300, 1000, 2000):
        # for every split and its results
        for split, result in results[n_classes].items():
            # stat results
            tp = pad_to_shape(result['tp:0'])
            fp = pad_to_shape(result['fp:0'])
            tn = pad_to_shape(result['tn:0'])
            fn = pad_to_shape(result['fn:0'])
            # calculate F1 score
            precision = tp / (tp + fp + 1e-6)
            recall = tp / (tp + fn + 1e-6)
            f1 = 2 * precision * recall / (precision + recall + 1e-6)
            # add columns to table
            n_classes = result['tp:0'].shape[0]
            columns += [np.around(x, decimals=4) for x in [tp, fp, tn, fn, precision, recall, f1]]
            headers += [f'{header}_{split}_{n_classes}' for header in ['tp', 'fp', 'tn', 'fn', 'precision', 'recall', 'f1']]
    # create a dataframe and write to csv file
    df = pd.DataFrame(data=np.array(columns).T, columns=headers)
    if os.path.isfile(CFG.WLASL_EVAL_METRICS_PATH):
        os.remove(CFG.WLASL_EVAL_METRICS_PATH)
    df.to_csv(CFG.WLASL_EVAL_METRICS_PATH, float_format='%.4f', sep=';')


if __name__ == '__main__':
    eval_wlasl(to_file=False, only_test=True)
    # eval_wlasl(to_file=True, only_test=False)
