class TextFilesReader(Reader):
    """
    Reads in one or more file and creates a dataset out of them.
        * If the input is a path to a single file we create a dataset with one document
        * If the input is a path to a directory we create a dataset where each file is one document
    The name of the file is the document ID.
    Reads in only files with .txt extension.
    When reading in each file we separate into parts by blank lines if there are any.
    """

    def __init__(self, path):
        self.path = path
        """path to a single file or a directory containing input files"""

    @staticmethod
    def __process_file(filename):
        document = Document()
        with open(filename) as file:
            part_id = 1
            for part in re.split('\n\n', file.read()):
                if part.strip():
                    document.parts['{}'.format(part_id)] = Part(part)
                    part_id += 1

        return os.path.split(filename)[-1], document

    def read(self):
        """
        :returns: nalaf.structures.data.Dataset
        """
        dataset = Dataset()
        if os.path.isdir(self.path):
            for filename in glob.glob(self.path + '/*.txt'):
                doc_id, doc = self.__process_file(filename)
                dataset.documents[doc_id] = doc
        else:
            if os.path.splitext(self.path)[-1] == '.txt':
                doc_id, doc = self.__process_file(self.path)
                dataset.documents[doc_id] = doc
            else:
                raise Exception('not a .txt file extension')

        return dataset class PMIDReader(Reader):
    """
    Reads in a single PMID or a list of PMIDs,
    downloads the associated articles (including title and abstract)
    and creates a dataset with one document per PMID.
    """
    def __init__(self, pmids):
        if hasattr(pmids, '__iter__'):
            self.pmids = pmids
        else:
            self.pmids = [pmids]

    def read(self):
        """
        :returns: nalaf.structures.data.Dataset
        """
        dataset = Dataset()
        with DownloadArticle() as da:
            for pmid, doc in da.download(self.pmids):
                dataset.documents[pmid] = doc
        return dataset class HTMLReader(Reader):
    """
    Reader for a local file system or hdfs of tagtog plain.html files.

    It reads either a single file or a directory (the contained .html's)
    """

    def __init__(self, path, whole_basename_as_docid=False, hdfs_url=None, hdfs_user=None):
        self.path = path
        """an html file or a directory containing .html files"""
        self.whole_basename_as_docid = whole_basename_as_docid

        self.hdfs_client = maybe_get_hdfs_client(hdfs_url, hdfs_user)

    def __read_directory_localfs(self):
        dataset = Dataset()

        filenames = glob.glob(str(self.path + "/**/*.html"), recursive=True) + glob.glob(str(self.path + "/**/*.xml"), recursive=True)
        for filename in filenames:
            dataset = self.__read_file_path_localfs(filename, dataset)

        return dataset

    def __read_file_path_localfs(self, filename, dataset=None):
        if dataset is None:
            dataset = Dataset()

        with open(filename, 'rb') as a_file:
            HTMLReader.read_file(a_file, filename, dataset, self.whole_basename_as_docid)

        return dataset

    def __read_directory_hdfs(self):
        dataset = Dataset()

        filenames = walk_hdfs_directory(self.hdfs_client, self.path, lambda fname: fname.endswith(".html") or fname.endswith(".xml"))
        for filename in filenames:
            dataset = self.__read_file_path_hdfs(filename, dataset)

        return dataset

    def __read_file_path_hdfs(self, filename, dataset=None):
        if dataset is None:
            dataset = Dataset()

        with self.hdfs_client.read(filename) as reader:
            HTMLReader.read_file(reader, filename, dataset, self.whole_basename_as_docid)

        return dataset

    @staticmethod
    def read_file(a_file, filename, dataset=None, whole_basename_as_docid=False):
        if dataset is None:
            dataset = Dataset()

        soup = BeautifulSoup(a_file, "html.parser")
        document = Document()

        for part in soup.find_all(id=re.compile('^s')):
            if re.match(r'^s[3-9]', part['id']):
                is_abstract = False
            else:
                is_abstract = True
            document.parts[part['id']] = Part(str(part.string), is_abstract=is_abstract)

        doc_id = os.path.basename(filename).replace('.plain.html', '').replace('.html', '').replace('.xml', '')
        if not whole_basename_as_docid and '-' in doc_id:
            doc_id = doc_id.split('-')[-1]

        dataset.documents[doc_id] = document

        return dataset

    def read(self):
        if self.hdfs_client is None:
            if os.path.isdir(self.path):
                return self.__read_directory_localfs()
            else:
                return self.__read_file_path_localfs(filename=self.path)

        else:
            if is_hdfs_directory(self.hdfs_client, self.path):
                return self.__read_directory_hdfs()
            else:
                return self.__read_file_path_hdfs(filename=self.path)

            return class StringReader(Reader):
    """
    Reads in a simple string and creates a dataset out of it.
    Useful for quick testing of model capabilities.

    The dataset contains a single Document with a single Part
    which contains the text given by the string.
    """

    def __init__(self, string):
        self.string = string
        """the string from which we form the dataset"""

    def read(self):
        """
        :returns: nalaf.structures.data.Dataset
        """
        part = Part(self.string)
        document = Document()
        dataset = Dataset()

        dataset.documents['doc_1'] = document
        document.parts['part_1'] = part

        return dataset class ConsoleWriter:
    """
    Writes the predicted annotations onto the console.
    """
    def __init__(self, ent1_class_id, ent2_class_id, color=True):
        self.color = self.__supports_color() if color else color
        self.protein_color_start = '\033[42m'
        self.mutation_color_start = '\033[41m'
        self.end_color = '\033[0m'
        self.ent1_class_id = ent1_class_id
        self.ent2_class_id = ent2_class_id


    def write(self, dataset):
        """
        :type dataset: nalaf.structures.data.Dataset()
        """
        for doc_id, doc in dataset.documents.items():
            print('DOCUMENT: {}'.format(doc_id))
            for part_id, part in doc.parts.items():
                print('PART {}'.format(part_id))
                self.___print_part(part)

    def ___print_part(self, part):
        if self.color:
            text = part.text
            total_offset = 0
            for ann in sorted(part.predicted_annotations, key=lambda x: x.offset):
                color = self.mutation_color_start if ann.class_id == self.ent2_class_id else self.protein_color_start
                text = text[:ann.offset+total_offset] + color + text[ann.offset+total_offset:]
                total_offset += 5
                text = text[:ann.offset+len(ann.text)+total_offset] + self.end_color + text[ann.offset+len(ann.text)+total_offset:]
                total_offset += 4
            print(text)
            print()
        else:
            padding = len(str(len(part.text)))
            print(part.text)
            print('ANNOTATIONS')
            for ann in sorted(part.predicted_annotations, key=lambda x: x.offset):
                if ann.class_id == self.ent2_class_id:
                    print('Mutation {0: <{pad}} {1: <{pad}} {2}'
                          .format(ann.offset, ann.offset+len(ann.text), ann.text, pad=padding))
                elif ann.class_id == self.ent1_class_id:
                    print('GGP      {0: <{pad}} {1: <{pad}} {2} {3}'
                          .format(ann.offset, ann.offset+len(ann.text), ann.text, ann.norms, pad=padding))
        print('RELATIONS')
        for rel in part.relations:
            print('{} ---> {}'.format(rel.text1, rel.text2))
        print()

    @staticmethod
    def __supports_color():
        """
        Returns True if the running system's terminal supports color, and False
        otherwise.
        """
        plat = sys.platform
        supported_platform = plat != 'Pocket PC' and (plat != 'win32' or 'ANSICON' in os.environ)
        is_a_tty = hasattr(sys.stdout, 'isatty') and sys.stdout.isatty()

        if not supported_platform or not is_a_tty:
            return False
        return True class TagTogFormat:
    """
    Ability to Export the dataset as Html + Ann.json database.
    """
    def __init__(self, dataset, use_predicted, to_save_to="resources/corpora/sample/anndoc", who="ml:nalaf", _annjson_folder="annjson", _html_folder="html", use_original_partids = True):
        """
        init function that does prepare annjson folder and html folder
        :param to_save_to: usually resources/corpora/[name of corpus]/anndoc/
        :type dataset: nalaf.structures.data.Dataset
        :param who:
        :param _annjson_folder:
        :param _html_folder:
        :return:
        """
        self.location = to_save_to
        """ root folder, that documents are saved into """
        self.data = dataset
        """ dataset param """
        self.who = who
        """ who parameter """
        self.use_original_partids = use_original_partids
        self.use_predicted = use_predicted

        # Possibility to use instance without writing files to disk
        if to_save_to:
            # check for root folder for files to save to
            if not os.path.isdir(self.location):
                print_verbose("mkdir", os.path.abspath(self.location))
                try:
                    os.makedirs(self.location)
                except FileExistsError:
                    pass

            # create subfolders if not existent
            # annjson folder
            self.annjson_path = os.path.join(self.location, _annjson_folder)
            """ subfolder where ann.json files are saved into """
            if not os.path.isdir(self.annjson_path):
                os.mkdir(self.annjson_path)
            # html folder
            self.html_folder = os.path.join(self.location, _html_folder)
            """ subfolder where html files are saved into """
            if not os.path.isdir(self.html_folder):
                os.mkdir(self.html_folder)

    def export(self, threshold_val):
        self.export_html()
        self.export_ann_json(threshold_val=threshold_val)

    def export_html(self):
        """
        Exporting Html files into folder with each html file being a document itself.
        Html files have sections and everything as if document was exported from tagtog.net itself.
        :return:
        """
        for docid, doc in self.data.documents.items():
            fname = os.path.join(self.html_folder, docid + ".html")

            with open(fname, 'wb') as f:

                # "tag" or "tag_attr" for their attributes

                html_attr = {
                    'id' : docid,
                    'data-origid' : docid,
                    'class' : 'anndoc',
                    'data-anndoc-version' : '2.0',
                    'lang' : '',
                    'xml:lang' : '',
                    'xmlns' : 'http://www.w3.org/1999/xhtml'
                }
                html = ET.Element('html', html_attr)

                head = ET.SubElement(html, 'head')

                # meta1 = ET.SubElement(head, 'meta', { 'charset' : 'UTF-8'} )
                # meta2 = ET.SubElement(head, 'meta', { 'name' : 'generator', 'content' : 'nalaf.utils.writers.TagTogFormat'} )
                # meta3 = ET.SubElement(head, 'meta', { 'name': 'dcterms.source', 'content' : 'http://www.ncbi.nlm.nih.gov/pubmed/' + docid } )  # deprecated maybe different sources

                title = ET.SubElement(head, 'title')
                title.text = docid

                body = ET.SubElement(html, 'body')

                article = ET.SubElement(body, 'article')

                section = ET.SubElement(article, 'section', { 'data-type' : 'article' } )
                div = ET.SubElement(section, 'div', { 'class' : 'content' } )

                for i, (partid, part) in enumerate(doc.parts.items()):
                    if not self.use_original_partids:
                        partid = "s1p{}".format(i + 1)
                    p = ET.SubElement(div, 'p', { 'id' : partid } )
                    p.text = part.text

                # print(ET.dump(html))
                # output = ET.tostring(html, encoding='UTF-8')
                f.write(ET.tostring(html, encoding='utf-8', method='html'))


    def get_single_ann_json(self, threshold_val=None, docid=None):
        if docid:
            doc = self.data.documents[docid]
        else:
            docid, doc = next(iter(self.data.documents.items()))

        partids = list(doc.parts.keys())
        if not self.use_original_partids:
            partids = []
            for i, (partid, part) in enumerate(doc.parts.items()):
                partid = "s1p{}".format(i + 1)
                partids.append(partid)

        # init empty json-object
        json_obj = {
            "annotatable": {
                "parts" : partids
            },
            "anncomplete": False,
            "sources": [
                {
                    "name": "ORIG",
                    "id": docid,
                    "url": ""
                }
                # each entry is a dict with "name", "id", "url"
            ],
            "metas": {
                # nothing important -->  # OPTIONAL add meta information from project
            },
            "entities": [
                # dict with "classId", "part", "offsets" (being [{"start","text"},...], confidence
            ],
            "relations": [
                # not important for here  # OPTIONAL get relations as well
            ]
        }

        for i, (partid, part) in enumerate(doc.parts.items()):
            if not self.use_original_partids:
                partid = "s1p{}".format(i + 1)

            for ann in (part.predicted_annotations if self.use_predicted else part.annotations):
                if threshold_val:
                    if ann.confidence >= threshold_val:
                        state = 'selected'
                    else:
                        state = 'pre-added'
                else:
                    state = ''

                fields = {}
                normalizations = {}

                for key, value in ann.norms.items():
                    if isinstance(value, str):
                        normalizations[key] = {
                            # "source": 'http://uniprot.org',
                            "recName": value,
                            # "confidence": 1  # todo discussion confidence from GNormPlus is not provided so just putting in here 1
                        }
                    else:
                        for substring in value:
                            normalizations[key + '_' + substring[0]] = {
                                # "source": 'http://uniprot.org',
                                "recName": substring,
                                # "confidence": 1  # todo discussion confidence from GNormPlus is not provided so just putting in here 1
                            }

                ent = {
                    "classId": ann.class_id,
                    "part": partid,
                    "offsets": [{"start": ann.offset, "text": ann.text}],
                    "confidence": {
                        "state": state,
                        "who": [
                            self.who
                        ],
                        "prob": ann.confidence
                    },
                    "fields": fields,
                    "normalizations": normalizations
                }
                json_obj['entities'].append(ent)

            for rel in doc.relations():
                ent_string_format = '{}|{},{}'
                ent_string1 = ent_string_format.format(partid, rel.start1, rel.start1 + len(rel.text1))
                ent_string2 = ent_string_format.format(partid, rel.start2, rel.start2 + len(rel.text2))
                relation = {
                    "classId": rel.class_id,
                    "directed": False,
                    "entities": [ent_string1, ent_string2],
                    "confidence": 1
                }
                json_obj['relations'].append(relation)

        return json_obj


    def export_ann_json(self, threshold_val=None):
        """
        Creates all Annotation files in the corresponding ann.json format.
        Description of ann.json-format: "https://github.com/tagtog/tagtog-doc/wiki/ann.json"
        :return:
        """
        for docid in self.data.documents.keys():
            fname = os.path.join(self.annjson_path, docid + ".ann.json")
            print_verbose(fname)
            with open(fname, 'w', encoding='utf-8') as f:
                json_obj = self.get_single_ann_json(threshold_val, docid)
                json.dump(json_obj, f) class PubTatorFormat:
    """
    Export into PubTatorFormat with or without annotations
    """
    def __init__(self, dataset, location="resources/corpora/idp4/pubtator.txt", no_annotations=True):
        self.location = location
        self.dataset = dataset
        self.no_annotations = no_annotations

    def export(self):
        with open(self.location, 'w', encoding='utf-8') as f:
            for pid, doc in self.dataset.documents.items():
                offset = 0
                f.write("{0}|t|{title}\n".format(pid, title=doc.get_title()))
                f.write("{0}|a|{text}\n".format(pid, text=doc.get_body()))
                for part in doc:
                    for ann in chain(part.annotations, part.predicted_annotations):
                        f.write("{0}\t{start}\t{end}\t{text}\t{classid}\n".format(pid, start=ann.offset,
                                                                                  end=ann.offset + len(ann.text),
                                                                                  text=ann.text, classid=ann.class_id))
                    offset += len(part.text)
                f.write("\n") class PrepareDatasetPipeline:
    """
    Prepares an instance of a dataset by executing modules in fixed order.
        * First executes the sentence splitter
        * Next executes the tokenizer
        * Finally executes each feature generator in the order they were provided

    :type splitter: nalaf.structures.data.Splitter
    :param splitter: the module responsible for splitting the text into sentences
    :type tokenizer: nalaf.structures.data.Tokenizer
    :param tokenizer: the module responsible for splitting the sentences into tokens
    :type feature_generators: collections.Iterable[FeatureGenerator]
    :param feature_generators: one or more modules responsible for generating features
    """

    def __init__(self, splitter=None, tokenizer=None, feature_generators=None):
        if not splitter:
            splitter = NLTKSplitter()
        if not tokenizer:
            tokenizer = TmVarTokenizer()
        if feature_generators is None:
            feature_generators = [SimpleFeatureGenerator(), PorterStemFeatureGenerator(),
                                  WindowFeatureGenerator((-3, -2, -1, 1, 2, 3), ['stem[0]'])]

        if isinstance(splitter, Splitter):
            self.splitter = splitter
        else:
            raise TypeError('not an instance that implements Splitter')

        if isinstance(tokenizer, Tokenizer):
            self.tokenizer = tokenizer
        else:
            raise TypeError('not an instance that implements Tokenizer')

        if hasattr(feature_generators, '__iter__'):
            for index, feature_generator in enumerate(feature_generators):
                if not isinstance(feature_generator, FeatureGenerator):
                    raise TypeError('not an instance that implements FeatureGenerator at index {}'.format(index))
            self.feature_generators = feature_generators
        elif isinstance(feature_generators, FeatureGenerator):
            self.feature_generators = [feature_generators]
        else:
            raise TypeError('not an instance or iterable of instances that implements FeatureGenerator')


    def execute(self, dataset):
        """
        :type dataset: nalaf.structures.data.Dataset()
        """

        self.splitter.split(dataset)
        self.tokenizer.tokenize(dataset)
        for feature_generator in self.feature_generators:
            print_verbose('Apply feature generator:', type(feature_generator))
            feature_generator.generate(dataset)


    def serialize(self, dataset, to_file=None):
        """
        :type dataset: nalaf.structures.data.Dataset()
        """

        types = [(type(self.splitter), self.splitter.__dict__),
                 (type(self.tokenizer), self.tokenizer.__dict__)]

        for feature_generator in self.feature_generators:
            types.append((type(feature_generator), feature_generator.__dict__))

        features = sorted(set(feature_name for token in dataset.tokens() for feature_name in token.features.keys()))

        from nalaf.utils.helpers import find_current_git_ref
        current_ref = find_current_git_ref()

        if to_file:
            with open(to_file, 'w') as file:
                file.write('git ref:{}'.format(current_ref))
                file.write('\nINSTANCES USED\n')
                file.writelines('\n'.join(repr(x) for x in types))
                file.write('\nFEATURES USED\n')
                file.writelines('\n'.join(repr(x) for x in features))

        return types, features, find_current_git_ref() class PyCRFSuite:

    def __init__(self, model_file=None):
        self.model_file = model_file

        if self.model_file is None:
            self.tagger = None
        else:
            self.tagger = pycrfsuite.Tagger()
            self.tagger.open(self.model_file)


    def close(self):
        if self.tagger is not None:
            self.tagger.close()


    def __del__(self):
        self.close()


    def __exit__(self, exc_type, exc_value, traceback):
        self.close()


    def annotate(self, corpus, class_id):
        """
        :type corpus: nalaf.structures.data.Dataset
        :type class_id: str ~ to annotate with
        """

        for sentence in corpus.sentences():
            labels = self.tagger.tag(pycrfsuite.ItemSequence(token.features for token in sentence))

            for token_index in range(len(sentence)):
                label = labels[token_index]
                try:
                    sentence[token_index].predicted_labels = [Label(label, self.tagger.marginal(label, token_index))]
                except Exception as e:
                    raise Exception("Exception when assining the predicted labels; likely a Multi-Thread problem", e)

        corpus.form_predicted_annotations(class_id)


    @staticmethod
    def train(data, model_file, params=None):
        """
        :type data: nalaf.structures.data.Dataset
        :type model_file: str ~ filename (from local file system) to save trained model to. If None, no model is saved.
        """

        trainer = pycrfsuite.Trainer()

        try:
            if params is not None:
                trainer.set_params(params)

            for sentence in data.sentences():
                trainer.append(pycrfsuite.ItemSequence([token.features for token in sentence]),
                               [token.original_labels[0].value for token in sentence])

            # The CRFSuite library handles the "pickling" of the file; saves the model here
            trainer.train(model_file)

        finally:
            trainer.clear()


    @staticmethod
    def tag(data, model_file, class_id):
        warnings.warn('Use non-static `annotate` instead', DeprecationWarning)

        """
        :type data: nalaf.structures.data.Dataset
        :type model_file: str
        """

        tagger = pycrfsuite.Tagger()

        try:
            tagger.open(model_file)

            for sentence in data.sentences():
                labels = tagger.tag(pycrfsuite.ItemSequence(token.features for token in sentence))

                for token_index in range(len(sentence)):
                    label = labels[token_index]
                    sentence[token_index].predicted_labels = [Label(label, tagger.marginal(label, token_index))]

            data.form_predicted_annotations(class_id)

        finally:
            tagger.close() class StubSameSentenceRelationExtractor(StubRelationExtractor):

    def __init__(self, entity1_class, entity2_class, relation_type, use_gold=True, use_pred=True):
        edge_generator = SentenceDistanceEdgeGenerator(entity1_class, entity2_class, relation_type, distance=0, use_gold=use_gold, use_pred=use_pred)
        super().__init__(edge_generator)


    def annotate(self, dataset):
        super().annotate(dataset) class Dataset:
    """
    Class representing a group of documents.
    Instances of this class are the main object that gets passed around and modified by different modules.

    :type documents: dict
    """

    def __init__(self):
        self.documents = OrderedDict()
        """
        documents the dataset consists of, encoded as a dictionary
        where the key (string) is the id of the document, for example PubMed id
        and the value is an instance of Document
        """


    def __len__(self):
        """
        the length (size) of a dataset equals to the number of documents it has
        """
        return len(self.documents)


    def __iter__(self):
        """
        when iterating through the dataset iterate through each document
        """
        for doc_id, document in self.documents.items():
            yield document


    def __contains__(self, item):
        return item in self.documents


    def parts(self):
        """
        helper functions that iterates through all parts
        that is each part of each document in the dataset

        :rtype: collections.Iterable[Part]
        """
        for document in self:
            for part in document:
                yield part


    def entities(self):
        """
        Yield all entities of the dataset.

        :rtype: collections.Iterable[Entity]
        """

        for part in self.parts():
            for annotation in part.annotations:
                yield annotation


    def annotations(self):
        warnings.warn('Use `self.entities` instead', DeprecationWarning)
        return self.entities()


    def predicted_entities(self):
        """
        Yield all predicted entities of the dataset.

        :rtype: collections.Iterable[Entity]
        """
        for part in self.parts():
            for annotation in part.predicted_annotations:
                yield annotation


    def predicted_annotations(self):
        warnings.warn('Use `self.predicted_entities` instead', DeprecationWarning)
        return self.predicted_entities()


    def relations(self):
        """
        Yield all relations of the Dataset.

        :rtype: collections.Iterable[Relation]
        """
        for part in self.parts():
            for rel in part.relations:
                yield rel


    def predicted_relations(self):
        """
        Yield all predicted relations of the dataset.

        :rtype: collections.Iterable[Relation]
        """
        for part in self.parts():
            for relation in part.predicted_relations:
                yield relation


    def plausible_relations_from_generated_edges(self):
        """
        Yield only the real relations that are obtainable from the corpus-generated edges.
        """
        for edge in self.edges():
            relation = edge.get_relation_if_is_real()
            if relation is not None:
                yield relation


    def compute_stats_relations_distances(self, relation_type, entity_map_fun=None, relation_accept_fun=None):
        """
        Returns a counter of the relationships distances.

        The relationships are mapped to unique strings as determined by entity_map_fun (see `map_relations`).

        The minimal distance of the mapped relations with same map key is used.
        """

        if next(self.predicted_relations(), None) is None:
            raise AssertionError("Relations for the corpus must have been predicted to fully exhaust the search")

        if entity_map_fun is None:
            entity_map_fun = Entity.__repr__

        counter_nums = Counter()

        for doc in self:
            # Group relationships at the document level (not part level, not corpus level)
            doc_relations = {}

            doc_relations = doc.map_relations(use_predicted=False, relation_type=relation_type, entity_map_fun=entity_map_fun)
            pred_doc_relations = doc.map_relations(use_predicted=True, relation_type=relation_type, entity_map_fun=entity_map_fun)

            for pred_key, pred_rels_with_distances in pred_doc_relations.items():

                if pred_key in doc_relations:
                    doc_relations[pred_key] += pred_rels_with_distances

                if relation_accept_fun is not None:
                    for real_key in doc_relations:
                        if relation_accept_fun(real_key, pred_key):
                            doc_relations[real_key] += pred_rels_with_distances

            for rel_key, rels_with_distances in doc_relations.items():
                rel, min_distance_for_unique_key = min(rels_with_distances, key=lambda reldist_tuple: reldist_tuple[1])
                counter_nums.update(["D" + str(min_distance_for_unique_key)])

        total = sum(counter_nums.values())

        counter_percts = Counter({key: (count / total) for (key, count) in counter_nums.items()})

        return (counter_nums, counter_percts)


    def sentences(self):
        """
        helper functions that iterates through all sentences
        that is each sentence of each part of each document in the dataset

        :rtype: collections.Iterable[list[Token]]
        """
        for part in self.parts():
            for sentence in part.sentences:
                yield sentence


    def tokens(self):
        """
        helper functions that iterates through all tokens
        that is each token of each sentence of each part of each document in the dataset

        :rtype: collections.Iterable[Token]
        """
        for sentence in self.sentences():
            for token in sentence:
                yield token


    def edges(self):
        """
        Yield all this Corpus's edges.

        :rtype: collections.Iterable[Edge]
        """
        for part in self.parts():
            for edge in part.edges:
                yield edge


    def label_edges(self):
        """
        label each edge with its REAL target (no prediction) - whether it is indeed a relation or not
        """
        for edge in self.edges():
            if edge.is_relation():
                edge.real_target = +1
            else:
                edge.real_target = -1


    def purge_false_relationships(self):
        """
        cleans false relationships by validating them
        :return:
        """
        for part in self.parts():
            part.relations[:] = [x for x in part.relations if x.validate_itself(part)]


    def partids_with_parts(self):
        """
        helper function that yields part id with part

        :rtype: collections.Iterable[(str, Part)]
        """
        for document in self:
            for part_id, part in document.key_value_parts():
                yield part_id, part


    def form_predicted_annotations(self, class_id, aggregator_function=arithmetic_mean):
        """
        Populates part.predicted_annotations with a list of Annotation objects
        based on the values of the field predicted_label for each token.

        One annotation is the chunk of the text (e.g. mutation mention)
        whose tokens have labels that are continuously not 'O'
        For example:
        ... O O O A D I S A O O O ...
        ... ----- annotation ---- ...
        here the text representing the tokens 'A, D, I, S, A' will be one predicted annotation (mention).
        Assumes that the 'O' label means outside of mention.

        Requires predicted_label[0].value for each token to be set.
        """
        # TODO
        warnings.warn('annotations actually means entities. This method and related attributes will soon be renamed')

        for part_id, part in self.partids_with_parts():
            for sentence in part.sentences:
                index = 0
                while index < len(sentence):
                    token = sentence[index]
                    confidence_values = []
                    if token.predicted_labels[0].value != 'O':
                        start = token.start
                        confidence_values.append(token.predicted_labels[0].confidence)
                        while index + 1 < len(sentence) \
                                and sentence[index + 1].predicted_labels[0].value not in ('O', 'B', 'A'):
                            token = sentence[index + 1]
                            confidence_values.append(token.predicted_labels[0].confidence)
                            index += 1
                        end = token.start + len(token.word)
                        confidence = aggregator_function(confidence_values)
                        part.predicted_annotations.append(Entity(class_id, start, part.text[start:end], confidence))
                    index += 1

        return self


    def form_predicted_relations(self):
        """
        Populates part.predicted_relations with a list of Relation objects
        based on the values of the field target for each edge.

        Each Relation object denotes a relationship between two entities
        of (usually) different classes. Each relation is given by a relation type.

        Requires edge.pred_target to be set for each edge.
        """

        for part in self.parts():
            for e in part.edges:

                if e.pred_target == +1:
                    r = e.get_potential_relation()
                    part.predicted_relations.append(r)

        return self


    def validate_entity_offsets(self):
        """
        Helper function to validate that the entities offsets match the entity text.
        Use it as a sanity check when writing or reading annotations external entities.
        """

        for docid, doc in self.documents.items():
            for partid, part in doc.parts.items():
                for e in chain(part.annotations, part.predicted_annotations):
                    readable_text = part.text[e.offset:e.offset + len(e.text)]

                    if not e.text == readable_text:
                        warnings.warn('the offsets ({} != {}) do not match in: {}/{}/{}'.format(e.text, part.text[e.offset:e.offset + len(e.text)], docid, partid, e))


    def generate_top_stats_array(self, class_id, top_nr=10, is_alpha_only=False):
        """
        An array for most occuring words.
        :param top_nr: how many top words are shown
        """
        # NOTE ambiguos words?

        raw_dict = {}

        for ann in self.entities():
            for word in ann.text.split(" "):
                lc_word = word.lower()
                if lc_word.isalpha() and ann.class_id == class_id:
                    if lc_word not in raw_dict:
                        raw_dict[lc_word] = 1
                    else:
                        raw_dict[lc_word] += 1

        # sort by highest number
        sort_dict = OrderedDict(sorted(raw_dict.items(), key=lambda x: x[1], reverse=True))
        print(json.dumps(sort_dict, indent=4))


    def clean_subclasses(self):
        """
        cleans all subclass by setting them to = False
        """
        for ann in self.entities():
            ann.subclass = False


    def get_size_chars(self):
        """
        :return: total number of chars in this dataset
        """
        return sum(doc.get_size() for doc in self.documents.values())


    def __repr__(self):
        def class_repr(class_id):
            return class_id + ": " + str(Counter(e.subclass for e in self.entities() if e.class_id == class_id))

        classes_repr = [class_repr(class_id) for class_id in {e.class_id for e in self.entities()}]

        return "Dataset({} documents and {} entities ({}))".format(len(self.documents), sum(1 for _ in self.entities()), classes_repr)


    def __str__(self):
        second_part = "\n".join(
            ["---DOCUMENT---\nDocument ID: '" + pmid + "'\n" + str(doc) for pmid, doc in self.documents.items()])
        return "----DATASET----\nNr of documents: " + str(len(self.documents)) + ', Nr of chars: ' + str(
            self.get_size_chars()) + '\n' + second_part


    def extend_dataset(self, other):
        """
        Does run on self and returns nothing. Extends the self-dataset with other-dataset.
        Each Document-ID that already exists in self-dataset gets skipped to add.
        :type other: nalaf.structures.data.Dataset
        """
        for key in other.documents:
            if key not in self.documents:
                self.documents[key] = other.documents[key]


    def prune_empty_parts(self):
        """
        deletes all the parts that contain no entities at all
        """
        for doc_id, doc in self.documents.items():
            part_ids_to_del = []
            for part_id, part in doc.parts.items():
                if len(part.annotations) == 0:  # or part.predict_entities ?
                    part_ids_to_del.append(part_id)
            for part_id in part_ids_to_del:
                del doc.parts[part_id]


    def prune_filtered_sentences(self, filterin=(lambda _: False), percent_to_keep=0):
        """
        Depends on labeler
        """
        empty_sentence = lambda s: all(t.original_labels[0].value == 'O' for t in s)

        for part in self.parts():
            tmp = []
            tmp_ = []
            for index, sentence in enumerate(part.sentences):
                do_use = not empty_sentence(sentence) or filterin(part.sentences_[index]) or random.uniform(0, 1) < percent_to_keep
                if do_use:
                    tmp.append(sentence)
                    tmp_.append(part.sentences_[index])
            part.sentences = tmp
            part.sentences_ = tmp_


    def prune_sentences(self, percent_to_keep=0):
        """
        * keep all sentences that contain at least one mention
        * keep a random selection of the rest of the sentences

        :param percent_to_keep: what percentage of the sentences with no mentions to keep
        :type percent_to_keep: float
        """
        for part in self.parts():
            # find which sentences have at least one mention
            sentences_have_ann = [any(sentence[0].start <= ann.offset < ann.offset + len(ann.text) <= sentence[-1].end
                                      for ann in part.annotations)  # pred_annotations ?
                                  for sentence in part.sentences]
            if any(sentences_have_ann):
                # choose a certain percentage of the ones that have no mention
                false_indices = [index for index in range(len(part.sentences)) if not sentences_have_ann[index]]
                chosen = random.sample(false_indices, round(percent_to_keep*len(false_indices)))

                # keep the sentence if it has a mention or it was chosen randomly
                part.sentences = [sentence for index, sentence in enumerate(part.sentences)
                                  if sentences_have_ann[index] or index in chosen]
            else:
                part.sentences = []

    def delete_subclass_annotations(self, subclasses, predicted=True):
        """
        Method does delete all annotations that have subclass.
        Will not delete anything if not specified that particular subclass.
        :param subclasses: one ore more subclasses to delete
        :param predicted: if False it will only consider Part.annotations array and not Part.pred_annotations
        """
        # TODO
        warnings.warn('annotations actually means entities. This method and related attributes will soon be renamed')

        # if it is not a list, create a list of 1 element
        if not hasattr(subclasses, '__iter__'):
            subclasses = [subclasses]

        for part in self.parts():
            part.annotations = [ann for ann in part.annotations if ann.subclass not in subclasses]
            if predicted:
                part.predicted_annotations = [ann for ann in part.predicted_annotations
                                              if ann.subclass not in subclasses]


    @staticmethod
    def _cv_kfold_split(keys, k, fold, validation_set=True):
        """
        Split keys (e.g. doc ids or indexes) into two groups: (training, evaluation)

        if validation_set is True,
            training = strictly training set
            evaluation = validation set
        otherwise
            training = training set + validation set
            evaluation = test set
        """
        total_size = len(keys)
        sub_size = round(total_size / k)

        def folds(k, fold):
            return [i % k for i in list(range(fold, fold + k))]

        subsamples = folds(k, fold)
        training = subsamples[0:k-2]
        validation = subsamples[k-2:k-1]
        test = subsamples[k-1:k]

        if validation_set:
            training = training
            evaluation = validation
        else:
            training = training + validation
            evaluation = test

        def create_keys_set(subsample_indexes):
            ret = []
            for sub in subsample_indexes:
                start = sub_size * sub
                end = (start + sub_size) if sub != (k-1) else total_size  # k-1 is the last subsample index
                ret += keys[start:end]
            return ret

        return (create_keys_set(training), create_keys_set(evaluation))


    @staticmethod
    def _cv_kfold_splits_randomize_keys(keys):
        random.seed(2727)
        random.shuffle(keys)
        return keys


    @staticmethod
    def _cv_kfold_splits_doc_keys_sets(doc_keys, k, validation_set):
        doc_keys = list(sorted(doc_keys))
        doc_keys = __class__._cv_kfold_splits_randomize_keys(doc_keys)

        for fold in range(k):
            training, evaluation = __class__._cv_kfold_split(doc_keys, k, fold, validation_set)
            yield training, evaluation


    def cv_kfold_splits(self, k, validation_set=True):

        def create_dataset(keys):
            ret = Dataset()
            for elem in keys:
                ret.documents[elem] = self.documents[elem]
            return ret

        for training, evaluation in __class__._cv_kfold_splits_doc_keys_sets(self.documents.keys(), k, validation_set):
            yield (create_dataset(training), create_dataset(evaluation))


    def fold_nr_split(self, n, fold_nr):
        """
        Sugar syntax for next(self.cv_split(n, fold_nr), None)
        """
        return next(self.cv_split(n, fold_nr), None)


    def cv_split(self, n=5, fold_nr=None):
        """
        Returns a generator of N tuples train & test random splits
        according to the standard N-fold cross validation scheme.
        If fold_nr is given, a generator of a single fold is returned (to read as next(g, None))
        Note that fold_nr is 0-indexed. That is, for n=5 the folds 0,1,2,3,4 exist.

        :param n: number of folds
        :type n: int

        :param fold_nr: optional, single fold number to return, 0-indexed.

        :return: a list of N train datasets and N test datasets
        :rtype: (list[nalaf.structures.data.Dataset], list[nalaf.structures.data.Dataset])
        """
        keys = list(sorted(self.documents.keys()))
        random.seed(2727)
        random.shuffle(keys)

        fold_size = int(len(keys) / n)

        def get_fold(fold_nr):
            """
            fold_nr starts in 0
            """
            start = fold_nr * fold_size
            end = start + fold_size
            test_keys = keys[start:end]
            train_keys = [key for key in keys if key not in test_keys]

            test = Dataset()
            for key in test_keys:
                test.documents[key] = self.documents[key]

            train = Dataset()
            for key in train_keys:
                train.documents[key] = self.documents[key]

            return train, test

        if fold_nr:
            assert(0 <= fold_nr < n)
            yield get_fold(fold_nr)
        else:
            for fold_nr in range(n):
                yield get_fold(fold_nr)


    def percentage_split(self, percentage=0.66):
        """
        Splits the dataset randomly into train and test dataset
        given the size of the train dataset in percentage.

        :param percentage: the size of the train dataset between 0.0 and 1.0
        :type percentage: float

        :return train dataset, test dataset
        :rtype: (nalaf.structures.data.Dataset, nalaf.structures.data.Dataset)
        """
        keys = list(sorted(self.documents.keys()))
        # 2727 is an arbitrary number when Alex was drunk one day, and it's just to have reliable order in data folds randomization
        random.seed(2727)
        random.shuffle(keys)

        len_train = int(len(keys) * percentage)
        train_keys = keys[:len_train]
        test_keys = keys[len_train:]

        train = Dataset()
        test = Dataset()

        for key in test_keys:
            test.documents[key] = self.documents[key]
        for key in train_keys:
            train.documents[key] = self.documents[key]

        return train, test


    def stratified_split(self, percentage=0.66):
        """
        Splits the dataset randomly into train and test dataset
        given the size of the train dataset in percentage.

        Additionally it tries to keep the distribution of entities
        in terms of subclass similar in both sets.

        :param percentage: the size of the train dataset between 0.0 and 1.0
        :type percentage: float

        :return train dataset, test dataset
        :rtype: (nalaf.structures.data.Dataset, nalaf.structures.data.Dataset)
        """
        from itertools import groupby
        train = Dataset()
        test = Dataset()

        strat = [(doc_id, Counter(ann.subclass for part in doc for ann in part.annotations))
                 for doc_id, doc in self.documents.items()]
        strat = sorted(strat, key=lambda x: (x[1].get(0, 0), x[1].get(1, 0), x[1].get(2, 0), x[0]))

        switch = 0
        for _, group in groupby(strat, key=lambda x: x[1]):
            group = list(group)
            if len(group) == 1:
                tmp = train if switch else test
                tmp.documents[group[0][0]] = self.documents[group[0][0]]
                switch = 1 - switch
            else:
                len_train = round(len(group) * percentage)
                random.seed(2727)
                random.shuffle(group)

                for key in group[:len_train]:
                    train.documents[key[0]] = self.documents[key[0]]
                for key in group[len_train:]:
                    test.documents[key[0]] = self.documents[key[0]]

        return train, test class AnnJsonAnnotationReader(AnnotationReader):
    """
    Reads the annotations from .ann.json format.

    Implements the abstract class Annotator.
    """

    def __init__(self, directory, read_only_class_id=None, delete_incomplete_docs=True, is_predicted=False, read_relations=False, whole_basename_as_docid=False, raise_exception_on_incosistencies=True, hdfs_url=None, hdfs_user=None):
        self.directory = directory
        """the directory containing *.ann.json files"""

        if read_only_class_id is not None and not isinstance(read_only_class_id, list):
            read_only_class_id = [read_only_class_id]
        self.read_only_class_id = read_only_class_id
        """whether to read in only entities with given class_id's (single id or list of). Otherwise if None, read all entities"""

        self.delete_incomplete_docs = delete_incomplete_docs
        """delete documents from the dataset that are not marked as 'anncomplete' provided the docs are not predicted"""
        self.is_predicted = is_predicted
        """whether the annotation is predicted or real, which determines where it will be saved"""
        self.read_relations = read_relations
        """whether relations should be read as well"""
        self.whole_basename_as_docid = whole_basename_as_docid
        self.raise_exception_on_incosistencies = raise_exception_on_incosistencies

        self.hdfs_client = maybe_get_hdfs_client(hdfs_url, hdfs_user)


    def annotate(self, dataset):
        """
        :type dataset: nalaf.structures.data.Dataset
        """
        read_docs = set()

        if self.hdfs_client is None:
            self.__read_files_localfs(dataset, read_docs)
        else:
            self.__read_files_hdfs(dataset, read_docs)

        # Delete docs with no ann.jsons
        docs_to_delete = set(dataset.documents.keys()) - read_docs
        for doc_id in docs_to_delete:
            del dataset.documents[doc_id]

        dataset.documents = OrderedDict((doc_id, doc) for doc_id, doc in dataset.documents.items())

        dataset.validate_entity_offsets()

        return dataset


    def __read_files_localfs(self, dataset, read_docs=None):
        if read_docs is None:
            read_docs = set()

        if not os.path.isdir(self.directory):
            filenames = [self.directory]
        else:
            filenames = glob.glob(str(self.directory + "/**/*.ann.json"), recursive=True)

        for filename in filenames:
            with open(filename, 'r', encoding="utf-8") as reader:
                doc_id = self.__read_annjson(reader, filename, dataset)
                if doc_id:
                    read_docs.add(doc_id)

        return read_docs


    def __read_files_hdfs(self, dataset, read_docs=None):
        if read_docs is None:
            read_docs = set()

        if not is_hdfs_directory(self.hdfs_client, self.directory):
            filenames = [self.directory]
        else:
            filenames = walk_hdfs_directory(self.hdfs_client, self.directory, lambda fname: fname.endswith(".ann.json"))

        for filename in filenames:
            with self.hdfs_client.read(filename, encoding="utf-8") as reader:
                doc_id = self.__read_annjson(reader, filename, dataset)
                read_docs.add(doc_id)

        return read_docs



    def __read_annjson(self, reader, filename, dataset):
        try:
            doc_id = os.path.basename(filename).replace('.ann.json', '').replace('.json', '')
            if not self.whole_basename_as_docid and '-' in doc_id:
                doc_id = doc_id.split('-')[-1]

            try:
                ann_json = json.load(reader)
            except JSONDecodeError:
                logging.exception("The annjson with docid={} seems malformed.".format(doc_id))
                return

            try:
                document = dataset.documents[doc_id]
            except Exception:
                logging.exception("The annjson with docid={} was not in the whole plain dataset.".format(doc_id))
                return doc_id

            if not (ann_json['anncomplete'] or self.is_predicted) and self.delete_incomplete_docs:
                del dataset.documents[doc_id]

            else:

                for e in ann_json['entities']:

                    if self.read_only_class_id is None or e['classId'] in self.read_only_class_id:

                        part = document.parts[e['part']]

                        try:
                            normalizations = {key: obj['source']['id'] for key, obj in e['normalizations'].items()}
                        except KeyError:
                            print_warning("The normalization is badly formatted: (docid={}) {}".format(doc_id, str(e['normalizations'])))
                            normalizations = None

                        entity = Entity(
                            e['classId'],
                            e['offsets'][0]['start'],
                            e['offsets'][0]['text'],
                            e['confidence']['prob'],
                            norms=normalizations)

                        if self.is_predicted:
                            part.predicted_annotations.append(entity)
                        else:
                            part.annotations.append(entity)

                if self.read_relations:
                    for relation in ann_json['relations']:
                        # Note: no distinction with predicted_relations yet

                        part = document.parts[relation['entities'][0].split('|')[0]]

                        e1_start = int(relation['entities'][0].split('|')[1].split(',')[0])
                        e2_start = int(relation['entities'][1].split('|')[1].split(',')[0])

                        rel_id = relation['classId']

                        e1 = part.get_entity(e1_start, use_pred=False, raise_exception_on_incosistencies=self.raise_exception_on_incosistencies)
                        e2 = part.get_entity(e2_start, use_pred=False, raise_exception_on_incosistencies=self.raise_exception_on_incosistencies)

                        if (not self.raise_exception_on_incosistencies and (e1 is None or e2 is None)):
                            continue

                        rel = Relation(rel_id, e1, e2)

                        part.relations.append(rel)

                # delete parts that are not annotatable
                annotatable_parts = set(ann_json['annotatable']['parts'])
                part_ids_to_del = []
                for part_id, part in document.parts.items():
                    if part_id not in annotatable_parts:
                        part_ids_to_del.append(part_id)
                for part_id in part_ids_to_del:
                    del document.parts[part_id]

            return doc_id

        except Exception as err:
            if self.raise_exception_on_incosistencies:
                raise err
            else:
                pass
import argparse
import os
import pkg_resources
from nalaf.utils.readers import TextFilesReader, PMIDReader, HTMLReader
from nalaf.utils.readers import StringReader
from nalaf.utils.writers import ConsoleWriter, TagTogFormat, PubTatorFormat
from nalaf.structures.dataset_pipelines import PrepareDatasetPipeline
from nalaf.learning.crfsuite import PyCRFSuite
from nalaf.learning.taggers import StubSameSentenceRelationExtractor
from nalaf.structures.data import Dataset
from nalaf.utils.annotation_readers import AnnJsonAnnotationReader




ENT1_CLASS_ID = 'e_x'
ENT2_CLASS_ID = 'e_y'
REL_ENT1_ENT2_CLASS_ID = 'r_z'
ENTREZ_GENE_ID = 'n_w'
UNIPROT_ID = 'n_v'


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='A simple demo for using the nalaf pipeline for prediction')

    parser.add_argument('--color', help='uses color for highlighting predictions if supported '
                                        'otherwise prints them in new line',
                        action='store_true', default=True, dest='color')
    parser.add_argument('--no_color', help='prints predictions in new line',
                        action='store_false', dest='color')

    parser.add_argument('-o', '--output_dir', help='write the output to the provided directory, '
                                                   'the format can be specified with the -f switch, '
                                                   'otherwise the output will be written to the standard console')
    parser.add_argument('-f', '--file_format', help='the format for writing the output to a directory',
                        choices=['ann.json', 'pubtator'], default='ann.json')

    group = parser.add_mutually_exclusive_group(required=True)

    group.add_argument('-s', '--string', help='string you want to predict for')
    group.add_argument('-d', '--dir_or_file', help='directory or file you want to predict for')
    group.add_argument('-p', '--pmids', nargs='+', help='a single PMID or a list of PMIDs separated by space')
    args = parser.parse_args()

    if args.string:
        dataset = StringReader(args.string).read()

    elif args.pmids:
        dataset = PMIDReader(args.pmids).read()

    elif os.path.exists(args.dir_or_file):
        file = args.dir_or_file

        file_extension = os.path.splitext(file)[1]

        if file_extension == ".txt":
            dataset = TextFilesReader(file).read()
        elif file_extension == ".html":

            dataset = HTMLReader(file, whole_basename_as_docid=True, hdfs_url=None, hdfs_user=None).read()