"""
Evaluate Interhand model and trained hand and face pose estimation models
"""
import os
# os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
import numpy as np
import matplotlib.pyplot as plt
from json import dump, load

from modules.config import CFG
from modules.metrics import InterhandMetric, FaceKeypointMetric, HandKeypointMetric

from modules.interhand_utils import init_dataset as init_hands_dataset
from modules.wflw_utils import init_dataset as init_faces_dataset
from modules.hands_utils import init_dataset as init_hands_dataset
from modules.interhand import InterhandEstimator
from modules.face_model import load_model as load_face_model
from modules.hands_model import load_model as load_hands_model


def evaluate_model(model, metric, dataset, residuals=False):
    """Evaluate general model with passed metric

    Args:
        model (tf.keras.Model): evaluated model
        metric (tf.keras.Metric): keras metric or a general metric object with update_state() method
        dataset (tf.keras.Dataset): iterable yielding x and y pairs
        residuals (bool, optional): flag wether to request residual logging from passed metric. Defaults to False.

    Returns:
        dict: output dictionary of metric results, if residuals are logged, return residuals as well
    """
    # iterate over dataset
    for x, y_true in dataset:
        # predict and update metric
        y_pred = model(x)
        metric.update_state(y_true, y_pred)
    # calculate metric results
    result = metric.result()
    # get residuals if passed
    if residuals:
        residuals = metric.get_residuals()
    # reset metric state and return
    metric.reset_state()
    return (result, residuals) if residuals else result


def evaluate_pretrained_interhand(handness_tholds, log_residuals):
    """Evaluate pretrained Interhand model

    Args:
        handness_tholds (list): list of classification thresholds
        log_residuals (bool): flag wether or not log and save metric residuals to file
    """
    # convert threshold to list if a float is passed
    handness_thresholds = handness_tholds if isinstance(handness_tholds, list) else [handness_tholds]
    # load testing set
    dataset = init_hands_dataset(
        CFG.INTERHAND_TFRECORD_PATHS['test'],
        CFG.INTERHAND_PARAMS.BATCH_SIZE,
        CFG.INTERHAND_PARAMS.IMG_SHAPE,
        CFG.INTERHAND_PARAMS.JOINT_HM_DIMS,
        CFG.INTERHAND_PARAMS.DEPTH_DIM,
        CFG.INTERHAND_PARAMS.ROOT_HM_DEPTH,
        CFG.INTERHAND_PARAMS.SIGMA,
        testing=True)
    # load model and weights
    model = InterhandEstimator(CFG.INTERHAND_PARAMS.IMG_SHAPE, CFG.INTERHAND_PARAMS.JOINT_HM_DIMS,
                               CFG.INTERHAND_PARAMS.DEPTH_DIM, CFG.INTERHAND_PARAMS.ROOT_HM_DEPTH)
    model.load_weights(CFG.INTERHAND_TF_PATH).expect_partial()
    # define metric
    metric = InterhandMetric(handness_thresholds, log_residuals=log_residuals)
    # evaluate model
    results, residuals = evaluate_model(model, metric, dataset, residuals=True)
    handness, depth, hand_joints, depth_joints = results
    print(f'Pretrained InterHand:\nHandness:\n{handness}\nDepth: {depth}\n2D Joints: {hand_joints}\nDepth Joints: {depth_joints}')
    # save residuals to file if requested
    if log_residuals:
        if os.path.isfile(CFG.INTERHAND_RESIDUALS_PATH):
            os.remove(CFG.INTERHAND_RESIDUALS_PATH)
        with open(CFG.INTERHAND_RESIDUALS_PATH, 'wt') as file:
            dump(residuals, file, indent=4)


def evaluate_residuals():
    """Evaluate logged Interhand residuals
    """
    # read residuals
    with open(CFG.INTERHAND_RESIDUALS_PATH, 'rt') as file:
        residuals = load(file)
    # calculate mean and std of tracked values
    joints_2d = np.array(residuals['joints_2d']) * 256
    mean_j_2d = np.mean(joints_2d)
    std_j_2d = np.std(joints_2d)
    joints_depth = np.array(residuals['joints_depth'])
    mean_j_depth = np.mean(joints_depth)
    std_j_depth = np.std(joints_depth)
    depth = np.array(residuals['depth'])
    mean_depth = np.mean(depth)
    std_depth = np.std(depth)
    # plot
    plt.figure()
    plt.subplot(1, 3, 1)
    plt.hist(joints_2d, bins=1000, density=True, range=(mean_j_2d-3*std_j_2d, mean_j_2d+3*std_j_2d))
    plt.title(f'2D Joints {mean_j_2d:.5f}±{std_j_2d:.5f}')
    plt.grid()
    plt.subplot(1, 3, 2)
    plt.hist(joints_depth, bins=1000, density=True, range=(mean_j_depth-3*std_j_depth, mean_j_depth+3*std_j_depth))
    plt.title(f'Depth Joints {mean_j_depth:.5f}±{std_j_depth:.5f}')
    plt.grid()
    plt.subplot(1, 3, 3)
    plt.hist(depth, bins=1000, density=True, range=(mean_depth-3*std_depth, mean_depth+3*std_depth))
    plt.title(f'Relative Depth {mean_depth:.5f}±{std_depth:.5f}')
    plt.grid()
    plt.show()


def evaluate_face_keypoints(backbone, decode_method='argmax'):
    """Evaluate trained face pose estimation model

    Args:
        backbone (str): model backbone configuration, see model implementation for possible inputs.
        decode_method (str, optional): heatmap decoding method, see model implementation for possible inputs. Defaults to 'argmax'.
    """
    # load dataset
    test_data = init_faces_dataset(
        CFG.WFLW_TFRECORD_PATHS['test'],
        CFG.WFLW_PARAMS.BATCH_SIZE,
        CFG.WFLW_PARAMS.IMG_SHAPE,
        CFG.WFLW_PARAMS.HM_DIMS,
        CFG.WFLW_PARAMS.SIGMA,
        testing=True
    )
    # load model
    model = load_face_model(CFG.WFLW_WEIGHTS[backbone], backbone, CFG.WFLW_PARAMS.IMG_SHAPE, CFG.WFLW_PARAMS.HM_DIMS,
                            imagenet=False, decode=decode_method, delta=CFG.WFLW_PARAMS.DELTA)
    # define metric
    metric = FaceKeypointMetric()
    # evaluate model
    overall_results, attrib_results = evaluate_model(model, metric, test_data)
    # scale results to input image plane size
    scale = CFG.WFLW_PARAMS.IMG_SHAPE[0]
    # print results
    attrib_results = ''.join(f'{name}: {scale*res:.2f}  ' for name, res in attrib_results.items())
    print(f'Face keypoint estimator with {backbone} backbone using {decode_method}: Overall: {scale*overall_results:.2f}  {attrib_results}')


def evaluate_hand_keypoints(backbone, decode_method='argmax'):
    """Evaluate trained hand pose estimation model

    Args:
        backbone (str): model backbone configuration, see model implementation for possible inputs.
        decode_method (str, optional): heatmap decoding method, see model implementation for possible inputs. Defaults to 'argmax'.
    """
    # load model
    model = load_hands_model(CFG.HANDS_WEIGHTS[backbone], backbone, CFG.HANDS_PARAMS.IMG_SHAPE, CFG.HANDS_PARAMS.HM_DIMS,
                             imagenet=False, decode=decode_method, delta=CFG.HANDS_PARAMS.DELTA)
    # evaluate model on different testing subsets
    datasets = ['test', 'coco_test', 'mwb_test', 'rhd_test']
    for dataset in datasets:
        # load dataset
        test_data = init_hands_dataset(
            CFG.HANDS_TFRECORD_PATHS[dataset],
            CFG.HANDS_PARAMS.BATCH_SIZE,
            CFG.HANDS_PARAMS.IMG_SHAPE,
            CFG.HANDS_PARAMS.HM_DIMS,
            CFG.HANDS_PARAMS.SIGMA,
            testing=True
        )
        # define metric
        metric = HandKeypointMetric([0.3])
        # evaluate model
        results = evaluate_model(model, metric, test_data)
        # scale results to input image plane size and print results
        scale = CFG.HANDS_PARAMS.IMG_SHAPE[0]
        print(f'Hand keypoint with {backbone} backbone estimator using {decode_method} on {dataset} subset')
        for result in results['handness']:
            print(f'\t- Handness @ {result["threshold"]:.2f}: Precision: {result["precision"]:.2f} Recall: {result["recall"]:.2f} F1: {result["f1"]:.2f}')
        print(f'\t- Keypoints: Single Hand: {scale*results["single_hand"]:.2f} Interacting: {scale*results["interacting"]:.2f}')
        print(f'\t- Area under ROC: {results["area_under_roc"]:.2f}')


if __name__ == '__main__':
    # evaluate_pretrained_interhand([0.9999], False)
    # evaluate_residuals()

    # evaluate_face_keypoints('efficientnet-b0')
    # evaluate_face_keypoints('efficientnet-b1')
    # evaluate_face_keypoints('efficientnet-b2')

    # evaluate_hand_keypoints('efficientnet-b0')
    # evaluate_hand_keypoints('efficientnet-b0', decode_method='compensated')
    # evaluate_hand_keypoints('efficientnet-b1')
    # evaluate_hand_keypoints('efficientnet-b1', decode_method='compensated')
    evaluate_hand_keypoints('efficientnet-b2')
    evaluate_hand_keypoints('efficientnet-b2', decode_method='compensated')
