import numpy as np
import random as rn
from Bio import AlignIO
import sys
import time
sys.path.append("src/")

from misc import *
from newick import *
from stateFreqs import *
from matrixBuilder import *
from evolver import *

# get dna/amino/codon lists
molecules = Genetics()

# read in tree
print "reading tree"
my_tree, flag_list  = readTree(file="trees/100.tre", show=False) # set True to print out the tree


# state frequencies
print "collecting state frequencies"
fgen = UserFreqs(by='amino', freqs = {'I':0.33, 'L':0.33, 'V':0.34} , save='stateFreqs.txt')
commonFreqs = fgen.getCodonFreqs()
fgen.save2file()


## temporary code for constructing multiple GY94 models. Will formalize in the future
numPart = 2 # number of partitions
kappa  = 4.5
omegas = [0.5, 1.8]
partLen = np.tile(100, numPart) # for now, equal size partitions
partitions = []
print "constructing models for", numPart, "partitions"
for i in range(numPart):
	# Define model object and its parameters. Build model matrix. Add tuple (partition length, model) to partitions list
	model = misc.Model()
	model.params = { "kappa": kappa, "omega": omegas[i], "stateFreqs": commonFreqs }
	m = GY94(model)
	model.Q = m.buildQ()
	partitions.append( (partLen[i], model) )

# Evolve
outfile = time.strftime("%m.%d_%H;%M;%S")+".fasta" # So, : is an illegal filename character
myEvolver = StaticEvolver(partitions = partitions, tree = my_tree)
myEvolver.sim_sub_tree(my_tree)
myEvolver.writeSequences()










######################### STUFF ###################
# parameters for rodrigue model
#model1.params={"nucMut":   {"AC":0.16, "AG":0.16, "AT":0.16, "CG":0.16, "CT":0.16, "GT":0.17}, 
#			   "nucFreqs": {'A':0.25, 'C':0.25, 'G':0.25, 'T':0.25},
#		       "aaVector": [0.1, 0.05, 0.1, 0.05, 0.025, 0.025, 0.025, 0.025, 0.05, 0.4, 0.005, 0.005, 0.005, 0.005, 0.005, 0.005, 0.005, 0.005, 0.05, 0.05]	
#			  }



















@@ -319,21 +319,8 @@
 tree
-, flag_list  
 = re
@@ -350,20 +350,8 @@
 tre%22
-, show=False
 ) # 
@@ -440,16 +440,17 @@
 encies%22%0A
+#
 fgen = U
@@ -489,18 +489,17 @@
 ':0.
-33
+2
 , 'L':0.
 33, 
@@ -498,18 +498,17 @@
 ':0.
-33
+2
 , 'V':0.
 34%7D 
@@ -507,10 +507,27 @@
 ':0.
-34
+2, 'C':0.2, 'F':0.2
 %7D , 
@@ -548,16 +548,84 @@
 qs.txt')
+ # these are 5 hydrophobic.%0Afgen = EqualFreqs(by='codon')#, file = '
 %0AcommonF
@@ -767,17 +767,18 @@
 mPart = 
-2
+ 1
  # numbe
@@ -802,17 +802,17 @@
 appa  = 
-4
+2
 .5%0Aomega
@@ -822,16 +822,30 @@
  %5B0.
-5, 1.8%5D%0A
+25%5D%0A%0A%0ApartLen = 1000 #
 part
@@ -1290,19 +1290,16 @@
 (partLen
-%5Bi%5D
 , model)
@@ -1311,16 +1311,33 @@
  Evolve%0A
+print %22evolving%22%0A
 outfile 
@@ -1548,16 +1548,271 @@
 ces()%0A%0A%0A
+# Write true rates%0Atruefile = %22truerates.txt%22%0Atruef = open(truefile, 'w')%0Atruef.write(%22position%5Ctomega%5Cn%22)%0Aposition = 1%0Afor i in range(numPart):%0A%09for l in range(partLen):%0A%09%09truef.write(str(position)+'%5Ct'+str(omegas%5Bi%5D)+'%5Cn')%0A%09%09position+=1%0Atruef.close()%0A%0A%0A
 %0A%0A%0A%0A%0A%0A%0A%0A

Edit to reflect newick.py changes