#!/usr/bin/env python

"""

Name: treemix_from_nex.py 

Author: Michael G. Harvey
Date: 12 May 2013

Description: Convert a nexus alignment of concatenated SNPs to the input file format for 
Treemix (Pickrell and Pritchard 2012).

Usage: 	python treemix_from_nex.py input_file output_directory number_of_populations \
			individuals_in_pop_1 individuals_in_pop_2 ... individuals_in_pop_n

Note that TreeMix uses only biallelic SNPs, so any sites that do not represent biallelic SNPs 
(including consideration of the outgroup), will not be used.

"""

import os
import sys
import argparse
from Bio import AlignIO


def get_args():
	parser = argparse.ArgumentParser(
			description="""Program description""")
	parser.add_argument(
			"in_file",
			type=str,
			help="""The input nexus file"""
		)
	parser.add_argument(
			"out_dir",
			type=str,
			help="""The output directory"""
		)
	parser.add_argument(
			"populations",
			type=int,
			help="""The number of populations"""
		)
	parser.add_argument(
			"pop_sizes",
			type=int,
			nargs='+',
			help="""The number of samples in each population"""
		)
	return parser.parse_args()


def main():
	args = get_args()
	alignment = AlignIO.read("{0}".format(args.in_file), "nexus")
	out = open("{0}treemix_file_out.txt".format(args.out_dir), 'wb')
	for x in range(args.populations):
		out.write("pop{0} ".format(x+1)) 
	out.write("outgroup\n")
	total_size = 0
	for w in xrange(alignment.get_alignment_length()):
		total_size = 0
		bases = alignment[:, w] 		
		uniqs = list()
		uniqs = list(set(bases))
		if len(uniqs) != 2:
			print "Skipping site {0} - not a biallelic SNP".format(w)	
		else:
			allele = bases[0]
			total_size = 0
			for x in range(args.populations):
				pop_size = args.pop_sizes[x]
				pop_bases = alignment[total_size:(total_size+pop_size), w]
				total_size += pop_size
				a = 0
				b = 0
				for pop_base in pop_bases:
					if pop_base == allele:
						a += 1
					else:
						b += 1
				out.write("{0},{1} ".format(a,b))
			ogs = alignment[(total_size):(total_size+2), w]
			a = 0
			b = 0
			for og in ogs:
				if og == allele:
					a += 1
				else:
					b += 1
			out.write("{0},{1}".format(a,b))
		out.write("\n")	
		out.flush()
	out.close()
	
if __name__ == '__main__':
    main()

@@ -48,17 +48,16 @@
 nex.py %0A
-%0A
 Author: 
@@ -478,58 +478,16 @@
 SNPs
- %0A(including consideration of the outgroup)
 , will 
+%0A
 not 
@@ -1331,16 +1331,8 @@
 te(%22
-outgroup
 %5Cn%22)
@@ -1490,16 +1490,157 @@
 bases))%0A
+%09%09nmuniqs = list()%0A%09%09for uniq in uniqs: # Remove sites with missing data%0A%09%09%09if uniq not in %5B'?', 'n', '?', 'N'%5D:%0A%09%09%09%09nmuniqs.append(uniq)%09%09%09%0A
 %09%09if len
@@ -1640,16 +1640,18 @@
 %09if len(
+nm
 uniqs) !
@@ -1738,18 +1738,45 @@
 lele
+1
  = 
-bases%5B0
+nmuniqs%5B0%5D%0A%09%09%09allele2 = nmuniqs%5B1
 %5D%0A%09%09
@@ -2024,24 +2024,25 @@
 se == allele
+1
 :%0A%09%09%09%09%09%09a +=
@@ -2055,159 +2055,19 @@
 %09%09el
-se:%0A%09%09%09%09%09%09b += 1%0A%09%09%09%09out.write(%22%7B0%7D,%7B1%7D %22.format(a,b))%0A%09%09%09ogs = alignment%5B(total_size):(total_size+2), w%5D%0A%09%09%09a = 0%0A%09%09%09b = 0%0A%09%09%09for og in ogs:%0A%09%09%09%09if og
+if pop_base
  == 
@@ -2076,32 +2076,12 @@
 lele
+2
 :%0A%09
-%09%09%09%09a += 1%0A%09%09%09%09else:%0A
 %09%09%09%09
@@ -2088,16 +2088,17 @@
 %09b += 1%0A
+%09
 %09%09%09out.w
@@ -2110,16 +2110,17 @@
 %22%7B0%7D,%7B1%7D
+ 
 %22.format
@@ -2211,9 +2211,8 @@
   main()
-%0A

update Treemix formatting script