"""
Visualize generated heatmaps of trained models
"""
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
import tensorflow as tf
import cv2
import numpy as np
import matplotlib.pyplot as plt

from modules.config import CFG
from modules import wflw_utils, face_model, hands_utils, hands_model


def wflw_heatmaps(backbone):
    """Visualize face keypoint estimation heatmaps

    Args:
        backbone (str): backbone (str): model backbone configuration, see model implementation for possible inputs
    """
    # load dataset
    test_data = wflw_utils.init_dataset(
        CFG.WFLW_TFRECORD_PATHS['test'],
        1,
        CFG.WFLW_PARAMS.IMG_SHAPE,
        CFG.WFLW_PARAMS.HM_DIMS,
        CFG.WFLW_PARAMS.SIGMA,
        testing=True
    )
    # load model
    model = face_model.load_model(CFG.WFLW_WEIGHTS[backbone], backbone, CFG.WFLW_PARAMS.IMG_SHAPE, CFG.WFLW_PARAMS.HM_DIMS, imagenet=False, decode=None)
    # load visualization color map
    cmap = plt.get_cmap('jet')
    # normalization function
    transorm = lambda x: cv2.cvtColor((np.delete(cmap(cv2.normalize(x, None, 0, 1, norm_type=cv2.NORM_MINMAX)), 3, 2) * 255).astype(np.uint8), cv2.COLOR_RGB2BGR)
    # padding function
    border = lambda x: cv2.copyMakeBorder(x, 1, 1, 1, 1, cv2.BORDER_CONSTANT, value=0)
    # viz grid number of columns
    step = 23
    # iterate over dataset
    for x, y_true in test_data:
        # forward pass
        heatmaps = model(x)
        # transform image to uint8 <0, 255>
        img = (x[0, ...].numpy() * 255).astype(np.uint8)
        # construct heatmap grid
        hms = tf.transpose(heatmaps[0, ...], (2, 0, 1)).numpy()
        hms_arr = np.concatenate([np.concatenate([border(transorm(h)) for h in hms[i:i+step, :, :]], axis=1) for i in range(0, 98-step+1, step)], axis=0)
        # visualize image and heatmaps
        cv2.imshow('img', img)
        cv2.imshow('heatmaps', hms_arr)
        if cv2.waitKey(0) == 27:
            break
    cv2.destroyAllWindows()


def hand_heatmaps(backbone):
    """Visualize hand keypoint estimation heatmaps

    Args:
        backbone (str): backbone (str): model backbone configuration, see model implementation for possible inputs
    """
    # load dataset
    test_data = hands_utils.init_dataset(
        CFG.HANDS_TFRECORD_PATHS['test'],
        1,
        CFG.HANDS_PARAMS.IMG_SHAPE,
        CFG.HANDS_PARAMS.HM_DIMS,
        CFG.HANDS_PARAMS.SIGMA,
        testing=True
    )
    # load model
    model = hands_model.load_model(CFG.HANDS_WEIGHTS[backbone], backbone, CFG.HANDS_PARAMS.IMG_SHAPE, CFG.HANDS_PARAMS.HM_DIMS, imagenet=False, decode=None)
    # load visualization color map
    cmap = plt.get_cmap('jet')
    # normalization function
    transorm = lambda x: cv2.cvtColor((np.delete(cmap(cv2.normalize(x, None, 0, 1, norm_type=cv2.NORM_MINMAX)), 3, 2) * 255).astype(np.uint8), cv2.COLOR_RGB2BGR)
    # padding function
    border = lambda x: cv2.copyMakeBorder(x, 1, 1, 1, 1, cv2.BORDER_CONSTANT, value=0)
    # viz grid number of columns
    step = 7
    # iterate over dataset
    for x, y_true in test_data:
        # forward pass
        handness, single_heatmaps, interacting_heatmaps = model(x)
        # transform image to uint8 <0, 255>
        img = (x[0, ...].numpy() * 255).astype(np.uint8)
        # construct heatmap grid for single and interacting hands separately
        hms = tf.transpose(single_heatmaps[0, ...], (2, 0, 1)).numpy()
        single_arr = np.concatenate([np.concatenate([border(transorm(h)) for h in hms[i:i+step, :, :]], axis=1) for i in range(0, 21-step+1, step)], axis=0)
        hms = tf.transpose(interacting_heatmaps[0, ...], (2, 0, 1)).numpy()
        interacting_arr = np.concatenate([np.concatenate([border(transorm(h)) for h in hms[i:i+step, :, :]], axis=1) for i in range(0, 42-step+1, step)], axis=0)
        # visualize image and heatmaps
        cv2.imshow('img', img)
        cv2.imshow('single heatmaps', single_arr)
        cv2.imshow('interacting heatmaps', interacting_arr)
        if cv2.waitKey(0) == 27:
            break
    cv2.destroyAllWindows()


if __name__ == '__main__':
    wflw_heatmaps('efficientnet-b1')
    hand_heatmaps('efficientnet-b1')
