import os
import warnings
import logging
import pickle
import logging
from importlib import resources

import pandas as pnd
import cobra

from ..commons.metrics import log_metrics
from ..commons.metrics import log_unbalances

from .gapfillutils import import_from_universe

from .polishing import remove_disconnected
from .polishing import remove_orphans
from .polishing import remove_sinks_demands

    
    
def load_input_universe(logger, universe):
    
    
    # built-in or fresh universe path
    path_to_unixml = None
    
    
    # use the buil-in one if no fresh universe wa provided
    if universe == '-':
        logger.info(f"No fresh universe provided: loading built-in whole universe...")
        with resources.path("gsrap.assets", f"universe.xml") as asset_path: 
            
            # this is the sbml to load
            path_to_unixml = asset_path
    
    
    # user provided a fresh universe
    else: 
        # check if files exist
        if os.path.isfile(universe) == False: 
            logger.error(f"Provided --universe doesn't exist: {universe}.")
            return 1
        
        # check extension
        if universe.endswith('.xml') == False:
            logger.error(f"Provided --universe must be in cobrapy-compatible SBML format (.xml extension).")
            return 1

        # this is the sbml to load
        logger.info(f"Loading provided fresh whole universe...")
        path_to_unixml = universe
        
        
    # Ignore eventual "https://identifiers.org/curator_notes/ does not conform to ..."
    with warnings.catch_warnings():  # temporarily suppress warnings for this block
        warnings.simplefilter("ignore")  # ignore all warnings

        # disable warnings
        cobra_logger = logging.getLogger("cobra.io.sbml")
        old_level = cobra_logger.level
        cobra_logger.setLevel(logging.ERROR)   

        # load the model
        universe = cobra.io.read_sbml_model(path_to_unixml)

        # restore original behaviour: 
        cobra_logger.setLevel(old_level)   
        
    
    # log main universe metrics:
    log_metrics(logger, universe)
    #log_unbalances(logger, universe)  # commented as it is 0 by design
        
        
    return universe



def prune_universe(logger, universe, taxon, transrule):
    
    
    # get available kingdoms and phyla
    avail_kingdoms = set()
    avail_phyla = set()
    for g in universe.genes:
        if 'kingdom' in g.annotation.keys():
            annot = g.annotation['kingdom']
            # saved as lists, resulting str after reading: now back to lists
            if type(annot) == str: annot = [annot]
            for i in annot:
                avail_kingdoms.add(i)
        if 'phylum' in g.annotation.keys():
            annot = g.annotation['phylum']
            # saved as lists, resulting str after reading: now back to lists
            if type(annot) == str: annot = [annot]
            for i in annot:
                avail_phyla.add(i)
                
    
    # check formatting
    if ':' not in taxon:
        logger.error(f"Format of provided --taxon is not compliant (missing ':').")
        return 1
    if len(taxon.split(':')) > 2:
        logger.error(f"Format of provided --taxon is not compliant (too many ':').")
        return 1
    level, name = taxon.split(':')
    allowed_levels = ['kingdom', 'phylum']
    if level not in allowed_levels:
        logger.error(f"Level '{level}' provided with --taxon is not allowed: please choose between {allowed_levels}.")
        return 1
    if level=='kingdom' and name not in avail_kingdoms:
        logger.error(f"Cannot find a taxon of level 'kingdom' named '{name}': please choose between {avail_kingdoms}.")
        return 1
    if level=='phylum' and name not in avail_phyla:
        logger.error(f"Cannot find a taxon of level 'phylum' named '{name}': please choose between {avail_phyla}.")
        return 1
        
    
    # create a copy of the universe
    logger.info(f"Pruning universe to represent taxon '{level}:{name}'...")
    pruned_universe = universe.copy()
    
    
    # remove sinks/demands
    n = remove_sinks_demands(logger, pruned_universe)   # none logger to avoid confusion.
        
        
    # mark orthologs for deletion
    to_remove = []
    for g in pruned_universe.genes:
        if level in g.annotation.keys():  # 'orphan' reactions remain!
            annot = g.annotation[level]
            # saved as lists, resulting str after reading: now back to lists
            if type(annot) == str: annot = [annot]
            if name not in annot:
                to_remove.append(g)
    logger.debug(f"Removing {len(to_remove)} non-taxon orthologs...")
           
        
    # delete marked orthologs
    # trick to avoid the WARNING "cobra/core/group.py:147: UserWarning: need to pass in a list" 
    # triggered when trying to remove reactions that are included in groups. 
    with warnings.catch_warnings():  # temporarily suppress warnings for this block
        warnings.simplefilter("ignore")  # ignore all warnings
        cobra_logger = logging.getLogger("cobra.util.solver")
        old_level = cobra_logger.level
        cobra_logger.setLevel(logging.ERROR)   

        cobra.manipulation.delete.remove_genes(pruned_universe, to_remove, remove_reactions=True)

        # restore original behaviour: 
        cobra_logger.setLevel(old_level) 
        
        
    # reintroduce transport reactions as orphans:
    if transrule == 'rescue': 
        rids_pruned_universe = [r.id for r in pruned_universe.reactions]
        for r in universe.reactions:
            if len(set([m.id.rsplit('_',1)[-1] for m in r.metabolites])) > 1:
                if r.id not in rids_pruned_universe:
                    logger.debug(f"Reintroducing trasport reaction '{r.id}'...")
                    import_from_universe(pruned_universe, universe, r.id, gpr='')  
                    
                    
    # remove universal orphans
    if transrule == 'rescue': keep_transport = True 
    elif transrule == 'ignore': keep_transport = False
    else: keep_transport = False
    n = remove_orphans(logger, pruned_universe, keep_transport)
                    
    
    # remove disconnected metabolites
    n = remove_disconnected(logger, pruned_universe)  # none logger to avoid confusion.
    
    
    # log main universe metrics:
    log_metrics(logger, pruned_universe)
    
    
    return pruned_universe



def load_input_eggnog(logger, eggnog):
    
    
    # load eggnog annotations
    df_eggnog = pnd.read_csv(eggnog, sep='\t', comment='#', header=None)
    df_eggnog.columns = 'query	seed_ortholog	evalue	score	eggNOG_OGs	max_annot_lvl	COG_category	Description	Preferred_name	GOs	EC	KEGG_ko	KEGG_Pathway	KEGG_Module	KEGG_Reaction	KEGG_rclass	BRITE	KEGG_TC	CAZy	BiGG_Reaction	PFAMs'.split('\t')
    df_eggnog = df_eggnog.set_index('query', drop=True, verify_integrity=True)
    
    
    return df_eggnog



def load_keggorg_like_eggnog(logger, keggorg, outdir):
    
    
    # load raw data, downloaded form kegg: 
    df_keggorg = pickle.load(open(os.path.join(outdir, f'{keggorg}.keggorg'), 'rb'))
    df_keggorg = df_keggorg.set_index('gid', drop=True, verify_integrity=True)


    # create an eggnog-like dataframe:
    df_eggnog_like = []   # list of dict future df
    for gid in df_keggorg.index:
        row_dict = {}

        row_dict['query'] = gid
        row_dict['PFAMs'] = ','.join(df_keggorg.loc[gid, 'Pfam']) if type(df_keggorg.loc[gid, 'Pfam'])==list else '-'
        row_dict['KEGG_ko'] = df_keggorg.loc[gid, 'ko'] if type(df_keggorg.loc[gid, 'ko'])==str else '-'

        df_eggnog_like.append(row_dict)
    df_eggnog_like = pnd.DataFrame.from_records(df_eggnog_like)


    # appen missing coluns and sort
    eggnog_columns = 'query	seed_ortholog	evalue	score	eggNOG_OGs	max_annot_lvl	COG_category	Description	Preferred_name	GOs	EC	KEGG_ko	KEGG_Pathway	KEGG_Module	KEGG_Reaction	KEGG_rclass	BRITE	KEGG_TC	CAZy	BiGG_Reaction	PFAMs'.split('\t')
    for c in eggnog_columns:  
        if c not in df_eggnog_like.columns:
            df_eggnog_like[c] = '-'
    df_eggnog_like = df_eggnog_like[eggnog_columns]


    # set the index like in eggnog
    df_eggnog_like = df_eggnog_like.set_index('query', drop=True, verify_integrity=True)
    return df_eggnog_like



def parse_eggnog(df_eggnog):
    
    
    # PART 1. get KO codes available
    gid_to_kos = {}
    ko_to_gids = {}
    for gid, kos in df_eggnog['KEGG_ko'].items():
        if kos == '-': 
            continue
            
        if gid not in gid_to_kos.keys(): 
            gid_to_kos[gid] = set()
            
        kos = kos.split(',')
        kos = [i.replace('ko:', '') for i in kos]
        for ko in kos: 
            if ko not in ko_to_gids.keys(): 
                ko_to_gids[ko] = set()
                
            # populate dictionaries
            ko_to_gids[ko].add(gid)
            gid_to_kos[gid].add(ko)

    
    return ko_to_gids, gid_to_kos



def get_modeled_kos(model):
    
    
    # get modeled KO ids:
    modeled_gid_to_ko = {}
    modeled_ko_to_gid = {}
    
    for g in model.genes:
        if g.id in ['orphan', 'spontaneous']: 
            continue
        corresponding_ko = g.annotation['ko']
        
        modeled_gid_to_ko[g.id] = corresponding_ko
        modeled_ko_to_gid[corresponding_ko] = g.id
        
    modeled_kos = list(modeled_gid_to_ko.values())
        
    return modeled_kos, modeled_gid_to_ko, modeled_ko_to_gid



def subtract_kos(logger, model, eggonog_ko_to_gids):
    
    
    modeled_kos, _, modeled_ko_to_gid = get_modeled_kos(model)
        
        
    to_remove = []  # genes to delete
    for ko in modeled_kos: 
        if ko not in eggonog_ko_to_gids.keys():
            gid_to_remove = modeled_ko_to_gid[ko]
            to_remove.append(model.genes.get_by_id(gid_to_remove))
            
    
    # remove also orphan reactions!
    to_remove.append(model.genes.get_by_id('orphan'))
    
    
    # delete marked genes!
    # trick to avoid the WARNING "cobra/core/group.py:147: UserWarning: need to pass in a list" 
    # triggered when trying to remove reactions that are included in groups. 
    with warnings.catch_warnings():  # temporarily suppress warnings for this block
        warnings.simplefilter("ignore")  # ignore all warnings
        cobra_logger = logging.getLogger("cobra.util.solver")
        old_level = cobra_logger.level
        cobra_logger.setLevel(logging.ERROR)   

        cobra.manipulation.delete.remove_genes(model, to_remove, remove_reactions=True)

        # restore original behaviour: 
        cobra_logger.setLevel(old_level)   
        
   
    logger.info(f"Found annotations for {len(model.genes)} orthologs modeled in universe.")
    return 0



def translate_remaining_kos(logger, model, eggonog_ko_to_gids):
    
    
    n_starting_orthologs = len(model.genes)
    _, modeled_gid_to_ko, _ = get_modeled_kos(model) 
    
    
    # iterate reactions:
    for r in model.reactions:

        gpr = r.gene_reaction_rule

        # force each gid to be surrounded by spaces: 
        gpr = ' ' + gpr.replace('(', ' ( ').replace(')', ' ) ') + ' '
        
        for gid in modeled_gid_to_ko.keys():
            if f' {gid} ' in gpr:
                
                new_gids = eggonog_ko_to_gids[modeled_gid_to_ko[gid]]
                gpr = gpr.replace(f' {gid} ', f' ({" or ".join(new_gids)}) ')       
            

        # remove spaces between parenthesis
        gpr = gpr.replace(' ( ', '(').replace(' ) ', ')')
        # remove spaces at the extremes: 
        gpr = gpr[1: -1]


        # New genes are introduced. Parethesis at the extremes are removed if not necessary. 
        r.gene_reaction_rule = gpr
        r.update_genes_from_gpr()
            
            
    # remaining old 'Cluster_'s need to removed.
    # remove if (1) hte ID starts with clusters AND (2) they are no more associated with any reaction
    to_remove = []
    for g in model.genes:
        
        if g.id in ['orphan', 'spontaneous']:
            continue
            
        if g.id in modeled_gid_to_ko.keys() and len(g.reactions)==0:
            to_remove.append(g)
            
    # warning suppression not needed here, as no reaction is actually removed.
    cobra.manipulation.delete.remove_genes(model, to_remove, remove_reactions=True)
    
        
    logger.info(f"Translated {n_starting_orthologs} orthologs to {len(model.genes)} strain-specific genes.")
    return 0
        
    

def restore_gene_annotations(logger, model, universe, eggonog_gid_to_kos):
    
    
    for g in model.genes:
        if g.id == 'spontaneous': 
            continue
            
        names = []    
        for ko in eggonog_gid_to_kos[g.id]:
            
            # get the corresponding universal gene:
            uni_g = None
            for ug in universe.genes:
                if 'ko' not in ug.annotation.keys():
                    continue
                if ug.annotation['ko']==ko:  # take the first (and only)
                    uni_g = ug
                    break
            if uni_g == None:  
                # The ko provided by eggnog-mapper is still not modeled in the universe.
                # Multiple ko are possible for each gene. Of these, only 1 could b modeled.
                continue
            
            
            # transfer annotations of this ko/universal gene:
            for key in uni_g.annotation.keys():
                if key == 'ko':
                    continue   # resulting models will loose links to kos.
                if key not in g.annotation:
                    g.annotation[key] = []
                items = uni_g.annotation[key]
                if type(items)==str:  items = [items]
                for i in items:
                    g.annotation[key].append(i)
                    
            # collect names
            names.append(uni_g.name)
        g.name = '; '.join(names)
        
        
        
def append_keggorg_gene_annots(logger, model, keggorg, outdir):
    

    # load raw data, downloaded form kegg: 
    df_keggorg = pickle.load(open(os.path.join(outdir, f'{keggorg}.keggorg'), 'rb'))
    df_keggorg = df_keggorg.set_index('gid', drop=True, verify_integrity=True)
    
    
    # KEGG can provide some useful (ie, used in Memote) gene annotations:
    for g in model.genes:
        if g.id in df_keggorg.index:
            
            g.annotation['kegg.genes'] = [keggorg + ':' + g.id]
            
            if 'NCBI-GeneID' in df_keggorg.columns:
                g.annotation['ncbigene'] = df_keggorg.loc[g.id, 'NCBI-GeneID'] if type(df_keggorg.loc[g.id, 'NCBI-GeneID'])==list else []
            if 'NCBI-ProteinID' in df_keggorg.columns:
                g.annotation['ncbiprotein'] = df_keggorg.loc[g.id, 'NCBI-ProteinID'] if type(df_keggorg.loc[g.id, 'NCBI-ProteinID'])==list else []
            if 'ASAP' in df_keggorg.columns:
                g.annotation['asap'] = df_keggorg.loc[g.id, 'ASAP'] if type(df_keggorg.loc[g.id, 'ASAP'])==list else []
            if 'UniProt' in df_keggorg.columns:
                g.annotation['uniprot'] = df_keggorg.loc[g.id, 'UniProt'] if type(df_keggorg.loc[g.id, 'UniProt'])==list else []
                
    logger.info("Added gene annotations retrieved from KEGG.")
                
                
                
            
