"""Code for identifying and anonymizing entities in NL and SQL."""

import copy
import json
from . import util

ENTITY_NAME = "ENTITY"
CONSTANT_NAME = "CONSTANT"
TIME_NAME = "TIME"
SEPARATOR = "#"


def timeval(string):
    """Returns the numeric version of a time.

    Inputs:
        string (str): String representing a time.

    Returns:
        String representing the absolute time.
    """
    if string.endswith("am") or string.endswith(
            "pm") and string[:-2].isdigit():
        numval = int(string[:-2])
        if len(string) == 3 or len(string) == 4:
            numval *= 100
        if string.endswith("pm"):
            numval += 1200
        return str(numval)
    return ""


def is_time(string):
    """Returns whether a string represents a time.

    Inputs:
        string (str): String to check.

    Returns:
        Whether the string represents a time.
    """
    if string.endswith("am") or string.endswith("pm"):
        if string[:-2].isdigit():
            return True

    return False


def deanonymize(sequence, ent_dict, key):
    """Deanonymizes a sequence.

    Inputs:
        sequence (list of str): List of tokens to deanonymize.
        ent_dict (dict str->(dict str->str)): Maps from tokens to the entity dictionary.
        key (str): The key to use, in this case either natural language or SQL.

    Returns:
        Deanonymized sequence of tokens.
    """
    new_sequence = []
    for token in sequence:
        if token in ent_dict:
            new_sequence.extend(ent_dict[token][key])
        else:
            new_sequence.append(token)

    return new_sequence


class Anonymizer:
    """Anonymization class for keeping track of entities in this domain and
       scripts for anonymizing/deanonymizing.

    Members:
        anonymization_map (list of dict (str->str)): Containing entities from
            the anonymization file.
        entity_types (list of str): All entities in the anonymization file.
        keys (set of str): Possible keys (types of text handled); in this case it should be
            one for natural language and another for SQL.
        entity_set (set of str): entity_types as a set.
    """
    def __init__(self, filename):
        self.anonymization_map = []
        self.entity_types = []
        self.keys = set()

        pairs = [json.loads(line) for line in open(filename).readlines()]
        for pair in pairs:
            for key in pair:
                if key != "type":
                    self.keys.add(key)
            self.anonymization_map.append(pair)
            if pair["type"] not in self.entity_types:
                self.entity_types.append(pair["type"])

        self.entity_types.append(ENTITY_NAME)
        self.entity_types.append(CONSTANT_NAME)
        self.entity_types.append(TIME_NAME)

        self.entity_set = set(self.entity_types)

    def get_entity_type_from_token(self, token):
        """Gets the type of an entity given an anonymized token.

        Inputs:
            token (str): The entity token.

        Returns:
            str, representing the type of the entity.
        """
        # these are in the pattern NAME:#, so just strip the thing after the
        # colon
        colon_loc = token.index(SEPARATOR)
        entity_type = token[:colon_loc]
        assert entity_type in self.entity_set

        return entity_type

    def is_anon_tok(self, token):
        """Returns whether a token is an anonymized token or not.

        Input:
            token (str): The token to check.

        Returns:
            bool, whether the token is an anonymized token.
        """
        return token.split(SEPARATOR)[0] in self.entity_set

    def get_anon_id(self, token):
        """Gets the entity index (unique ID) for a token.

        Input:
            token (str): The token to get the index from.

        Returns:
            int, the token ID if it is an anonymized token; otherwise -1.
        """
        if self.is_anon_tok(token):
            return self.entity_types.index(token.split(SEPARATOR)[0])
        else:
            return -1

    def anonymize(self,
                  sequence,
                  tok_to_entity_dict,
                  key,
                  add_new_anon_toks=False):
        """Anonymizes a sequence.

        Inputs:
            sequence (list of str): Sequence to anonymize.
            tok_to_entity_dict (dict): Existing dictionary mapping from anonymized
                tokens to entities.
            key (str): Which kind of text this is (natural language or SQL)
            add_new_anon_toks (bool): Whether to add new entities to tok_to_entity_dict.

        Returns:
            list of str, the anonymized sequence.
        """
        # Sort the token-tok-entity dict by the length of the modality.
        sorted_dict = sorted(tok_to_entity_dict.items(),
                             key=lambda k: len(k[1][key]))[::-1]

        anonymized_sequence = copy.deepcopy(sequence)

        if add_new_anon_toks:
            type_counts = {}
            for entity_type in self.entity_types:
                type_counts[entity_type] = 0
            for token in tok_to_entity_dict:
                entity_type = self.get_entity_type_from_token(token)
                type_counts[entity_type] += 1

        # First find occurrences of things in the anonymization dictionary.
        for token, modalities in sorted_dict:
            our_modality = modalities[key]

            # Check if this key's version of the anonymized thing is in our
            # sequence.
            while util.subsequence(our_modality, anonymized_sequence):
                found = False
                for startidx in range(
                        len(anonymized_sequence) - len(our_modality) + 1):
                    if anonymized_sequence[startidx:startidx +
                                           len(our_modality)] == our_modality:
                        anonymized_sequence = anonymized_sequence[:startidx] + [
                            token] + anonymized_sequence[startidx + len(our_modality):]
                        found = True
                        break
                assert found, "Thought " \
                    + str(our_modality) + " was in [" \
                    + str(anonymized_sequence) + "] but could not find it"

        # Now add new keys if they are present.
        if add_new_anon_toks:

            # For every span in the sequence, check whether it is in the anon map
            # for this modality
            sorted_anon_map = sorted(self.anonymization_map,
                                     key=lambda k: len(k[key]))[::-1]

            for pair in sorted_anon_map:
                our_modality = pair[key]

                token_type = pair["type"]
                new_token = token_type + SEPARATOR + \
                    str(type_counts[token_type])

                while util.subsequence(our_modality, anonymized_sequence):
                    found = False
                    for startidx in range(
                            len(anonymized_sequence) - len(our_modality) + 1):
                        if anonymized_sequence[startidx:startidx + \
                            len(our_modality)] == our_modality:
                            if new_token not in tok_to_entity_dict:
                                type_counts[token_type] += 1
                                tok_to_entity_dict[new_token] = pair

                            anonymized_sequence = anonymized_sequence[:startidx] + [
                                new_token] + anonymized_sequence[startidx + len(our_modality):]
                            found = True
                            break
                    assert found, "Thought " \
                        + str(our_modality) + " was in [" \
                        + str(anonymized_sequence) + "] but could not find it"

            # Also replace integers with constants
            for index, token in enumerate(anonymized_sequence):
                if token.isdigit() or is_time(token):
                    if token.isdigit():
                        entity_type = CONSTANT_NAME
                        value = new_token
                    if is_time(token):
                        entity_type = TIME_NAME
                        value = timeval(token)

                    # First try to find the constant in the entity dictionary already,
                    # and get the name if it's found.
                    new_token = ""
                    new_dict = {}
                    found = False
                    for entity, value in tok_to_entity_dict.items():
                        if value[key][0] == token:
                            new_token = entity
                            new_dict = value
                            found = True
                            break

                    if not found:
                        new_token = entity_type + SEPARATOR + \
                            str(type_counts[entity_type])
                        new_dict = {}
                        for tempkey in self.keys:
                            new_dict[tempkey] = [token]

                        tok_to_entity_dict[new_token] = new_dict
                        type_counts[entity_type] += 1

                    anonymized_sequence[index] = new_token

        return anonymized_sequence
