"""
NLM CellLink example code:
Count various statistics on BioC-XML documents (including number of documents;
 number of passages; number of mentions, unique mentions, identifiers, and
 unique identifiers for each entity type).

Run: python characterize.py <path_1> ... <path_n>
For example: python characterize.py train.xml val.xml
"""

import re
import sys
import os
import statistics

from bioc import biocxml
from collections import Counter

MENTION_TYPE_KEY = "type"
IDENTIFIER_KEY = "identifier"


def tokenize(text):
    return re.findall(r'\b\w+\b', text.replace("_", " ").lower())


def parse_ellipsis_element(element):
    if element is None:
        return (None, [])
    element = element.strip()
    if element == "None" or element == "-":
        return (None, [])
    end_paren_index = element.find(")")
    accession_index = end_paren_index + 1 if end_paren_index >= 0 else -1
    qualifier_text = element[:accession_index] if accession_index >= 0 else None
    element_identifier_text = element[accession_index:] if accession_index >= 0 else element
    element_identifiers = element_identifier_text.split(",")
    return (qualifier_text, element_identifiers)


def parse_identifier_list(identifier_list_text):
    if identifier_list_text is None:
        return []
    return [parse_ellipsis_element(element) for element in identifier_list_text.split(";")]


class BCC_docid:
    def __init__(self, docid, pmid, pmcid):
        self.docid = docid
        self.pmid = pmid
        self.pmcid = pmcid


class BCC_passage:
    def __init__(self, passage_id, passage_text):
        self.id = passage_id
        self.text = passage_text


class BCC_annotation:
    def __init__(self, mention_text, mention_type, identifier):
        self.mention_text = mention_text
        self.mention_type = mention_type
        self.identifier = identifier


class BioCCharacterizer:
    def __init__(self):
        self.document_ids = list()
        self.passages = list()
        self.annotations = list()
    
    def process_path(self, pathname, extension = None):
        if os.path.isdir(pathname):
            self.process_dir(pathname, extension)
        elif os.path.isfile(pathname):
            self.process_file(pathname)
        else:
            raise RuntimeError(f"Path is not a directory or normal file: {pathname}")

    def process_dir(self, dirname, extension = None):
        dir = os.listdir(dirname)
        for item in dir:
            filename = dirname + "/" + item
            if os.path.isfile(filename) and (extension is None or filename.endswith(extension)):
                self.process_file(filename)

    def process_file(self, filename, extension = None):
        if not extension is None and not filename.endswith(extension):
            print("Ignoring file " + filename)
            return
        print("Processing file " + filename)
        with open(filename, "r") as fp:
            collection = biocxml.load(fp)
        for document in collection.documents:
            header = document.passages[0]
            pmid = header.infons.get("article-id_pmid")
            pmcid = header.infons.get("article-id_pmc")
            self.document_ids.append(BCC_docid(document.id, pmid, pmcid))
            for passage_index, passage in enumerate(document.passages):
                passage_id = passage.infons.get("passage_id", (document.id, passage_index))
                self.passages.append(BCC_passage(passage_id, passage.text))
                for annotation in passage.annotations:
                    self.annotations.append(BCC_annotation(annotation.text, 
                                                           annotation.infons.get(MENTION_TYPE_KEY),
                                                           annotation.infons.get(IDENTIFIER_KEY)))

def ID_is_exact(identifier):
    if "(skos:related)" in identifier:
        return False
    elif "-" in identifier:
        return False
    elif "(skos:exact)" in identifier:
        return True
    elif "None" == identifier:
        return False
    else:
        raise Exception("Identifier qualifiers not formatted as expected.")


def main():
    characterizer = BioCCharacterizer()

    for path in sys.argv[1:]:
        characterizer.process_path(path, ".xml")

    print("DOCUMENTS")
    print("Number of documents: {}".format(len(characterizer.document_ids)))
    print("Number of unique PMIDs: {}".format(len({docid.pmid for docid in characterizer.document_ids if not docid.pmid is None})))
    print("Number of unique PMCIDs: {}".format(len({docid.pmcid for docid in characterizer.document_ids if not docid.pmcid is None})))
    print()
    print("PASSAGES")
    print("Number of passages: {}".format(len(characterizer.passages)))
    print("Number of unique passage texts: {}".format(len(set(passage.text for passage in characterizer.passages))))
    print()
    print("TOKENS")
    passage_tokens = [token for passage in characterizer.passages for token in tokenize(passage.text)]
    print("Number of tokens: {}".format(len(passage_tokens)))
    print("Number of unique tokens: {}".format(len(set(passage_tokens))))
    print()
    print("ANNOTATIONS")
    print("Number of annotations: {}".format(len(characterizer.annotations)))

    print("Cumulative count of unique mentions")
    mention_counts = Counter([annotation.mention_text for annotation in characterizer.annotations])
    max_count = max(mention_counts.values())
    count_counts = Counter(mention_counts.values())
    cumulative = 0
    for c in range(1, max_count + 1):
        cc = count_counts.get(c, 0)
        cumulative += cc
        #print(f"{c}\t{cc}\t{cumulative}")
    types2info = dict()
    types2info["ALL"] = {
        "mentions": list(),
        "tokens": list(),
        "identifiers": list(),
        "mentions_not_exact": list()
    }
    for annotation in characterizer.annotations:
        if not annotation.mention_type in types2info:
            types2info[annotation.mention_type] = {
                "mentions": list(),
                "tokens": list(),
                "identifiers": list(),
                "mentions_not_exact": list()
            }
        types2info["ALL"]["mentions"].append(annotation.mention_text)
        types2info[annotation.mention_type]["mentions"].append(annotation.mention_text)
        tokens = tokenize(annotation.mention_text)
        types2info["ALL"]["tokens"].extend(tokens)
        types2info[annotation.mention_type]["tokens"].extend(tokens)
        parsed_identifiers = parse_identifier_list(annotation.identifier)
        if not ID_is_exact(annotation.identifier):
            types2info["ALL"]["mentions_not_exact"].append(annotation.mention_text)
            types2info[annotation.mention_type]["mentions_not_exact"].append(annotation.mention_text)
        flattened_identifiers = [identifier for _qualifier, identifiers in parsed_identifiers for identifier in identifiers if not identifier is None]
        types2info["ALL"]["identifiers"].extend(flattened_identifiers)
        types2info[annotation.mention_type]["identifiers"].extend(flattened_identifiers)
    for type, info in types2info.items():
        print("Type \"{}\"".format(type))
        print("\tNumber of annotations: {}".format(len(info["mentions"])))
        print("\tNumber of unique mentions: {}".format(len(set(info["mentions"]))))
        annotation_lengths = [len(mention) for mention in info["mentions"]]
        print("\tMention lengths: mean={:.2f} median={}".format(statistics.mean(annotation_lengths), statistics.median(annotation_lengths)))
        print("\tNumber of total tokens: {}".format(len(info["tokens"])))
        print("\tNumber of unique tokens: {}".format(len(set(info["tokens"]))))
        print("\tNumber of total identifiers: {}".format(len(info["identifiers"])))
        print("\tNumber of unique identifiers: {}".format(len(set(info["identifiers"]))))
        print("\tNumber of unique mentions without an exact identifier: {}".format(len(set(info["mentions_not_exact"]))))


    print("Done.")

if __name__ == "__main__":
    main()