"""
Utilities for model inference
"""
import tensorflow as tf
import cv2
import numpy as np

from .config import CFG
from .interhand import InterhandEstimator
from .face_model import load_model as load_face_model
from .hands_model import load_model as load_hands_model
from .estimation_utils import load_skeleton
from .interhand_utils import preprocess_image as interhand_preprocess_image
from .wflw_utils import preprocess_image as faces_preprocess_image
from .hands_utils import preprocess_image as hands_preprocess_image


class ExportModel:
    """Export model for inference
    """
    def __init__(self, model):
        """s:
            model (tf.keras.Model): model to export
        """
        self.model = model
    
    @tf.function(experimental_relax_shapes=True)
    def f(self, x):
        """Inference call

        Args:
            x (tf.Tensor): input tensor

        Returns:
            any: model outputs
        """
        return self.model(x)


def draw_pose(img, joints, handness, parents=None):
    """Draw hand pose onto canvas image

    Args:
        img (np.ndarray): BGR image
        joints (np.ndarray): array of keypoints of N subjects [N, 42, 2]
        handness (list): list of tuples of two hand presence indicators, [[left hand, right hand], ...]
        parents (list, optional): parent list defining hand skeleton. Defaults to None.

    Returns:
        np.ndarray: output image
    """
    # repeat validity to match the number of keypoints of each hand
    validity = np.repeat(handness, 21, axis=1)
    # for every subject
    for j, (jnts, val) in enumerate(zip(joints, validity)):
        # for every keypoint
        for i, pnt in enumerate(jnts):
            if not val[i]:
                continue
            color = (202, 183, 42) if j == 0 else (102, 215, 254) # (202, 183, 42) if i < 21 else (102, 215, 254)
            # draw lines to illustrate skeleton
            if parents is not None and parents[i] != -1 and val[parents[i]]:
                img = cv2.line(img, tuple(jnts[parents[i], :2].astype(np.int32)), tuple(pnt[:2].astype(np.int32)), color, 1)
            # draw keypoint
            img = cv2.circle(img, tuple(pnt[:2].astype(np.int32)), 3, color, -1)
    return img


def draw_keypoints(img, keypoints):
    """Draw skeleton-less keypoints onto canvas image

    Args:
        img (np.ndarray): BGR image
        keypoints (np.ndarray): array of M keypoints of N subjects [N, M, 2]

    Returns:
        np.ndarray: output image
    """
    # for every subject
    for kpnts in keypoints:
        # for every keypoint
        for pnt in kpnts:
            img = cv2.circle(img, tuple(pnt.astype(np.int32)), 1, (209, 96, 61), -1)
    return img


def preprocess_interhand(image, bboxes, center_format=False):
    """Preprocess image for inference on Interhand model

    Args:
        image (np.ndarray): input BGR image
        bboxes (np.ndarray): array of input boxes defining ROIs in input image to be passed to the model [N, 4]
        center_format (bool, optional): flag if passed boxes are in center format [x_center, y_center, width, height],
                                        if False - boxes are expected in format [x_min, y_min, width, height]. Defaults to False.

    Returns:
        tf.Tensor, list: tensor of input images, list of inverse transformations - one for every input bounding box
    """
    imgs = list()
    transforms = list()
    # for every input bounding box
    for bbox in bboxes:
        # preprocess image and get inverse transform
        img, _, inv_trans = interhand_preprocess_image(image, bbox, CFG.INTERHAND_PARAMS.IMG_SHAPE, center_format=center_format)
        # convert to tf.Tensor and scale to <0, 1>
        img = tf.expand_dims(img.astype(np.float32), axis=0) / 255
        imgs.append(img)
        transforms.append(inv_trans)
    # if no boxes are passed
    if len(imgs) == 0:
        return None, None
    imgs = tf.concat(imgs, axis=0)
    return imgs, transforms


def preprocess_hands(image, bboxes, center_format=False):
    """Preprocess image for inference on hand pose estimation model

    Args:
        image (np.ndarray): input BGR image
        bboxes (np.ndarray): array of input boxes defining ROIs in input image to be passed to the model [N, 4]
        center_format (bool, optional): flag if passed boxes are in center format [x_center, y_center, width, height],
                                        if False - boxes are expected in format [x_min, y_min, width, height]. Defaults to False.

    Returns:
        tf.Tensor, list: tensor of input images, list of inverse transformations - one for every input bounding box
    """
    imgs = list()
    transforms = list()
    # for every input bounding box
    for bbox in bboxes:
        # preprocess image and get inverse transform
        img, _, inv_trans = hands_preprocess_image(image, bbox, CFG.HANDS_PARAMS.IMG_SHAPE, center_format=center_format, scale=1.5)
        # convert to tf.Tensor and scale to <0, 1>
        img = tf.expand_dims(img.astype(np.float32), axis=0) / 255
        imgs.append(img)
        transforms.append(inv_trans)
    # if no boxes are passed
    if len(imgs) == 0:
        return None, None
    imgs = tf.concat(imgs, axis=0)
    return imgs, transforms


def preprocess_faces(image, bboxes, center_format=False):
    """Preprocess image for inference on face pose estimation model

    Args:
        image (np.ndarray): input BGR image
        bboxes (np.ndarray): array of input boxes defining ROIs in input image to be passed to the model [N, 4]
        center_format (bool, optional): flag if passed boxes are in center format [x_center, y_center, width, height],
                                        if False - boxes are expected in format [x_min, y_min, width, height]. Defaults to False.

    Returns:
        tf.Tensor, list: tensor of input images, list of inverse transformations - one for every input bounding box
    """
    imgs = list()
    transforms = list()
    # for every input bounding box
    for bbox in bboxes:
        # preprocess image and get inverse transform
        img, _, inv_trans = faces_preprocess_image(image, bbox, CFG.INTERHAND_PARAMS.IMG_SHAPE, center_format, scale=1.5)
        # convert to tf.Tensor and scale to <0, 1>
        img = tf.expand_dims(img.astype(np.float32), axis=0) / 255
        imgs.append(img)
        transforms.append(inv_trans)
    # if no boxes are passed
    if len(imgs) == 0:
        return None, None
    imgs = tf.concat(imgs, axis=0)
    return imgs, transforms


def postprocess_interhand(predictions, transforms=None):
    """Postprocess interhand pose estimation outputs

    Args:
        predictions (list): list of predictions [handedness, joints, relative_depth]
        transforms (list, optional): list of inverse transformations for 2D keypoints to project them back
                                     to original image plane, if None - don't project. Defaults to None.

    Returns:
        np.ndarray, np.ndarray, float: handedness, joints, relative depth of hands
    """
    # convert to numpy
    handness = predictions[0].numpy()
    joints = predictions[1].numpy()
    relative_depth = predictions[2].numpy()
    # scale keypoints from <0, 1> to image plane
    joints[:, :, :2] = joints[:, :, :2] * CFG.INTERHAND_PARAMS.IMG_SHAPE[1::-1]
    # if transformations are passed
    if transforms is not None:
        # transform predicted keypoints to original image space
        for jnts, tform in zip(joints, transforms):
            jnts[:, :2] = np.dot(tform, np.concatenate((jnts[:, :2], np.ones((jnts.shape[0], 1))), axis=1).T).T
    return handness, joints, relative_depth


def postprocess_hands(predictions, transforms=None):
    """Postprocess hand pose estimation outputs

    Args:
        predictions (list): list of predictions [handedness, keypoints]
        transforms (list, optional): list of inverse transformations for 2D keypoints to project them back
                                     to original image plane, if None - don't project. Defaults to None.

    Returns:
        np.ndarray, np.ndarray: handedness - presence indicator of left and right hands, keypoints - transformed if transforms are passed
    """
    # convert to numpy and scale keypoints from <0, 1> to image plane
    handness = predictions[0].numpy() > CFG.HANDS_PARAMS.HANDNESS_THOLD
    keypoints_single_hand = predictions[1].numpy() * CFG.HANDS_PARAMS.IMG_SHAPE[1::-1]
    keypoints_interacting = predictions[2].numpy() * CFG.HANDS_PARAMS.IMG_SHAPE[1::-1]
    # filter keypoints based on classification results
    keypoints = [inter if hnd else single for hnd, single, inter in zip(handness, keypoints_single_hand, keypoints_interacting)]
    # if transformations are passed
    if transforms is not None:
        tformed_keypoints = list()
        # transform predicted keypoints to original image space
        for kpts, tform in zip(keypoints, transforms):
            tformed_keypoints.append(np.dot(tform, np.concatenate((kpts, np.ones((kpts.shape[0], 1))), axis=1).T).T)
        keypoints = tformed_keypoints
    # convert interacting hand indicator to right and left hand presence indicators
    handness = np.concatenate((np.ones_like(handness, dtype=np.bool), handness), axis=1)
    return handness, keypoints


def postprocess_faces(keypoints, transforms=None):
    """Postprocess hand pose estimation outputs

    Args:
        keypoints (tf.Tensor): predicted keypoints
        transforms (list, optional): list of inverse transformations for 2D keypoints to project them back
                                     to original image plane, if None - don't project. Defaults to None.

    Returns:
        np.ndarray: keypoints - transformed if transforms are passed
    """
    # convert to numpy and scale keypoints from <0, 1> to image plane
    keypoints = keypoints.numpy() * CFG.WFLW_PARAMS.IMG_SHAPE[1::-1]
    # if transformations are passed
    if transforms is not None:
        # transform predicted keypoints to original image space
        for kpts, tform in zip(keypoints, transforms):
            kpts[...] = np.dot(tform, np.concatenate((kpts, np.ones((kpts.shape[0], 1))), axis=1).T).T
    return keypoints


def init_interhand_estimator():
    """Initialize Interhand hand pose estimator for inference

    Returns:
        ExportModel, list: exported model for inference, list of parent keypoints representing hand skeleton
    """
    # load model, build it and load 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.build((1, *CFG.INTERHAND_PARAMS.IMG_SHAPE, 3))
    model.load_weights(CFG.INTERHAND_TF_PATH).expect_partial()
    # export for inference, make first forward pass on random inputs
    model = ExportModel(model)
    model.f(np.random.rand(1, *CFG.INTERHAND_PARAMS.IMG_SHAPE, 3))
    # load hand skeleton
    skeleton = load_skeleton(CFG.INTERHAND_SKELETON_PATH)
    return model, skeleton


def init_hand_estimator(backbone):
    """Initialize hand pose estimator for inference

    Args:
        backbone (str): backbone feature extractor, check model's implementation to see expected values

    Returns:
        ExportModel, list: exported model for inference, list of parent keypoints representing hand skeleton
    """
    # load model, build it and load weights
    model = load_hands_model(CFG.HANDS_WEIGHTS[backbone], backbone, CFG.HANDS_PARAMS.IMG_SHAPE, CFG.HANDS_PARAMS.HM_DIMS,
                             imagenet=False, decode='argmax', delta=CFG.HANDS_PARAMS.DELTA)
    model.build((1, *CFG.HANDS_PARAMS.IMG_SHAPE, 3))
    # export for inference, make first forward pass on random inputs
    model = ExportModel(model)
    model.f(np.random.rand(1, *CFG.HANDS_PARAMS.IMG_SHAPE, 3))
    # load hand skeleton
    skeleton = load_skeleton(CFG.COCO_SKELETON_PATH)
    return model, skeleton


def init_face_estimator(backbone):
    """Initialize face pose estimator for inference

    Args:
        backbone (str): backbone feature extractor, check model's implementation to see expected values

    Returns:
        ExportModel: exported model for inference
    """
    # load model, build it and load weights
    model = load_face_model(CFG.WFLW_WEIGHTS[backbone], backbone, CFG.WFLW_PARAMS.IMG_SHAPE, CFG.WFLW_PARAMS.HM_DIMS, decode='argmax')
    model.build((1, *CFG.WFLW_PARAMS.IMG_SHAPE, 3))
    model.load_weights(CFG.WFLW_WEIGHTS[backbone]).expect_partial()
    # export for inference, make first forward pass on random inputs
    model = ExportModel(model)
    model.f(np.random.rand(1, *CFG.WFLW_PARAMS.IMG_SHAPE, 3))
    return model


def estimate_interhand(model, image, bboxes, center_format=False):
    """Infer Interhand estimator, runs preprocessing, inference and postprocessing steps

    Args:
        model (ExportModel): Interhand model exported for inference
        image (np.ndarray): input image
        bboxes (np.ndarray): array of bounding boxes of hands
        center_format (bool, optional): flag if passed boxes are in center format [x_center, y_center, width, height],
                                        if False - boxes are expected in format [x_min, y_min, width, height]. Defaults to False.

    Returns:
        np.ndarray, np.ndarray, float: handedness, joints, relative depth of hands
    """
    imgs, transforms = preprocess_interhand(image, bboxes, center_format=center_format)
    # return empty arrays if no boxes are passed
    if imgs is None:
        return np.zeros((0, 2)), np.zeros((0, 0, 3)), np.zeros((0, ))
    output = model.f(imgs)
    handness, joints, relative_depth = postprocess_interhand(output, transforms)
    return handness, joints, relative_depth


def estimate_hands(model, image, bboxes, center_format=False):
    """Infer hand pose estimator, runs preprocessing, inference and postprocessing steps

    Args:
        model (ExportModel): hand pose estimation model exported for inference
        image (np.ndarray): input image
        bboxes (np.ndarray): array of bounding boxes of hands
        center_format (bool, optional): flag if passed boxes are in center format [x_center, y_center, width, height],
                                        if False - boxes are expected in format [x_min, y_min, width, height]. Defaults to False.

    Returns:
        np.ndarray, np.ndarray: handedness - presence indicator of left and right hands, keypoints
    """
    imgs, transforms = preprocess_hands(image, bboxes, center_format)
    # return empty arrays if no boxes are passed
    if imgs is None:
        return np.zeros((0, 2)), np.zeros((0, 0, 2))
    output = model.f(imgs)
    handness, keypoints = postprocess_hands(output, transforms)
    return handness, keypoints


def estimate_face(model, image, bboxes, center_format=False):
    """Infer face pose estimator, runs preprocessing, inference and postprocessing steps

    Args:
        model (ExportModel): face pose estimation model exported for inference
        image (np.ndarray): input image
        bboxes (np.ndarray): array of bounding boxes of faces
        center_format (bool, optional): flag if passed boxes are in center format [x_center, y_center, width, height],
                                        if False - boxes are expected in format [x_min, y_min, width, height]. Defaults to False.

    Returns:
        np.ndarray: keypoints
    """
    imgs, transforms = preprocess_faces(image, bboxes, center_format)
    # return empty array if no boxes are passed
    if imgs is None:
        return np.zeros((0, 0, 2))
    output = model.f(imgs)
    keypoints = postprocess_faces(output, transforms)
    return keypoints
