import gzip
import logging
import os
import pandas as pd

from utils.databases.superclasses.fasta_db_superclass import FastaDatabase


logger = logging.getLogger(__name__)
logger.addHandler(logging.NullHandler())


class NcbiNrDatabase(FastaDatabase):
    name = "ncbi_nr"
    valid_types = [name]
    mandatory_keys = FastaDatabase.mandatory_keys[:]
    required_db_types = ["ncbi_taxdump", "ncbi_accession2taxid"]

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def get_list_of_acc2tax_files(self, *args, **kwargs):
        return self.get_helper_db_by_type('ncbi_accession2taxid').get_all_acc2tax_files(*args, **kwargs)

    def get_taxdump_map_file(self):
        logger.warning("Deprecated: Use 'get_tax2annot_map_file' instead")
        return self.get_tax2annot_map_file()

    def get_db_map_file(self):
        return self.get_acc2tax_map_file()

    def get_acc2tax_map_file(self):
        return self.join_with_db_path(os.path.join("processed", "acc2tax.map.gz"))

    def get_tax2annot_map_file(self, *args, **kwargs):
        return self.join_with_db_path(os.path.join("processed", "tax2annot.map.gz"), *args, **kwargs)

    def get_taxdump_file(self):
        return self.get_helper_db_by_type('ncbi_taxdump').get_taxdump_tar_gz_file()

    @staticmethod
    def create_acc2tax_map(taxdump_in, l_acc2taxid_in, nr_out, nr_missing_out):
        """
        Create table of two columns containing accessions from any file in l_acc2taxid_in, whose associated TaxID
         is present in taxdump_in
        :param taxdump_in:
        :param l_acc2taxid_in:
        :param nr_out:
        :param nr_missing_out:
        :return:
        """
        nr_header = ['#accession', 'taxid']
        # get tax IDs from taxdump.map
        taxids = pd.read_csv(taxdump_in, sep='\t', skiprows=1, usecols=[0], names=['taxid'], dtype=int)
        # set is smaller and performs better
        taxids = set(taxids['taxid'])

        # get acc2taxid from protacc2taxid, accv=accession version
        lst_chunks_acc2taxid = []
        for f in l_acc2taxid_in:
            lst_chunks_acc2taxid.append(
                pd.read_csv(
                    f, sep='\t', skiprows=1, names=['acc', 'acc-version', 'taxid', 'gi'],
                    usecols=['acc-version', 'taxid'], dtype={'acc-version': str, 'taxid': int}, chunksize=10 ** 7
                )
            )

        with gzip.open(nr_out, 'wt') as out_map_handle, \
                gzip.open(nr_missing_out, 'wt') as out_handle_protacc_only:
            for h in [out_map_handle, out_handle_protacc_only]:
                h.write('\t'.join(nr_header) + '\n')
            for chunk_accessor in lst_chunks_acc2taxid:
                for chunk_acc2taxid in chunk_accessor:
                    # boolean series of taxids from protacc2taxid present in taxdump.map
                    s_bool_acctaxid_in_taxdump = chunk_acc2taxid['taxid'].isin(taxids)
                    shared_tax = chunk_acc2taxid[s_bool_acctaxid_in_taxdump]
                    protacc_only = chunk_acc2taxid[~s_bool_acctaxid_in_taxdump]
                    shared_tax.to_csv(out_map_handle, sep='\t', header=False, index=False)
                    protacc_only.to_csv(out_handle_protacc_only, sep='\t', header=False, index=False)
