import spacy
from scispacy.abbreviation import AbbreviationDetector
from scispacy.hyponym_detector import HyponymDetector
from scispacy.linking import EntityLinker
from negspacy.negation import Negex

def get_abbreviations(model, text):
    """
    returns a list of tuples in the form (abbreviation, expanded form), each element being a str
    """

    # logging
    print(f"Identifying abbrevations using {model}")
    partial_input = '\n'.join(text.split('\n')[:5])
    print(f"Input text (truncated): {partial_input}\n...")

    # abbreviation detection with scispacy
    nlp = spacy.load(model)
    nlp.add_pipe("abbreviation_detector")
    doc = nlp(text)
    abbreviations = [(abrv.text, abrv._.long_form.text) for abrv in doc._.abbreviations]

    return abbreviations

def get_hyponyms(model, text):
    """
    returns a list of tuples in the form (hearst_pattern, entity_1, entity_2, ...), each element being a str
    """

    # logging
    print(f"Extracting hyponyms using {model}")
    partial_input = '\n'.join(text.split('\n')[:5])
    print(f"Input text (truncated): {partial_input}\n...")

    # hyponym detection with scispacy
    nlp = spacy.load(model)
    nlp.add_pipe("hyponym_detector", last=True, config={"extended": True})
    doc = nlp(text)
    hearst_patterns = [tuple([str(element) for element in pattern]) for pattern in doc._.hearst_patterns]

    return hearst_patterns

def get_linked_entities(model, text):
    """
    returns a dictionary in the form {named entity: list of strings each describing one piece of linked information}
    """

    # logging
    print(f"Entity linking using {model}")
    partial_input = '\n'.join(text.split('\n')[:5])
    print(f"Input text (truncated): {partial_input}\n...")

    # entity linking with scispacy
    output = {}

    nlp = spacy.load(model)
    nlp.add_pipe("scispacy_linker", config={"resolve_abbreviations": True, "linker_name": "umls"})
    doc = nlp(text)

    ents = doc.ents
    linker = nlp.get_pipe("scispacy_linker")

    for entity in ents:
        cur = []
        for umls_ent in entity._.kb_ents:
            cur.append(str(linker.kb.cui_to_entity[umls_ent[0]]))
        output[entity] = cur

    return output

def get_named_entities(model, text):
    """
    returns a list of strings, each string is an identified named entity
    """

    # logging
    print(f"Extracting named entities using {model}")
    partial_input = '\n'.join(text.split('\n')[:5])
    print(f"Input text (truncated): {partial_input}\n...")

    # named recognition with scispacy
    nlp = spacy.load(model)
    doc = nlp(text)
    named_entities = [str(ent) for ent in doc.ents]

    return named_entities


def get_negation_entities(model, text):
    """
    returns a list of pairs, default model is "en_core_web_sm"
    Negspacy is a spaCy pipeline component that evaluates whether Named Entities are negated in text.
    Example:
    >> test = get_negation_entities("en_core_web_sm","She does not like Steve Jobs but likes Apple products.")
    >> print (test)
    [(True, 'Steve Jobs'), (False, 'Apple')]
    """

    # logging
    print(f"Extracting whether Named Entities are negated using {model}")
    partial_input = '\n'.join(text.split('\n')[:5])
    print(f"Input text (truncated): {partial_input}\n...")

    # named recognition with scispacy
    nlp = spacy.load(model)
    nlp.add_pipe("negex", config={"ent_types":["PERSON","ORG","NORP","GPE"]})
    doc = nlp(text)
    pairs = [(ent._.negex,ent.text) for ent in doc.ents]

    return pairs