import alignment_util
import alignment
from Bio import SeqIO
from collections import OrderedDict
from math import log
import pickle
from PhyloTree import *
from PhyloTree_util import *
import PhyloTreeAlgs

pickle_output = open('/home/kumara3/CSE620K/proj6/CoV.dist','wb') 
newick_tree_obj = open('/home/kumara3/CSE620K/proj6/CoV.tree','wb')

def readfasta(filename):                                                    ##reads downloaded fasta file###
    store_seq = {} 
    with open (filename) as fh:
        for seq_info in SeqIO.parse(fh, "fasta") :
            store_seq[seq_info.description] = seq_info.seq
    return store_seq
    fh.close()

def find_distance(sequence_dictionary):
    ## Takes dictionary generated by readfasta and calculates distance matrix###
    popped_list = []
    p_distance = 0
    evol_distance = 0
    distance_matrix= {}
    
    copied_dict = sequence_dictionary.copy()                                 ###Making a copy of sequence_dictionary##
    for key in copied_dict:
        
        
        if key not in popped_list:
            seq1 = copied_dict[key]
            popped_list.append(key)   
        
        for i in sequence_dictionary:
            seq2 = sequence_dictionary[i]
            s = alignment_util.readScoringMatrix("/home/kumara3/CSE620K/proj6/Blosum62.txt")
            output_key_i = alignment.SmithWatermanAffine(seq1,seq2,s,7,1)  ###SmithWatermanAffine for calculating the alignment between pair of sequences###
            output = [it for it in output_key_i]                            ## alignment stored in a list##
            count_m = 0                                                     ## count for number of matches in sequences##
            count_mis= 0                                                    ## count for number of mismatches in sequence##
           
            X = [x for x in output[1:2]]                                    ## sequence 1 from alignment
            Y = [y for y in output[2:3]]                                    ## sequence 2 from alignment
            X_string = ''.join(X)
            Y_string = ''.join(Y)
            size = max(len(X_string), len(Y_string))                              
                    
            
            for k in range(0,size):
                if X_string[k] == Y_string[k]:
                    count_m += 1
                elif X_string[k] != Y_string[k]:
                    count_mis += 1
                else:
                    if X_string[k] == '-' or Y_string[k] == '-':
                        print "IGNORE"
            
            p_distance = count_mis/float(count_m+count_mis)                 ## calculating the p distance
            evol_distance = -0.75*log(float(1-4*p_distance/3))              ## calculating distance for evolutionary matrix
            distance_matrix.setdefault(key,{}).setdefault(i,evol_distance)  ## creating and filing the distance matrix 
            if key == i:
                distance_matrix[key][i] = 0 
    return distance_matrix                                                

dictionary_seq = readfasta('/home/kumara3/CSE620K/proj6/CoV.fa')
distance_matrix_to_pickle = find_distance(dictionary_seq)
phylotree_object = PhyloTreeAlgs.NJ(distance_matrix_to_pickle)
pickle.dump(phylotree_object,newick_tree_obj)
pickle.dump(distance_matrix_to_pickle,pickle_output)
pickle_output.close()
newick_tree_obj.close()

