"""
NLM CellLink example code:
Identify the most common cell phenotype mention in a BioC-XML file
 (such as the CellLink train set) and print where it can be found.

To run, replace filepath.
"""

import bioc
import collections

filepath = "train.xml" # edit here
cell_pheno_mentions = collections.defaultdict(list)

# load the bioc-xml file
with open(filepath, 'r', encoding='utf-8') as readfp:
    bioc_collection = bioc.load(readfp)

# iterate through each annotation
for bioc_document in bioc_collection.documents:
    for bioc_passage in bioc_document.passages:
        for bioc_annotation in bioc_passage.annotations:
            
            # filter for cell phenotypes
            if bioc_annotation.infons['type'] == 'cell_phenotype':
                
                # store metadata from the bioc_passage and bioc_annotation
                metadata = {
                    'pmid': bioc_passage.infons['article-id_pmid'],
                    'pmc': bioc_passage.infons.get('article-id_pmc'), # not all passages will have a PMCID
                    'passage_index': bioc_passage.infons['passage_id'].split('_')[1],
                    'annotation_offset': bioc_annotation.locations[0].offset
                    }
                
                # save the mention text and metadata
                cell_pheno_mentions[bioc_annotation.text].append(metadata)



most_frequent_mention = max(cell_pheno_mentions, key = lambda mention: len(cell_pheno_mentions[mention]))
frequency = len(cell_pheno_mentions[most_frequent_mention])
example_metadata = cell_pheno_mentions[most_frequent_mention][0]

print(f"The most frequent cell phenotype mention in the input bioc-xml is '{most_frequent_mention}',",
      f"which appears {frequency} times, including in article PMID {example_metadata['pmid']}",
      f"(PMC{example_metadata['pmc']}), at the {example_metadata['annotation_offset']}st character",
      f"of passage #{example_metadata['passage_index']}.")

# Expected output:
# "The most frequent cell phenotype mention in the input bioc-xml is 'macrophages', which appears 109 times, including in article PMID 32429316 (PMC7279007), at the 231st character of passage #28."