from annotations import Entity, Relation

from typing import List, Dict, Union, Tuple, Callable, Optional

import warnings

import numpy

class HealthRecord:

    """

    Objects that represent a single electronic health record

    """

    def __init__(self, record_id: str = "1", text_path: Optional[str] = None,

                 ann_path: Optional[str] = None,

                 text: Optional[str] = None,

                 tokenizer: Callable[[str], List[str]] = None,

                 is_bert_tokenizer: bool = True,

                 is_training: bool = True) -> None:

        """

        Initializes a health record object

        Parameters

        ----------

        record_id : int

            A unique ID for the record.

        text_path : str

            Path for the ehr record txt file.

        ann_path : str, optional

            Path for the annotation file. The default is None.

        text: str

            If text_path is not specified, the actual text for the

            record

        tokenizer: Callable[[str], List[str]], optional

            The tokenizer function to use. The default is None.

        is_bert_tokenizer: bool

            If the tokenizer is a BERT-based wordpiece tokenizer.

            The default is False.

        is_training : bool, optional

            Specifies if the record is a training example.

            The default is True.

        """

        if is_training and ann_path is None:

            raise AttributeError("Annotation path needs to be "

                                 "specified for training example.")

        if text_path is None and text is None:

            raise AttributeError("Either text or text path must be "

                                 "specified.")

        self.record_id = record_id

        self.is_training = is_training

        if text_path is not None:

            self.text = self._read_ehr(text_path)

        else:

            self.text = text

        self.char_to_token_map: List[int] = []

        self.token_to_char_map: List[int] = []

        self.tokenizer = None

        self.is_bert_tokenizer = is_bert_tokenizer

        self.elmo = None

        self.set_tokenizer(tokenizer)

        self.split_idx = None

        if ann_path is not None:

            annotations = self._extract_annotations(ann_path)

            self.entities, self.relations = annotations

        else:

            self.entities = None

            self.relations = None

    @staticmethod

    def _read_ehr(path: str) -> str:

        """

        Internal function to read EHR data.

        Parameters

        ----------

        path : str

            Path for EHR record.

        Returns

        -------

        str

            EHR record as a string.

        """

        f = open(path)

        raw_data = f.read()

        f.close()

        return raw_data

    @staticmethod

    def _extract_annotations(path: str) \

            -> Tuple[Dict[str, Entity], Dict[str, Relation]]:

        """

        Internal function that extracts entities and relations

        as a dictionary from an annotation file.

        Parameters

        ----------

        path : str

            Path for the ann file.

        Returns

        -------

        Tuple[Dict[str, Entity], Dict[str, Relation]]

            Entities and relations.

        """

        f = open(path)

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

        f.close()

        entities = {}

        relations = {}

        # Relations with entities that haven't been processed yet

        relation_backlog = []

        for line in raw_data:

            if line.startswith('#'):

                continue

            line = line.split('\t')

            # Remove empty strings from list

            line = list(filter(None, line))

            if not line or not line[0]:

                continue

            if line[0][0] == 'T':

                assert len(line) == 3

                idx = 0

                # Find the end of first word, which is the entity type

                for idx in range(len(line[1])):

                    if line[1][idx] == ' ':

                        break

                char_ranges = line[1][idx + 1:]

                # Get all character ranges, separated by ;

                char_ranges = [r.split() for r in char_ranges.split(';')]

                # Create an Entity object

                ent = Entity(entity_id=line[0],

                             entity_type=line[1][:idx])

                r = [char_ranges[0][0], char_ranges[-1][1]]

                r = list(map(int, r))

                ent.set_range(r)

                ent.set_text(line[2])

                entities[line[0]] = ent

            elif line[0][0] == 'R':

                assert len(line) == 2

                rel_details = line[1].split(' ')

                entity1 = rel_details[1].split(':')[-1]

                entity2 = rel_details[2].split(':')[-1]

                if entity1 in entities and entity2 in entities:

                    rel = Relation(relation_id=line[0],

                                   relation_type=rel_details[0],

                                   arg1=entities[entity1],

                                   arg2=entities[entity2])

                    relations[line[0]] = rel

                else:

                    # If the entities aren't processed yet, 

                    # add them to backlog to process later

                    relation_backlog.append([line[0], rel_details[0],

                                             entity1, entity2])

            else:

                # If the annotation is not a relation or entity, warn user

                msg = f"Invalid annotation encountered: {line}, File: {path}"

                warnings.warn(msg)

        for r in relation_backlog:

            rel = Relation(relation_id=r[0], relation_type=r[1],

                           arg1=entities[r[2]], arg2=entities[r[3]])

            relations[r[0]] = rel

        return entities, relations

    def _compute_tokens(self) -> None:

        """

        Computes the tokens and character <-> token index mappings

        for EHR text data.

        """

        self.tokens = list(map(lambda x: str(x), self.tokenizer(self.text)))

        char_to_token_map = []

        token_to_char_map = []

        j = 0

        k = 0

        for i in range(len(self.tokens)):

            # For BioBERT, a split within a word is denoted by ##

            if self.is_bert_tokenizer and self.tokens[i].startswith("##"):

                k += 2

            # Characters that are discarded from tokenization

            while self.text[j].lower() != self.tokens[i][k].lower():

                char_to_token_map.append(char_to_token_map[-1])

                j += 1

            # For SciSpacy, if there are multiple spaces, it removes

            # one and keeps the rest

            if self.text[j] == ' ' and self.text[j + 1] == ' ':

                char_to_token_map.append(char_to_token_map[-1])

                j += 1

            token_start_idx = j

            # Go over each letter in token and original text

            while k < len(self.tokens[i]):

                if self.text[j].lower() == self.tokens[i][k].lower():

                    char_to_token_map.append(i)

                    j += 1

                    k += 1

                else:

                    msg = f"Error computing token to char map. ID: {self.record_id}"

                    raise Exception(msg)

            token_end_idx = j

            token_to_char_map.append((token_start_idx, token_end_idx))

            k = 0

        # Characters at the end which are discarded by tokenizer

        while j < len(self.text):

            char_to_token_map.append(char_to_token_map[-1])

            j += 1

        assert len(char_to_token_map) == len(self.text)

        assert len(token_to_char_map) == len(self.tokens)

        self.char_to_token_map = char_to_token_map

        self.token_to_char_map = token_to_char_map

    def get_tokens(self) -> List[str]:

        """

        Returns the tokens.

        Returns

        -------

        List[str]

            List of tokens.

        """

        if self.tokenizer is None:

            raise AttributeError("Tokenizer not set.")

        return self.tokens

    def set_tokenizer(self, tokenizer: Callable[[str], List[str]]) \

            -> None:

        """

        Set the tokenizer for the object.

        Parameters

        ----------

        tokenizer : Callable[[str], List[str]]

            The tokenizer function to use.

        """

        self.tokenizer = tokenizer

        if tokenizer is not None:

            self._compute_tokens()

    def get_token_idx(self, char_idx: int) -> int:

        """

        Returns the token index from character index.

        Parameters

        ----------

        char_idx : int

            Character index.

        Returns

        -------

        int

            Token index.

        """

        if self.tokenizer is None:

            raise AttributeError("Tokenizer not set.")

        token_idx = self.char_to_token_map[char_idx]

        return token_idx

    def get_char_idx(self, token_idx: int) -> int:

        """

        Returns the index for the first character of the specified

        token index.

        Parameters

        ----------

        token_idx : int

            Token index.

        Returns

        -------

        int

            Character index.

        """

        if self.tokenizer is None:

            raise AttributeError("Tokenizer not set.")

        char_idx = self.token_to_char_map[token_idx]

        return char_idx

    def get_labels(self) -> List[str]:

        """

        Get token labels in IOB format.

        Returns

        -------

        List[str]

            Labels.

        """

        if self.tokenizer is None:

            raise AttributeError("No tokens found. Set tokenizer first.")

        ent_label_map = {'Drug': 'DRUG', 'Strength': 'STR', 'Duration': 'DUR',

                         'Route': 'ROU', 'Form': 'FOR', 'ADE': 'ADE', 'Dosage': 'DOS',

                         'Reason': 'REA', 'Frequency': 'FRE'}

        labels = ['O'] * len(self.tokens)

        for ent in self.entities.values():

            start_idx = self.get_token_idx(ent.range[0])

            end_idx = self.get_token_idx(ent.range[1])

            for idx in range(start_idx, end_idx + 1):

                if idx == start_idx:

                    labels[idx] = 'B-' + ent_label_map[ent.name]

                else:

                    labels[idx] = 'I-' + ent_label_map[ent.name]

        return labels

    def get_split_points(self, max_len: int = 510,

                         new_line_ind: List[str] = None,

                         sent_end_ind: List[str] = None) -> List[int]:

        """

        Get the splitting points for tokens.

        > It includes as many paragraphs as it can within the

        max_len - 2 token limit. (2 less because BERT needs

                                  to add 2 special tokens)

        > If it can't find a single complete paragraph,

        it will split on the last verifiable new line that

        starts with a new sentence.

        > If it can't find that as well, it splits on token max_len - 2.

        Parameters

        ----------

        max_len : int, optional

            Maximum number tokens in one example. The default is 510

            for BERT.

        new_line_ind : List[str], optional

            New line indicators. Strings other than numbers.

            The default is ['[', '#', '-', '>', ' '].

        sent_end_ind : List[str], optional

            Sentence end indicators. The default is ['.', '?', '!'].

        Returns

        -------

        List[int]

            Splitting indices, includes the first and last index.

            Need to add 1 to the end indices if accessing

            with list splicing.

        """

        if new_line_ind is None:

            new_line_ind = ['[', '#', '-', '>', ' ']

        if sent_end_ind is None:

            sent_end_ind = ['.', '?', '!']

        split_idx = [0]

        last_par_end_idx = 0

        last_line_end_idx = 0

        for i in range(len(self.text)):

            curr_counter = self.get_token_idx(i) - split_idx[-1]

            if curr_counter >= max_len:

                # If not even a single paragraph has ended

                if last_par_end_idx == 0 and last_line_end_idx != 0:

                    split_idx.append(last_line_end_idx)

                elif last_par_end_idx != 0:

                    split_idx.append(last_par_end_idx)

                else:

                    split_idx.append(self.get_token_idx(i))

                last_par_end_idx = 0

                last_line_end_idx = 0

            if i < len(self.text) - 2 and self.text[i] == '\n':

                if self.text[i + 1] == '\n':

                    last_par_end_idx = self.get_token_idx(i - 1)

                if self.text[i + 1] == '.' or self.text[i + 1] == '*':

                    last_par_end_idx = self.get_token_idx(i + 1)

                if self.text[i + 1] in new_line_ind or \

                        self.text[i + 1].isdigit() or \

                        self.text[i - 1] in sent_end_ind:

                    last_line_end_idx = self.get_token_idx(i)

        split_idx.append(len(self.tokens))

        self.split_idx = split_idx

        return self.split_idx

    def get_annotations(self) -> Dict[str, Union[list, dict]]:

        """

        Get entities and relations in a dictionary.

        Entities are referenced with the key 'entities'

        and relations with 'relations'

        Returns

        -------

        Dict[Dict[str, Entity], Dict[str, Relation]]

            Entities and relations.

        """

        if self.entities is None or self.relations is None:

            raise AttributeError("Annotations not available")

        return {'entities': self.entities, 'relations': self.relations}

    def get_entities(self) -> Dict[str, Entity]:

        """

        Get the entities.

        Returns

        -------

        Dict[str, Entity]

            Entity ID: Entity object.

        """

        if self.entities is None:

            raise AttributeError("Entities not set")

        return self.entities

    def get_relations(self) -> Dict[str, Relation]:

        """

        Get the entity relations.

        Returns

        -------

        Dict[str, Relation]

            Relation ID: Relation Object.

        """

        if self.relations is None:

            raise AttributeError("Relations not set")

        return self.relations

    def _compute_elmo_embeddings(self) -> None:

        """

        Computes the Elmo embeddings for each token in EHR text data.

        """

        # noinspection PyUnresolvedReferences

        elmo_embeddings = self.elmo.embed_sentence(self.tokens)[-1]

        self.elmo_embeddings = elmo_embeddings

    def set_elmo_embedder(self, elmo: Callable[[str], numpy.ndarray]) -> None:

        """

        Set Elmo embedder for object.

        Parameters

        ----------

        elmo :

            The Elmo embedder to use.

        """

        self.elmo = elmo

        if elmo is not None:

            self._compute_elmo_embeddings()

    def get_elmo_embeddings(self) -> numpy.ndarray:

        """

        Get the elmo embeddings.

        Returns

        -------

        List[int]:

            Elmo embeddings for each word

        """

        if self.elmo_embeddings is None:

            raise AttributeError("Elmo embeddings not set")

        return self.elmo_embeddings