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

from .model import SignRecognitionModel
from .dataset_utils import load_json, keypoints_to_input_space
from .config import CFG


class ExportModel:
    """Export model for inference
    """
    def __init__(self, model):
        """
        Args:
            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)


class SignLanguageClassifier:
    """Sign language classifier class
    """
    def __init__(self, n_classes, frames_per_sign):
        """
        Args:
            n_classes (int): number of predicted classes, one of (100, 300, 1000, 2000)
            frames_per_sign (int): model sequence length, number of frames inputted to the model
        """
        self.n_classes = n_classes
        self.frames_per_sign = frames_per_sign
        # load class map and inter it to ID: name
        class_map = load_json(CFG.WLASL_DATASET_PATHS['class_map'])
        self.class_map = {val: key for key, val in class_map.items()}
        # load model and its weights, build it and export for inference
        input_shape = (self.frames_per_sign, 1*98*2+2*21*2)
        self.model = SignRecognitionModel(input_shape, n_classes)
        self.model.build((1, *input_shape))
        self.model.load_weights(CFG.WLASL_WEIGHTS[n_classes]).expect_partial()
        self.model = ExportModel(self.model)
        self.model.f(np.random.rand(1, *input_shape))
        # list of sequences of keypoints, kept for inference
        self.hand_keypoints = list()
        self.face_keypoints = list()
        self.classes = list()

    def update(self, hand_keypoints, face_keypoints):
        """Update sequences with new predictions

        Args:
            hand_keypoints (np.ndarray): predicted hand keypoints
            face_keypoints (np.ndarray): predicted face keypoints
        """
        # ignore if no keypoints are passed
        if len(self.hand_keypoints) == 0 and hand_keypoints.size == 0:
            return
        # take only two hands and one face objects
        hand_keypoints = hand_keypoints.reshape(-1, 21, 2)[:2, ...]
        face_keypoints = face_keypoints.reshape(-1, 98, 2)[:1, ...]
        # pad with zeros to match 2 hands and 1 face
        hand_keypoints = np.pad(hand_keypoints, ((0, 2-hand_keypoints.shape[0]), (0, 0), (0, 0)))
        face_keypoints = np.pad(face_keypoints, ((0, 1-face_keypoints.shape[0]), (0, 0), (0, 0)))
        # add to list
        self.hand_keypoints.append(hand_keypoints)
        self.face_keypoints.append(face_keypoints)
        # remove frames older then sequence length
        self.hand_keypoints = self.hand_keypoints[-self.frames_per_sign:]
        self.face_keypoints = self.face_keypoints[-self.frames_per_sign:]
        # if all keypoints are zeros, reset list
        if not np.any(np.array(self.hand_keypoints)):
            self.hand_keypoints = list()
            self.face_keypoints = list()

    def preprocess(self):
        """Preprocess keypoints for inference

        Returns:
            np.ndarray: sequence of feature vectors for inference
        """
        hand_keypoints, face_keypoints = keypoints_to_input_space(self.hand_keypoints, self.face_keypoints)
        hand_keypoints = hand_keypoints.reshape(1, -1, 2*21*2)
        face_keypoints = face_keypoints.reshape(1, -1, 1*98*2)
        keypoints = np.concatenate((hand_keypoints, face_keypoints), axis=-1)
        return keypoints

    def postprocess(self, prediction, n_output=1):
        """Postprocess classification output

        Args:
            prediction (tf.Tensor): model prediction
            n_output (int, optional): number of top classification outputs. Defaults to 1.

        Returns:
            list, list: lists of predicted class names and their probabilities
        """
        self.classes.append(prediction[0, ...])
        # filter class predictions as a mean from previous 5 time steps
        self.classes = self.classes[-5:]
        classes = np.mean(self.classes, axis=0)
        # sort by probabilities and keep only n for output
        sorted_indices = np.argsort(classes)[-n_output:]
        # get class names and probabilities from indices
        probabilities = [classes[i] for i in reversed(sorted_indices)]
        classes = [self.class_map[i] for i in reversed(sorted_indices)]
        return classes, probabilities

    def classify(self, hand_keypoints, face_keypoints, n_output=1):
        """Classify from predicted hand and face keypoints in current time step

        Args:
            hand_keypoints (np.ndarray): hand keypoints predicted in current time step
            face_keypoints (np.ndarray): face keypoints predicted in current time step
            n_output (int, optional): number of classes to output. Defaults to 1.

        Returns:
            list, list: lists of predicted class names and their probabilities
        """
        # update state with keypoints
        self.update(hand_keypoints, face_keypoints)
        # output nones if sequence length is not reached
        if len(self.hand_keypoints) < self.frames_per_sign:
            return n_output * ['None'], n_output * [0.]
        # preprocess, predict and postprocess
        preprocessed = self.preprocess()
        out = self.model.f(preprocessed)
        classes, probabilities = self.postprocess(out, n_output=n_output)
        return classes, probabilities


def draw_classes(image, classes, probs):
    """Draw predicted classes onto image canvas

    Args:
        image (np.ndarray): BGR image
        classes (list): list of class names (strings)
        probs (list): list of predicted class probabilities
    """
    # starting point
    point = [10, 20]
    for cls, prob in zip(classes, probs):
        # draw text
        cv2.putText(image, f'{prob:.2f} {cls}', point, cv2.FONT_HERSHEY_DUPLEX, 1, (255, 0, 0), 2)
        # next line
        point[1] += 30
