from easy_ko.load_genomes import generate_all_data
from analyze_gene import *

from pygenomeviz import GenomeViz

def get_fig(name, val):
    gene, df, info = val['A_allele'], val['df_gene_analysis'], val['fig_info']
    if gene['strand'] == '+':
        strand = +1
    else:
        strand = -1


    end = int(df['end'].max() + 10)
    start = int(df['start'].min() - 10)

    gv = GenomeViz()
    track = gv.add_feature_track(name, end)

    track.add_feature(gene['start'], gene['end'], strand, plotstyle='bigbox', fc='none', lw=5.0, zorder=+1)

    for i, row in df.iterrows(): 
        name = row['part name']
        if not(pd.isna(row['start'])):
            s = row['start']
            e = row['end']
            text = dict(rotation=45, vpos="top")
            # todo: change zorder to make primers on top of gRNA
            if 'gRNA' in name:
                name = 'gRNA'
                c = 'black'
                track.add_feature(s, e, label=name, plotstyle='bigbox', fc=c)
            else:
                if 'F_orf' in name:
                    name = 'F_orf'
                    c = 'blue'
                    strand = +1
                elif 'R_orf' in name:
                    name = 'R_orf'
                    c = 'blue'
                    strand = -1
                    text=dict(rotation=-45, vpos="bottom")
                elif 'F_5flank' in name:
                    name = 'F_5flank'
                    c = 'orange'
                    strand = +1
                elif 'R_3flank' in name:
                    name = 'R_3flank'
                    c = 'orange'
                    strand = -1
                    text=dict(rotation=-45, vpos="bottom")

                if not('SATFLP' in name):
                    track.add_feature(s, e, strand, label=name, fc=c, text_kws=text) # can change arrow_shaft_ratio
    
    backup_guides = info['backup_guides']

    for i in range(len(backup_guides)):
        guide = backup_guides[i]
        s = guide['crRNA_start']
        e = s + 20
        name = f'backup gRNA {str(i+1)}'
        c = 'gray'
        track.add_feature(s, e, label=name, plotstyle='bigbox', fc=c)


    fig = gv.plotfig()
    track.ax.set_xlim(start, end)

    return fig

def visualize_figs(fig_data):
    figs = []
    for gene_name, val in fig_data.items():
        figs.append(get_fig(gene_name, val))
    return figs

# you can use a python notebook to call this function
def main(input_genes_file, GTF_file, assembly_file, PCR_components_file, black_list_file, bar = 0):
    data = generate_all_data(input_genes_file,GTF_file,assembly_file, PCR_components_file, black_list_file)

    df_output = pd.DataFrame()
    fig_data = {}
    input_genes_df = data.input_genes_df
    print(input_genes_df) #eek
    for i, row in input_genes_df.iterrows():
        gene_id, gene_name = str(row["systematic name"]).strip(), str(row["gene name"]).strip() # strip whitespace
        print(f'Analyzing {gene_name}...')
        df_gene_analysis, A_allele, fig_info = analyze_gene(data, gene_id, gene_name)
        df_output = pd.concat([df_output, df_gene_analysis], ignore_index=True)
        fig_data[gene_name] = {'A_allele': A_allele, 'df_gene_analysis': df_gene_analysis, 'fig_info': fig_info}
        if bar:
            bar.progress(round((i+1)/len(input_genes_df)*100))

    return df_output, fig_data
    data.genome_GTF_df #GTF file
    data.genome_assembly_dict # assembly file
    data.input_genes_df # list of genes to query
    data.gene_info # info for a given gene
    data.black_list_df # list of recently duplicated genes