#!/usr/bin/python3
#Tom Delomas - 2018
#Read vcf file and select loci that match pattern expected under various sex determination systems
#Must give the script the file path to a tab delimited text file with sample names in one column, phenotype (male or female) in teh second column, and each sample on its own line
#sample names must match column names of vcf file
##candidate_XY.txt - females all homozygous for one allele, males all heterozygous
##candidate_WZ.txt - males all homozygous for one allele, females all heterozygous
#written specifically for vcf's generated by STACKS, other vcf files that are structured differently may or may not function with this script
#in terminal, call as:		 python vcf_sex_markers.py file/path/to/your_vcf_file.vcf file/path/to/sample_phenotypes.txt

import sys, re
import pandas as pd

def Filter(data, het_gamet, hom_gamet, suffix, version):
	output_file = open('candidate_' + suffix + '.txt', 'w')
	output_file.write('Locus\tNum_fem\tNum_male\tTotal_geno\n')
	#loop through loci
	for i in range(0, data.shape[0], 1):

		#test that all hom_gamet are homozygous for the same allele
		zero_hom = 0
		one_hom = 0
		het = 0
		fail_hom = False
		fail_het = False
		for hom_sample in hom_gamet:
			if '0/1' in data.loc[:,hom_sample].iloc[i] or '1/0' in data.loc[:,hom_sample].iloc[i]:
				fail_hom = True
				break
			elif '0/0' in data.loc[:,hom_sample].iloc[i]:
				zero_hom += 1
			elif '1/1' in data.loc[:,hom_sample].iloc[i]:
				one_hom += 1
			if one_hom != 0 and zero_hom != 0:
				fail_hom = True
				break
		if fail_hom:
			continue
		for het_sample in het_gamet:			#test that all het_gamet are hetereozygous
			if '0/1' in data.loc[:,het_sample].iloc[i] or '1/0' in data.loc[:,het_sample].iloc[i]:
				het += 1
			elif '0/0' in data.loc[:,het_sample].iloc[i] or '1/1' in data.loc[:,het_sample].iloc[i]:
				fail_hom = True
				break
		if fail_hom:
			continue
		#if pass male and female output locus name, number of females genotyped, number of males genotyped, and total number of fish genotyped
		if suffix == 'XY':
			if version == 1:
				output_file.write(str(data.loc[:,'ID'].iloc[i]) + '\t' + str(zero_hom + one_hom) + '\t' + str(het) + '\t' + str(zero_hom + one_hom + het) + '\n')
			else:
				output_file.write(str(data.loc[:,'#CHROM'].iloc[i]) + '_' + str(data.loc[:,'POS'].iloc[i]) + '\t' + str(zero_hom + one_hom) + '\t' + str(het) + '\t' + str(zero_hom + one_hom + het) + '\n')
		else:
			if version == 1:
				output_file.write(str(data.loc[:,'ID'].iloc[i]) + '\t' + str(het) + '\t' + str(zero_hom + one_hom) + '\t' + str(zero_hom + one_hom + het) + '\n')
			else:
				output_file.write(str(data.loc[:,'#CHROM'].iloc[i]) + '_' + str(data.loc[:,'POS'].iloc[i]) + '\t' + str(het) + '\t' + str(zero_hom + one_hom) + '\t' + str(zero_hom + one_hom + het) + '\n')

	output_file.close()

def Main():
	#find header row of vcf file
	vcf = sys.argv[1]
	header = 0
	version = 2		#assume Stacks v2
	for line in open(vcf, 'r'):
		if 'Stacks v1' in line:		#check if Stacks v1
			version = 1
		if '#CHROM' in line:
			break
		header += 1
		
	data = pd.read_csv(open(vcf, 'r'), sep='\t', header=header, index_col=False)	#read in vcf ignoring metainformation lines
	
	#generate lists of male and female sample names
	females, males = [], []
	try:
		for line in open(sys.argv[2], 'r'):
			separated = re.split('[\t\n]', line)
			if separated[1].lower() == 'male':
				males.append(separated[0])
			elif separated[1].lower() == 'female':
				females.append(separated[0])
	except:
		print('cannot open supplied sample_list')
		return
	print('Read in ', str(len(males)), ' males and ', str(len(females)), ' females.')

	####first assuming XX/XY
	print('Searching for candidates under XX-XY (male heterogamety)')
	Filter(data, males, females, 'XY', version)
	
	####next assuming ZZ/WZ	
	print('Searching for candidates under ZZ-WZ (female heterogamety)')
	Filter(data, females, males, 'WZ', version)

Main()
			 	 
			 
			 
			 