# This code was developed and authored by Jerzy Twarowski in Malkova Lab at the University of Iowa 

# Contact: jerzymateusz-twarowski@uiowa.edu, tvarovski1@gmail.com

## Contents:

## extractSeqFromFastaToList

## validateSquence

## compl

## rev_compl

## unique

## createChrList

## dataFrameImport

import logging

def extractSeqFromFastaToList(fasta_file_path: str) -> list:

    '''Extracts sequence information from a FASTA file and returns it as a nested list.

    Args:

        fasta_file_path (str): The path to the FASTA file.

    Returns:

        list: A nested list of the form [[title_0, sequence_0], [title_1, sequence_1], ...].

            Each element of the list is a list containing the title and sequence of a sequence in the FASTA file.

    '''

    fasta_file = open(fasta_file_path, 'r')

    contents = fasta_file.readlines()

    fasta_file.close()

    fasta_list=[]

    for i in contents:

        if '>' in i:

            fasta_list.append([i.strip('>').strip(),''])

        else:

            fasta_list[-1][1] = fasta_list[-1][1]+i.strip()

    logging.info(f"Extraction of sequence information from {fasta_file_path} finished.")

    return fasta_list

def validateSquence(sequence: str) -> bool:

    '''Checks if a DNA sequence contains only valid nucleotides.

    Args:

        sequence (str): A DNA sequence to be validated.

    Returns:

        bool: True if the sequence contains only valid nucleotides, False otherwise.

            If the sequence contains non-canonical nucleotides, a log message is generated.

            If the sequence contains 'N's, a warning log message is generated.

            The log messages are written to the default logger.

    '''

    bases = "ATGCatgcN"

    for i in sequence:

        if i not in bases:

            logging.info(f"Sequence doesn't contain cannonical nucleotides: {i}")

            return(False)

        if i == "N":

            logging.warning(f"Warning, sequence contains 'N's")

    return(True)

def compl(base: str) -> str:

    """

    Returns the complementary base for a given DNA base.

    Args:

        base (str): A single character representing a DNA base. Must be one of 'A', 'T', 'G', 'C', or '-'.

    Returns:

        str: The complementary base for the given DNA base. Returns '-' if the input is '-'.

    Raises:

        ValueError: If the input is not one of 'A', 'T', 'G', 'C', or '-'.

    """

    if base == "A":

        return('T')

    elif base == "T":

        return('A')

    elif base == "G":

        return('C')

    elif base == "C":

        return('G')

    elif base == "-":

        return('-')

    else:

        logging.error(f"Invalid base: {base}")

        raise ValueError(f"Invalid base: {base}")

def rev_compl(seq: str) -> str:

    """

    Returns the reverse complement of a given DNA sequence.

    Args:

        seq (str): A string representing a DNA sequence. Must only contain characters 'A', 'T', 'G', 'C', or '-'.

    Returns:

        str: The reverse complement of the given DNA sequence.

    Raises:

        ValueError: If the input sequence contains characters other than 'A', 'T', 'G', 'C', or '-' (Through compl())

    """

    new_seq = ""

    for base in seq:

        new_base = compl(base)

        new_seq = new_base + new_seq

    return(new_seq)

def unique(list1: list) -> list:

    """

    Returns a new list containing only the unique elements of the input list.

    Args:

        list1 (list): A list of elements.

    Returns:

        list: A new list containing only the unique elements of the input list.

    Examples:

        >>> unique([1, 2, 3, 2, 1])

        [1, 2, 3]

        >>> unique(['a', 'b', 'c', 'b', 'a'])

        ['a', 'b', 'c']

    """

    list_set = set(list1) 

    unique_list = (list(list_set))

    return unique_list

def createChrList(chrNum: int) -> list:

    """

    Creates a list of chromosome names.

    Args:

        chrNum (int): The number of chromosomes to include in the list.

    Returns:

        list: A list of chromosome names, where each name is a string of the form "chrX", where X is the chromosome number.

    Examples:

        >>> createChrList(3)

        ['chr1', 'chr2', 'chr3']

        >>> createChrList(5)

        ['chr1', 'chr2', 'chr3', 'chr4', 'chr5']

    """

    chrList = []

    for i in range(chrNum):

        chrList.append(f"chr{i+1}")

    return(chrList)

def dataFrameImport(directory: str) -> tuple:

    """

    Imports all TSV files in a directory as Pandas dataframes.

    Args:

        directory (str): The path to the directory containing the TSV files.

    Returns:

        tuple: A tuple containing two elements:

            - A list of Pandas dataframes, where each dataframe corresponds to a TSV file in the directory.

            - A list of sample names, where each name is a string representing the name of the corresponding TSV file.

    """

    import os

    import pandas as pd

    frames_list_samples = []

    sample_names = []

    counter=1

    for filename in os.listdir(directory):

            if filename.endswith(".tsv"):

                    sample_names.append(filename[:10].strip())

                    path = os.path.join(directory, filename)

                    globals()[f"sample{counter}"] = pd.read_csv(path, sep='\t')

                    print(f'dataframe {sample_names[-1]} has {len(globals()[f"sample{counter}"])} total rows')

                    frames_list_samples.append(globals()[f"sample{counter}"])

                    counter+=1

    logging.info(f"found {len(frames_list_samples)} samples in {directory}")

    logging.info(len(sample_names), sample_names)

    return(frames_list_samples, sample_names)

def pullGenomicContext(list_of_positions: list, fasta_file_path: str, context_flank: int = 5) -> list:

    """

    Extracts genomic context of a specified length around each position in a list of positions.

    Args:

        list_of_positions (list): A list of positions, where each position is a list of two elements: the chromosome name and the position.

        fasta_file_path (str): The path to the FASTA file containing the reference genome.

        context_flank (int, optional): The length of the context to extract on either side of each position. Defaults to 5.

    Returns:

        list: A list of genomic contexts, where each context is a list of four elements:

            - The sequence of bases to the left of the position.

            - The query base at the position.

            - The sequence of bases to the right of the position.

            - The chromosome name.

    """

    fasta_list = extractSeqFromFastaToList(fasta_file_path)

    context_list = []

    for i in fasta_list:

        logging.debug(f"extracting context from {i[0]}. It has {len(i[1])} bases.")

        sequence = i[1]

        for j in list_of_positions:

            #check if it is the same chromosome

            if j[0] != i[0]:

                continue

            #set j to be the position

            j = j[1]

            #find the position in the fasta list

            #remember that the fasta sequences start with 1, not 0, so we need to subtract 1 from the position

            try:

                query_base = sequence[j-1]

            except:

                logging.warning(f"position {j} not found in {i[0]}. Skipping...")

                continue

            #extract the context

            context_left = sequence[j-1-context_flank:j-1]

            context_right = sequence[j:j+context_flank]

            context_list.append([context_left, query_base, context_right, i[0]])

    return(context_list)

def drawGenomicContext(context_list: list, show: bool = False, **kwargs: dict) -> None:

    '''

    Draws a sequence logo of a list of genomic contexts.

    Args:

        context_list (list): A list of genomic contexts, where each context is a list of four elements:

            - The sequence of bases to the left of the position.

            - The query base at the position.

            - The sequence of bases to the right of the position.

            - The chromosome name.

        show (bool, optional): Whether to show the plot. Defaults to False.

        **kwargs (dict): Additional keyword arguments to pass to the function. Supported arguments include:

            - save_path (str): The path to save the plot to.

    Returns:

        None

    Raises:

        ImportError: If the required packages matplotlib.pyplot and pandas are not installed.

    '''

    import matplotlib.pyplot as plt

    import pandas as pd

    #create a dataframe from the context list

    df = pd.DataFrame(context_list, columns=['context_left', '0', 'context_right', 'chromosome'])

    #split the context into columns by base

    split_left = df['context_left'].apply(lambda x: pd.Series(list(x)))

    split_right = df['context_right'].apply(lambda x: pd.Series(list(x)))

    #rename the columns such that they are numbered. e.g. -3,-2,-1, 0, 1, 2, 3

    split_left.columns = [f"-{i}" for i in range(len(split_left.columns), 0, -1)]

    split_right.columns = [f"{i+1}" for i in range(len(split_right.columns))]

    #make a new dataframe with the context split into columns, query still in the middle

    df_context = pd.concat([split_left, df['0'], split_right], axis=1)

    df_context_freq = df_context.apply(pd.value_counts).fillna(0)              #get the frequencies of each base in each column

    df_context_freq = df_context_freq.div(df_context_freq.sum(axis=0), axis=1) #normalize the frequencies

    df_context_freq = df_context_freq.transpose()

    df_context_freq = df_context_freq*100                                      #set the frequencies to %

    #set color scheme for the bases so that they are consistent

    colors = {

                'A':'tab:green', 

                'T':'tab:red',

                'G':'tab:orange',

                'C':'tab:blue',

                'N':'tab:gray',

                }

    #plot the dataframe as a stacked barplot

    column_count = len(df_context_freq.columns) #count the number of columns

    df_context_freq.plot.bar(stacked=True, figsize=(column_count+3,4), color = list(colors.values()))

    plt.legend(loc='center left', bbox_to_anchor=(1.0, 0.5))

    plt.ylabel('Frequency (%)')

    plt.xlabel('Relative Position')

    if 'save_path' in kwargs:

        plt.savefig(kwargs['save_path'], bbox_inches='tight', dpi=300)

        logging.info(f"saved plot to {kwargs['save_path']}")

    if show == True:

        plt.show()

    plt.close()