import_igc_data.py

#!/usr/bin/env python
import argparse
import logging
import sys
from itertools import islice
from requests.exceptions import HTTPError

from bioapi import MetageneDBCatalogGeneAPI, MetageneDBCatalogTaxonomyAPI
from slugify import slugify

from metagenedb.common.utils.parsers import IGCLineParser

_LOGGER = logging.getLogger(__name__)


class ImportIGCGenes(object):
    METAGENEDB_GENE_API = MetageneDBCatalogGeneAPI
    METAGENEDB_TAXONOMY_API = MetageneDBCatalogTaxonomyAPI

    PHYLUM_COL = 'taxo_phylum'
    GENUS_COL = 'taxo_genus'
    SELECTED_KEYS = ['gene_id', 'length', 'kegg_ko', PHYLUM_COL, GENUS_COL]

    def __init__(self, annotation_file, url, skip_tax=False, skip_functions=False):
        self.annotation_file = annotation_file
        self.url = url
        self.metagenedb_gene_api = self.METAGENEDB_GENE_API(base_url=self.url)
        self.metagenedb_taxonomy_api = self.METAGENEDB_TAXONOMY_API(base_url=self.url)
        self.processed_genes = 0
        self.skipped_genes = 0
        # Skip some insertion if specified in script options
        self.skip_tax = skip_tax
        self.skip_functions = skip_functions

    def _parse_gene(self, raw_line, selected_keys=SELECTED_KEYS):
        """
        Use IGCLineParser and return selected keys
        """
        gene_parser = IGCLineParser()
        all_dict = gene_parser.gene(raw_line)
        selected_dict = {k: v for k, v in all_dict.items() if k in selected_keys}
        return selected_dict

    def _select_taxonomy(self, taxonomy_dict, unknown_val='unknown'):
        """
        Select the taxonomy to be assigned for the gene.
        genus has priority on phylum. If both unknow, remove the taxonomy key
        """
        phylum = taxonomy_dict.pop(self.PHYLUM_COL)
        genus = taxonomy_dict.pop(self.GENUS_COL)
        resp_dict = {}
        if genus != unknown_val:
            resp_dict = self.metagenedb_taxonomy_api.get_all(params={'name': genus, 'rank': 'genus'})
            if len(resp_dict['results']) > 1:
                _LOGGER.warning(f"More than 1 result found for genus {genus}. First result is kept.")
        elif phylum != unknown_val:
            resp_dict = self.metagenedb_taxonomy_api.get_all(params={'name': phylum, 'rank': 'phylum'})
            if len(resp_dict['results']) > 1:
                _LOGGER.warning(f"More than 1 result found for phylum {phylum}. First result is kept.")
        if resp_dict:
            taxonomy_dict.update(
                {'taxonomy': resp_dict['results'][0]['tax_id']}
            )
        return taxonomy_dict

    def _clean_gene(self, gene_dict):
        gene_dict['gene_name'] = gene_dict['gene_id']
        gene_dict['gene_id'] = slugify(gene_dict['gene_id'])
        gene_dict['functions'] = [{'function_id': gene_dict.pop('kegg_ko')}]
        if self.skip_tax:
            gene_dict.pop('taxonomy')
        if self.skip_functions:
            gene_dict.pop('functions')
        return gene_dict

    def _upsert_gene(self, gene_dict):
        clean_gene_dict = self._clean_gene(gene_dict)
        try:
            gene_id = clean_gene_dict['gene_id']
            self.metagenedb_gene_api.get(gene_id)  # Try to get obj to check if it exists
            self.metagenedb_gene_api.put(clean_gene_dict, entry_id=gene_id)
        except HTTPError:
            self.metagenedb_gene_api.post(clean_gene_dict)

    def _insert_gene_list(self, chunk_genes):
        for gene_line in chunk_genes:
            gene_dict = self._parse_gene(gene_line)
            gene_dict_with_taxo = self._select_taxonomy(gene_dict)
            try:
                self._upsert_gene(gene_dict_with_taxo)
            except HTTPError as e:
                self.skipped_genes += 1
                _LOGGER.warning(f"{e.response.json()} for gene_id: {gene_dict.get('gene_id')}. Insertion skipped.")

    def load_annotation_file_to_db_in_chunks(self, chunk_size=100000):
        with open(self.annotation_file, 'r') as file:
            while True:
                chunk_genes = list(islice(file, chunk_size))
                if not chunk_genes:
                    break
                self.processed_genes += len(chunk_genes)
                self._insert_gene_list(chunk_genes)
                _LOGGER.info(f"{self.processed_genes} genes inserted/updated so far...")
        _LOGGER.info(f"[DONE] {self.processed_genes} genes inserted/updated.")
        _LOGGER.info(f"[DONE] {self.skipped_genes} genes skipped.")


def parse_arguments():
    """
    Defines parser.
    """
    parser = argparse.ArgumentParser(description='Populate database from a given IGC annotation file.')
    # Common arguments for analysis and annotations
    parser.add_argument('annotation', help='IGC annotation file')
    parser.add_argument('--url', help='base URL of the instance.', default='http://localhost/')
    parser.add_argument('--skip_taxonomy', action='store_true', help='Skip taxonomy information from genes.')
    parser.add_argument('--skip_functions', action='store_true', help='Skip functions information from genes.')

    try:
        return parser.parse_args()
    except SystemExit:
        sys.exit(1)


def run():
    args = parse_arguments()
    import_igc_genes = ImportIGCGenes(args.annotation, args.url,
                                      skip_tax=args.skip_taxonomy, skip_functions=args.skip_functions)
    import_igc_genes.load_annotation_file_to_db_in_chunks()


if __name__ == "__main__":
    run()