import warnings
import nltk
from nltk.translate.bleu_score import SmoothingFunction
from nltk.corpus import stopwords
from nltk.stem.snowball import SnowballStemmer
from evaluate.VqaMedEvaluatorBase import VqaMedEvaluatorBase
import string

class BleuEvaluator(VqaMedEvaluatorBase):
    """"""

    # Used for NLTK
    remove_stopwords = True
    stemming = True
    case_sensitive = False

    def __init__(self, predictions, ground_truth):
        super(BleuEvaluator, self).__init__(predictions, ground_truth)

    def get_name(self):
        return 'bleu'

    def evaluate(self):
        """
       This is the only method that will be called by the framework
       Parameter 'submission_file_path': Path of the submitted runfile
       returns a _result_object that can contain up to 2 different scores
       """
        predictions = self.predictions
        ground_truth = self.ground_truth

        # Compute second score
        bleu = self._compute_bleu(predictions, ground_truth)
        return bleu

    def _compute_bleu(self, predictions, ground_truth):
        '''
    Compute and return the secondary 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
    '''
        # Hide warnings
        warnings.filterwarnings('ignore')
        # English Stopwords
        stops = set(stopwords.words("english"))
        # Stemming
        stemmer = SnowballStemmer("english")

        # Remove punctuation from string
        translator = str.maketrans('', '', string.punctuation)

        candidate_pairs = self._readresult(predictions)

        gt_pairs = self._readresult(ground_truth)

        # Define max score and current score
        max_score = len(gt_pairs)
        current_score = 0

        i = 0
        for image_key in candidate_pairs:

            # Get candidate and GT caption
            candidate_caption = candidate_pairs[image_key]
            gt_caption = gt_pairs[image_key]

            # Optional - Go to lowercase
            if not BleuEvaluator.case_sensitive:
                candidate_caption = candidate_caption.lower()
                gt_caption = gt_caption.lower()

            # Split caption into individual words (remove punctuation)
            candidate_words = nltk.tokenize.word_tokenize(candidate_caption.translate(translator))
            gt_words = nltk.tokenize.word_tokenize(gt_caption.translate(translator))

            # Optional - Remove stopwords
            if BleuEvaluator.remove_stopwords:
                candidate_words = [word for word in candidate_words if word.lower() not in stops]
                gt_words = [word for word in gt_words if word.lower() not in stops]

            # Optional - Apply stemming
            if BleuEvaluator.stemming:
                candidate_words = [stemmer.stem(word) for word in candidate_words]
                gt_words = [stemmer.stem(word) for word in gt_words]

            # Calculate BLEU score for the current caption

            try:
                # If both the GT and candidate are empty, assign a score of 1 for this caption
                if len(gt_words) == 0 and len(candidate_words) == 0:
                    bleu_score = 1
                # Calculate the BLEU score
                else:
                    bleu_score = nltk.translate.bleu_score.sentence_bleu([gt_words], candidate_words,
                                                                         smoothing_function=SmoothingFunction().method0)

                # Increase calculated score
                current_score += bleu_score
            # Handle problematic cases where BLEU score calculation is impossible
            except ZeroDivisionError:
                msg = f'Problem with {gt_words} {candidate_words}'
                print(msg)
                # raise Exception(msg)
        return current_score / max_score

    def _readresult(self, tuples):
        pairs = {}
        for row in tuples:
            pairs[row.q_id] = row.answer
        return pairs
