import pandas as pd
import numpy as np

"""helper functions for CRISPR primer design"""


def rename_attributes(att, named_attribute = None):
    """function to interpret attributes column"""
    
    if "=" in att:
        # assume this means GFF formatting
        dict_attributes = {
            i.split("=")[0]: i.split("=")[1]
            for i in att.split(";")
            if len(i.split("=")) > 1
        }
        
        # Try to extract the ID, Name, or Parent in that order
        ID = None


        if named_attribute:
            try:
                ID = dict_attributes.get(named_attribute)
                return ID
            except Exception as e:
                print(f"Error extracting {named_attribute} from annotations file: {e}")
        try:
            ID = dict_attributes.get("ID") or dict_attributes.get("Name") or dict_attributes.get("Parent")
        except Exception as e:
            print(f"Error extracting ID from annotations file: {e}")


    else:
        # assume this means GTF formatting
        dict_attributes = {i.split(" ")[0].strip('''"'''):i.split(" ")[1].strip('''"''').strip(''';"''') for i in att.split("; ") if len(i.split(" ")) > 1}

        # Try to extract the gene_id, transcript_id, or gene_name in that order
        ID = None
        if named_attribute:
            try:
                ID = dict_attributes.get(named_attribute)
                return ID
            except Exception as e:
                print(f"Error extracting {named_attribute} from annotations file: {e}")
        

        try:
            ID = dict_attributes.get("gene_id") or dict_attributes.get("transcript_id") or dict_attributes.get("gene_name")
        except Exception as e:
            print(f"Error extracting ID from annotations file: {e}")

    return ID


def load_GTF(GTF_file):
    """function to read GTF file and attributes"""
    dfGTF = pd.read_csv(GTF_file, sep = "\t", header = None)

    GTF_columns = ['seqname', 'source', 'feature', 'start', 'end', 'score', 'strand', 'frame', 'attribute']
    dfGTF.rename(columns = dict(zip(dfGTF.columns,GTF_columns )), inplace=True)

    # float to int
    dfGTF['start'] = dfGTF['start'].astype(int)
    dfGTF['end'] = dfGTF['end'].astype(int)

    # make sure start is always less than end
    dfGTF['start'], dfGTF['end'] = np.where(dfGTF['start'] > dfGTF['end'], dfGTF['end'], dfGTF['start']), np.where(dfGTF['start'] > dfGTF['end'], dfGTF['start'], dfGTF['end'])

    # add columns for gene_id and gene_biotype (e.g. CDS, tRNA, rRNA)
    dfGTF["gene_id"] = dfGTF.apply(
        lambda row: rename_attributes(row["attribute"], named_attribute = "gene_id"),
        axis=1
        )
    dfGTF["gene_biotype"] = dfGTF.apply(
        lambda row: rename_attributes(row["attribute"], named_attribute = "gene_biotype"),
        axis=1
        )

    # use only genes that code for proteins (CDS)
    dfGTF = dfGTF[dfGTF["feature"]=="CDS"].copy()
    
    # this column isn't needed anymore:
    dfGTF = dfGTF.drop(columns='feature')
    
    return dfGTF

def get_gene_info(dfGTF, gene_id, columns): # returns start and end as first two
    """function to get gene info from gtf file"""
    filtered_df = dfGTF[dfGTF["gene_id"]==gene_id][['start','end'] + columns]

    # to account for splices (still count entire gene)
    locs = []
    for i, row in filtered_df.iterrows():
        locs += [row['start'], row['end']]
    
    gene_info = filtered_df.iloc[0].copy()
    # print(gene_info)

    start = min(locs)
    end = max(locs)
    gene_info['start'], gene_info['end'] = start, end
    # print(gene_info)
    return gene_info