from itertools import groupby
import numpy as np


class EntityClassifier:
    def __init__(self):
        pass

    def _get_grouped_by_length(self, entities):
        sorted_by_len = sorted(entities, key=lambda entity: len(entity.get_span()), reverse=True)

        entities_by_length = {}
        for length, group in groupby(sorted_by_len, lambda entity: len(entity.get_span())):
            entities = list(group)
            entities_by_length[length] = entities

        return entities_by_length

    def _filter_max_length(self, entities):
        entities_by_length = self._get_grouped_by_length(entities)
        max_length = max(list(entities_by_length.keys()))

        return entities_by_length[max_length]

    def _select_max_prior(self, entities):
        priors = [entity.get_prior() for entity in entities]
        return entities[np.argmax(priors)]

    def _get_casing_difference(self, word1, original):
        difference = 0
        for w1, w2 in zip(word1, original):
            if w1 != w2:
                difference += 1

        return difference

    def _filter_most_similar(self, entities):
        similarities = np.array(
            [self._get_casing_difference(entity.get_span().text, entity.get_original_alias()) for entity in entities])

        min_indices = np.where(similarities == similarities.min())[0].tolist()

        return [entities[i] for i in min_indices]

    def __call__(self, entities):
        filtered_by_length = self._filter_max_length(entities)
        filtered_by_casing = self._filter_most_similar(filtered_by_length)

        return self._select_max_prior(filtered_by_casing)