import sys
import gym
import yaml
import math
import random
import logging
import numpy as np
from os import path
import pandas as pd

import Levenshtein as lev

from src.abstract.environment import AgentEnv
from src.vision.vision_agent import VisionAgent
from src.finger.finger_agent import FingerAgent
from src.proofread.proofread_agent import ProofreadAgent
from src.display.touchscreendevice import TouchScreenDevice
from src.utilities.utils import distance

class SupervisorEnvironment(AgentEnv):

    def __init__(self, layout_config, agent_params, train):

        self.logger = logging.getLogger(__name__)
        self.config_file = None
        if path.exists(path.join('configs', layout_config)):
            with open(path.join("configs", layout_config), 'r') as file:
                self.config_file = yaml.load(file, Loader=yaml.FullLoader)
                self.logger.info("Device Configurations loaded.")
        else:
            self.logger.error("File doesn't exist: Failed to load %s file under configs folder." % layout_config)
            sys.exit(0)

        if self.config_file:
            self.device = TouchScreenDevice(self.config_file['layout_file'], self.config_file['config'])
            self.user_distance = self.device.device_params['user_distance']
        else:
            self.device = None

        self.vision_agent = VisionAgent(layout_config, agent_params['vision'])
        self.finger_agent = FingerAgent(layout_config, agent_params['finger'], 0, False)
        self.proofread_agent = ProofreadAgent(layout_config, agent_params['proofread'])

        self.actions = agent_params['supervisor']['action_type']
        self.sat_desired = agent_params['supervisor']['sat_desired']
        self.found_reward = agent_params['supervisor']['reward']
        self.agent_id = 0
        self.eye_loc = None
        self.prev_eye_loc = None
        self.finger_loc = None
        self.belief = None
        self.finger_q = 0.0
        self.proofread_q = 0.0
        self.is_error = None
        self.sentence_to_type = None
        self.key_found = False
        self.typed = None
        self.eye_on_keyboard = None
        self.is_terminal = False
        self.n_chars = 0
        self.line = ''
        self.typed_detailed = ""
        self.eye_model_time = 0
        self.finger_model_time = 0
        self.eye_on_kb_time = 0
        self.action_space = gym.spaces.discrete.Discrete(len(self.actions))
        self.mt = 0
        self.model_time = 0
        self.steps = 0
        self.n_fixations_freq = 0
        self.gaze_shift_kb_to_txt_freq = 0
        self.gaze_shift_txt_to_kb_freq = 0
        self.n_back_space_freq = 0
        self.immediate_backspace_freq = 0
        self.delay_backspace_freq = 0
        self.finger_travel_dist = 0
        self.saccade_time = 0
        self.encoding_time = 0
        self.chunk = 0
        self.eye_test_data = []
        self.finger_test_data = []
        self.sentence_test_data = []
        self.fixation_duration = []
        self.eye_viz_log = []
        self.finger_viz_log = []
        self.typing_viz_log = []
        self.chunk_length = []
        self.train = train
        self.belief_state = None
        if not self.train:
            self.eye_test_data.append(["model time", "eyeloc x", "eyeloc y", "action x", "action y", "type"])
            self.finger_test_data.append(["model time", "fingerloc x", "fingerloc y", "action x", "action y", "type"])
            self.eye_viz_log.append(['subject_id', 'sentence_n', 'trialtime', 'x', 'y'])
            self.finger_viz_log.append(['subject_id', 'block', 'sentence_n', 'trialtime', 'x', 'y'])
            self.typing_viz_log.append(['subject_id', 'sentence_n', 'trialtime', 'event'])

            data = pd.read_csv(path.join("data", "user_data", agent_params['supervisor']['corpus']), sep=",")
            # assuming the 1st column will always be sentence id
            self.sentences_id = data.iloc[:, 0].values
            self.sentences = list(data.iloc[:, 1].values)
            self.sentences_bkp = self.sentences.copy()
        else:
            data = pd.read_csv(path.join("data", "user_data", "train_corpus.csv"), sep=",")
            self.sentences = list(data.iloc[:, 0].values)

        self.vision_agent.agent.load()
        self.finger_agent.load()
        self.proofread_agent.agent.load()

    def update_model_time(self, delta):
        """
        Function to update model runtime.
        :param delta: time to increment in ms.
        """
        self.model_time += delta

    def make_eye_movement(self, char):
        """
        Function to perform eye movement from current position to target char.
        :param char: target character to move eyes to.
        :return movement time in seconds.
        """

        self.prev_eye_loc = self.eye_loc
        # movement values are in seconds. But, for visualisation we use ms. Converting everything here.
        (mt_enc, mt_exec, mt_enc_l), self.mt, self.eye_loc, _, moved = self.vision_agent.type_char(char, self.eye_loc)

        if moved:
            self.n_fixations_freq += 1

        self.fixation_duration.append(self.eye_model_time - self.saccade_time)

        self.eye_model_time += (self.mt * 1000)

        self.saccade_time = round(self.eye_model_time - (mt_enc_l * 1000), 4)

        if not self.train:
            self.eye_test_data.append(
                [round(self.eye_model_time - (mt_enc_l * 1000) - (mt_exec * 1000) + 50, 4), self.prev_eye_loc[0],
                 self.prev_eye_loc[1], "", "", 'encoding'])
            self.eye_test_data.append(
                [round(self.eye_model_time - (mt_enc_l * 1000), 4), self.prev_eye_loc[0], self.prev_eye_loc[1],
                 self.eye_loc[0], self.eye_loc[1], 'saccade'])
            if mt_enc_l > 0:
                self.eye_test_data.append(
                    [round(self.eye_model_time, 4), self.eye_loc[0], self.eye_loc[1], "", "", 'late encoding'])

            self.eye_viz_log.append([self.agent_id,
                                     self.sentences_id[self.sentences_bkp.index(self.sentence_to_type[:-1])],
                                     round(self.eye_model_time, 4), self.eye_loc[0], self.eye_loc[1]])

        return self.mt

    def make_finger_movement(self, char, sigma_desired, eye_time):
        """
         Function to perform finger movement from current position to target char.
        """
        finger_time = 0
        self.mt, _, _, self.finger_q = self.finger_agent.type_char(char, sigma_desired, self.eye_on_keyboard)
        self.finger_loc = self.finger_agent.env.finger_location
        self.finger_travel_dist += self.finger_agent.env.dist
        finger_time += self.mt

        # Keep iterating the finger model until a peck is performed and appending the finger data
        while self.finger_agent.env.action_type == self.finger_agent.env.actions[0]:
            self.finger_model_time += (self.mt * 1000)
            if not self.train:
                self.finger_test_data.append(
                    [round(self.finger_model_time, 4), self.finger_loc[0], self.finger_loc[1], "", "", "move"])
                self.finger_viz_log.append([self.agent_id,
                                            self.sentences_id[self.sentences_bkp.index(self.sentence_to_type[:-1])],
                                            2, round(self.finger_model_time, 4), self.finger_loc[0],
                                            self.finger_loc[1]])

            self.mt, _, _, self.finger_q = self.finger_agent.type_char(char, sigma_desired, self.eye_on_keyboard)
            self.finger_loc = self.finger_agent.env.finger_location
            self.finger_travel_dist += self.finger_agent.env.dist
            finger_time += self.mt

        self.finger_model_time += (self.mt * 1000)
        if not self.train:
            self.finger_test_data.append(
                [round(self.finger_model_time, 4), self.finger_loc[0], self.finger_loc[1], self.finger_loc[0],
                 self.finger_loc[1], "peck"])
            self.finger_viz_log.append([self.agent_id,
                                        self.sentences_id[self.sentences_bkp.index(self.sentence_to_type[:-1])],
                                        2, round(self.finger_model_time, 4), self.finger_loc[0],
                                        self.finger_loc[1]])
            self.typing_viz_log.append([self.agent_id,
                                        self.sentences_id[self.sentences_bkp.index(self.sentence_to_type[:-1])],
                                        round(self.finger_model_time, 4), 'pressed key'])

            # finger waits for eye and eye waits for finger. Next action taken when both have reacted for a target.
            if self.finger_model_time > self.eye_model_time:
                self.eye_on_keyboard = True if not tuple(self.eye_loc) in self.proofread_agent.env.proof_locs else False
                if self.eye_on_keyboard:
                    self.eye_on_kb_time += (self.finger_model_time - self.eye_model_time)
                self.eye_model_time = self.finger_model_time
                self.eye_test_data.append(
                    [round(self.eye_model_time, 4), self.eye_loc[0], self.eye_loc[1], "", "", 'wait'])
            elif self.finger_model_time < self.eye_model_time:
                self.finger_model_time = self.eye_model_time
                self.finger_test_data.append(
                    [round(self.finger_model_time, 4), self.finger_loc[0], self.finger_loc[1], "", "", "wait"])

        # See what is typed and update
        letter = self.device.get_character(self.finger_loc[0], self.finger_loc[1])

        if letter == '<':
            self.typed = self.typed[:-1]
            self.typed_detailed += '<'
        elif letter != '>':
            self.typed += letter
            self.chunk += 1
            self.typed_detailed += letter
        else:
            # pressed enter key. this is terminal state.
            self.typed += letter
            self.typed_detailed += letter
            self.is_terminal = True

    def step(self, action):
        """
        Perform a single step in the environment.
        Args:
            action: int value representing, 0: type or 1: proof read.

        Returns:
            state: current finger and proofread agent q values.
            reward: scalar reward value for taking action.
            terminal: if episode is done or not.
            info: dictionary of episode info.
        """
        self.steps += 1
        if self.sentence_to_type.startswith(self.typed):
            # Updates what is left to type
            self.line = self.sentence_to_type.replace(self.typed, '', 1)

        char = self.line[0]

        if action < len(self.sat_desired):
            # Typing action was selected.
            sigma_desired = self.sat_desired[action]
            self.logger.debug("Typing character %s, with desired SAT: %.2f" % (char, sigma_desired))

            # EYE MOVEMENT
            eye_time = self.make_eye_movement(char)

            # FINGER MOVEMENT
            self.eye_on_keyboard = True if not tuple(self.eye_loc) in self.proofread_agent.env.proof_locs else False
            if self.eye_on_keyboard:
                self.eye_on_kb_time += (eye_time * 1000)
            self.make_finger_movement(char, sigma_desired, eye_time)

            if not self.is_error:
                self.is_error = not self.finger_agent.env.hit

                # code block hard-coded for now. This is for intelligent text entry.
                # TODO: clean this up and move to device.
                char_coord = self.device.get_coordinate(char)
                x = random.choice(char_coord[0])
                y = random.choice(char_coord[1])
                dist = distance([x, y], self.finger_loc)

                # choose the correct correction rate. default 0.0 which means don't correct.
                # Also, only correct chars that are 1 key away from target.
                if random.random() < 0.0 and dist <= 1.5:
                    self.typed = self.typed[:-1]
                    self.typed += char
                    self.is_error = False
                    if char == '>':
                        self.is_terminal = True

            # update proofread
            self.update_proofread(self.proofread_agent.env.error_probability_list[self.finger_agent.env.sat_true])

        else:
            # EYE MOVEMENT FOR PROOFREADING
            sigma_desired = self.sat_desired[(action - len(self.sat_desired))]

            self.logger.debug("Proofreading, with desired SAT: %.2f" % sigma_desired)

            self.chunk_length.append(self.chunk)
            self.chunk = 0

            # Check if eyes are on keyboard
            self.eye_on_keyboard = True if not tuple(self.eye_loc) in self.proofread_agent.env.proof_locs else False

            self.prev_eye_loc = self.eye_loc
            (mt_enc, mt_exec, mt_enc_l), self.mt, self.eye_loc = self.proofread_agent.proofread_text(self.eye_loc)

            if tuple(self.eye_loc) in self.proofread_agent.env.proof_locs and self.eye_on_keyboard:
                self.gaze_shift_kb_to_txt_freq += 1
            self.n_fixations_freq += 1

            self.eye_model_time += (self.mt * 1000)
            if not self.train and self.finger_model_time < self.eye_model_time:
                self.finger_model_time = self.eye_model_time
                self.finger_test_data.append(
                    [round(self.finger_model_time, 4), self.finger_loc[0], self.finger_loc[1], "", "", "wait"])

            if not self.train:
                self.eye_test_data.append(
                    [round(self.eye_model_time - mt_enc_l - mt_exec + 50, 4), self.prev_eye_loc[0],
                     self.prev_eye_loc[1], "", "", 'encoding'])
                self.eye_test_data.append(
                    [round(self.eye_model_time - mt_enc_l, 4), self.prev_eye_loc[0], self.prev_eye_loc[1],
                     self.eye_loc[0], self.eye_loc[1], 'saccade'])
                if mt_enc_l > 0:
                    self.eye_test_data.append(
                        [round(self.eye_model_time, 4), self.eye_loc[0], self.eye_loc[1], "", "",
                         'late encoding'])
                self.eye_viz_log.append([self.agent_id,
                                         self.sentences_id[self.sentences_bkp.index(self.sentence_to_type[:-1])],
                                         round(self.eye_model_time, 4), self.eye_loc[0], self.eye_loc[1]])

            # Once proofread. Error prob reset to lowest value.
            self.proofread_agent.env.reset_error_prob()
            self.proofread_agent.env.set_belief()
            self.proofread_q = self.proofread_agent.get_q_value()

            # If errors found on proofreading.
            if self.is_error:
                if len(self.typed) > len(self.sentence_to_type):
                    indexes = [i for i in range(len(self.sentence_to_type)) if
                               self.typed[i] != self.sentence_to_type[i]]
                else:
                    indexes = [i for i in range(len(self.typed)) if self.typed[i] != self.sentence_to_type[i]]
                if len(indexes) > 0:
                    err_idx = indexes[0]
                else:
                    err_idx = len(self.typed)

                # Calculating required back spaces to remove errors
                n_bs = len(self.typed) - err_idx

                self.n_back_space_freq += n_bs

                if n_bs == 1:
                    self.immediate_backspace_freq += 1
                if n_bs > 1:
                    self.delay_backspace_freq += 1

                # For number of required backspaces
                for bs in range(n_bs):
                    char = '<'
                    # EYE MOVEMENT
                    eye_time = self.make_eye_movement(char)
                    if self.eye_on_keyboard:
                        self.eye_on_kb_time += (eye_time * 1000)

                    # FINGER MOVEMENT
                    self.eye_on_keyboard = True if not tuple(
                        self.eye_loc) in self.proofread_agent.env.proof_locs else False
                    self.make_finger_movement(char, sigma_desired, eye_time)

                # Errors corrected after proofreading
                self.is_error = False

        total_mt = (self.eye_model_time + self.finger_model_time) / 1000.0
        reward = self.reward(action, total_mt)

        self.set_belief()

        if self.is_terminal and not self.train:
            # log sentence level data.
            self.logger.debug("typed: %s" % self.typed)
            index = self.sentences_bkp.index(self.sentence_to_type[:-1])
            corrected = 0
            uncorrected = 0
            if self.typed == self.sentence_to_type and self.n_back_space_freq > 0:
                corrected = 1

            if not (self.typed == self.sentence_to_type):
                uncorrected = 1

            wpm = ((len(self.typed) - 1) / 5.0) / (self.finger_model_time / 60000.0)
            err_lev_dist = lev.distance(self.sentence_to_type, self.typed)
            proportion_gaze_kb = self.eye_on_kb_time / self.eye_model_time
            iki = self.finger_model_time / len(self.typed_detailed)
            c_ln = np.mean(self.chunk_length) if len(self.chunk_length) > 0 else 0.0
            line = [self.sentences_id[index], self.agent_id, self.sentence_to_type, wpm, err_lev_dist,
                    self.gaze_shift_kb_to_txt_freq, self.n_back_space_freq, self.immediate_backspace_freq,
                    self.delay_backspace_freq, proportion_gaze_kb, self.n_fixations_freq,
                    self.finger_travel_dist, iki, corrected, uncorrected, np.mean(self.fixation_duration), c_ln]
            self.sentence_test_data.append(line)

        return self.belief_state, reward, self.is_terminal, {}

    def reset(self):
        """
        Function to be called on start of a trial. It resets the environment
        and sets the initial belief state.
        :return: current belief state.
        """
        self.logger.debug("Resetting Environment for start of new trial.")
        self.steps = 0
        self.eye_loc = self.device.start()
        self.logger.debug("Eye initialised to location: {%d, %d}" % (self.eye_loc[0], self.eye_loc[1]))
        self.finger_loc = self.device.start()
        self.logger.debug("Finger initialised to location: {%d, %d}" % (self.finger_loc[0], self.finger_loc[1]))
        self.vision_agent.env.reset()
        self.finger_agent.env.reset()
        self.proofread_agent.env.reset()
        self.finger_agent.env.finger_location = self.finger_loc
        self.vision_agent.env.eye_location = self.eye_loc
        self.proofread_agent.env.eye_location = self.eye_loc
        self.mt = 0
        self.key_found = False
        self.is_terminal = False
        self.eye_on_keyboard = True
        self.typed = ""
        self.typed_detailed = ""
        self.finger_q = 0.0
        self.eye_model_time = 0
        self.finger_model_time = 0
        self.eye_on_kb_time = 0
        self.gaze_shift_kb_to_txt_freq = 0
        self.gaze_shift_txt_to_kb_freq = 0
        self.n_back_space_freq = 0
        self.immediate_backspace_freq = 0
        self.delay_backspace_freq = 0
        self.n_fixations_freq = 0
        self.finger_travel_dist = 0
        self.saccade_time = 0
        self.encoding_time = 0
        self.chunk = 0
        self.fixation_duration.clear()
        self.chunk_length.clear()

        self.proofread_agent.env.set_belief()
        self.proofread_q = self.proofread_agent.get_q_value()
        if self.train:
            self.sentence_to_type = random.choice(self.sentences)
        else:
            self.sentence_to_type = self.sentences.pop(0)

        self.sentence_to_type += '>'
        self.logger.debug('typing: %s' % self.sentence_to_type)
        self.line = self.sentence_to_type
        self.n_chars = len(self.sentence_to_type)
        if not self.train:
            # Initial data appended for both eye and finger
            self.eye_test_data.append(
                [round(self.eye_model_time, 4), self.eye_loc[0], self.eye_loc[1], "", "", "start"])

            self.finger_test_data.append(
                [round(self.finger_model_time, 4), self.finger_loc[0], self.finger_loc[1], "", "", "start"])

        self.set_belief()
        return self.belief_state

    def reward(self, action, movement_time):
        """
        Function for calculating R(a) = total character in sentence - movement time - levenshtein distance.
        :param action: tuple for finger movement action taken by agent.
        :param movement_time: movement time in seconds for taking action.
        :return: reward: float value to denote goodness of action and action type
        """

        if self.steps == 100 and not self.is_terminal:
            # reached max length.
            self.is_terminal = True
            return np.clip(-movement_time - lev.distance(self.sentence_to_type, self.typed_detailed), -20, 20)
        elif not self.is_terminal:
            return 0.0
        else:
            return np.clip(self.n_chars - movement_time - lev.distance(self.sentence_to_type, self.typed_detailed),
                           -20, 20)

    def render(self, mode='human'):
        pass

    def set_belief(self):
        """
        Function to update belief state.
        """
        self.belief_state = np.asarray([self.finger_q, self.proofread_q]).astype(np.float32)
        self.logger.debug("current belief state is {%s}" % str(self.belief_state))

    def update_proofread(self, error_chance):
        """
        Function to update proof read belief.
        """
        obs_prob = self.proofread_agent.env.observation_probability

        if self.is_error:
            obs_error = obs_prob
        else:
            obs_error = 1 - obs_prob

        self.proofread_agent.env.error_prob = self.proofread_agent.env.update_error_belief(obs_error,
                                                                                self.proofread_agent.env.error_prob,
                                                                                error_chance)

        self.logger.debug("Updated error probability to %.2f" % self.proofread_agent.env.error_prob)
        self.proofread_agent.env.set_belief()
        self.proofread_q = self.proofread_agent.get_q_value()


# Custom transformer models for Joint NLU.
# last edited: 10.2.2021
# SP
import transformers
from transformers import (
    AutoConfig,
    AutoModel,
    AutoTokenizer,
    RobertaModel,
    Trainer,
    TrainingArguments,
    XLMRobertaModel,
)

from transformers.models.roberta.modeling_roberta import (
    RobertaPreTrainedModel,
    RobertaClassificationHead,
)
from transformers.modeling_outputs import (
    SequenceClassifierOutput,
    TokenClassifierOutput,
)
from transformers import PreTrainedModel
import math
import torch
import torch.nn as nn
from torch.nn import CrossEntropyLoss, MSELoss
import numpy as np
from preprocessing.util import *

class IntentClassifier(nn.Module):
    "classifier head for intent detection"

    def __init__(self, input_dim, num_intent_labels, dropout_rate=0.01):
        super(IntentClassifier, self).__init__()
        self.dropout = nn.Dropout(dropout_rate)
        self.dense = nn.Linear(input_dim, input_dim * 2)
        self.activation = nn.Tanh()
        self.linear = nn.Linear(input_dim * 2, num_intent_labels)

    def forward(self, x):
        x = self.dense(x)
        x = self.activation(x)
        x = self.dropout(x)
        return self.linear(x)

class SimpleIntentClassifier(nn.Module):
    "vanilla dense layer for ablation study"
    def __init__(self, input_dim, num_intent_labels, dropout_rate=0.01):
        super(SimpleIntentClassifier, self).__init__()
        self.dropout = nn.Dropout(dropout_rate)
        self.dense = nn.Linear(input_dim, input_dim * 2)
        self.activation = nn.Tanh()
        self.linear = nn.Linear(input_dim * 2, num_intent_labels)

    def forward(self, x):
        x = self.dense(x)
        x = self.activation(x)
        x = self.dropout(x)
        return self.linear(x)

class CustomClassificationHead(nn.Module):
    "deprecated XLM-R intent classification head"

    def __init__(self, num_labels, dropout=0.01, hidden_size=768):
        super(CustomClassificationHead, self).__init__()

        self.dense = nn.Linear(hidden_size, hidden_size)
        self.dropout = nn.Dropout(dropout)
        self.out_proj = nn.Linear(hidden_size, num_labels)

    def forward(self, features, **kwargs):

        x = features[:, 0, :]
        x = self.dropout(x)
        x = self.dense(x)
        x = torch.tanh(x)
        x = self.dropout(x)
        x = self.out_proj(x)
        return x

class SlotClassifier(nn.Module):
    "token classification layer"

    def __init__(self, input_dim, num_slot_labels, dropout_rate=0.01):
        super(SlotClassifier, self).__init__()
        self.dropout = nn.Dropout(dropout_rate)
        self.linear = nn.Linear(input_dim, num_slot_labels)

    def forward(self, x):
        x = self.dropout(x)
        return self.linear(x)

class JointClassifier(RobertaPreTrainedModel):
    "XLM-R Joint Classifier. Based on the Architecture described by Chen 2019"

    def __init__(self, config, num_intents=12, num_slots=31, return_dict=True):
        # bug
        # num intents not passing through when loading model
        super(JointClassifier, self).__init__(config)

        self.num_labels = config.num_labels

        self.num_intent_labels = num_intents
        self.num_slot_labels = num_slots

        self.roberta = XLMRobertaModel(config, add_pooling_layer=True)
        self.intent_clf = IntentClassifier(768, self.num_intent_labels)
        self.slot_clf = SlotClassifier(768, self.num_slot_labels)

        self.return_dict = return_dict

        #initial lize weights. Uses inherited Roberta Init
        self.init_weights()

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        intent_labels=None,
        slot_labels=None,
    ):
        """main forward function for the NN"""

        # pass inputs and attention masks into XLM Roberta
        outputs = self.roberta(
            input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids
        )

        # Hidden Layer output
        sequence_output = outputs[0]

        # pooler output
        cls_output = outputs[1]
        intent_logits = self.intent_clf(cls_output)
        slot_logits = self.slot_clf(sequence_output)
        total_loss = 0.0

        # if label is not empty
        if slot_labels is not None:

            slot_cross_ent = nn.CrossEntropyLoss()

            # apply attention masks to padding tokens
            if attention_mask is not None:
                active_loss = attention_mask.view(-1) == 1
                active_logits = slot_logits.view(-1, self.num_slot_labels)[active_loss]
                active_labels = slot_labels.view(-1)[active_loss]
                # calculate x-entropy
                slot_loss = slot_cross_ent(active_logits, active_labels)

            else:
                slot_loss = slot_cross_ent(
                    slot_logits.view(-1, self.num_slot_labels), slot_labels.view(-1)
                )

            # print(slot_loss)
            total_loss += slot_loss

        if intent_labels is not None:
            intent_cross_ent = nn.CrossEntropyLoss()
            intent_loss = intent_cross_ent(
                intent_logits.view(-1, self.num_intent_labels),
                intent_labels.view(-1),
            )
            total_loss += intent_loss

        outputs = ((intent_logits, slot_logits),) + outputs[
            2:
        ]  # add hidden states and attention if they are here
        outputs = (total_loss,) + outputs

        # return as dictionary
        if self.return_dict:
            return {"loss": total_loss, "intents": intent_logits, "slots": slot_logits}

        # tuple -> loss, intent logits, slot logits
        return outputs

def config_init(model_name):
    """simple wrapper for autoconfig"""
    conf = AutoConfig.from_pretrained(model_name)
    return conf

def tokenizer_init(model_name):
    """simple wrapper for auto tokenizer"""
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    return tokenizer

def main(data):
    pass

if __name__ == "__main__":
    main()


import numpy as np

from ..colors import Color
from .widget import Widget, overlapping_region
from .widget_data_structures import Point, Size, Rect

class _Root(Widget):
    """
    Root widget. Meant to be instantiated by the `App` class. Renders to terminal.
    """
    def __init__(self, app, env_out, default_char, default_color: Color):
        self._app = app
        self.env_out = env_out
        self.default_char = default_char
        self.default_color = default_color
        self.children = [ ]

        self.resize(env_out.get_size())

    def resize(self, dim: Size):
        """
        Resize canvas. Last render is erased.
        """
        self.env_out.erase_screen()
        self.env_out.flush()

        self._dim = dim

        self._last_canvas = np.full(dim, self.default_char, dtype=object)
        self._last_colors = np.full((*dim, 6), self.default_color, dtype=np.uint8)

        self.canvas = np.full_like(self._last_canvas, "><")  # "><" will guarantee an entire screen redraw.
        self.colors = self._last_colors.copy()

        # Buffer arrays to re-use in the `render` method:
        self._char_diffs = np.zeros_like(self.canvas, dtype=np.bool8)
        self._color_diffs = np.zeros_like(self.colors, dtype=np.bool8)
        self._reduced_color_diffs = np.zeros_like(self.canvas, dtype=np.bool8)

        for child in self.children:
            child.update_geometry()

    @property
    def top(self):
        return 0

    @property
    def left(self):
        return 0

    @property
    def pos(self):
        return Point(0, 0)

    @property
    def absolute_pos(self):
        return Point(0, 0)

    @property
    def is_transparent(self):
        return False

    @property
    def is_visible(self):
        return True

    @property
    def parent(self):
        return None

    @property
    def root(self):
        return self

    @property
    def app(self):
        return self._app

    def absolute_to_relative_coords(self, coord):
        return coord

    def render(self):
        """
        Paint canvas. Render to terminal.
        """
        # Swap canvas with last render:
        self.canvas, self._last_canvas = self._last_canvas, self.canvas
        self.colors, self._last_colors = self._last_colors, self.colors

        # Bring arrays into locals:
        canvas = self.canvas
        colors = self.colors

        char_diffs = self._char_diffs
        color_diffs = self._color_diffs
        reduced_color_diffs = self._reduced_color_diffs

        env_out = self.env_out
        write = env_out._buffer.append

        # Erase canvas:
        canvas[:] = self.default_char
        colors[:, :] = self.default_color

        overlap = overlapping_region
        height, width = canvas.shape
        rect = Rect(
            0,
            0,
            height,
            width,
            height,
            width,
        )

        for child in self.children:
            if region := overlap(rect, child):
                dest_slice, child_rect = region
                child.render(canvas[dest_slice], colors[dest_slice], child_rect)

        # Find differences between current render and last render:
        # (This is optimized version of `(last_canvas != canvas) | np.any(last_colors != colors, axis=-1)`
        # that re-uses buffers instead of creating new arrays.)
        np.not_equal(self._last_canvas, canvas, out=char_diffs)
        np.not_equal(self._last_colors, colors, out=color_diffs)
        np.any(color_diffs, axis=-1, out=reduced_color_diffs)
        np.logical_or(char_diffs, reduced_color_diffs, out=char_diffs)

        write("\x1b[?25l")  # Hide cursor

        ys, xs = np.nonzero(char_diffs)
        for y, x, color, char in zip(ys, xs, colors[ys, xs], canvas[ys, xs]):
            # The escape codes for moving the cursor and setting the color concatenated:
            write("\x1b[{};{}H\x1b[0;38;2;{};{};{};48;2;{};{};{}m{}".format(y + 1, x + 1, *color, char))

        write("\x1b[0m")  # Reset attributes
        env_out.flush()

    def dispatch_press(self, key_press):
        """
        Dispatch key press to descendants until handled.
        """
        any(widget.dispatch_press(key_press) for widget in reversed(self.children))

    def dispatch_click(self, mouse_event):
        """
        Dispatch mouse event to descendents until handled.
        """
        any(widget.dispatch_click(mouse_event) for widget in reversed(self.children))


from __future__ import print_function
import time
import cv2
import argparse
import torch
import torch.utils.data
from torch import nn, optim
from torch.nn import functional as F
from torchvision import datasets, transforms
from torchvision.utils import save_image
from gridworld.algorithms.composite_dataset import CompositeDataset
import numpy as np
import copy
from gridworld.algorithms.composite_models import NaiveModel
from tensorboardX import SummaryWriter
import os
import time
parser = argparse.ArgumentParser(description='VAE MNIST Example')
parser.add_argument('--batch-size', type=int, default=256, metavar='N',
                    help='input batch size for training (default: 128)')
parser.add_argument('--epochs', type=int, default=3000, metavar='N',
                    help='number of epochs to train (default: 10)')
parser.add_argument('--no-cuda', action='store_true', default=False,
                    help='enables CUDA training')
parser.add_argument('--seed', type=int, default=1, metavar='S',
                    help='random seed (default: 1)')
parser.add_argument('--log-interval', type=int, default=100, metavar='N',
                    help='how many batches to wait before logging training status')
parser.add_argument('--tasks', type=str, nargs="*", default='EatBreadPolicy',
                    help='which task to run')
parser.add_argument('--train_size', type=int, default=None,
                    help='which task to run')
parser.add_argument('--test_size', type=int, default=1000,
                    help='which task to run')
parser.add_argument('--resume', type=str, default=None,
                    help='A model to load')
parser.add_argument('--feature-dim', type=int, default=256)
parser.add_argument('-G', action='store_true', default=False,
                    help='on GCP')

parser.add_argument('--letter', type=str, default='A',
                    help='A model to load')
args = parser.parse_args()
args.cuda = not args.no_cuda and torch.cuda.is_available()

torch.manual_seed(args.seed)
print("tasks", args.tasks)
device = torch.device("cuda" if args.cuda else "cpu")
def worker_init_fn(worker_id):
    np.random.seed()
    st0 = np.random.get_state()[1][0]

kwargs = {'num_workers': 4, 'pin_memory': False} if args.cuda else {}
print("train_loader")

data_dir = 'data/Oct_4tasks_images/'

train_loader = torch.utils.data.DataLoader(
    CompositeDataset(directory=data_dir,
                             train=True, size=args.train_size),
    batch_size=args.batch_size, shuffle=True, worker_init_fn=worker_init_fn,**kwargs)
print("test_loader")
test_loader = torch.utils.data.DataLoader(
    CompositeDataset(directory=data_dir,
                             train=False, size=args.test_size),
    batch_size=args.batch_size, shuffle=True,worker_init_fn=worker_init_fn, **kwargs)
args.train_size = len(train_loader.dataset)
start_time = time.time()
result_list = []
model = NaiveModel(device,task_embedding_dim=args.feature_dim).to(device)

optimizer = optim.Adam(model.parameters(), lr=1e-3)
print("done model")

def accuracy(logpi, label):
    max_index = logpi.max(dim = 1)[1]
    return (max_index == label).sum()

def v3_loss_function(logpi, label):
    return torch.nn.functional.cross_entropy(logpi, label)

def process_data(data):
    image = data['image'].to(device)
    post_image = data['post_image'].to(device)
    pre_image = data['pre_image'].to(device)
    first_image = data['init_image'].to(device)
    final_image = data['final_image'].to(device)
    action =  data['action'].to(device)

    return first_image, image, final_image,  pre_image, post_image, action#, (pos_x, pos_y)

def train(epoch):
    model.train()
    train_loss = 0
    train_acc = 0
    num_batch = 0
    feat_loss_total = 0
    for batch_idx, data in enumerate(train_loader):
        first_image, image,final_image, pre_image, post_image, action = process_data(data)
        optimizer.zero_grad()
        logpi =model.forward(first_image, image, pre_image, post_image)
        loss = v3_loss_function(logpi, action)
        #import pdb;pdb.()
        acc = accuracy(logpi,action).item()
        num_batch +=len(image)
        loss.backward()
        train_loss += loss.item()
        train_acc += acc
        optimizer.step()

        if batch_idx % args.log_interval == 0:

            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.3f}'.format(
                epoch, batch_idx * len(image), len(train_loader.dataset),
                100. * batch_idx / len(train_loader),
                loss.item() ,
            ),
            )
    writer.add_scalar('Accuracy/train', train_acc/num_batch, epoch )
    writer.add_scalar('Loss/train',train_loss / num_batch, epoch )
    writer.add_scalar('Regloss/train',feat_loss_total/num_batch, epoch)
    print('====> Epoch: {} Average loss: {:.4f} Accuracy: {:.4f} FeatLoss: {:.4f} '.format(
          epoch, train_loss / num_batch, train_acc/num_batch, feat_loss_total/num_batch))

def test(epoch):
    model.eval()
    test_loss = 0
    test_acc = 0
    num_batch = 0
    total_feat_loss = 0
    with torch.no_grad():
        for i, data in enumerate(test_loader):
            first_image, image, final_image,  pre_image, post_image, action = process_data(data)
            logpi =model.forward(first_image, image, pre_image, post_image)
            loss = v3_loss_function(logpi, action)
            acc = accuracy(logpi,action).item()
            test_loss += loss.item()
            test_acc += acc
            num_batch += len(image)

    writer.add_scalar('Loss/test', test_loss/num_batch, epoch )
    writer.add_scalar('Regloss/test', total_feat_loss/num_batch, epoch)
    writer.add_scalar('Accuracy/test', test_acc/num_batch, epoch )
    print('====> Test set loss: {:.4f} Accuracy: {:.4f} FeatLoss: {:.4f} '.format(test_loss/num_batch, test_acc/num_batch, total_feat_loss/num_batch))

def save_model(model, epoch, path):
    state_dict = model.state_dict()
    opt_dict = optimizer.state_dict()

    torch.save((state_dict, epoch, len(train_loader.dataset), opt_dict), path)

if __name__ == "__main__":
    filename = 'naive_4tasks'+'_images_feat'+str(args.feature_dim)+'_'+args.letter+'_{:06d}'.format(args.train_size)
    writer = SummaryWriter(comment=filename)
    basedirectory = 'octresults/'+filename
    print("Tensorboard logdir",  writer.logdir)
    if not os.path.isdir(basedirectory):
        os.mkdir(basedirectory)
    if args.resume is not None:
        model_path = args.resume
        #start_epoch = int(model_path[-5:-3])
        state_dict, start_epoch, _, opt_dict = torch.load(model_path)
        model.load_state_dict(state_dict)
        optimizer.load_state_dict(opt_dict)
    else:
        start_epoch = 0
        test(0)
        path_name =  basedirectory+'/epoch{:05d}.pt'.format( 0)
        print('saved',path_name)
        save_model(model, 0, path_name)
    #test(0)
    for epoch in range(start_epoch+1, args.epochs + 1):
        print("starting train")
        t = time.time()
        writer.add_scalar("Time/time", t, epoch)
        train(epoch)
        test(epoch)
        if epoch % 10 == 0:
            path_name =  basedirectory+'/epoch{:05d}.pt'.format(epoch)
            print('saved',path_name)
            save_model(model, epoch,  path_name)
        t2 = time.time()
        writer.add_scalar("Time/delta", t2-t, epoch)


import numpy as np
import torch
import cv2

def  referto_skeleton_(pose:np.ndarray,ref_idx:int):
    #inplace opt
    pose = pose.reshape(len(pose),-1,2)
    ref_x = pose[:,ref_idx,0] #.copy()
    ref_y = pose[:,ref_idx,1]
    for i in range(len(pose)):
        pose[i,:,0]-= ref_x[i]
        pose[i,:,1]-= ref_y[i]
    return pose
def  normalize_referto_skeleton_(pose:np.ndarray,ref_idx:int):
    #inplace opt
    pose = pose.reshape(len(pose),-1,2)
    ref_x = pose[:,ref_idx,0] #.copy()
    ref_y = pose[:,ref_idx,1]
    for i in range(len(pose)):
        pose[i,:,0]-= ref_x[i]
        pose[i,:,1]-= ref_y[i]

        xlen = np.max(pose[i,:,0]) - np.min(pose[i,:,0])
        ylen = np.max(pose[i,:,1]) - np.min(pose[i,:,1])

        if(xlen==0): pose[i,:,0]=0
        else:pose[i,:,0] /= xlen
        if(ylen==0): pose[i,:,1]=0
        else:pose[i,:,1] /= ylen
    return pose

def normalize_min_(pose:np.ndarray):
    pose = pose.reshape(len(pose),-1,2)
    for i in range(len(pose)):
        xmin = np.min(pose[i,:,0])
        ymin = np.min(pose[i,:,1])
        xlen = np.max(pose[i,:,0]) - xmin
        ylen = np.max(pose[i,:,1]) - ymin

        if(xlen==0): pose[i,:,0]=0
        else:
            pose[i,:,0] -= xmin
            pose[i,:,0] /= xlen
        if(ylen==0): pose[i,:,1]=0
        else:
            pose[i,:,1] -= ymin
            pose[i,:,1] /= ylen
    return pose

def normalize_min_seq_(pose:np.ndarray):
    pose = pose.reshape(len(pose),-1,2)
    xmin = np.min(pose[:,:,0])
    ymin = np.min(pose[:,:,1])
    xlen = np.max(pose[:,:,0]) - xmin
    ylen = np.max(pose[:,:,1]) - ymin
    if(xlen==0): pose[:,:,0]=0
    else:
        pose[:,:,0] -= xmin
        pose[:,:,0] /= xlen
    if(ylen==0): pose[:,:,1]=0
    else:
        pose[:,:,1] -= ymin
        pose[:,:,1] /= ylen
    return pose

def single_normalize_min_(pose:np.ndarray):
    pose = pose.reshape(-1,2)
    xmin = np.min(pose[:,0])
    ymin = np.min(pose[:,1])
    xlen = np.max(pose[:,0]) - xmin
    ylen = np.max(pose[:,1]) - ymin
    if(xlen==0): pose[:,0]=0
    else:
        pose[:,0] -= xmin
        pose[:,0] /= xlen
    if(ylen==0): pose[:,:,1]=0
    else:
        pose[:,1] -= ymin
        pose[:,1] /= ylen
    return pose

#fork from alphapose.utils.vis
RED = (0, 0, 255)
GREEN = (0, 255, 0)
BLUE = (255, 0, 0)
CYAN = (255, 255, 0)
YELLOW = (0, 255, 255)
ORANGE = (0, 165, 255)
PURPLE = (255, 0, 255)
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
DEFAULT_FONT = cv2.FONT_HERSHEY_SIMPLEX
def vis_frame_fast(frame, im_res, add_bbox=False, format='coco'):
    '''
    frame: frame image
    im_res: im_res of predictions
    format: coco or mpii

    return rendered image
    '''
    if format == 'coco':
        l_pair = [
            (0, 1), (0, 2), (1, 3), (2, 4),  # Head
            (5, 6), (5, 7), (7, 9), (6, 8), (8, 10),
            (17, 11), (17, 12),  # Body
            (11, 13), (12, 14), (13, 15), (14, 16)
        ]
        p_color = [(0, 255, 255), (0, 191, 255), (0, 255, 102), (0, 77, 255), (0, 255, 0),  # Nose, LEye, REye, LEar, REar
                   (77, 255, 255), (77, 255, 204), (77, 204, 255), (191, 255, 77), (77, 191, 255), (191, 255, 77),  # LShoulder, RShoulder, LElbow, RElbow, LWrist, RWrist
                   (204, 77, 255), (77, 255, 204), (191, 77, 255), (77, 255, 191), (127, 77, 255), (77, 255, 127), (0, 255, 255)]  # LHip, RHip, LKnee, Rknee, LAnkle, RAnkle, Neck
        line_color = [(0, 215, 255), (0, 255, 204), (0, 134, 255), (0, 255, 50),
                      (77, 255, 222), (77, 196, 255), (77, 135, 255), (191, 255, 77), (77, 255, 77),
                      (77, 222, 255), (255, 156, 127),
                      (0, 127, 255), (255, 127, 77), (0, 77, 255), (255, 77, 36)]
    elif format == 'mpii':
        l_pair = [
            (8, 9), (11, 12), (11, 10), (2, 1), (1, 0),
            (13, 14), (14, 15), (3, 4), (4, 5),
            (8, 7), (7, 6), (6, 2), (6, 3), (8, 12), (8, 13)
        ]
        p_color = [PURPLE, BLUE, BLUE, RED, RED, BLUE, BLUE, RED, RED, PURPLE, PURPLE, PURPLE, RED, RED, BLUE, BLUE]
    else:
        NotImplementedError

    # im_name = os.path.basename(im_res['imgname'])
    img = frame.copy()
    height, width = img.shape[:2]
    for index,human in enumerate(im_res['result']):
        part_line = {}
        kp_preds = human['keypoints']
        kp_scores = human['kp_score']
        kp_preds = torch.cat((kp_preds, torch.unsqueeze((kp_preds[5, :] + kp_preds[6, :]) / 2, 0)))
        kp_scores = torch.cat((kp_scores, torch.unsqueeze((kp_scores[5, :] + kp_scores[6, :]) / 2, 0)))
        # Draw bboxes
        if add_bbox:
            if 'box' in human.keys():
                bbox = human['box']
                bbox = [bbox[0], bbox[0]+bbox[2], bbox[1], bbox[1]+bbox[3]]
            else:
                from PoseFlow.poseflow_infer import get_box
                keypoints = []
                for n in range(kp_scores.shape[0]):
                    keypoints.append(float(kp_preds[n, 0]))
                    keypoints.append(float(kp_preds[n, 1]))
                    keypoints.append(float(kp_scores[n]))
                bbox = get_box(keypoints, height, width)
            # color = get_color_fast(int(abs(human['idx'][0])))
            cv2.rectangle(img, (int(bbox[0]), int(bbox[2])), (int(bbox[1]), int(bbox[3])), BLUE, 2)
            # Draw indexes of humans
            if 'idx' in human.keys():
                cv2.putText(img, ''.join(str(index)), (int(bbox[0])+26, int((bbox[2] + 46))), DEFAULT_FONT, 2, RED, 2)
        # Draw keypoints
        for n in range(kp_scores.shape[0]):
            if kp_scores[n] <= 0.35:
                continue
            cor_x, cor_y = int(kp_preds[n, 0]), int(kp_preds[n, 1])
            part_line[n] = (cor_x, cor_y)
            cv2.circle(img, (cor_x, cor_y), 3, p_color[n], -1)
        # Draw limbs
        for i, (start_p, end_p) in enumerate(l_pair):
            if start_p in part_line and end_p in part_line:
                start_xy = part_line[start_p]
                end_xy = part_line[end_p]
                cv2.line(img, start_xy, end_xy, line_color[i], 2 * (kp_scores[start_p] + kp_scores[end_p]) + 1)
    return img


import logging
from tflite_runtime.interpreter import Interpreter
import azure.functions as func
import time
import cv2
import numpy as np
import base64

face_cascade = cv2.CascadeClassifier('HttpTrigger1/haarcascade_frontalface_default.xml')
interpreter = Interpreter("HttpTrigger1/converted_model_opt.tflite")
interpreter.allocate_tensors()

def main(req: func.HttpRequest) -> func.HttpResponse:
    name = req.params.get('flag')
    try:
        if name:
            image = req.get_body()  # raw data with base64 encoding
            decoded_data = base64.b64decode(image)
            np_data = np.fromstring(decoded_data,np.uint8)
            frame = cv2.imdecode(np_data,cv2.IMREAD_UNCHANGED)
            if(frame.shape[1] > frame.shape[0]):
                width = int(640)
                height = int(480)
                dim = (width, height)
                frame = cv2.resize(frame, dim, interpolation = cv2.INTER_AREA)
            else:
                width = int(480)
                height = int(640)
                dim = (width, height)
                frame = cv2.resize(frame, dim, interpolation = cv2.INTER_AREA)
            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            faces = face_cascade.detectMultiScale(gray, 1.045,5, minSize=(100, 100))
            marker = 0
            predict=0
            # Loop face detection
            for (x,y,w,h) in faces:
                if(w>109): # Detecting only the closest faces
                    # Taking the processing times to get the FPS
                    start_time = time.time()
                    # Processing the image to go through the TFlite model
                    raw_data = []
                    roi_gray = gray[y:y+h, x:x+w]
                    roi_gray = cv2.resize(roi_gray, (64,64), interpolation = cv2.INTER_AREA)
                    roi_gray = roi_gray/255.0
                    raw_data.append(roi_gray)
                    raw_data = np.array(raw_data, dtype="float32")
                    raw_data = np.expand_dims(raw_data, axis=3)
                    # Passing the image through the neural network.
                    input_index = interpreter.get_input_details()[0]["index"]
                    output_index = interpreter.get_output_details()[0]["index"]
                    interpreter.set_tensor(input_index, raw_data)
                    interpreter.invoke()
                    predictions = interpreter.get_tensor(output_index)
                    predict=np.argmax(predictions[0])
                    if(predict):
                        predict="Facemask ON"
                    else:
                        predict="Facemask OFF"
                    #print("FPS: ", 1.0 / (time.time() - start_time))
                    return func.HttpResponse(f"{predict}")
                else:
                    return func.HttpResponse(f"No Face Detected")
        else:
            return func.HttpResponse(f"Flag Error")
    except:
        return func.HttpResponse(f"Request Error")


#!/usr/bin/env python

import rospy
import tf
import sys
import math
import numpy
from geometry_msgs.msg import PointStamped
from deictic_gestures.srv import LookAt, CanLookAt
from deictic_gestures.msg import LookAtStatus
from pepper_resources_synchronizer_msgs.srv import MetaStateMachineRegister, MetaStateMachineRegisterRequest
from pepper_resources_synchronizer_msgs.msg import SubStateMachine_pepper_arm_manager_msgs, SubStateMachine_pepper_head_manager_msgs, SubStateMachine_pepper_base_manager_msgs
from pepper_head_manager_msgs.msg import StateMachineStatePrioritizedPoint as HeadStatePoint
from pepper_base_manager_msgs.msg import StateMachineStatePrioritizedAngle as BaseStateAngle
from resource_management_msgs.msg import StateMachineTransition, MessagePriority
from resource_synchronizer_msgs.msg import MetaStateMachinesStatus

from tf.transformations import translation_matrix, quaternion_matrix, quaternion_from_euler

LOOK_AT_MAX_SPEED = 0.3
LOOK_AT_SPEED = 0.1
LOOK_AT_MAX_ANGLE = math.radians(80)
MIN_DIST_MOVE = 0.1

HEAD_MANAGER_NAMES = ["pepper_head_manager"]
BASE_MANAGER_NAMES = ["pepper_base_manager"]

def center_radians(a):
    while a >= math.pi:
        a -= 2*math.pi
    while a < -math.pi:
        a += 2*math.pi
    return a

def transformation_matrix(t, q):
    translation_mat = translation_matrix(t)
    rotation_mat = quaternion_matrix(q)
    return numpy.dot(translation_mat, rotation_mat)

class LookAtSrv(object):
    def __init__(self):

        self.services = {"look_at": rospy.Service('/deictic_gestures/look_at', LookAt, self.handle_look_at),
                         "can_look_at": rospy.Service('/deictic_gestures/can_look_at', CanLookAt, self.handle_can_look_at)}

        self.tfListener = tf.TransformListener()
        self.parameters = {"fixed_frame": rospy.get_param("global_frame_id", "/map"),
                           "robot_footprint": rospy.get_param("footprint_frame_id", "/base_footprint"),
                           "look_at_max_speed": rospy.get_param("look_at_max_speed", LOOK_AT_MAX_SPEED)}

        self.publishers = {
            "result_point": rospy.Publisher('/deictic_gestures/looking_point', PointStamped, queue_size=5)}

        self.pointing_ids = []
        self.current_state = LookAtStatus.IDLE
        self.fsm_status = rospy.Subscriber("/pepper_resources_synchronizer/state_machine_status",
                                           MetaStateMachinesStatus, self.on_fsm_status)
        self.status_pub = rospy.Publisher("/deictic_gestures/look_at/status", LookAtStatus, queue_size=5, latch=True)

        self.resource_synchronizer = rospy.ServiceProxy("/pepper_resources_synchronizer/state_machines_register", MetaStateMachineRegister)

        self.current_lookat_frame = None
        self.current_lookat_point = None

    def on_fsm_status(self, msg):
        """

        :param msg:
        :type msg: MetaStateMachinesStatus
        :return:
        """
        for id in self.pointing_ids:
            if msg.id == id:
                rospy.loginfo("LOOK AT: state machine status")
                for i in range(len(msg.resource)):
                    if msg.resource[i] in HEAD_MANAGER_NAMES and msg.state_name[i] == "look" and self.current_state != LookAtStatus.LOOK:
                        self.current_state = LookAtStatus.LOOK
                        st = LookAtStatus()
                        st.status = LookAtStatus.LOOK
                        self.status_pub.publish(st)
                        rospy.loginfo("LOOK AT: status LOOK")
                    if msg.resource[i] in BASE_MANAGER_NAMES and msg.state_name[i] == "base_turn" and self.current_state != LookAtStatus.ROTATE:
                        self.current_state = LookAtStatus.ROTATE
                        st = LookAtStatus()
                        st.status = LookAtStatus.ROTATE
                        self.status_pub.publish(st)
                        rospy.loginfo("LOOK AT: status ROTATE")
                    if msg.resource[i] in HEAD_MANAGER_NAMES and msg.state_name[i] == "" and self.current_state != LookAtStatus.IDLE:
                        self.current_state = LookAtStatus.IDLE
                        st = LookAtStatus()
                        st.status = LookAtStatus.FINISHED
                        self.status_pub.publish(st)
                        rospy.loginfo("LOOK AT: status FINISHED")
                        st.status = LookAtStatus.IDLE
                        self.status_pub.publish(st)
                        self.pointing_ids.remove(id)
                        rospy.loginfo("LOOK AT: status IDLE")

    def distance(self, point):
        x = point[0]-self.current_lookat_point[0]
        y = point[1]-self.current_lookat_point[1]
        z = point[2]-self.current_lookat_point[2]
        return math.sqrt(x*x+y*y+z*z)

    def handle_can_look_at(self, req):
        try:
            if self.tfListener.canTransform("/torso", req.point.header.frame_id, rospy.Time()):
                (translation, rotation) = self.tfListener.lookupTransform('/torso', req.point.header.frame_id, rospy.Time())
                t = transformation_matrix(translation, rotation)
                p = numpy.atleast_2d([req.point.point.x, req.point.point.y, req.point.point.z, 1]).transpose()
                new_p = numpy.dot(t, p)
                angle = math.atan2(new_p[1, 0], new_p[0, 0])
                if math.fabs(angle) > LOOK_AT_MAX_ANGLE:
                    if angle > 0:
                        angle += 0.1
                    else:
                        angle -= 0.1
                    return False, angle
                else:
                    return True, 0
        except Exception as e:
            rospy.logerr("[look_at_srv] Exception occurred :" + str(e))
            return False, 0

    def handle_look_at(self, req, ):
        self.parameters["look_at_max_speed"] = rospy.get_param("look_at_max_speed", LOOK_AT_MAX_SPEED)
        rospy.loginfo("Look at request received")
        try:
            if self.tfListener.canTransform("/torso", req.point.header.frame_id, rospy.Time()):
                (translation, rotation) = self.tfListener.lookupTransform('/torso', req.point.header.frame_id, rospy.Time())
                t = transformation_matrix(translation, rotation)
                p = numpy.atleast_2d([req.point.point.x, req.point.point.y, req.point.point.z, 1]).transpose()
                new_p = numpy.dot(t, p)
                to_move = False
                dist = None
                speed_up = False
                if self.current_lookat_point:
                    dist = self.distance([new_p[0, 0], new_p[1, 0], new_p[2, 0]])
                    if dist > MIN_DIST_MOVE or self.current_lookat_frame != req.point.header.frame_id:
                        to_move = True

                else:
                    to_move = True
                if dist is not None:
                    speed_up = True if dist > 0.35 else False

                if to_move:
                    (translation, rotation) = self.tfListener.lookupTransform('torso', req.point.header.frame_id,
                                                                              rospy.Time())
                    t = transformation_matrix(translation, rotation)
                    p = numpy.atleast_2d([req.point.point.x, req.point.point.y, req.point.point.z, 1]).transpose()
                    new_p = numpy.dot(t, p)
                    #self.tracker.stopTracker()
                    #self.publishers["result_point"].publish()
                    #rospy.logwarn("lookat")

                    if speed_up :
                        look_at_speed = LOOK_AT_MAX_SPEED
                    else:
                        look_at_speed = LOOK_AT_SPEED
                    self.current_lookat_point = [new_p[0, 0], new_p[1, 0], new_p[2, 0]]

                    angle = math.atan2(new_p[1, 0], new_p[0, 0])
                    angle_diff = 0.0
                    if math.fabs(angle) > LOOK_AT_MAX_ANGLE:
                        if not req.with_base:
                            return False
                        angle_diff = min(center_radians(angle - LOOK_AT_MAX_ANGLE), center_radians(angle + LOOK_AT_MAX_ANGLE), key=lambda x: abs(x))

                    if angle_diff != 0:
                        rospy.loginfo("Will turn: " + str(math.degrees(angle_diff)))

                    after_rot = transformation_matrix([0, 0, 0], quaternion_from_euler(0, 0, -angle_diff))
                    new_p = numpy.dot(after_rot, new_p)
                    rospy.loginfo("Point after rotation: " + str(new_p))

                    mfsm = MetaStateMachineRegisterRequest()
                    headfsm = SubStateMachine_pepper_head_manager_msgs()
                    basefsm = SubStateMachine_pepper_base_manager_msgs()
                    head_wait_turned = HeadStatePoint()
                    head_wait_turned.header.id = "_head_wait_turned"
                    base_turn = BaseStateAngle()
                    base_turn.data = angle_diff
                    base_turn.header.id = "base_turn"
                    head_moving = HeadStatePoint()
                    head_moving.header.id = "head_moving"
                    head_moving.data.header.frame_id = "torso"
                    head_moving.data.header.stamp = rospy.Time.now()
                    head_moving.data.point.x, head_moving.data.point.y, head_moving.data.point.z = new_p[0, 0], new_p[1, 0], new_p[2, 0]
                    head_look = HeadStatePoint()
                    head_look.header.id = "look"
                    head_look.data.header.frame_id = "torso"
                    head_look.data.header.stamp = rospy.Time.now()
                    head_look.data.point.x, head_look.data.point.y, head_look.data.point.z = new_p[0, 0], new_p[1, 0], new_p[2, 0]
                    base_idle = BaseStateAngle()
                    base_idle.header.id = "_base_idle"
                    head_end = HeadStatePoint()
                    head_end.header.id = "head_end"
                    head_end.data = head_look.data
                    t_head_wait_turn = StateMachineTransition()
                    t_head_wait_turn.next_state = "head_moving"
                    t_head_wait_turn.end_condition.duration = rospy.Duration(-1)
                    t_head_wait_turn.end_condition.timeout = rospy.Duration(-1)
                    t_head_wait_turn.end_condition.regex_end_condition.append("__synchro__turned")
                    head_wait_turned.header.transitions.append(t_head_wait_turn)
                    t_base_turn = StateMachineTransition()
                    t_base_turn.next_state = "_base_idle"
                    t_base_turn.end_condition.timeout = rospy.Duration(-1)
                    t_base_turn.end_condition.duration = rospy.Duration(-1)
                    t_base_turn.end_condition.regex_end_condition.append("__synchro__turned")
                    base_turn.header.transitions.append(t_base_turn)
                    t_head_moving = StateMachineTransition()
                    t_head_moving.end_condition.duration = rospy.Duration(-1)
                    t_head_moving.end_condition.timeout = rospy.Duration(-1)
                    t_head_moving.end_condition.regex_end_condition.append("__done__")
                    t_head_moving.next_state = "look"
                    head_moving.header.transitions.append(t_head_moving)
                    t_head_look = StateMachineTransition()
                    t_head_look.end_condition.duration = rospy.Duration(-1)
                    t_head_look.end_condition.timeout = rospy.Duration(-1)
                    t_head_look.end_condition.regex_end_condition.append("human_perceived")
                    t_head_look.next_state = "head_end"
                    t_head_look_fail = StateMachineTransition()
                    t_head_look_fail.end_condition.duration = rospy.Duration(60)
                    t_head_look_fail.end_condition.timeout = rospy.Duration(-1)
                    t_head_look_fail.end_condition.regex_end_condition.append("stop_look_at")
                    t_head_look_fail.next_state = "head_end"
                    head_look.header.transitions.append(t_head_look)
                    head_look.header.transitions.append(t_head_look_fail)
                    t_head_end = StateMachineTransition()
                    t_head_end.next_state = "end"
                    t_head_end.end_condition.timeout = rospy.Duration(-1)
                    t_head_end.end_condition.duration = rospy.Duration(-1)
                    t_head_end.end_condition.regex_end_condition.append("__synchro__end")
                    head_end.header.transitions.append(t_head_end)
                    t_base_end = StateMachineTransition()
                    t_base_end.next_state = "end"
                    t_base_end.end_condition.duration = rospy.Duration(-1)
                    t_base_end.end_condition.timeout = rospy.Duration(-1)
                    t_base_end.end_condition.regex_end_condition.append("__synchro__end")
                    base_idle.header.transitions.append(t_base_end)

                    headfsm.state_machine.states_PrioritizedPoint.append(head_wait_turned)
                    headfsm.state_machine.states_PrioritizedPoint.append(head_moving)
                    headfsm.state_machine.states_PrioritizedPoint.append(head_look)
                    headfsm.state_machine.states_PrioritizedPoint.append(head_end)
                    basefsm.state_machine.states_PrioritizedAngle.append(base_turn)
                    basefsm.state_machine.states_PrioritizedAngle.append(base_idle)
                    if angle_diff != 0.0:
                        headfsm.header.initial_state = "_head_wait_turned"
                        basefsm.header.initial_state = "base_turn"
                    else:
                        headfsm.header.initial_state = "head_moving"
                        basefsm.header.initial_state = "_base_idle"
                    headfsm.header.timeout = rospy.Duration(-1)
                    headfsm.header.begin_dead_line = rospy.Time.now() + rospy.Duration(15)
                    basefsm.header.begin_dead_line = rospy.Time.now() + rospy.Duration(15)
                    basefsm.header.timeout = rospy.Duration(-1)

                    mfsm.state_machine_pepper_head_manager = headfsm
                    mfsm.state_machine_pepper_base_manager = basefsm
                    mfsm.header.timeout = rospy.Duration(-1)
                    mfsm.header.begin_dead_line = rospy.Time.now() + rospy.Duration(15)
                    mfsm.header.priority.value = MessagePriority.URGENT

                    ret = self.resource_synchronizer.call(mfsm)
                    self.pointing_ids.append(ret.id)
                    self.current_state = LookAtStatus.IDLE

                return True
            else:
                if self.current_lookat_frame is not None:
                    if req.point.header.frame_id == self.current_lookat_frame:
                        self.end_predicate(self.world.timeline, "isLookingAt", "robot", object_name=self.current_lookat_frame)
                        self.current_lookat_frame = None
                return False
        except Exception as e:
            rospy.logerr("[look_at_srv] Exception occurred :" + str(e))
            return False

if __name__ == '__main__':
    sys.argv = [arg for arg in sys.argv if "__name" not in arg and "__log" not in arg]
    sys.argc = len(sys.argv)

    rospy.init_node('look_at_srv')
    l = LookAtSrv()
    rospy.spin()
    exit(0)


import mxnet as mx
import numpy as np
import random
import cv2
import Augment

class LPRIter(mx.io.DataIter):
    def __init__(self, path_imgrec="", batch_size=128, data_shape=(3, 94, 24), path_imglist="",
                 label_width=8, mean_img="", augment = True,
                 **kwargs):
        super(LPRIter, self).__init__()
        self.batch_size = batch_size
        self.data_shape = data_shape

        if augment:
            self.rec = mx.io.ImageRecordIter(
                path_imgrec=path_imgrec,
                path_imglist=path_imglist,
                label_width=label_width,
                mean_img=mean_img,
                data_name='data',
                label_name='label',
                data_shape=(3,24,94),
                batch_size=batch_size,
                # max_aspect_ratio    = 0.25,
                min_aspect_ratio    = None,
                brightness          = 0.1,
                contrast            = 0.1,
                saturation          = 0.1,
                random_h            = 9,
                # random_s            = 10,
                # random_l            = 10,
                max_rotate_angle    = 5,
                max_shear_ratio     = 0.1,
                preprocess_threads=8,
                shuffle=True,
                num_parts=1,
                part_index=0
            )
        else:
            self.rec = mx.io.ImageRecordIter(
                path_imgrec=path_imgrec,
                path_imglist=path_imglist,
                label_width=label_width,
                mean_img=mean_img,
                data_name='data',
                label_name='label',
                data_shape=(3,24,94),
                batch_size=batch_size,
                preprocess_threads=8,
                shuffle=True,
                num_parts=1,
                part_index=0
            )

        self.provide_data = [('data', (batch_size,data_shape[0], data_shape[1], data_shape[2]))]
        self.provide_label = [('label', (self.batch_size, label_width))]

        self.reset()

    def reset(self):
        self.rec.reset()

    def iter_next(self):
        return self._get_batch()

    def next(self):
        if self.iter_next():
            return self._batch
        else:
            raise StopIteration

    def _get_batch(self):
        self._batch = self.rec.next()
        if not self._batch:
            return False
        # modify label
        label = self._batch.label[0].asnumpy()

        for i in range(len(label)):
            #if i==0:
                #print(label[i])
            for j in range(len(label[i])):
                if label[i][j] == 78:
                    label[i][j] = 0
                else:
                    label[i][j] += 1
        self._batch.label = [mx.nd.array(label)]

        buf = self._batch.data[0].asnumpy()
        data = np.zeros((self.batch_size,self.data_shape[0],self.data_shape[1],self.data_shape[2]))

        for i in range(len(buf)):
            im = buf[i] #
            augment_type = random.randint(0,2)
            ori_im = im.transpose(1,2,0).astype(np.uint8)
            im = cv2.imread("/opt/data/plate/im.jpg")
            # if augment_type == 0:
            # im_Distort = Augment.GenerateDistort(im, 12)
            # cv2.imwrite("/opt/data/plate/im_Distort.jpg", im_Distort)
            # if augment_type == 1:
            im_Stretch = Augment.GenerateStretch(im, 12)
            cv2.imwrite("/opt/data/plate/im_Stretch.jpg", im_Stretch)
            # if augment_type == 2:
            im_Perspective = Augment.GeneratePerspective(im)
            cv2.imwrite("/opt/data/plate/im_Perspective.jpg", im_Perspective)
            data[i]= im.transpose((0,2,1))
            # cv2.imwrite("/opt/data/plate/im.jpg", im.transpose((1,2,0)))

            '''
            if i==0:
                print(data[i].shape)
                #cv2.imshow("e",data[i])
                cv2.imwrite("res3.bmp",data[i].transpose((1,2,0)))
                cv2.waitKey(0)'''
        self._batch.data = [mx.nd.array(data)]

        return True

if __name__ == "__main__":

    batch_size = 128
    train = LPRIter(
        path_imgrec         = '/opt/data/plate/rec/train.rec',
        path_imglist        = '/opt/data/plate/rec/resized_plate_train.txt',
        label_width         = 8,
        mean_img            = '/opt/data/plate/rec/mean_plate.bin',
        data_shape          = (3,94,24),
        batch_size          = batch_size,
        augment             = True
    )

    for epoch in range(0, 1):
        data_iter = iter(train)
        end_of_batch = False
        next_data_batch = next(data_iter)
        data_batch = next_data_batch


# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""This module provides functions to calculate error rate in different level.
e.g. wer for word-level, cer for char-level.
"""
from itertools import groupby

import editdistance
import numpy as np

__all__ = ['word_errors', 'char_errors', 'wer', 'cer', "ErrorCalculator"]

def _levenshtein_distance(ref, hyp):
    """Levenshtein distance is a string metric for measuring the difference
    between two sequences. Informally, the levenshtein disctance is defined as
    the minimum number of single-character edits (substitutions, insertions or
    deletions) required to change one word into the other. We can naturally
    extend the edits to word level when calculate levenshtein disctance for
    two sentences.
    """
    m = len(ref)
    n = len(hyp)

    # special case
    if ref == hyp:
        return 0
    if m == 0:
        return n
    if n == 0:
        return m

    if m < n:
        ref, hyp = hyp, ref
        m, n = n, m

    # use O(min(m, n)) space
    distance = np.zeros((2, n + 1), dtype=np.int32)

    # initialize distance matrix
    for j in range(n + 1):
        distance[0][j] = j

    # calculate levenshtein distance
    for i in range(1, m + 1):
        prev_row_idx = (i - 1) % 2
        cur_row_idx = i % 2
        distance[cur_row_idx][0] = i
        for j in range(1, n + 1):
            if ref[i - 1] == hyp[j - 1]:
                distance[cur_row_idx][j] = distance[prev_row_idx][j - 1]
            else:
                s_num = distance[prev_row_idx][j - 1] + 1
                i_num = distance[cur_row_idx][j - 1] + 1
                d_num = distance[prev_row_idx][j] + 1
                distance[cur_row_idx][j] = min(s_num, i_num, d_num)

    return distance[m % 2][n]

def word_errors(reference, hypothesis, ignore_case=False, delimiter=' '):
    """Compute the levenshtein distance between reference sequence and
    hypothesis sequence in word-level.

    :param reference: The reference sentence.
    :type reference: str
    :param hypothesis: The hypothesis sentence.
    :type hypothesis: str
    :param ignore_case: Whether case-sensitive or not.
    :type ignore_case: bool
    :param delimiter: Delimiter of input sentences.
    :type delimiter: char
    :return: Levenshtein distance and word number of reference sentence.
    :rtype: list
    """
    if ignore_case:
        reference = reference.lower()
        hypothesis = hypothesis.lower()

    ref_words = list(filter(None, reference.split(delimiter)))
    hyp_words = list(filter(None, hypothesis.split(delimiter)))

    edit_distance = _levenshtein_distance(ref_words, hyp_words)
    # `editdistance.eavl precision` less than `_levenshtein_distance`
    # edit_distance = editdistance.eval(ref_words, hyp_words)
    return float(edit_distance), len(ref_words)

def char_errors(reference, hypothesis, ignore_case=False, remove_space=False):
    """Compute the levenshtein distance between reference sequence and
    hypothesis sequence in char-level.

    :param reference: The reference sentence.
    :type reference: str
    :param hypothesis: The hypothesis sentence.
    :type hypothesis: str
    :param ignore_case: Whether case-sensitive or not.
    :type ignore_case: bool
    :param remove_space: Whether remove internal space characters
    :type remove_space: bool
    :return: Levenshtein distance and length of reference sentence.
    :rtype: list
    """
    if ignore_case:
        reference = reference.lower()
        hypothesis = hypothesis.lower()

    join_char = ' '
    if remove_space:
        join_char = ''

    reference = join_char.join(list(filter(None, reference.split(' '))))
    hypothesis = join_char.join(list(filter(None, hypothesis.split(' '))))

    edit_distance = _levenshtein_distance(reference, hypothesis)
    # `editdistance.eavl precision` less than `_levenshtein_distance`
    # edit_distance = editdistance.eval(reference, hypothesis)
    return float(edit_distance), len(reference)

def wer(reference, hypothesis, ignore_case=False, delimiter=' '):
    """Calculate word error rate (WER). WER compares reference text and
    hypothesis text in word-level. WER is defined as:

    .. math::
        WER = (Sw + Dw + Iw) / Nw

    where

    .. code-block:: text

        Sw is the number of words subsituted,
        Dw is the number of words deleted,
        Iw is the number of words inserted,
        Nw is the number of words in the reference

    We can use levenshtein distance to calculate WER. Please draw an attention
    that empty items will be removed when splitting sentences by delimiter.

    :param reference: The reference sentence.
    :type reference: str
    :param hypothesis: The hypothesis sentence.
    :type hypothesis: str
    :param ignore_case: Whether case-sensitive or not.
    :type ignore_case: bool
    :param delimiter: Delimiter of input sentences.
    :type delimiter: char
    :return: Word error rate.
    :rtype: float
    :raises ValueError: If word number of reference is zero.
    """
    edit_distance, ref_len = word_errors(reference, hypothesis, ignore_case,
                                         delimiter)

    if ref_len == 0:
        raise ValueError("Reference's word number should be greater than 0.")

    wer = float(edit_distance) / ref_len
    return wer

def cer(reference, hypothesis, ignore_case=False, remove_space=False):
    """Calculate charactor error rate (CER). CER compares reference text and
    hypothesis text in char-level. CER is defined as:

    .. math::
        CER = (Sc + Dc + Ic) / Nc

    where

    .. code-block:: text

        Sc is the number of characters substituted,
        Dc is the number of characters deleted,
        Ic is the number of characters inserted
        Nc is the number of characters in the reference

    We can use levenshtein distance to calculate CER. Chinese input should be
    encoded to unicode. Please draw an attention that the leading and tailing
    space characters will be truncated and multiple consecutive space
    characters in a sentence will be replaced by one space character.

    :param reference: The reference sentence.
    :type reference: str
    :param hypothesis: The hypothesis sentence.
    :type hypothesis: str
    :param ignore_case: Whether case-sensitive or not.
    :type ignore_case: bool
    :param remove_space: Whether remove internal space characters
    :type remove_space: bool
    :return: Character error rate.
    :rtype: float
    :raises ValueError: If the reference length is zero.
    """
    edit_distance, ref_len = char_errors(reference, hypothesis, ignore_case,
                                         remove_space)

    if ref_len == 0:
        raise ValueError("Length of reference should be greater than 0.")

    cer = float(edit_distance) / ref_len
    return cer

class ErrorCalculator():
    """Calculate CER and WER for E2E_ASR and CTC models during training.

    :param y_hats: numpy array with predicted text
    :param y_pads: numpy array with true (target) text
    :param char_list: List[str]
    :param sym_space: <space>
    :param sym_blank: <blank>
    :return:
    """

    def __init__(self,
                 char_list,
                 sym_space,
                 sym_blank,
                 report_cer=False,
                 report_wer=False):
        """Construct an ErrorCalculator object."""
        super().__init__()

        self.report_cer = report_cer
        self.report_wer = report_wer

        self.char_list = char_list
        self.space = sym_space
        self.blank = sym_blank
        self.idx_blank = self.char_list.index(self.blank)
        if self.space in self.char_list:
            self.idx_space = self.char_list.index(self.space)
        else:
            self.idx_space = None

    def __call__(self, ys_hat, ys_pad, is_ctc=False):
        """Calculate sentence-level WER/CER score.

        :param paddle.Tensor ys_hat: prediction (batch, seqlen)
        :param paddle.Tensor ys_pad: reference (batch, seqlen)
        :param bool is_ctc: calculate CER score for CTC
        :return: sentence-level WER score
        :rtype float
        :return: sentence-level CER score
        :rtype float
        """
        cer, wer = None, None
        if is_ctc:
            return self.calculate_cer_ctc(ys_hat, ys_pad)
        elif not self.report_cer and not self.report_wer:
            return cer, wer

        seqs_hat, seqs_true = self.convert_to_char(ys_hat, ys_pad)
        if self.report_cer:
            cer = self.calculate_cer(seqs_hat, seqs_true)

        if self.report_wer:
            wer = self.calculate_wer(seqs_hat, seqs_true)
        return cer, wer

    def calculate_cer_ctc(self, ys_hat, ys_pad):
        """Calculate sentence-level CER score for CTC.

        :param paddle.Tensor ys_hat: prediction (batch, seqlen)
        :param paddle.Tensor ys_pad: reference (batch, seqlen)
        :return: average sentence-level CER score
        :rtype float
        """
        cers, char_ref_lens = [], []
        for i, y in enumerate(ys_hat):
            y_hat = [x[0] for x in groupby(y)]
            y_true = ys_pad[i]
            seq_hat, seq_true = [], []
            for idx in y_hat:
                idx = int(idx)
                if idx != -1 and idx != self.idx_blank and idx != self.idx_space:
                    seq_hat.append(self.char_list[int(idx)])

            for idx in y_true:
                idx = int(idx)
                if idx != -1 and idx != self.idx_blank and idx != self.idx_space:
                    seq_true.append(self.char_list[int(idx)])

            hyp_chars = "".join(seq_hat)
            ref_chars = "".join(seq_true)
            if len(ref_chars) > 0:
                cers.append(editdistance.eval(hyp_chars, ref_chars))
                char_ref_lens.append(len(ref_chars))

        cer_ctc = float(sum(cers)) / sum(char_ref_lens) if cers else None
        return cer_ctc

    def convert_to_char(self, ys_hat, ys_pad):
        """Convert index to character.

        :param paddle.Tensor seqs_hat: prediction (batch, seqlen)
        :param paddle.Tensor seqs_true: reference (batch, seqlen)
        :return: token list of prediction
        :rtype list
        :return: token list of reference
        :rtype list
        """
        seqs_hat, seqs_true = [], []
        for i, y_hat in enumerate(ys_hat):
            y_true = ys_pad[i]
            eos_true = np.where(y_true == -1)[0]
            ymax = eos_true[0] if len(eos_true) > 0 else len(y_true)
            # NOTE: padding index (-1) in y_true is used to pad y_hat
            seq_hat = [self.char_list[int(idx)] for idx in y_hat[:ymax]]
            seq_true = [
                self.char_list[int(idx)] for idx in y_true if int(idx) != -1
            ]
            seq_hat_text = "".join(seq_hat).replace(self.space, " ")
            seq_hat_text = seq_hat_text.replace(self.blank, "")
            seq_true_text = "".join(seq_true).replace(self.space, " ")
            seqs_hat.append(seq_hat_text)
            seqs_true.append(seq_true_text)
        return seqs_hat, seqs_true

    def calculate_cer(self, seqs_hat, seqs_true):
        """Calculate sentence-level CER score.

        :param list seqs_hat: prediction
        :param list seqs_true: reference
        :return: average sentence-level CER score
        :rtype float
        """
        char_eds, char_ref_lens = [], []
        for i, seq_hat_text in enumerate(seqs_hat):
            seq_true_text = seqs_true[i]
            hyp_chars = seq_hat_text.replace(" ", "")
            ref_chars = seq_true_text.replace(" ", "")
            char_eds.append(editdistance.eval(hyp_chars, ref_chars))
            char_ref_lens.append(len(ref_chars))
        return float(sum(char_eds)) / sum(char_ref_lens)

    def calculate_wer(self, seqs_hat, seqs_true):
        """Calculate sentence-level WER score.

        :param list seqs_hat: prediction
        :param list seqs_true: reference
        :return: average sentence-level WER score
        :rtype float
        """
        word_eds, word_ref_lens = [], []
        for i, seq_hat_text in enumerate(seqs_hat):
            seq_true_text = seqs_true[i]
            hyp_words = seq_hat_text.split()
            ref_words = seq_true_text.split()
            word_eds.append(editdistance.eval(hyp_words, ref_words))
            word_ref_lens.append(len(ref_words))
        return float(sum(word_eds)) / sum(word_ref_lens)


"""
The purpose of this module is to implement the database of faces and information.

We will implement 4 functions :

    - Add(picture, person_name, picture_file_name) to add a person in the database.
    - Remove(link, hardRemove = True) to remove a person from the database. the hardRemove parameter defines whether we remove physically the information from the disk.
    - GetImage(name) to retrieve a picture of a person knowing their name.
    - GetProfile(name) to retrieve the profile of a person knowing their name.

Our implementation is based on an array of the form

    [(encodings, name, path_to_image, profile)]

for every stored image.
"""

###############################################################################
# Imports.
###############################################################################

# Utilitary packages.
import os
import shutil
import re
from os.path import join, split
import os.path

# Packages used for image processing and numeric computing.
import numpy as np
import facialRecognition
import cv2

###############################################################################
# Main content of the program.
###############################################################################

class database:
    """
    A class for the management of the database.
    """
    def __init__(self, file_name = join('Database','London_database')):
        """
        Initialization of the class.

        :param file_name: Filename to the numpy file containing the encoding
        information. It must be an array of the form [(encodings, name, path_to_image, profile)]
        """
        # Initialize constructor.
        self.file_name = file_name

        # Open file.
        print('Loading faces from file ' + self.file_name)
        self.table_faces = np.load(self.file_name + '.npy')
        print('Loaded ' + str(len(self.table_faces)) + ' faces.')

        # Get path to images folder.
        self.folder_name_images = join(split(self.file_name)[0], split(split(self.table_faces[0][2])[0])[0])

        # Initialize algorithm for facial recognition.
        self.facial_recognition = facialRecognition.faceComparator()

        # Default profile value.
        self.DEFAULT_PROFILE = 'No Arup People profile'

    def add(self, frame, face_name, file_name, check_name = True):
        """
        Attempts to update the database adding picture in path 'self.folder_name_images/face_name/file_name.jpg'.
        It computes and returns 2 booleans : name_already_exists and one_face_detected,
        and the database is updated iff (one_face_detected && !name_already_exists).

        :param frame: The image.
        :param face_name: The name of the person in the image.
        :param file_name: The filename we want to give to the image.
        :param check_name: Boolean to decide whether we check if the name is already in the database.
        :return: Two booleans, name_already_exists and one_face_detected, whose value show whether the database was successfully updated or not, and the correponding problem if not.
        """
        # Initialise validity booleans.
        name_already_exists = False
        one_face_detected = True

        # Compute encodings, and input validity.
        encodings = self.facial_recognition.face_encodings(frame)
        if len(encodings) != 1:
            one_face_detected = False

        if check_name and one_face_detected:
            try:
                # Create dedicated folder and add picture.
                os.mkdir(join(self.folder_name_images, face_name))
            except Exception:
                name_already_exists = True

        if one_face_detected and not name_already_exists:
            # Save data
            link = join(self.folder_name_images, face_name, file_name) + '.jpg'
            self.table_faces = np.append(self.table_faces, [(encodings[0], face_name, link, self.DEFAULT_PROFILE)], axis = 0)
            cv2.imwrite(join(self.folder_name_images, face_name, file_name) + '.jpg', frame)
            np.save(self.file_name, self.table_faces)
        # Return results.
        return name_already_exists, one_face_detected

    def remove(self, link, hard_remove = True):
        """
        Remove the corresponding link image from the database.
        If hard_remove, we also delete physically the image from the computer.

        :param link: The link to the image we want to erase.
        :param hard_remove: Parameter to decide whether or not we physically erase the image from the computer.
        """
        # Remove the corresponding list to table_faces.
        self.table_faces = list(filter(lambda x : x[2] != link, self.table_faces))
        # Try to remove from folder if required.
        if hard_remove:
            try:
                os.remove(link)
                # We even delete the folder if it is empty.
                link_to_folder = os.path.split(link)[0]
                if len(os.listdir(link_to_folder)) == 0:
                    shutil.rmtree(link_to_folder)
            except Exception:
                print('Hard remove did not succeed.')
        # Actualize database.
        np.save(self.file_name, self.table_faces)

    def getImage(self, name):
        """
        Returns the image corresponding to the name.

        :param name: The considered name.
        :return: The corresponding image. Returns None if no image is found (even though it should not happen).
        """
        try:
            # Get link to image from the table of encodings.
            link = [path_to_image for (encoding, user_name, path_to_image, profile) in self.table_faces if user_name == name][0]
            # Return cv2 image.
            return cv2.imread(link)
        except:
            # Return default value.
            return None

    def getProfile(self, name):
        """
        Returns the profile corresponding to the name.

        :param name: The considered name.
        :return: A string corresponding to the profile. Either precomputed profile or 'Unable to match description with profile' if name is not in the database (even though it should not happen).
        """
        try:
            # Get profile from the table of encodings.
            return [profile for (encoding, user_name, path_to_image, profile) in self.table_faces if user_name == name][0]
        except:
            # Return default value.
            return self.DEFAULT_PROFILE