# 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:
## findpairings
## findFrequencies
## plot
## runOligoFreqAnalysis

import pandas as pd
from .base import extractSeqFromFastaToList
from .base import validateSquence

def findpairings(sequence: str, pairing_length: int) -> list:
    '''
    Find all oligonucleotides of a given length in a DNA sequence.

    Args:
        sequence (str): the DNA sequence to search in.
        pairing_length (int): the length of the oligonucleotides to find.

    Returns:
        list: a list of all oligonucleotides of length 'pairing_length' found in 'sequence'.
    '''

    pairing_list = []
    for i in range(len(sequence)):
        next_chunk = sequence.upper()[i:i+pairing_length]
        if len(next_chunk) == pairing_length:
            pairing_list.append(next_chunk)
    return(pairing_list)

def findFrequencies(alist: list) -> dict:
    '''
    Count the frequency of each item in a list.

    Args:
        alist (list): the list to count the frequency of.

    Returns:
        dict: a dictionary where the keys are the items in the list and the values are their frequency.
    '''

    sequence_dict = {}

    for i in alist:
        if i in sequence_dict:
            sequence_dict[i]+=1
        else:
            sequence_dict[i] = 1
    return(sequence_dict)

def plot(oligo_dict: dict, title: str) -> None:
    '''
    Plot a histogram of the frequency of each item in a dictionary.

    Args:
        oligo_dict (dict): the dictionary to plot the frequency of.
        title (str): the title of the plot.

    Returns:
        None
    '''
    freq_list = list(oligo_dict.items())
    freq_list.sort(key=lambda x:x[1], reverse=True)
    df = pd.DataFrame(freq_list, columns=['oligo', 'frequency'])
    total_oligos = df.frequency.sum()
    df['frequency'] = df.frequency / total_oligos
    df.plot(kind='bar', x='oligo', figsize=(15,8), title=title+" Oligonucleotide Frequency", xlabel='Oligonucleotide', ylabel='Frequency (%)')

def runOligoFreqAnalysis(file_path: str) -> None:
    '''
    Run oligonucleotide frequency analysis on all sequences in a FASTA file.

    Args:
        file_path (str): the path to the FASTA file to analyze.

    Returns:
        None
    '''
    import matplotlib.pyplot as plt
    fa_seq_list = extractSeqFromFastaToList(file_path)

    for fa_seq in fa_seq_list:
        sequence = fa_seq[1]
        title = fa_seq[0]

        if not validateSquence(sequence):
            print(f'skipping "{title}", sequence validation failed')
            continue

        pairing_list_1 = findpairings(sequence,2)
        seqdeck_1 = findFrequencies(pairing_list_1)

        pairing_list_2 = findpairings(sequence,1)
        seqdeck_2 = findFrequencies(pairing_list_2)

        plot(seqdeck_1, title)
        plot(seqdeck_2, title)
        plt.show()


if __name__ == "__main__":
    import sys
    file_path = sys.argv[1]
    runOligoFreqAnalysis(file_path)