from evaluate.VqaMedEvaluatorBase import VqaMedEvaluatorBase
from nltk.corpus import wordnet as wn
from scipy import spatial


class WbssEvaluator(VqaMedEvaluatorBase):
    """"""

    def __init__(self, predictions, ground_truth):
        super().__init__(predictions, ground_truth)
        # Used for WUPS
        self.word_pair_dict = {}

    def get_name(self):
        return 'wbss'

    def evaluate(self):
        return self._compute_wbss(self.predictions, self.ground_truth)

    def _compute_wbss(self, predictions, ground_truth):
        """
    Compute and return the primary score
    Parameter 'predictions' : predictions object generated by the load_predictions method
    NO VALIDATION OF THE RUNFILE SHOULD BE IMPLEMENTED HERE
    We assume that the predictions in the parameter are valid
    Validation should be handled in the load_predictions method
        :param predictions:
        :param ground_truth:
        :return:
    """
        count = 0
        totalscore_wbss = 0.0
        for tuple1, tuple2 in zip(ground_truth, predictions):
            QID1 = tuple1.q_id
            QID2 = tuple2.q_id
            imageID1 = tuple1.image_id
            imageID2 = tuple2.image_id
            ans1 = tuple1.answer
            ans2 = tuple2.answer
            assert (QID1 == QID2)
            assert (imageID1 == imageID2)

            count += 1
            QID = QID1

            if ans1 == ans2:
                score_wbss = 1.0
            elif ans2.strip() == "":  # Added by Ivan (Handle case of empty answer)
                score_wbss = 0
            else:
                score_wbss = self._calculateWBSS(ans1, ans2)

            totalscore_wbss += score_wbss

        return totalscore_wbss / float(count)

    def _calculateWBSS(self, S1, S2):
        if S1 is None or S2 is None:
            return 0.0
        dictionary = self._constructDict(S1.split(), S2.split())
        vector1 = self._getVector_wordnet(S1, dictionary)
        vector2 = self._getVector_wordnet(S2, dictionary)
        cos_similarity = self._calculateCosineSimilarity(vector1, vector2)
        return cos_similarity

    def _constructDict(self, list1, list2):
        return list(set(list1 + list2))

    def _getVector_wordnet(self, S, dictionary):
        vector = [0.0] * len(dictionary)
        for index, word in enumerate(dictionary):
            # update score for vector[index]
            for wordinS in S.split():
                if wordinS == word:
                    score = 1.0
                else:
                    score = self._wups_score(word, wordinS)
                if score > vector[index]:
                    vector[index] = score
        return vector

    def _wups_score(self, word1, word2):
        score = self._wup_measure(word1, word2)
        return score

    def _wup_measure(self, a, b, similarity_threshold=0.925, debug=False):
        """
        Returns Wu-Palmer similarity score.
        More specifically, it computes:
            max_{x \in interp(a)} max_{y \in interp(b)} wup(x,y)
            where interp is a 'interpretation field'
        """
        if debug: print('Original', a, b)
        # if word_pair_dict.has_key(a+','+b):
        if a + ',' + b in self.word_pair_dict.keys():
            return self.word_pair_dict[a + ',' + b]

        def get_semantic_field(a):
            return wn.synsets(a, pos=wn.NOUN)

        if a == b: return 1.0

        interp_a = get_semantic_field(a)
        interp_b = get_semantic_field(b)
        if debug: print(interp_a)

        if interp_a == [] or interp_b == []:
            return 0.0

        if debug: print('Stem', a, b)
        global_max = 0.0
        for x in interp_a:
            for y in interp_b:
                local_score = x.wup_similarity(y)
                if debug: print('Local', local_score)
                if local_score > global_max:
                    global_max = local_score
        if debug: print('Global', global_max)

        # we need to use the semantic fields and therefore we downweight
        # unless the score is high which indicates both are synonyms
        if global_max < similarity_threshold:
            interp_weight = 0.1
        else:
            interp_weight = 1.0

        final_score = global_max * interp_weight
        self.word_pair_dict[a + ',' + b] = final_score
        return final_score

    def _calculateCosineSimilarity(self, vector1, vector2):
        return 1 - spatial.distance.cosine(vector1, vector2)


def wbss_score(y_true: iter, y_pred: iter) -> float:
    evaluator = WbssEvaluator(predictions=y_pred, ground_truth=y_true)
    return evaluator.evaluate()
