import easy_ko.helper_functions as helper

from dataclasses import dataclass, field
import pandas as pd
from Bio import SeqIO
import pandas as pd

@dataclass
class SharedData:
    genome_GTF_df: pd.DataFrame = None
    genome_assembly_dict: dict = field(default_factory=dict)
    gene_info: str = None
    black_list_df: pd.DataFrame = None
    input_genes_df: pd.DataFrame = None
    components: pd.DataFrame = None



def read_input_excel(data: SharedData, input_genes_file):
    """Store entire Excel sheet with input genes"""
    data.input_genes_df = pd.read_excel(input_genes_file).rename(columns=str.strip).rename(columns=str.lower)
    
def read_primer_components(data: SharedData, PCR_components_file):
    """Store entire Excel sheet with input genes"""
    df_components = pd.read_excel(PCR_components_file).rename(columns=str.strip).rename(columns=str.lower)
    df_components['part sequence'] = df_components['part sequence'].str.upper()
    data.components = df_components



def read_GTF(data: SharedData, GTF_file):
    """Store the entire GTF as a DataFrame"""
    # use helper function to load GTF
    dfGTF = helper.load_GTF(GTF_file)
    data.genome_GTF_df = dfGTF
    



def read_genome_assembly(data: SharedData, assembly_file):
    """Store sequences in a dictionary {id: seq}"""
    data.genome_assembly_dict = {record.id: record.seq.upper() for record in SeqIO.parse(assembly_file, "fasta")}


def read_black_list(data: SharedData, black_list_file ):

    """
    Black_list is list of recent duplicated genes. 
    Can be used to refine gRNA search when gene of interest
    has a close sequence match in the genome
    """
    data.black_list_df = pd.read_csv(black_list_file).rename(columns=str.strip).rename(columns=str.lower)
    

# ----- orchestrator function -----
def generate_all_data(input_genes_file,
                      GTF_file,
                      assembly_file, 
                      PCR_components_file,
                      black_list_file ="C_albicans_SC5314_A22_current_orf_genomic_all_v_all.csv"
                      ) -> SharedData:
    
    data = SharedData()  
      
    # these now standardize excel and csv columns to lowercase and apply strip() to remove whitespace
    read_input_excel(data, input_genes_file)
    read_GTF(data, GTF_file)
    read_genome_assembly(data, assembly_file)
    read_primer_components(data, PCR_components_file)
    
    if black_list_file:
        read_black_list(data, black_list_file)
    
    return data


