from .f1 import f1_score_with_precision_recall
from .utils.tools import *
import joblib
from scipy.sparse import hstack
import numpy as np
import os
import contractions
import requests


class PEDANT:
    def __init__(self): 
        current_dir = os.path.dirname(__file__) 
        model_dir = os.path.join(current_dir, 'classifier') 
        model_path = os.path.join(model_dir, 'lr_classifier.pkl') 
        vectorizer_path = os.path.join(model_dir, 'tf-idf_vectorizer.pkl') 

        clf_url = 'https://github.com/zli12321/qa_metrics/raw/master/qa_metrics/classifier/lr_classifier' 
        vectorizer_url = 'https://github.com/zli12321/qa_metrics/raw/master/qa_metrics/classifier/tf-idf_vectorizer' 
        # if not os.path.exists(model_path) or file_needs_update(clf_url, model_path) == True: 
        if True:
            if not os.path.exists(model_dir): 
                os.makedirs(model_dir)
            try:
                # print('Downloaded model path: ', model_path) 
                download_link(model_path, clf_url, 'PANDA model') 
                download_link(vectorizer_path, vectorizer_url, 'PANDA evaluation model tokenizer') 
            except:
                pass
        

        self.model = joblib.load(model_path)
        self.tokenizer = joblib.load(vectorizer_path)

    def download_latest_model(self):
        current_dir = os.path.dirname(__file__) 
        model_dir = os.path.join(current_dir, 'classifier') 
        model_path = os.path.join(model_dir, 'lr_classifier.pkl') 
        vectorizer_path = os.path.join(model_dir, 'tf-idf_vectorizer.pkl') 

        if not os.path.exists(model_dir): 
            os.makedirs(model_dir) 

        clf_url = 'https://github.com/zli12321/qa_metrics/raw/master/qa_metrics/classifier/lr_classifier' 
        vectorizer_url = 'https://github.com/zli12321/qa_metrics/raw/master/qa_metrics/classifier/tf-idf_vectorizer'
        print('Downloaded model path: ', model_path) 
        download_link(model_path, clf_url, 'PANDA model') 
        download_link(vectorizer_path, vectorizer_url, 'PANDA evaluation model tokenizer') 
        self.model = joblib.load(model_path) 
        self.tokenizer = joblib.load(vectorizer_path)

    '''
    Return the confidence score between the reference and candidate answers. 
    reference, candidate, and question are strings.
    '''
    def get_score(self, reference, candidate, question):
        reference = remove_punctuation(normalize_answer(str(reference)))
        candidate = remove_punctuation(normalize_answer(str(candidate)))

        if reference in candidate:
            return 1.0

        input_texts = []
        f1_scores, precisions, recalls = [], [], []
        curr_input_text = "[CLS] " + candidate + " [SEP] " + reference + " [SEP] " + remove_punctuation(normalize_answer(str(question))) + " [SEP]"
        input_texts.append(curr_input_text)
        f, p, r = calculate_f1_score_with_precision(str(candidate), str(reference))
       
        f1_scores.append(f)
        precisions.append(p)
        recalls.append(r)

        f1_scores = np.array(f1_scores).reshape(-1, 1)
        precisions = np.array(precisions).reshape(-1, 1)
        recalls = np.array(recalls).reshape(-1, 1)

        texts = self.tokenizer.transform(input_texts)

        '''
        Concatenate text features with f1 features
        '''
        features = hstack([texts, f1_scores, precisions, recalls])
        pred_probas = self.model.predict_proba(features)

        return pred_probas[0][0]


    '''
    Returns the classifier confidence score for the candidate answer matching judgment if the reference and candidate answers
    are lists. The reference and candidate answers can lists of strings or just strings. The question is a string.
    '''
    def get_scores(self, reference, candidate, question):
        # Calculate the F1 score between the referee and candidate
        confidence_scores = {}
        if isinstance(reference, list) and isinstance(candidate, list):
            references = reference
            candidates = candidate

            
            for reference in references:
                for candidate in candidates:
                    score = self.get_score(reference, candidate, question)
                    confidence_scores[reference] = {}
                    confidence_scores[reference][candidate] = score
                            
            return confidence_scores

            # for candidate in candidates:
            #     input_texts = []
            #     f1_scores, precisions, recalls = [], [], []
            #     for reference in references:

            #         input_texts.append("[CLS] " + candidate + " [SEP] " + reference + " [SEP] " + question + " [SEP]")
            #         f1_results = f1_score_with_precision_recall(reference, candidate)
            #         f, p, r = f1_results['f1'], f1_results['precision'], f1_results['recall']
            #         f1_scores.append(f)
            #         precisions.append(p)
            #         recalls.append(r)

            #     f1_scores = np.array(f1_scores).reshape(-1, 1)
            #     precisions = np.array(precisions).reshape(-1, 1)
            #     recalls = np.array(recalls).reshape(-1, 1)

            #     texts = self.tokenizer.transform(input_texts)

            #     '''
            #     Concatenate text features with f1 features
            #     '''
            #     features = hstack([texts, f1_scores, precisions, recalls])
            #     pred_probas = self.model.predict_proba(features)

            #     for i in range(len(pred_probas)):
            #         if references[i] not in confidece_scores:
            #             confidece_scores[references[i]] = {}
            #         confidece_scores[references[i]][candidate] = pred_probas[i][0]
                            
            # return confidece_scores
        elif isinstance(reference, list):
            references = reference
            for reference in references:
                score = self.get_score(reference, candidate, question)
                confidence_scores[reference] = {}
                confidence_scores[reference][candidate] = score
                            
            return confidence_scores

            #     input_texts.append("[CLS] " + candidate + " [SEP] " + reference + " [SEP] " + question + " [SEP]")
            #     f1_results = f1_score_with_precision_recall(reference, candidate)
            #     f, p, r = f1_results['f1'], f1_results['precision'], f1_results['recall']
            #     f1_scores.append(f)
            #     precisions.append(p)
            #     recalls.append(r)

            # f1_scores = np.array(f1_scores).reshape(-1, 1)
            # precisions = np.array(precisions).reshape(-1, 1)
            # recalls = np.array(recalls).reshape(-1, 1)

            # texts = self.tokenizer.transform(input_texts)

            # '''
            # Concatenate text features with f1 features
            # '''
            # features = hstack([texts, f1_scores, precisions, recalls])
            # pred_probas = self.model.predict_proba(features)

            
            # for i in range(len(pred_probas)):
            #     confidece_scores[references[i]] = {}
            #     confidece_scores[references[i]][candidate] = pred_probas[i][0]
                            
            # return confidece_scores
        elif isinstance(candidate, list):
            candidates = candidate
            for candidate in candidates:
                score = self.get_score(reference, candidate, question)
                confidence_scores[reference] = {}
                confidence_scores[reference][candidate] = score
                            
            return confidence_scores
        else:
            confidence_scores[reference] = {}
            confidence_scores[reference][candidate] = self.get_score(reference, candidate, question)

            return confidence_scores

    '''
    Given a list of reference, candidate, and a question, return the pair with the highest confidence score.
    '''
    def get_highest_score(self, reference, candidate, question):
        confidence_scores = self.get_scores(reference, candidate, question)

        max_score = -1
        max_pair = (None, None)

        for reference, candidates in confidence_scores.items():
            for candidate, score in candidates.items():
                if score > max_score:
                    max_score = score
                    max_pair = (reference, candidate)

        return max_pair, max_score

    '''
    Input your reference and candidate answers, and the question.
    Return True if the candidate answer is deemed correct. Else, False.
    '''
    def evaluate(self, reference, candidate, question):
        # Calculate the F1 score between the referee and candidate
        judgment = False
        if isinstance(reference, list) and isinstance(candidate, list):
            references = [remove_punctuation(normalize_answer(str(ele))) for ele in reference]
            candidates = [remove_punctuation(normalize_answer(str(ele))) for ele in candidate]
            question = remove_punctuation(normalize_answer(str(question)))

            for candidate in candidates:
                if judgment == False:
                    input_texts = []
                    f1_scores, precisions, recalls = [], [], []
                    for reference in references:
                        if reference in candidate:
                            return True
                        input_texts.append("[CLS] " + candidate + " [SEP] " + reference + " [SEP] " + question + " [SEP]")
                        f1_results = f1_score_with_precision_recall(reference, candidate)
                        f, p, r = f1_results['f1'], f1_results['precision'], f1_results['recall']
                        f1_scores.append(f)
                        precisions.append(p)
                        recalls.append(r)

                        f1_scores = np.array(f1_scores).reshape(-1, 1)
                        precisions = np.array(precisions).reshape(-1, 1)
                        recalls = np.array(recalls).reshape(-1, 1)

                        texts = self.tokenizer.transform(input_texts)

                    '''
                    Concatenate text features with f1 features
                    '''
                    features = hstack([texts, f1_scores, precisions, recalls])
                    preds = self.model.predict(features)

                    if "correct" in preds:
                        judgment = True
                            
            return judgment
        elif isinstance(reference, list):
            references = [remove_punctuation(normalize_answer(str(ele))) for ele in reference]
            candidate = remove_punctuation(normalize_answer(str(candidate)))
            input_texts, tagged_texts = [], []
            f1_scores, precisions, recalls = [], [], []
            for reference in references:
                if reference in candidate:
                    return True
                curr_input_text = "[CLS] " + candidate + " [SEP] " + reference + " [SEP] " + contractions.fix(normalize_answer(str(question))) + " [SEP]"
                input_texts.append(curr_input_text)
                f, p, r = calculate_f1_score_with_precision(str(candidate), str(reference))
                f1_scores.append(f)
                precisions.append(p)
                recalls.append(r)

            f1_scores = np.array(f1_scores).reshape(-1, 1)
            precisions = np.array(precisions).reshape(-1, 1)
            recalls = np.array(recalls).reshape(-1, 1)

            texts = self.tokenizer.transform(input_texts)

            result = hstack([texts, f1_scores, precisions, recalls])
            output = self.model.predict(result)


            return "correct" in output
        elif isinstance(candidate, list):
            candidates = [remove_punctuation(normalize_answer(str(ele))) for ele in candidate]
            reference = remove_punctuation(normalize_answer(str(reference)))
            question = remove_punctuation(normalize_answer(str(question)))

            input_texts = []
            f1_scores, precisions, recalls = [], [], []
            for candidate in candidates:
                if reference in candidate:
                    return True
                input_texts.append("[CLS] " + candidate + " [SEP] " + reference + " [SEP] " + question + " [SEP]")
                f1_results = f1_score_with_precision_recall(reference, candidate)
                f, p, r = f1_results['f1'], f1_results['precision'], f1_results['recall']
                f1_scores.append(f)
                precisions.append(p)
                recalls.append(r)

            f1_scores = np.array(f1_scores).reshape(-1, 1)
            precisions = np.array(precisions).reshape(-1, 1)
            recalls = np.array(recalls).reshape(-1, 1)

            texts = self.tokenizer.transform(input_texts)

            '''
            Concatenate text features with f1 features
            '''
            features = hstack([texts, f1_scores, precisions, recalls])
            preds = self.model.predict(features)

            if "correct" in preds:
                judgment = True
                            
            return judgment
        else:
            reference = remove_punctuation(normalize_answer(str(reference)))
            candidate = remove_punctuation(normalize_answer(str(candidate)))
            question = remove_punctuation(normalize_answer(str(question)))

            if reference in candidate:
                return True

            input_texts = []
            f1_scores, precisions, recalls = [], [], []
            input_texts.append("[CLS] " + candidate + " [SEP] " + reference + " [SEP] " + question + " [SEP]")
            f1_results = f1_score_with_precision_recall(reference, candidate)
            f, p, r = f1_results['f1'], f1_results['precision'], f1_results['recall']
            f1_scores.append(f)
            precisions.append(p)
            recalls.append(r)

            f1_scores = np.array(f1_scores).reshape(-1, 1)
            precisions = np.array(precisions).reshape(-1, 1)
            recalls = np.array(recalls).reshape(-1, 1)

            texts = self.tokenizer.transform(input_texts)

            '''
            Concatenate text features with f1 features
            '''
            features = hstack([texts, f1_scores, precisions, recalls])
            preds = self.model.predict(features)

            if "correct" in preds:
                judgment = True
                            
            return judgment