from flair.data import Sentence

from flair.models import SequenceTagger, TextClassifier

from flair.tokenization import SciSpacyTokenizer

from transformers import pipeline, TextClassificationPipeline, AutoTokenizer, TFBertForTokenClassification, BertForSequenceClassification, AutoModelForSequenceClassification

from transformers.trainer import Trainer, TrainingArguments

from stqdm import stqdm

from allennlp.predictors.predictor import Predictor

import os

import fitz

import streamlit

import wikipedia

import nltk.data

import os

os.environ["TOKENIZERS_PARALLELISM"] = "false"

class InferenceADE:  

    ''' Voting classifier using 3 different models for ADE detection '''

    def __init__(self, pipeline_scibert, pipeline_biolink, model_hunflair):

        self.f1_score_biolink = 0.96   # not real f1 scores for now

        self.f1_score_scibert = 0.80

        self.f1_score_hunflair = 0.90

        self.pipeline_scibert = pipeline_scibert

        self.pipeline_biolink = pipeline_biolink

        self.model_hunflair = model_hunflair

    def __call__(self, sentence):

        result_bert = self.pipeline_scibert(sentence)[0]

        result_biolink = self.pipeline_biolink(sentence)[0]

        s = Sentence(sentence)

        self.model_hunflair.predict(s)

        result_hunflair = s.labels[0].to_dict()

        if result_bert['label'] == 'LABEL_0':

            pred_scibert = [result_bert['score'], 1-result_bert['score']]

        elif result_bert['label'] == 'LABEL_1':

            pred_scibert = [1-result_bert['score'], result_bert['score']]

        if result_biolink['label'] == 'LABEL_0':

            pred_biolink = [result_biolink['score'], 1-result_biolink['score']]

        elif result_biolink['label'] == 'LABEL_1':

            pred_biolink = [1-result_biolink['score'], result_biolink['score']]

        if result_hunflair['value'] == '0':

            pred_hunflair = [result_hunflair['confidence'], 1-result_hunflair['confidence']]

        elif result_hunflair['value'] == '1':

            pred_hunflair = [1-result_hunflair['confidence'], result_hunflair['confidence']]

        # voting classifier

        weighted_average_1 = float((self.f1_score_biolink * pred_biolink[0] + self.f1_score_scibert * pred_scibert[0] + self.f1_score_hunflair * pred_hunflair[0]) / (self.f1_score_biolink + self.f1_score_scibert + self.f1_score_hunflair))

        weighted_average_2 = float((self.f1_score_biolink * pred_biolink[1] + self.f1_score_scibert * pred_scibert[1] + self.f1_score_hunflair * pred_hunflair[1]) / (self.f1_score_biolink + self.f1_score_scibert + self.f1_score_hunflair))

        return [weighted_average_1, weighted_average_2]

def extraction(filename: str, choices: list[bool], use_streamlit: bool = True):

    ''' Takes as input the name of a pdf file and extract the wanted entities from it 

        Outputs a dictionary with the entities' names as keys and the list of entities as values

    '''

    tokenizer_split_sentences = nltk.data.load('tokenizers/punkt/english.pickle')

    root = './NER-Medical-Document/processed_files/' + filename[:-4] + '/'

    results = {}

    models = {}

    limit = 20 # maximum number of files to process, used for testing if the pdf file is too big

    ind = 0

    # only add the necessary models to the dictionary 'models' (avoid unnecessary loading)

    model_url = 'https://storage.googleapis.com/allennlp-public-models/coref-spanbert-large-2021.03.10.tar.gz'

    predictor = Predictor.from_path(model_url)

    if choices[0]:

        tagger_chemicals: SequenceTagger = SequenceTagger.load('./NER-Medical-Document/training_results/best-model.pt')

        models[0] = tagger_chemicals

    if choices[1]:

        tagger_diseases: SequenceTagger = SequenceTagger.load('hunflair-disease')

        models[1] = tagger_diseases

    if choices[2]:

        tagger_dates = SequenceTagger.load("flair/ner-english-ontonotes-fast")

        models[2] = tagger_dates

    if choices[3]:

        # 3 different methods for ADE detection

        method = 2

        if method == 1:

            # Use a token classification model (classify a token as an ADE, not a sentence)

            model_adverse_name = "abhibisht89/spanbert-large-cased-finetuned-ade_corpus_v2" # model name from huggingface.co/models

            model_adverse = TFBertForTokenClassification.from_pretrained(model_adverse_name, from_pt=True)

            tokenizer_adverse = AutoTokenizer.from_pretrained(model_adverse_name)

            models[3] = pipeline("token-classification", model = model_adverse, tokenizer = tokenizer_adverse, grouped_entities=True)

        elif method == 2:

            # Sentence classification: use HunFlair model + negation detection

            tokenizer_neg = AutoTokenizer.from_pretrained("bvanaken/clinical-assertion-negation-bert")

            model_neg = AutoModelForSequenceClassification.from_pretrained("bvanaken/clinical-assertion-negation-bert")

            pipeline_neg = TextClassificationPipeline(model=model_neg, tokenizer=tokenizer_neg)

            model_hunflair = TextClassifier.load('./NER-Medical-Document/training_results/flair_bert/best-model.pt')

            models[3] = model_hunflair

        elif method == 3:

            # Sentence classification: use the InferenceADE class to design a voting classifier

            model_scibert_name = 'NER-Medical-Document/training_results/scibert_scivocab_uncased'

            model_scibert = BertForSequenceClassification.from_pretrained(model_scibert_name)

            tokenizer_scibert = AutoTokenizer.from_pretrained('allenai/scibert_scivocab_uncased')

            pipeline_scibert = pipeline("text-classification", model = model_scibert, tokenizer = tokenizer_scibert)

            model_biolink_name = 'NER_Medical-Document/training_results/BioLinkBERT-base'

            model_biolink = BertForSequenceClassification.from_pretrained(model_biolink_name)

            tokenizer_biolink = AutoTokenizer.from_pretrained('michiyasunaga/BioLinkBERT-base')

            pipeline_biolink = pipeline("text-classification", model = model_biolink, tokenizer = tokenizer_biolink)

            model_hunflair = TextClassifier.load('./NER-Medical-Document/training_results/flair_bert/best-model.pt')

            models[3] = InferenceADE(pipeline_scibert, pipeline_biolink, model_hunflair)

    if choices[4]:

        tagger_doses = SequenceTagger.load("flair/ner-english-ontonotes-fast")

        models[4] = tagger_doses

    dic = {0: 'Chemicals', 1: 'Diseases', 2: 'Dates', 3: 'Adverse effects', 4: 'Doses'}

    info = 'Extracting '

    for j, c in enumerate(choices):

        if c:

            info += dic[j].lower() + ' and '

    streamlit.write(info[:-5] + ' entities...')

    local_results_chemicals = []

    local_results_diseases = []

    local_results_dates = []

    local_results_adverse = []

    local_results_doses = []

    dic_doses_chemicals = {}

    for i, file in enumerate(os.listdir(root)):

        if file.endswith('.txt') and ind < limit:

            ind += 1

            with open(root+file, 'r') as f:

                paragraphs = f.read().split('\n\n')

                sentences = []

                for p in paragraphs:

                    sentences.extend(tokenizer_split_sentences.tokenize(p))

                    sentences[-1] = sentences[-1]+ '\n\n'                

                if not use_streamlit:

                    size_window = 3

                    list_coref = [sentences[i] for i in range(size_window)]

                    try:

                        prediction = predictor.predict(document=' '.join(list_coref))

                    except:

                        pass

                for s in stqdm(range(len(sentences))):

                    sentence_ = sentences[s]

                    ### Coreference resolution ###

                    if not use_streamlit:

                        if s >= size_window:

                            list_coref.append(sentence_)

                            list_coref.pop(0)

                            try:

                                prediction = predictor.predict(document=' '.join(list_coref))

                                transformed_chunk = predictor.coref_resolved(' '.join(list_coref))

                                paragraphs2 = transformed_chunk.split('\n\n')

                                for p in paragraphs2:

                                    sentences2 = tokenizer_split_sentences.tokenize(p)

                                sentence_transformed = sentences2[-1]

                            except:

                                sentence_transformed = sentence_

                                pass

                        else:

                            sentence_transformed = sentence_

                    sentence_ = sentence_.replace('\n', ' ')

                    if len(sentence_) >= 4:

                        print(sentence_)

                        sentence = Sentence(sentence_, use_tokenizer=SciSpacyTokenizer())

                        for j, c in enumerate(choices):

                            if c:

                                tagger = models[j]

                                if dic[j] == 'Adverse effects':

                                    if method == 1:

                                        result = tagger(sentence_)

                                        if result != []:

                                            for entity in tagger(sentence_):

                                                if entity['entity_group'] == 'ADR':

                                                    local_results_dates.append(entity['word'])

                                    elif method == 2:

                                        sentence = Sentence(sentence_, use_tokenizer=SciSpacyTokenizer())  # create new instance of Sentence

                                        tagger.predict(sentence)

                                        result = sentence.labels[0].to_dict()

                                        if result['value'] == '1':

                                            if not pipeline_neg(sentence_)[0]['label'] == 'ABSENT':

                                                print('DETECTED')

                                                local_results_adverse.append(sentence_)

                                    elif method == 3:

                                        result = tagger(sentence_)

                                        print(result)

                                        if result[1] > 0.5:

                                            local_results_adverse.append(sentence_)

                                        else:

                                            models[0].predict(sentence)

                                            found = False

                                            for annotation_layer in sentence.annotation_layers.keys():

                                                for entity in sentence.get_spans(annotation_layer):

                                                        found = True

                                                        sentence_2 = sentence_.replace(entity.text, 'aspirin')

                                            if found: 

                                                result = tagger(sentence_2)

                                                print(result)

                                                if result[1] > 0.5:

                                                    local_results_adverse.append(sentence_)

                                else: 

                                    tagger.predict(sentence)

                                    for annotation_layer in sentence.annotation_layers.keys():

                                        for entity in sentence.get_spans(annotation_layer):

                                            if dic[j] == 'Chemicals':

                                                local_results_chemicals.append(entity.text)

                                                detected_chemicals = True

                                                entity_chemical = entity.text

                                            elif dic[j] == 'Diseases':

                                                local_results_diseases.append(entity.text)

                                            elif dic[j] == 'Dates':

                                                if entity.tag == 'DATE':

                                                    local_results_dates.append(entity.text)

                                            elif dic[j] == 'Doses':

                                                if entity.tag == 'QUANTITY':

                                                    local_results_doses.append(entity.text)

                                                    if detected_chemicals:

                                                        print('YES')

                                                        print(detected_chemicals)

                                                        print(sentence)

                                                        dic_doses_chemicals[entity.text] = entity_chemical

                                                    else:

                                                        dic_doses_chemicals[entity.text] = 'unknown'

                        detected_chemicals = False

    for j, c in enumerate(choices):

        if c:

            if dic[j] == 'Chemicals':

                # next line is to avoid detecting some characters as 'drugs' (happened sometimes)

                local_results_chemicals = [x for x in local_results_chemicals if x not in ['(', ')', '[', ']', '{', '}', ' ', '']]

                results[dic[j]] = list(set(local_results_chemicals))

            elif dic[j] == 'Diseases':

                results[dic[j]] = list(set(local_results_diseases))

            elif dic[j] == 'Dates':

                results[dic[j]] = list(set(local_results_dates))

            elif dic[j] == 'Adverse effects':

                results[dic[j]] = list(set(local_results_adverse))

            elif dic[j] == 'Doses':

                results[dic[j]] = list(set(local_results_doses))

    streamlit.write('Done!')

    return results

def higlight(filename: str, choices: list[bool]):

    ''' Highlight the entities chosen in the pdf file whose name is 'filename' '''

    root = './NER-Medical-Document/processed_files/' + filename[:-4] + '/'

    pdf_input = fitz.open(root+filename)

    results = extraction(filename, choices)

    streamlit.write('Highlighting entities...')

    text_instances = {}

    # go through all the pages of the chosen pdf file

    for page in pdf_input:

        # search all the occurences of the entities in the page

        for name, entities in results.items():

            text_instances[name] = [page.search_for(text) for text in entities]

        for name, instances in text_instances.items():

            add_definition = False

            if name == 'Chemicals':

                color = (1, 1, 0)

                add_definition = True

            elif name == 'Diseases':

                color = (0, 1, 0)

                add_definition = True

            elif name == 'Dates':

                color = (0, 0.7, 1)

                add_definition = False

            elif name == 'Adverse effects':

                color = (1, 0, 0)

                add_definition = False

            elif name == 'Doses':

                color = (1, 0, 1)

                add_definition = False

            # highlight each occurence of the entity in the page

            for i, inst in enumerate(instances):

                for x in inst:

                    # handle the case where an entity should not be highlighted (see README): the idea is too check the surrounding characters to 

                    # detect if the occurence of the entity is part of another word or not

                    # check the typical distance between 2 letters in the word (because it depends on the font size)

                    dist_letters = (x[2]-x[0])/len(results[name][i])

                    # draw a larger rectangle to check the surrounding characters

                    rect_larger = fitz.Rect(x[0]-dist_letters, x[1], x[2]+dist_letters, x[3])

                    non_accepted_chars = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'

                    word = page.get_textbox(rect_larger).lower()

                    for sub in results[name][i].split():

                        word = word.replace(sub.lower(), '')

                    keep = True

                    for l in word:

                        if l in non_accepted_chars:

                            keep = False

                    if not keep:

                        continue     # ignore the occurence of the entity if it is part of another word

                    annot = page.add_highlight_annot(x)

                    annot.set_colors({"stroke": color})

                    annot.set_opacity(0.4)

                    if add_definition:

                        try:

                            annot.set_popup(x)

                            info = annot.info

                            info["title"] = "Definition"

                            if name == 'Chemicals':

                                info["content"] = wikipedia.summary(results[name][i] +  ' (drug)').split('.')[0]

                            else:

                                info["content"] = wikipedia.summary(results[name][i] + f' ({name.lower()[:-1]})').split('.')[0]

                            annot.set_info(info)

                        except:

                            pass

                    annot.update()

    if os.path.exists(root+filename[:-4]+'_highlighted.pdf'):

            os.remove(root+filename[:-4]+'_highlighted.pdf')

    pdf_input.save(root+filename[:-4]+'_highlighted.pdf')

    streamlit.write('Done!')

if __name__ == '__main__':

    print(extraction('0.txt', [True, True, True, True, True], use_streamlit=False))