# 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:
## wordsInSequence
## findComplexity
## movingWindow
## createLogFile
## createSequenceList
## createDataFrameFromLogFile

import pandas as pd
import ast
from pandas import DataFrame
from .base import unique

def wordsInSequence(sequence: str, treeLevel: int) -> list:
    """
    Extracts all possible words of a specified length from a genetic sequence.

    Args:
        sequence (str): The genetic sequence to extract words from.
        treeLevel (int): The length of the words to extract.

    Returns:
        list: A list of words, where each word is a string of length treeLevel.

    Examples:
        >>> wordsInSequence("ATCGATCG", 3)
        ['ATC', 'TCG', 'CGA', 'GAT', 'ATC', 'TCG']

        >>> wordsInSequence("ATCGATCG", 4)
        ['ATCG', 'TCGA', 'CGAT', 'GATC', 'ATCG']
    """
    length = len(sequence)
    max_possible_words_in_sequence = length-(treeLevel-1)
    wordList = []
    for i in range(max_possible_words_in_sequence):
        wordList.append(sequence[0+i:treeLevel+i])
    return wordList

def findComplexity(sequence, treeLevel, complexity_threshold):
    # takes a genetic sequence and calculates the linguistic complexity 
    # for that sequence for word length treeLevel

    wordList = wordsInSequence(sequence,treeLevel)
    wordList = unique(wordList)

    if len(sequence) > 4**treeLevel:
        complexity = len(wordList)/(4**treeLevel)
    else:
        complexity = len(wordList)/(len(sequence)-(treeLevel-1))

    if (complexity < complexity_threshold) & (complexity > 0.2):
        print("Complexity at tree level " + str(treeLevel) + " is " + str(complexity) + " for sequence: "+str(sequence))

    return ([complexity, treeLevel, sequence])

def movingWindow(sequence: str, treeLevel: int = 8, chunkSize: int = 20, complexity_threshold: float = 0.2, showGraph: bool = False) -> list:
    """
    Calculates the linguistic complexity scores for windows of a specified size and word lengths from 1 to a specified level.

    Args:
        sequence (str): The genetic sequence to calculate complexity scores for.
        treeLevel (int, optional): The maximum length of words to use in the complexity calculation. Defaults to 8.
        chunkSize (int, optional): The size of the windows to use in the complexity calculation. Defaults to 20.
        complexity_threshold (float, optional): The threshold below which a window is considered to have low complexity. Defaults to 0.2.
        showGraph (bool, optional): Whether to show a graph of the complexity scores. Defaults to False.

    Returns:
        list: A list containing the lowest complexity score window in the format [Boolean, complexity, treeLevel], where Boolean denotes if the complexity score is lower (True) than complexity_threshold.

    Raises:
        ValueError: If the sequence is too short to calculate complexity scores.

    Examples:
        >>> movingWindow("ATCGATCGATCGATCGATCGATCGATCGATCGATCGATCGATCGATCGATCGATCGATCGATCG", treeLevel=4, chunkSize=10, complexity_threshold=0.1, showGraph=True)
        [True, 0.0, 4]

        >>> movingWindow("ATCGATCGATCGATCGATCGATCGATCGATCGATCGATCGATCGATCGATCGATCGATCGATCG", treeLevel=10, chunkSize=5, complexity_threshold=0.5, showGraph=True)
        [False, 0.5, 10]
    """

    complexityList=[]

    for i in range(1, treeLevel+1):
        for chunk in range(len(sequence)-(chunkSize-1)):

            chunkSeq = sequence[chunk:(chunk+chunkSize)]
            complexityList.append(findComplexity(chunkSeq, i, complexity_threshold))

    try:
        lowest=min(complexityList, key=lambda x: x[0])

    except ValueError:
        print(f"ERROR with sequence: {sequence}. It is probably too short len={len(sequence)}. Skipping it.")
        print("complexityList: ", complexityList)
        lowest = [0,0,0]

    #print(complexityList)

    lowLvl=lowest[1]

    if showGraph:

        df = DataFrame(complexityList,columns=['complexity','tree','sequence'])
        df = df[['complexity','tree']]
        ax = df[df.tree == lowLvl].filter(items=["complexity"]).plot(fontsize=15, grid=True, figsize=(20,8))
        ax.axhline(complexity_threshold, color="red", linestyle="--")
        ax.axhline(1, color="green", linestyle="--")
        ax.set_ylim(ymin=0)
        ax.set_title(str(lowest[2]))

    if lowest[0] < complexity_threshold:
        return([True, lowest[0], lowest[1]])
        #return([True, lowest])
    else:
        return([False, lowest[0], lowest[1]])
        #return([False, lowest])

def createLogFile(seqlist: list, filename: str, complexity_threshold: float = 0.2, chunkSize: int = 20) -> None:
    """
    Creates a log file containing the lowest linguistic complexity score for each sequence in a list.

    Args:
        seqlist (list): A list of genetic sequences to calculate complexity scores for.
        filename (str): The name of the output file to write the log to.
        complexity_threshold (float, optional): The threshold below which a sequence is considered to have low complexity. Defaults to 0.2.
        chunkSize (int, optional): The size of the windows to use in the complexity calculation. Defaults to 20.

    Returns:
        None
    """
    outputf = open(filename, "w")
    truemmbir = 0
    for i in seqlist:
        seqScore = movingWindow(i, complexity_threshold = complexity_threshold, chunkSize = chunkSize)
        if seqScore[0] == False:
            truemmbir += 1
        outputf.write(str(seqScore[1])+"\n")
    print("sequences above threshold:", truemmbir)

    outputf.close()

def createSequenceList(filename: str) -> list:
    """
    Extracts genetic sequences from a file and returns them as a list.

    Args:
        filename (str): The name of the file to extract sequences from.

    Returns:
        list: A list of genetic sequences, where each sequence is a string.
    """
    seqfile = open(filename, "r")
    lines = seqfile.readlines()
    seqlist=[]
    for line in lines:
        if "bir:" in line:
            if len(line) > 100:
                seqlist.append(line.strip("\n").strip("bir:-"))
    seqlist = list(dict.fromkeys(seqlist))
    print(len(seqlist))

    longseq=0
    for i in seqlist:
        if len(i) > 300:
            longseq+=1
    print(longseq)

    return seqlist

def createDataFrameFromLogFile(logfile_path: str) -> DataFrame:
    """
    Creates a pandas DataFrame from a log file containing linguistic complexity scores.

    Args:
        logfile (str): The name of the log file to create the DataFrame from.

    Returns:
        DataFrame: A pandas DataFrame containing the complexity scores, where each row represents a sequence and contains the complexity score, the word length, and the sequence itself.
    """
    logfile = open(logfile_path, "r")
    lines = logfile.readlines()
    loglist=[]
    for line in lines:
        if "[" in line:
                loglist.append(line.strip("\n"))

    for i in range(len(loglist)):
        loglist[i] = ast.literal_eval(loglist[i]) 
    print("total sequences:", len(loglist))

    df = DataFrame(loglist,columns=['complexity', 'level', 'sequence'])
    
    return df