from utils.databases.superclasses.fasta_db_superclass import FastaDatabase
from utils.databases.superclasses.database_superclasses import Database
import pandas as pd
import gzip
import os


class UniprotBaseDatabase(FastaDatabase):
    mandatory_keys = FastaDatabase.mandatory_keys + ['acc_regexp']
    required_db_types = ["uniprot_tax"]

    def get_list_of_acc2tax_files(self, *args, **kwargs):
        return [self.get_helper_db_by_type('accession2taxid').get_prot_acc2tax_file(*args, **kwargs)]

    def get_acc2tax_map_file(self):
        return self.get_helper_db_by_type('uniprot_tax').get_acc2tax_map_file()

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

    def get_acc_regex_pattern(self):
        return self._config['acc_regexp']

    def get_tax2annot_map_file(self, *args, **kwargs):
        return self.get_helper_db_by_type('uniprot_tax').get_tax2annot_map_file(*args, **kwargs)


class UniprotTremblDatabase(UniprotBaseDatabase):
    name = 'uniprot_tr'
    valid_types = [name]


class UniprotSwissprotDatabase(UniprotBaseDatabase):
    name = "uniprot_sp"
    valid_types = [name]


class UniprotCompleteDatabase(UniprotBaseDatabase):
    name = 'uniprot_complete'
    valid_types = [name]
    required_db_types = UniprotBaseDatabase.required_db_types + ["uniprot_sp", "uniprot_tr"]

    def get_fasta_gz_file_list(self, *args, **kwargs):
        l_files = self.get_helper_db_by_type('uniprot_sp').get_fasta_gz_file_list(*args, **kwargs) + \
                  self.get_helper_db_by_type('uniprot_tr').get_fasta_gz_file_list(*args, **kwargs)
        return l_files


class UniprotTaxDatabase(Database):
    name = 'uniprot_tax'
    valid_types = [name]
    mandatory_keys = Database.mandatory_keys + ['files']

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

    def get_idmap(self):
        return self.join_with_db_path(self._config["files"]["idmap"])

    def get_taxmap_txt(self, *args, **kwargs):
        return self.join_with_db_path(self._config["files"]["uniprot_tax_txt"], *args, **kwargs)

    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)

    @staticmethod
    def create_acc2tax_map(uniprot_tax_in, idmapping_in, sharedtax_out, protacc_only_out):
        col_names = ['#accession', 'taxid']
        # get taxa from uniprot_tax.txt
        uniprot_taxids = pd.read_csv(
            uniprot_tax_in, header=0, sep='\t', usecols=[0], names=['taxid'], dtype=int
        )
        # get idmapping.dat.ncbi_taxid.gz
        chunks_acc_taxid = pd.read_csv(
            idmapping_in, sep='\t', names=col_names, dtype={'#accession': str, 'taxid': int}, chunksize=10 ** 7
        )

        with gzip.open(sharedtax_out, 'wt') as out_map_handle, \
                gzip.open(protacc_only_out, 'wt') as out_handle_unip_only:
            for h in [out_map_handle, out_handle_unip_only]:
                h.write('\t'.join(col_names) + '\n')
            for chunk_acc_to_taxid in chunks_acc_taxid:
                # boolean series of taxids from idmapping present in uniprot_tax
                s_bool_acctaxid_in_uptaxids = chunk_acc_to_taxid['taxid'].isin(uniprot_taxids['taxid'])
                shared_tax = chunk_acc_to_taxid[s_bool_acctaxid_in_uptaxids]
                acc_in_up_only = chunk_acc_to_taxid[~s_bool_acctaxid_in_uptaxids]
                shared_tax.to_csv(out_map_handle, sep='\t', header=False, index=False)
                acc_in_up_only.to_csv(out_handle_unip_only, sep='\t', header=False, index=False)
