"""
| **Abbreviation** | **Description** |
|------------------|-----------------|
| O                | Outside of a named entity
| B-MIS            | Beginning of a miscellaneous entity right after another miscellaneous entity
| I-MIS            | Miscellaneous entity
| B-PER            | Beginning of a person’s name right after another person’s name
| I-PER            | Person’s name
| B-ORG            | Beginning of an organization right after another organization
| I-ORG            | Organization
| B-LOC            | Beginning of a location right after another location
| I-LOC            | Location
"""

from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline

from enum import Enum


class DictKey(Enum):
    ENTITY = 'entity'
    SCORE = 'score'
    INDEX = 'index'
    WORD = 'word'
    START = 'start'
    END = 'end'


class NER:
    def __init__(self, text_to_analyse):
        """
        The Constructor for the Named Entity Recognition class.
        :param text_to_analyse: The text in which to find named entities.
        """
        self.tokenizer = AutoTokenizer.from_pretrained("dslim/bert-base-NER")

        self.model = AutoModelForTokenClassification.from_pretrained("dslim/bert-base-NER")

        self.nlp = pipeline("ner", model=self.model, tokenizer=self.tokenizer, grouped_entities=True)
        if self.nlp is None:
            raise ValueError("Unable to load pipeline from DSLIM BERT model")

        self.text_to_analyse = text_to_analyse
        self.results = self.nlp(text_to_analyse)
        self.all_entities = self.get_list_of_entities()
        self.unique_entities = self.unique_entities()
        self.markdown = None
        self.markdown_text = None

    def get_entity_value(self, key: DictKey, item_index):
        """
        Extracts the value for a specific key (as an Enum) from a specific dictionary item in the list.
        :param key: DictKey Enum representing the key for which the value is required.
        :param item_index: Index of the item in the list to process.
        :return: Value for the given key in the specified dictionary item, or None if key is not found.
        """
        if item_index < len(self.results):
            return self.results[item_index].get(key.value)
        else:
            raise ValueError("The supplied list index is out of bounds")

    def get_list_of_entities(self):
        """
        Returns a list of all entities in the original text, in the order they appear. There may be repeated
        entities in this list.
        :return: A list of all entities in the original text.
        """
        # create a list where each item is the value of word from each of the dictionaries in self.results
        return [item.get(DictKey.WORD.value) for item in self.results]

    def entity_markdown(self):
        """
        Convert a string to markdown format and change the color of specified substrings to red.
        """
        self.markdown = self.text_to_analyse

        for substring in self.get_list_of_entities():
            self.markdown = self.markdown.replace(substring, f'<span style = "color:red;">{substring}</span>')

        self.markdown_text = self.markdown.replace('\n', '  \n')  # Two spaces at the end of line for markdown new line

    def unique_entities(self):
        """
        Return a list of all unique entities in the original text.
        :return: A list of unique entities.
        """
        unique_set = set()  # Sets are faster than lists for checking membership

        # Create a new list to store the unique strings in order
        unique_list = []

        for string in self.all_entities:
            if string not in unique_set:
                unique_set.add(string)
                unique_list.append(string)

        return unique_list