from typing import Literal
from aldepyde.data import get_store


def hydropathy(sequence: str, scale: Literal['kd'] = 'kd', window: int = 0, window_average: bool = False) -> list:
    hydropathies = []
    den = len(sequence) if window_average\
    else 1
    if window < 0:
        return [sum([get_store().hydropathy(aa, scale=scale) for aa in sequence])/den]
    for i in range(window, len(sequence) - window, 1):
        residue_window = sequence[i - window:i + window + 1]
        hydropathies.append(sum([get_store().hydropathy(aa, scale=scale) for aa in residue_window])/(len(residue_window) if window_average
                                                                                                    else 1))
    if len(hydropathies) == 0:
        return [sum([get_store().hydropathy(aa, scale=scale) for aa in sequence])/den]
    return hydropathies

def helical_prop(sequence: str, window: int = 0, window_average: bool = False) -> list:
    propensities = []
    den = len(sequence) if window_average\
    else 1
    if window < 0:
        return [sum([get_store().helical_prop(aa) for aa in sequence])/den]
    for i in range(window, len(sequence) - window, 1):
        residue_window = sequence[i - window:i + window + 1]
        propensities.append(sum([get_store().helical_prop(aa) for aa in residue_window])/(len(residue_window) if window_average
                                                                                        else 1))
    if len(propensities) == 0:
        return [sum([get_store().helical_prop(aa) for aa in sequence])/den]
    return propensities

# TODO Implement this algorithm
# https://bio.libretexts.org/Courses/Ouachita_Baptist_University/Reyna_Cell_Biology/02%3A_2-(T2-first_lecture)_Protein_Structure/2.05%3A_G._Predicting_Protein_Properties_From_Sequences/2.5.02%3A_G2._Prediction_of_Secondary_Structure

def chou_fasman():
    pass


def smith_waterman(sequence_1:str, sequence_2:str, match:int=1, gap:int=-2, mismatch:int=-1, minimum=0, recurse:bool=True) -> tuple[float, tuple[str, str]]:
    matrix = [[0] * (len(sequence_1) + 1)]
    pointer_matrix = [[[(-1,-1)]] * (len(sequence_1) + 1)]
    max_value = -float('inf')
    max_coords = list()
    coords = (-1, -1)
    for i in range(1, len(sequence_2) + 1):
        matrix.append(list([0]))
        pointer_matrix.append(list([[(-1, -1)]]))
        for v in range(1, len(sequence_1) + 1):
            u = (matrix[i-1][v] + gap, 'u')
            l = (matrix[i][v-1] + gap, 'l')
            d = (matrix[i-1][v-1] + match if sequence_1[v-1] == sequence_2[i-1] else matrix[i-1][v-1] + mismatch, 'd')
            s = (0, 's')
            scores = [u, l, d, s]
            max_sub_score = max([x[0] for x in scores])
            directions = []
            val = -1
            for value_dir in scores:
                if value_dir[0] == max_sub_score:
                    val = max_sub_score
                    if value_dir[1] == 'd':
                        directions.append((i-1, v-1))
                    elif value_dir[1] == 'u':
                        directions.append((i-1,v))
                    elif value_dir[1] == 'l':
                        directions.append((i,v-1))
            pointer_matrix[-1].append(directions)
            matrix[-1].append(val)
            if val > max_value:
                max_value = val
                max_coords = list()
            if val == max_value:
                max_coords.append((i,v))

    #traceback
    # TODO Make recursion optional
    import copy
    results = []
    def pointer_helper(al1, al2, mc):
        # print(pointer_matrix[mc[0]][mc[1]])
        curr_1 = al1
        curr_2 = al2
        for new_point in pointer_matrix[mc[0]][mc[1]]:
            c1 = sequence_1[mc[1] - 1] if mc[1] - 1 == new_point[1] else '-'
            c2 = sequence_2[mc[0] - 1] if mc[0] - 1 == new_point[0] else '-'

            al1.append(c1)
            al2.append(c2)
            # al1.append(sequence_1[mc[1] - 1]) if mc[1] - 1 == new_point[1] else al1.append('-')
            # al2.append(sequence_2[mc[0] - 1]) if mc[0] - 1 == new_point[0] else al2.append('-')
            if matrix[new_point[0]][new_point[1]] != 0:
                pointer_helper(al1, al2, new_point)
                # al1, al2 = pointer_helper(copy.deepcopy(curr_1), copy.deepcopy(curr_2), new_point)
            else:
                results.append(("".join(reversed(al1)), "".join(reversed(al2))))
            al1.pop()
            al2.pop()
        # return al1, al2

    alignments = list()
    for mc in max_coords:
        pointer_helper(list(), list(), mc)
    for result in results:
        alignments.append((result[0], result[1]))
    print(alignments)
    return max_value, alignments
