"""Contains several classes for parsing text against the UMLS"""

from umls_tools import errors as err, objects
from warnings import warn

# Used for printline debugging
import pprint

# Used to set the environmat variables
import os
import time
import re
import itertools
import math
import subprocess


warn(
    "The parse module is deprecated as the code is very old"
    ", this may be updated in the future",
    DeprecationWarning,
    2,
)

java_path = "/usr/bin/java"
os.environ["JAVAHOME"] = java_path

pp = pprint.PrettyPrinter()


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class RunGeniaTagger(object):
    """A wrapper around the geniatagger command line program. This replaces
    the geniatagger.python module which is broken.

    Parameters
    ----------
    path_to_tagger : `str`
        The full path to the geniatagger command line program.
    """

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, path_to_tagger):
        self._path_to_tagger = path_to_tagger
        self._dir_to_tagger = os.path.dirname(path_to_tagger)
        self._tagger = subprocess.Popen(
            "./" + os.path.basename(path_to_tagger),
            cwd=self._dir_to_tagger,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
        )

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def parse(self, text):
        """Parse some text with the geniatagger.

        Parameters
        ----------
        text : `str`
            The text to parse.

        Returns
        -------
        results : `list` of `tuple`
            The geniatagger results.
        """
        results = list()

        for oneline in text.split("\n"):
            self._tagger.stdin.write(oneline + "\n")
            while True:
                r = self._tagger.stdout.readline()[:-1]
                if not r:
                    break
                results.append(tuple(r.split("\t")))
        return results


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class GeniaTagger(object):
    """A wrapper for the geniatagger, this initialises the tagger and parses
    text into a phrase object.
    """

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self):
        """
        Initialise
        """
        self.tagger = RunGeniaTagger(os.environ["GENIATAGGER"])

        # This will display the raw text of the tast text string to be parsed
        self._last_text = None

        # Let the genia tagger load
        time.sleep(7)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def parse(self, text):
        """
        Parse some text with the GeniaTagger

        Parameters
        ----------

        text : :obj:`str`
            The text to parse with the genia tagger
        """

        # Set the last text value
        self._last_text = text

        # Parse the text
        phrase = self._parse(text)
        return phrase
        # pp.pprint(phrase)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _parse(self, text):
        """
        The private method actually does the parsing

        Parameters
        ----------

        text : :obj:`str`
            The text to parse with the genia tagger
        """

        chunks = []
        # Parse the phrase with the GENIA tagger
        for c, i in enumerate(self.tagger.parse(text)):
            # print(i)
            # Create a token
            token = Token(i[0], i[1], i[2], i[4])

            # This tells us if we are at the start of a chunk
            chunk_spec = i[3]

            # B=NP, B-PP are the start of chunks
            if chunk_spec.startswith("B-"):
                chunks.append(Chunk(token, chunk_spec=chunk_spec, start=c))
            elif chunk_spec.startswith("I-"):
                # If the token is "within" the chunk then add it to the
                # last chunk. Note that we assume that the tokens are in order
                chunks[-1].add_token(token)
            else:
                # Anything else we will just treat as a new chunk
                chunks.append(Chunk(token, chunk_spec=chunk_spec, start=c))

        return Phrase(*chunks)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def last_text(self):
        return self._last_text


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
# Classes representing the text components                                    @
# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class Phrase(object):
    def __init__(self, *chunks):
        self._chunks = []
        self.add_chunks(*chunks)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __str__(self):
        """
        Called when str(object) is called
        """
        line = []

        # Build all the strings from the tokens and then join with spaces. Note
        # that this is NOT a reverse of tokenisation and will NOT be exact
        for i in self._chunks:
            line.append(str(i))

        return " ".join(line)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __repr__(self):
        """
        Called when Chunk is PrettyPrinted
        """
        return self.report()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __eq__(self, other):
        """
        Determine equality in the phrase. We compare start, end and text.

        Parameters
        ----------
        other : `object`
            Object to compare with this phrase.

        Returns
        -------
        result : `bool` or `NotImplementedType`
            ``True`` when start/end/text all match, otherwise ``False``.
            Returns ``NotImplemented`` for non-``Phrase`` objects.
        """
        if not isinstance(other, Phrase):
            return NotImplemented

        if (
            self.start == other.start
            and self.end == other.end
            and str(self) == str(other)
        ):
            return True

        return False

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __getitem__(self, index):
        """
        make an iterator
        """
        return self._chunks[index]

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __len__(self):
        """
        The length of the chunk is the text length
        """
        return len(str(self))

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def start(self):
        """
        Return the start position of the chunk (i.e. the position of the first
        token)
        """

        try:
            return self._chunks[0].start
        except IndexError:
            return None

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def end(self):
        """
        Return the start position of the chunk (i.e. the position of the first
        token)
        """
        try:
            return self._chunks[-1].end
        except IndexError:
            return None

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def report(self, level=0):
        """
        print out everything nicely, useful for debugging
        """

        line = [
            "%sPhrase:%i|%i|%i: %s"
            % ("\t" * level, self.start, self.end, len(self), str(self))
        ]
        line.append(self._report_chunks((level + 1)))

        return "\n".join(line)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def add_chunks(self, *chunks):
        for i in chunks:
            self._add_chunk(i)

        self._sort_chunks()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def chunks(self):
        """
        generator for the chunks
        """
        for i in self._chunks:
            yield i

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def tokens(self):
        """
        generator for the chunks
        """
        for c in self._chunks:
            for t in c.tokens():
                yield t

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _add_chunk(self, chunk):
        if isinstance(chunk, Chunk) is False:
            raise err.GeniaParserError(
                "only instances of Chunk can be added to %s"
                % self.__class__.__name__
            )

        for i in self._chunks:
            if i.overlap(chunk) > 0:
                raise err.GeniaParserError(
                    "incoming chunk overlaps present chunk"
                )
        self._chunks.append(chunk)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _sort_chunks(self):
        """
        Perform a sort on the chunks to make sure they are in order
        """
        self._chunks.sort(key=lambda x: x.start)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _report_chunks(self, level):
        """
        Report all the tokens in the chunk

        Returns
        -------

        str
            A line joined on "\n"
        """

        # Build a line and then join it when built
        line = []
        for i in self._chunks:
            # We get the tokens to tab indent 1 tab
            line.append(i.report(level=level))

        return "\n".join(line)


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class Chunk(object):
    """
    A class representing a chunk partial parse of a sentence
    """

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, *args, **kwargs):
        """
        tokens: The tokens that make up the chunk
        """

        # Store the **kwargs
        # TODO: Clean up the chunk spec (remove B-/I-)
        self._chunk_spec = self._parse_chunk_spec(kwargs.pop("chunk_spec", "O"))
        self._start = kwargs.pop("start", 0)
        self.report_ngrams = kwargs.pop("report_ngrams", False)

        # Store the token objects
        self._tokens = []

        # Add all the tokens to the chunk
        for i in args:
            self.add_token(i)

        self.ngrams = []

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __str__(self):
        """
        Called when str(object) is called
        """
        line = []

        # Build all the strings from the tokens and then join with spaces. Note
        # that this is NOT a reverse of tokenisation and will NOT be exact
        for i in self._tokens:
            line.append(str(i))

        return " ".join(line)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __repr__(self):
        """
        Called when Chunk is PrettyPrinted
        """
        return self.report()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __eq__(self, other):
        """
        Determine equality in the chunk. We compare start, end and text.

        Parameters
        ----------
        other : `object`
            Object to compare with this chunk.

        Returns
        -------
        result : `bool` or `NotImplementedType`
            ``True`` when start/end/text all match, otherwise ``False``.
            Returns ``NotImplemented`` for non-``Chunk`` objects.
        """
        if not isinstance(other, Chunk):
            return NotImplemented

        if (
            self.start == other.start
            and self.end == other.end
            and str(self) == str(other)
        ):
            return True

        return False

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __getitem__(self, index):
        """
        make an iterator
        """
        return self._tokens[index]

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __len__(self):
        """
        The length of the chunk is the text length
        """
        return len(str(self))

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def tokens(self):
        """
        generator for the chunks
        """
        for t in self._tokens:
            yield t

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def ntokens(self):
        return len(self._tokens)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def start(self):
        """
        Return the start position of the chunk (i.e. the position of the first
        token)
        """
        return self._start
        # return self._tokens[0].phrase_position

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def end(self):
        """
        Return the start position of the chunk (i.e. the position of the first
        token)
        """
        return self._start + (len(self._tokens) - 1)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def chunk_spec(self):
        """
        Return what the chunk represents, i.e. is it a noun phrase verb phrase
        etc
        """
        return self._chunk_spec

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def is_symbol(self):
        """Is the chunk comprised of all symbols i.e. ``\\w``."""
        return all([i.is_symbol for i in self._tokens])

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def overlap(self, other):
        """
        test to see if the coordinates of 1 chunk overlap the coordinates of
        other chunk

        Parameters
        ----------

        other : :obj:`Chunk`
            A chunk to test for overlapping coordinates

        Returns
        -------

        int
            The length of the overlap
        """

        return max(0, min(self.end, other.end) - max(self.start, other.start))

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def report(self, level=0):
        """
        print out everything nicely, useful for debugging
        """

        line = [
            "%sChunk:%i|%i|%i|%s: %s"
            % (
                "\t" * level,
                self.start,
                self.end,
                len(self),
                self._chunk_spec,
                str(self),
            )
        ]
        line.append(self._report_tokens((level + 1)))

        return "\n".join(line)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def summary(self):
        """
        Returns all the toke summaries in the form (NE text/POS)
        """
        line = []

        for i in self._tokens:
            line.append(i.summary())

        return " ".join(line)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def as_tuple(self):
        """
        Get the chunk and token infomation as a tuple of tuples

        Returns
        -------

        Returns a tuple with all the token info in it
        """
        line = []

        for i in self._tokens:
            line.append(i.as_tuple())

        return tuple(line)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def add_token(self, token):
        """
        Add a token to the chunk
        """

        # If the token is already in the chunk then we do not continue as it
        # is probably a callback
        if token in self._tokens:
            return

        # Otherwise add the token
        self._tokens.append(token)

        # Do a call back to make sure the token "knows" it has been added to a
        # chunk
        token.add_chunk(self)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _report_tokens(self, level):
        """
        Report all the tokens in the chunk

        Returns
        -------

        str
            A line joined on "\n"
        """

        # Build a line and then join it when built
        line = []
        for i in self._tokens:
            # We get the tokens to tab indent 1 tab
            line.append(i.report(level=level))

        return "\n".join(line)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _parse_chunk_spec(self, chunk_spec):
        """
        Cleans up the chunk spec by removing I- or B-

        Parameters
        ----------

        chunk_spec : :obj:`str`
            the chunk spec to clean up

        Returns
        -------

        str
            The cleaned chunk spec
        """

        return re.sub(
            r"[BI]-([A-Z]+)",
            r"\1",
            chunk_spec,
            flags=re.IGNORECASE,
        )

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def make_ngram(self, n):
        """
        Make n-grams of n size
        """

        if n > len(self._tokens):
            raise IndexError("n > len(tokens)")

        i = 0
        while i + n <= len(self._tokens):
            yield self._init_chunk(
                self._tokens[i : i + n],
                start=self.start + i,
                chunk_spec="NGRAM{0}_{1}".format(n, self.chunk_spec),
            )
            i += 1

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def all_ngrams(self):
        """Create chunks consiting of all possible ngrams of the chunk"""
        for ngram_size in range(len(self._tokens) - 1, 0, -1):
            for n in self.make_ngram(ngram_size):
                yield n

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _init_chunk(self, token_list, **kwargs):
        """
        Initialise the chunk, used when creating ngrams
        """
        return self.__class__(*token_list, **kwargs)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    # def _get_ngram_str(self):
    #     """Return a report string detailing the length n-gram of the chunk
    #     and the string of the chunk.
    #     """
    #     return "%i-gram: %s | " % (len(self.tokens),str(self))

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    # def _report_ngrams(self):
    #     """Return report strings for ngrams
    #     """
    #     # Build a line and then join it when built
    #     line = []

    #     # In addition to the tokens displayed by the chunk display all the
    #     # CUI's and cui's of n-grams
    #     line.append(self._get_ngram_str())

    #     # Now loop through any n-gram strings that the CUI has
    #     for k in self.ngrams:
    #         line.append(k._report_ngrams())

    #     return "\n".join(line)


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class Token(object):
    """
    A class representing a token of text
    """

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, text, stem, part_of_speech, entity, parent_chunk=None):
        """
        A token of text that has been parsed with the GeniaTagger

        Parameters
        ----------

        text : :obj:`str`
            The text that forms the token
        stem : :obj:`str`
            A stemmed form of the text
        pos : :obj:`str`
            Part-of-speech tag for the token
        entity : :obj:`str`
            Any mamed entities that have been asigned by GeniaTagger
        parent_chuck : :obj:`Chunk` (optional)
            If the Token is associated with a chunk then add it here, the token
            will determine it's position in the chunk
        """

        self._text = text
        self._stem = stem
        self._part_of_speech = part_of_speech
        self._entity = entity

        # This is the tokens position within the chunk
        # initialised at None and will stay that have if no chunk is added
        self._chunk_pos = None

        # Thi is the tokens position within the whole phrase
        self._phrase_pos = None

        # We start with no parent chunks
        self._parent_chunk = None

        # Add any parent chunks to the token
        self._add_parent_chunk(parent_chunk)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __str__(self):
        """
        When str is called we return the token text
        """
        return self.text

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __repr__(self):
        """
        When str is called we return the token text
        """
        return self.report()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __len__(self):
        """
        Define what happens if len(Token) is used. It returns the length of the
        text
        """
        return len(self.text)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __getitem__(self, index):
        """
        make an iterator
        """
        return self._text[index]

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def text(self):
        """
        Return the text of the token
        """
        return self._text

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def stem(self):
        """
        Return the stem of the token
        """
        return self._stem

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def part_of_speech(self):
        """
        Return the part of speech tag of the token
        """
        return self._part_of_speech

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def entity(self):
        """
        Return any named entities of the token
        """
        return self._entity

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def chunk_position(self):
        """
        Return token position in the parent chunk
        """
        return self._chunk_position

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def phrase_position(self):
        """
        Return token position in the whole phrase
        """
        return self._phrase_position

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def is_symbol(self):
        """Is the chunk comprised of all symbols i.e. ``\\w``."""
        return bool(re.match(r"\w+", str(self)))

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def parent_chuck(self):
        """Get the parent chunk."""
        return self._parent_chunk

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def add_chunk(self, chunk):
        """
        Add a parent chunk to the token. Note that the token can only have one
        parent chunk. Adding a parent chunk to a token that already has one
        will result in an error.

        Parameters
        ----------

        chunk : :obj:`Chunk`
            The parent chunk to associate with the token

        Raises
        ------

        GeniaParserError
            If the token already has a chunk of chunk is None
        """

        if chunk is None:
            raise err.GeniaParserError("no chunk to add")

        # Now do the adding
        self._add_parent_chunk(chunk)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _add_parent_chunk(self, chunk):
        """
        Add a parent chunk to the token. Note that the token can only have one
        parent chunk. Adding a parent chunk to a token that already has one
        will result in an error.

        Parameters
        ----------

        chunk : :obj:`Chunk`
            The parent chunk to associate with the token

        Raises
        ------

        GeniaParserError
            If the token already has a chunk
        """

        if chunk is None:
            return

        if self._parent_chunk is not None and self._parent_chunk != chunk:
            raise err.GeniaParserError("a token can only have 1 parent chunk")

        # Now we make sure the token is bound to the chunk
        chunk.add_token(self)

        # Now loop through all the tokens in the chunk
        for c, i in enumerate(chunk):
            if i == self:
                self._chunk_pos = c
                self._phrase_pos = c + chunk.start

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def report(self, level=0):
        """
        Nice printing, useful for debugging

        Parameters
        ----------

        level : :obj:`int`
            The number of tag stops to place at the start of the report
        """
        return "%sToken:%s|%s: %-15s%-15s%-7s%-7s" % (
            ("\t" * level),
            str(self._phrase_pos),
            str(self._chunk_pos),
            self._text,
            self._stem,
            self._part_of_speech,
            self._entity,
        )

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def summary(self):
        """
        Returns the token in the form (NE text/POS)
        """
        return "(%s %s/%s)" % (self._entity, self._text, self._part_of_speech)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def as_tuple(self):
        """
        Returns a tuple with all the token info in it
        """
        return (self._text, self._part_of_speech, self._entity)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def get_stem_token(self):
        """
        Create a new Token object where the text == the stem
        """

        # Create a token, this has no parent chunk otherwise the parent chunk
        # would get 2 copies of the token
        return Token(
            self.stem,
            self.stem,
            self.part_of_speech,
            self.entity,
            parent_chunk=None,
        )


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class CuiPhraseNgramMapper(object):
    """
    Attempts to map phrases (as returned by the GeniaTagger) to CUIs
    """

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, umls_query, phrase):
        """
        Takes a phrase and a UMLS query object
        """
        self._umls_query = umls_query
        self._phrase = phrase

        # Process the phrase
        self._process_phrase()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __repr__(self):
        return_str = []

        for c, i in enumerate(self.chunk_mappers):
            return_str.append("\nCHUNK [{0}]".format(c))
            return_str.append(str(i))

        return "\n".join(return_str)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def phrase_map(self):
        """
        Yield the annotated phrase with different permutations of CUI mappings
        """

        # Loop through all the chunk mappers and get them to yield their
        # mappings
        # Each chunk mapper will contain a chunk and cuis that map to the
        # chunk if the whole chunk does not map to a CUI then ngrams
        # (sub-chunks) are created and tested against UMLS.
        mapped_phrase = []
        for cm in self.chunk_mappers:
            # A call to chunk map will generate all available CUI mappings
            # for that chunk and ngrams of that chunk
            sub_chunk_maps = []

            for scm in cm.chunk_map():
                sub_chunk_maps.append(scm)

            mapped_phrase.append(sub_chunk_maps)
            # for mapping in cm.chunk_map():
            #     mapping.

        # The mapped phrase will be a list of lists where each sub list
        # contains the mapped sub-chunk permutations
        for i in itertools.product(*mapped_phrase):
            # for c in i:
            yield (sum([s[0] for s in i]), i)
            # yield i

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _process_phrase(self):
        """Process the phrase"""
        self.chunk_mappers = []

        # Loop through all the chunks and attempt to map them
        for c, chunk in enumerate(self._phrase.chunks()):
            # print("CHUNK {0}".format(c))
            # pp.pprint(chunk)
            # Store all the chunk mappers
            self.chunk_mappers.append(
                CuiChunkNgramMapper(self._umls_query, chunk)
            )


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class CuiChunkNgramMapper(object):
    """Attempts to map a chunk to CUIs"""

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, umls_query, chunk):
        """
        Takes the umls query object and a chunk
        """
        self._umls_query = umls_query
        self._chunk = chunk

        # Search the chunk against the UMLS
        self._search_chunk()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __repr__(self):
        return_str = []

        # If there are no mappings then report it
        if len(self._chunk_mappings) == 0:
            return "No Mappings!"

        # Loop through all the mappings
        for c, i in enumerate(self._chunk_mappings):
            return_str.append("SUB_MATCH [{0}]".format(c + 1))
            return_str.append(str(i))

        return "\n".join(return_str)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _search_chunk(self):
        """
        Search the text of the chunk against the UMLS
        """

        # Will store the mappings between chunks and CUIs or ngrams of chunks
        # and CUIs
        self._chunk_mappings = []

        # Check to see if it needs masking from a CUI search
        try:
            self._chunk_mappings.append(self._mask_in(self._chunk))
            return
        except RuntimeError:
            pass

        # First attempt to match the whole chunk to a CUI
        cui = self._umls_query.match_term(str(self._chunk))

        # If we have a match then we store it and do not process any ngrams
        if len(cui) > 0:
            self._chunk_mappings.append(
                ChunkMappingResult(
                    self._chunk, cui, self._chunk, db=self._umls_query
                )
            )
        else:
            try:
                # Search for all the STEMs and text from the chunk
                self._chunk_mappings.extend(self._stem_matches(self._chunk))
            except RuntimeError:
                # The stems have no matches
                self._search_ngram_chunks()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _search_ngram_chunks(self):
        """
        Search for CUI matches on ngrams and all their STEMs if needed
        """
        # Single length chunks will not produce ngrams so we store them
        # so they do not drop out
        if self._chunk.ntokens == 1:
            try:
                # Check to see if it needs masking from a CUI search
                self._chunk_mappings.append(self._mask_in(self._chunk))
                return
            except RuntimeError:
                # Store the NGRAM irrespective if it has a match
                self._chunk_mappings.append(
                    ChunkMappingResult(
                        self._chunk, [], self._chunk, db=self._umls_query
                    )
                )
        else:
            # Generate all the ngrams of the chunk and store any results
            for ngram in self._chunk.all_ngrams():
                try:
                    # Check to see if it needs masking from a CUI search
                    self._chunk_mappings.append(self._mask_in(ngram))
                    continue
                except RuntimeError:
                    pass

                cui = self._umls_query.match_term(str(ngram))

                if len(cui) == 0:
                    try:
                        self._chunk_mappings.extend(self._stem_matches(ngram))
                        continue
                    except RuntimeError:
                        pass

                # Store the NGRAM irrespective if it has a match
                self._chunk_mappings.append(
                    ChunkMappingResult(
                        ngram, cui, self._chunk, db=self._umls_query
                    )
                )

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _mask_in(self, chunk):
        """
        Loop through the chunk mapping results, if they are length == 1 and are
        a IN, CC, TO (POS), create dummy CUIs with POS tags as the semantic
        type
        """

        if chunk.ntokens > 1:
            raise RuntimeError("can only mask chunks with single tokens")

        # Now loop through the tokens, there will only be 1 but it is a
        # generator
        for token in chunk.tokens():
            if token.part_of_speech in ["IN", "TO", "CC"]:
                return ChunkMappingResult(
                    chunk,
                    [
                        objects.Cui(
                            "C000000",
                            token.part_of_speech,
                            stys=[
                                objects.Sty(
                                    token.part_of_speech,
                                    term=token.part_of_speech,
                                )
                            ],
                        )
                    ],
                    self._chunk,
                    db=self._umls_query,
                )
            elif token.text in "({[]})":
                return ChunkMappingResult(
                    chunk,
                    [
                        objects.Cui(
                            "C000000",
                            "BRAKET",
                            stys=[objects.Sty("BRAKET", term="BRACKET")],
                        )
                    ],
                    self._chunk,
                    db=self._umls_query,
                )
            elif token.text in ".,?!":
                return ChunkMappingResult(
                    chunk,
                    [
                        objects.Cui(
                            "C000000",
                            "PUNCT",
                            stys=[objects.Sty("PUNCT", term="PUNCT")],
                        )
                    ],
                    self._chunk,
                    db=self._umls_query,
                )

            # If we get here then we have not matched anything so we raise a
            # RuntimeError
            raise RuntimeError("nothing to mask")

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _stem_matches(self, chunk):
        """
        Produce all permutations of stemmed and unstemmed tokens for searching
        against the UMLS
        """

        # Holds all the stemmed and unstemmed tokens in the chunk
        combinations = []

        # Loop through all the tokens in the chunk
        for i in chunk.tokens():
            token_stems = []

            token_stems.append(i)
            # if i.stem not in token_stems:
            token_stems.append(i.get_stem_token())
            combinations.append(token_stems)

        # This will hold all the unique CUIs that are generated from the
        # searches of the STEMs and TEXT from the tokens
        all_matches = []
        cui_check = set()
        for i in itertools.product(*combinations):
            # Create a chunk for the TEXT/STEM combinations
            # text_to_search = " ".join(i)
            stem_chunk = Chunk(
                *i,
                start=chunk.start,
                chunk_spec="STEM_{0}".format(chunk.chunk_spec),
            )
            cuis = self._umls_query.match_term(str(stem_chunk))

            # For each stem_chunk we only store CUIs that have not been
            # seen in any other stem_chunks. If there are no unique CUIs
            # then we do not keep the stem_chunk
            uniq_cuis = []
            for c in cuis:
                if c.cui not in cui_check:
                    cui_check.add(c.cui)
                    uniq_cuis.append(c)

            if len(uniq_cuis) > 0:
                all_matches.append(
                    ChunkMappingResult(
                        stem_chunk, uniq_cuis, chunk, db=self._umls_query
                    )
                )

        # Raise an error if stems do not find anything
        if len(all_matches) == 0:
            raise RuntimeError("stems found no CUIs")

        return all_matches

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def chunk_map(self):
        """
        Yield the annotated phrase with different permutations of CUI mappings
        """

        # First sort the chunk mappings based on start positions
        # self._chunk_mappings.sort(key=lambda x: x.chunk.start)

        # Convert the chunk mappings to a dick with the keys as the start
        # positions and the values as a list of the chunk mappings with
        # those start positions
        cm = {}
        for m in self._chunk_mappings:
            cm.setdefault(m.chunk.start, [])
            cm[m.chunk.start].append(m)

        # Start position, different chunks will have different start positions
        start_pos = min(cm.keys())

        # We build all the sub-chunk permutations starting from the
        # sub-chunks that occur at the begining of the chunk
        # all_sub_chunks = []
        for i in cm[start_pos]:
            # Seq is a list of lists. Each sub-list is a representation of the
            # chunk interms of sub-chunk combinations
            for sequence in self._build_chunk_sequences([i], cm):
                yield (sum([s.weighted_ic for s in sequence]), sequence)
                # yield sequence

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _build_chunk_sequences(self, chunk_seq, chunk_map):
        """
        Recursively build sequences of chunks
        """

        all_seq = []

        # Loop through all the mappings for chunks that start after
        # the one we are processing
        try:
            for cm in chunk_map[chunk_seq[-1].chunk.end + 1]:
                all_seq.extend(
                    self._build_chunk_sequences(chunk_seq + [cm], chunk_map)
                )
        except KeyError:
            all_seq.append(chunk_seq)

        return all_seq


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class ChunkMappingResult(object):
    """
    Holds chunks and CUI mappings
    """

    TOTAL_STR = 43538830.0

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, chunk, cuis, parent_chunk, db=None):
        """
        Takes a chunk and a CUI. CUIs should be a list
        """
        self._chunk = chunk
        self._cuis = cuis
        self._parent_chunk = parent_chunk
        self._db = db
        self._set_ic()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __repr__(self):
        chunk_str = ["<{0}: {1}>".format(self._weighted_ic, str(self._chunk))]
        # chunk_str.append("\nCUIs\n=====")

        cui_data = []

        if len(self._cuis) == 0:
            cui_data.append("(NA,NA,NA)")

        for i in self._cuis:
            cui_data.append(
                "({0},{1},{2})".format(
                    i.cui, i.term, "|".join([s.term for s in i.stys])
                )
            )
        chunk_str.append("".join(cui_data))

        return "({0})".format(" ".join(chunk_str))

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def chunk(self):
        return self._chunk

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _set_ic(self):
        """
        Calculate the information content for the chunk with the CUI result
        """

        # As all the CUI matches should be the same for the chunk we just need
        # the IC of the first one
        try:
            # A call to generate an index error
            self._cuis[0].term
            # pp.pprint(self._db.genia_string(self._cuis[0].term))

            # loop through all the tokens in the chunk and get their
            # information content
            ics = []
            # TODO: I should really search for the IC of all the CUIs and pick
            # TODO: the highest
            for i in self._chunk.tokens():
                try:
                    ics.append(
                        -1.0
                        * math.log(
                            float(len(self._db.genia_string(i.text)))
                            / ChunkMappingResult.TOTAL_STR
                        )
                    )
                except TypeError:
                    ics.append(0.0)
            self._ic = sum(ics)
        except IndexError:
            # No database or no CUIs
            self._ic = 0.0

        # We weight the IC by the proportion of the parent chunk that the match
        # occupies
        self._weighted_ic = (
            float(self._chunk.ntokens) / float(self._parent_chunk.ntokens)
        ) * self._ic

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def ic(self):
        return self._ic

    @property
    def cuis(self):
        return self._cuis

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def weighted_ic(self):
        return self._weighted_ic

    # # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    # def search_chunk(db, chunk):
    #     """
    #     Search a chunk against the UMLS
    #     """
    # cui = db.match_term(str(chunk))
    # print(chunk)
    # print(cui)


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class SemanticTypeTagger(object):
    """
    Maps semantic types back to TAGs that will perform a similar job to POS tag
    these will allow for parsing with the NLTK RegexParser to pull out entities
    """

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, mappings):
        """
        Mapper will be a CUI mapper of some description and meppings is a
        dict of mappings
        """
        # self.mapper = mapper
        self.mappings = mappings

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def get_tags(self, phrase_mapping):
        """
        yield the tags for each phrase, starting with the highest scoring
        phrase
        """

        all_cui_comb = []

        # Loop through all the mapped chunks in the phrase
        for chunk_map_group in phrase_mapping[1]:
            for chunk_map in chunk_map_group[1]:
                # print(chunk_map)
                chunk_cui = []
                try:
                    # Force an index error
                    chunk_map.cuis[0]

                    for c in chunk_map.cuis:
                        for s in c.stys:
                            try:
                                chunk_cui.append(
                                    (
                                        str(chunk_map.chunk),
                                        self.__class__.TAGS[s.term],
                                        c.cui,
                                    )
                                )
                            except KeyError:
                                chunk_cui.append(
                                    (str(chunk_map.chunk), "UNKNOWN", c.cui)
                                )
                except IndexError:
                    chunk_cui.append((str(chunk_map.chunk), "UNKNOWN", None))

                all_cui_comb.append(chunk_cui)

        for i in itertools.product(*all_cui_comb):
            yield i


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class GwasCatSematicTypeTagger(SemanticTypeTagger):
    """
    Semantic type tagger with tags relevent to the GWas Catalogue
    """

    TAGS = {
        "Quantitative Concept": "MEASURE",
        "Qualitative Concept": "MEASURE",
        "Functional Concept": "MEASURE",
        "Laboratory Procedure": "PROCEDURE",
        "Clinical Attribute": "CLINICAL",
        "Diagnostic Procedure": "PROCEDURE",
        "Finding": "RESULT",
        "Age Group": "COHORT",
        "Body Part, Organ, or Organ Component": "TISSUE",
        "Tissue": "TISSUE",
        "Disease or Syndrome": "DISEASE",
        "Neoplastic Process": "DISEASE",
        "Sign or Symptom": "DISEASE",
        "Pharmacologic Substance": "DRUG",
        "Therapeutic or Preventive Procedure": "THERAPY",
        "Enzyme": "PROTEINS",
        "Organism Function": "ORG_FUNCTION",
        "Biologically Active Substance": "BIOMARKER",
        "Organic Chemical": "CHEMICAL",
        "TO": "TO",
        "BRACKET": "BRACKET",
    }

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self):
        """
        Initialise, the mapper is built in
        """
        super(GwasCatSematicTypeTagger, self).__init__(self.__class__.TAGS)
