from utils.databases.superclasses.database_superclasses import Database

cog_funcats_dict = {'A': ('information storage and processing', 'RNA processing and modification'),
                    'B': ('information storage and processing',
                          'Chromatin structure and dynamics'),
                    'C': ('metabolism', 'Energy production and conversion'),
                    'D': ('cellular processes and signaling',
                          'Cell cycle control, cell division, chromosome partitioning'),
                    'E': ('metabolism', 'Amino acid transport and metabolism'),
                    'F': ('metabolism', 'Nucleotide transport and metabolism'),
                    'G': ('metabolism', 'Carbohydrate transport and metabolism'),
                    'H': ('metabolism', 'Coenzyme transport and metabolism'),
                    'I': ('metabolism', 'Lipid transport and metabolism'),
                    'J': ('information storage and processing',
                          'Translation, ribosomal structure and biogenesis'),
                    'K': ('information storage and processing', 'Transcription'),
                    'L': ('information storage and processing',
                          'Replication, recombination and repair'),
                    'M': ('cellular processes and signaling',
                          'Cell wall/membrane/envelope biogenesis'),
                    'N': ('cellular processes and signaling', 'Cell motility'),
                    'O': ('cellular processes and signaling',
                          'Posttranslational modification, protein turnover, chaperones'),
                    'P': ('metabolism', 'Inorganic ion transport and metabolism'),
                    'Q': ('metabolism',
                          'Secondary metabolites biosynthesis, transport and catabolism'),
                    'R': ('poorly characterized', 'General function prediction only'),
                    'S': ('poorly characterized', 'Function unknown'),
                    'T': ('cellular processes and signaling', 'Signal transduction mechanisms'),
                    'U': ('cellular processes and signaling',
                          'Intracellular trafficking, secretion, and vesicular transport'),
                    'V': ('cellular processes and signaling', 'Defense mechanisms'),
                    'W': ('cellular processes and signaling', 'Extracellular structures'),
                    'Y': ('cellular processes and signaling', 'Nuclear structure'),
                    'Z': ('cellular processes and signaling', 'Cytoskeleton')}


class EggnogDatabaseV4(Database):
    name = "eggnog_v4"
    valid_types = [name]
    mandatory_keys = Database.mandatory_keys[:]

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

    def get_og_annotations_file(self, *args, **kwargs):
        return self.join_with_db_path(self._config["files"]['og_annotations'], *args, **kwargs)

    def get_fun_cats_file(self, *args, **kwargs):
        return self.join_with_db_path(self._config["files"]['fun_cats'], *args, **kwargs)

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


class EggnogDatabaseV5(Database):
    name = "eggnog_v5"
    valid_types = [name]
    mandatory_keys = Database.mandatory_keys[:]

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

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

    @staticmethod
    def write_task_map(task_map, emapper_result, source_task_string,
                       unclassified_annotation_string):
        headline = "#acc\tsource\tmain_role\tsubrole\tog\tdesc\tseed_ortholog\n"
        with open(task_map, 'w') as outhandle:
            outhandle.write(headline)
            with open(emapper_result, 'r') as result_handle:
                for line in result_handle:
                    if line.startswith('#'):
                        continue
                    fields = line.strip().split("\t")
                    acc = fields[0]
                    seed_orthologue = fields[1]
                    eggnog_ogs = fields[4]
                    cog_categories = fields[6]
                    description = fields[7]
                    # eggnog_ogs: OG@taxonID|taxonName,OG@taxonID|taxonName (for complete lineage starting with root)
                    # take OG from most specific lineage item
                    og = eggnog_ogs.split(',')[-1].split('@')[0]
                    for cog_category in cog_categories:
                        if cog_category == "-":
                            mainrole = subrole = "-"
                        else:
                            mainrole, subrole = cog_funcats_dict[cog_category]
                        lst_annotations_raw = [mainrole, subrole, og, description, seed_orthologue]
                        lst_annotations = [unclassified_annotation_string if x == "-" else x for x in
                                           lst_annotations_raw]
                        outhandle.write("\t".join([acc, source_task_string, *lst_annotations]) + "\n")
