import nwalign3 as nw
import itertools
import numpy as np

# Needleman–Wunsch algorithm based insetion technique
class NW_INSERTION:

    def __init__(self, gap_open_weight = -5, gap_extend_weight = -2, matrix = 'PAM250'):
        self.gap_open_weight = gap_open_weight
        self.gap_extend_weight = gap_extend_weight
        self.matrix = matrix

    # compute nw score of two sequence using nwalign3 library method.
    def get_nw_score(self, sequence_ref, sequence):
        return nw.score_alignment(sequence_ref, sequence, 
        gap_open=self.gap_extend_weight,gap_extend=self.gap_extend_weight, matrix=self.matrix)

    # This method take one sequence and an insertion item as parameter and return the 
    # the sequence with best insertion based on highest nw score.
    def get_best_nw_based_insertion(self, target, insertion_item):
        best_inserted_seq = target
        best_score = 0
        for i in range(0, len(target)-len(insertion_item)+1):
            seq = list(target)
            for j in range(0, len(insertion_item)):
                seq[i+j] = insertion_item[j]
            seq = ''.join(seq)
            if i == 0:
                best_score = self.get_nw_score(target, seq)
                best_inserted_seq = seq
                print(seq, best_score)
            else:
                score = self.get_nw_score(target, seq)
                print(seq, score)
                if score > best_score:
                    best_score = score
                    best_inserted_seq = seq
                
        return best_inserted_seq

    # This method take one sequence and insertion items as parameter and return the 
    # the sequence with best insertion, worst insertion, best nw score, worst nw score and average nw score.
    # it consider all permutation of the insertions if keep_order is set to True.
    # otherwise it will keep the order of the insertion items.
    def get_best_nw_based_insertions(self, target, insertion_items, keep_order = False, verbose=False):
        if not keep_order:
            insertion_items = list(itertools.permutations(insertion_items))
        else:
            insertion_items = [tuple(insertion_items)]

        best_inserted_seq = target
        worst_inserted_seq = target
        best_score = 0
        worst_score = 0
        all_scores = []
        first_time = True
        for item_list in insertion_items:
            # put all the items on top of target and at the begining of the target
            # then first take the last item and push it towards the end 
            # then do the same for the second last item and so on until the last item
            # it is similar to moving the parts of curtles which were all at the left and
            # and we need to move it towards end one by one.

            # putting all at left 
            next = 0
            seq = list(target)
            next_list = []
            for item in item_list:
                next_list.append(next)
                for n in item:
                    seq[next] = n
                    next += 1
            seq = ''.join(seq)
            score = self.get_nw_score(target, seq)
            all_scores.append(score)
            if first_time:
                best_inserted_seq = seq
                worst_inserted_seq = seq
                best_score = score
                worst_score = score
                first_time = False
            if score > best_score:
                best_inserted_seq = seq
                best_score = score
            if score < worst_score:
                worst_inserted_seq = seq
                worst_score = score
            if verbose:
                print('------------')
                print(item_list)
                print(seq, score)  

            # moving to write one by one starting by the last one 
            last = len(seq)
            for next, item in zip(next_list[::-1], item_list[::-1]):
                if verbose:
                    print(next, last, item)
                for i in range(next+1, last-len(item)+1):
                    seq = list(seq)
                    # put original from target
                    for j in range(next, i):
                        seq[j] = target[j]
                    # moving item by one place
                    j = i
                    for n in item:
                        seq[j] = n
                        j += 1
                    seq = ''.join(seq)
                    score = self.get_nw_score(target, seq)
                    all_scores.append(score)
                    if score > best_score:
                        best_inserted_seq = seq
                        best_score = score
                    if score < worst_score:
                        worst_inserted_seq = seq
                        worst_score = score
                    if verbose:
                        print(seq, score)
                last = last-len(item)+1      

        return best_inserted_seq, worst_inserted_seq, np.max(all_scores), np.min(all_scores), np.mean(all_scores)

    # This method take one sequence as host and insertion items as parameter and return the
    # the sequence with best insertion/ worst insertion, best nw score / worst nw score
    # (based on the parameter value 'return_best') and average nw score .
    # The default value is true to return sequence with best and value false to return sequence with worst insertions.
    # It consider greedy approach of insertions. So, First fragment will be moved over the host and fixed at best/worst
    # place. Then the second fragment will find its best/worst possible place leaving out the places hold by the first
    # fragment. The process will continue for all the fragments. Any position will not be used between two already
    # placed fragments if the length is shorter than the fragment length to avoid overriding of previous fragments.

    def get_nw_based_insertions_greedy(self, target, insertion_items, return_best=True, verbose=False):
        final_seq = target  # keeping the final sequence will all fragments inserted (with best or worst insertions)
        all_scores = []  # keeping all the scores calculated during the process.
        # Inside locked_positions, value 1 means the position is locked by previously inserted fragments
        locked_positions = [0 for i in range(0,len(target))]
        if verbose:
            print('start with:', target)
            print('item to insert:')
            for fragment in insertion_items:
                print(fragment)
        for fragment in insertion_items:
            first_time = True
            best_loc = -1 # keeping the final location for this fragment (for best or worst insertion)
            best_seq = '' # final sequence after the fragment is inserted depending on best or worst insertion.
            best_score = 0  # keeping the final score for the fragment(for best or worst insertion)
            # move item over every position of host/target and skip the locked positions.
            for i in range(0, (len(target)-len(fragment)+1)):
                # if fragment will overlapping with previously inserted fragments then do nothing
                if 1 in locked_positions[i:(i+len(fragment))]:
                    continue
                # put the new fragment in place
                seq = final_seq[:i]+fragment+final_seq[i+len(fragment):]
                # calculate the overall nw score with initial host sequence
                score = self.get_nw_score(target, seq)
                # if the new fragment is considered for an available position.
                # then consider the position is best possible insertion position and score as best score
                # to start with the process for the fragment
                if first_time:
                    first_time = False
                    best_score = score
                    best_loc = i
                    best_seq = seq
                    continue
                # if this is not the first ever position to consider
                # and we want to return a sequence with best(nw) insertion
                if return_best and score > best_score:
                    best_score = score
                    best_loc = i
                    best_seq = seq
                    continue
                # if this is not the first ever position to consider
                # and we want to return a sequence with worst(nw) insertion
                if not return_best and score < best_score:
                    best_score = score
                    best_loc = i
                    best_seq = seq
                    continue
            if best_loc == -1:
                print("Error: Did not find space to insert a fragment ", fragment, "in host ", target)
                print("Last insertions: ", final_seq)
                exit()
            if verbose:
                print()
                print('best insertion location and score: ',best_loc, best_score)
                print(final_seq)
                print(locked_positions)
            final_seq = best_seq
            all_scores.append(best_score)
            # locking locations for stopping further use
            for i in range(best_loc, (best_loc+len(fragment))):
                locked_positions[i] = 1
            if verbose:
                print(final_seq)
                print(locked_positions)
                print()
        if return_best:
            return final_seq, np.max(all_scores), np.mean(all_scores)
        else:
            return final_seq, np.min(all_scores), np.mean(all_scores)


