#!/usr/bin/env python
# Aylwyn Scally 2012-2018
#TODO: handle ./. output e.g. from vcf-merge

import sys
import optparse
#import os
import re
import os.path
from numpy import median
import logging
from logging import error, warning, info, debug, critical
import gzip
from string import maketrans
#import difflib

loglevel = logging.INFO

# defaults
VCF_FIXEDCOLS = 9
PSMC_BINSIZE = 100
PSMC_N_RATIO = 0.9
sampsep = '-'

#	fout.write('\t'.join([str(x) for x in args]) + '\n')

def genotype(al1, al2):
	return(''.join(sorted([al1, al2])))

def hapsitetypes(n):
# haploid allele combinations
	for gt in ['0', '1']:
		if n > 1:
			for cc in hapsitetypes(n-1):
				yield(''.join((gt, cc)))
		else:
			yield(gt)


def dipsitetypes(n):
# diploid genotype combinations
# to get segregating sites drop first and last elements of returned vals
	for gt in ['00', '01', '11']:
		if n > 1:
			for cc in dipsitetypes(n-1):
				yield('-'.join((gt, cc)))
		else:
			yield(gt)

def phases(gt):
# diploid allele combinations
	yield(gt[0] + gt[2])
	if gt[1] != '|' and gt[0] != gt[2]:
#		yield(gt[::-1])
		yield(gt[2] + gt[0])

def phasecombs(gts, sep = ''):
# different phase combinations of diploid genotypes in gts 
	if len(gts) > 1:
		for cc in phasecombs(gts[1:], sep):
			for ph in phases(gts[0]):
				yield(sep.join((ph + cc)))
	else:
		for ph in phases(gts[0]):
			yield(ph)

class FastaStream(object):
	def __init__(self, fh):
		self.fh = fh
		self.ccount = 0

	def write(self, c):
		self.fh.write(c)
		self.ccount += len(c)
		if self.ccount >= 60:
			self.fh.write('\n')
			self.ccount = 0

	def newrec(self, name):
		self.fh.write('\n>%s\n' % str(name))
		self.ccount = 0


#test
tgt = ['A/B', 'C|D', 'E/F']
assert ','.join(phasecombs(tgt)) == 'ABCDEF,BACDEF,ABCDFE,BACDFE'
tgt = ['A/B', 'C/D', 'E/F']
assert ','.join(phasecombs(tgt)) == 'ABCDEF,BACDEF,ABDCEF,BADCEF,ABCDFE,BACDFE,ABDCFE,BADCFE'

#p = optparse.OptionParser(usage = '%prog [file] [--segsites] [--diversity] [-c cond_sample]')
p = optparse.OptionParser()
p.add_option('-H', '--header', action='store_true', default = False)
p.add_option('--debug', action='store_true', default = False)
p.add_option('--depths', action='store_true', default = False, help = 'include read depths in rates')
p.add_option('--depthsonly', action='store_true', default = False, help = 'output per-individual read depths only')
p.add_option('--no_calls', action='store_true', default = False, help = 'suppress output of calls')
p.add_option('--indels', action='store_true', default = False, help = 'include indels')
p.add_option('--segsites', action='store_true', default = False, help = 'output segregating site flag')
p.add_option('--mediandep', action='store_true', default = False, help = 'output median read depth')
p.add_option('--diversity', action='store_true', default = False, help = 'output diversity statistics')
p.add_option('--Fst', action='store_true', default = False, help = 'output Fst (Weir & Cockerham)')
p.add_option('-c', '--condsamp', default = '', help = 'condition on het in sample CONDSAMP')
p.add_option('--rates', action='store_true', default = False, help = 'output per-individual het and hom-non-ref rates')
p.add_option('--nosites', action='store_true', default = False, help = 'suppress per-site output')
p.add_option('--vars', action='store_true', default = False, help = 'output variant sites only')
p.add_option('--segsep', action='store_true', default = False, help = 'output seg site separations (e.g. for input to msmc)')
p.add_option('--psmcfa', action='store_true', default = False, help = 'output psmcfa (for input to psmc; NOTE psmc only works for two chrs)')
p.add_option('--replacecalls', default='', help = 'replace calls at sites listed in REPLACECALLS, a vcf.gz file of replacement records (e.g. phased SNP calls). NOTE: no checking is done to ensure samples match')
p.add_option('--risk', default='', help = 'calculate risk scores based on allele risk scores in RISK; format: chr, pos, rsID, allele, score; assumes additive model')
p.add_option('--effect_dir', action='store_true', default= False, help = 'use only the risk score sign (+/- 1), not its magnitude')
p.add_option('--callmask', default='', help = 'bed.gz file of uncallable regions')
p.add_option('--pops', default='', help = 'file of sample population assignments (line format: sample pop)')
p.add_option('--aims', action='store_true', default = False, help = 'output AIM sites (requires --pops)')
p.add_option('--alleles', action='store_true', default = False, help = 'output alleles')
p.add_option('--pseudodip', action='store_true', default = False, help = 'create pseudodiploids from consecutive pairs of input samples (assumed haploid so exclude hets)')
p.add_option('--haploid', action='store_true', default = False, help = 'haploid input calls')

opt, args = p.parse_args()

if opt.depthsonly:
	opt.depths = True

#if opt.ratesonly:
#	opt.rates = True

if opt.debug:
	loglevel = logging.DEBUG
logging.basicConfig(format = '%(module)s:%(lineno)d:%(levelname)s: %(message)s', level = loglevel)

if args:
	fin = open(args[0])
else:
	fin = sys.stdin
fout = sys.stdout

if opt.replacecalls:
	info('Reading replacement calls from %s' % opt.replacecalls)
	reprecord = {}
	for line in gzip.open(opt.replacecalls):
		if line.startswith('#'):
			continue
		tok = line.split()
		reprecord[tok[0]+tok[1]] = tok[VCF_FIXEDCOLS-1:]

if opt.risk:
	info('Reading risk rsIDs, alleles and scores from %s' % opt.risk)
	SNPscore = {}
	SNPallele = {}
	for line in open(opt.risk):
		if line.startswith('#'):
			continue
		tok = line.split()
		SNPscore[tok[2]] = float(tok[4])
		if opt.effect_dir:
			SNPscore[tok[2]] /= abs(float(tok[4]))
		SNPallele[tok[2]] = tok[3]
	sampsep = '\t'

if opt.callmask:
	info('Reading uncallable regions from %s' % opt.callmask)
	callmask = {}
	if opt.callmask.endswith('.gz'):
		cmf = gzip.open(opt.callmask)
	else:
		cmf = open(opt.callmask)
	for line in cmf:
		if line.startswith('#'):
			continue
		tok = line.split()
		callmask[int(tok[1]) + 1] = int(tok[2])

if opt.pops:
	info('Reading sample population assignments from %s' % opt.pops)
	samppop = {}
	for line in open(opt.pops):
		if line.startswith('#'):
			continue
		tok = line.split()
		samppop[tok[0]] = tok[1]


	if opt.aims:
		opt.vars = True
		fout.write('\t'.join(['#chr', 'pos'] + poplist + ['AIM_pop']) + '\n')

#	for line in fin:
#		if line.startswith('#'):
#			sys.stdout.write(line)
#			continue
#		tok = line.split()
##		print(tok[0]+tok[1])
#		if tok[0]+tok[1] in reprecord:
##			print(tok[0]+tok[1])
#			tok[VCF_FIXEDCOLS-1:] = reprecord[tok[0]+tok[1]]
#		sys.stdout.write('\t'.join(tok) + '\n')
#	sys.exit(0)

for line in fin:
#	if not opt.noheader:
#		sys.stdout.write(line)
	if line.startswith('#CHROM'):
		tok = line.split()
		nsamp = len(tok) - VCF_FIXEDCOLS
		sampnames = tok[VCF_FIXEDCOLS:]
		sampn = dict([(x[1], int(x[0])) for x in enumerate(sampnames)])

		popsampnums = {}
		if opt.pops: # fill poplist with pops actually present
			poplist = []
			for samp in sampnames:
				if not samp in samppop: # look for matching prefix
					warning('%s not found in %s' % (samp, opt.pops))
					prefixes = [os.path.commonprefix([s, samp]) for s in samppop.keys()]
					longpref = max(prefixes, key=len)
					if len(longpref) > 6: #ARBITRARY HACK TO FIND BEST MATCHING SAMPLE
						match = samppop.keys()[prefixes.index(longpref)]
						samppop[samp] = samppop[match]
						warning('Using %s\t%s' % (match, samppop[match]))

				if not samppop[samp] in poplist: # error if samp not in opt.pops
					poplist.append(samppop[samp])
					popsampnums[samppop[samp]] = []
				popsampnums[samppop[samp]].append(sampn[samp])

		else:
			poplist = ['all_samples']
			popsampnums['all_samples'] = range(nsamp)

		if opt.condsamp:
			for i, name in enumerate(sampnames):
				if name.find(condsamp) >= 0:
					condcol = i + VCF_FIXEDCOLS
#				   print(i, condsamp, tok[condcol])
					break
		break

if opt.risk:
	totalrisk = [0 for x in range(nsamp)]

if opt.rates:
	if opt.depths:
		depcount = [0 for i in range(nsamp)]
	hetcount = [0 for i in range(nsamp)]
	hnrcount = [0 for i in range(nsamp)]
	nsites = 0
	firstcoord = ''

if opt.psmcfa:
	faout = FastaStream(sys.stdout)

#if opt.header:
#	if opt.pops:
#		fout.write('#VCF_POPULATIONS:\t%s\n' % ('\t'.join(poplist)))
#	if not opt.no_calls:
#		fout.write('\t'.join(['#SAMPLES'] + sampnames) + '\n')
if opt.header: #column headers
	header = []
	if opt.rates:
		header += ['#SAMPLES']
	else:
		header += ['CHR', 'POS', 'ID']
	if opt.pops:# and opt.rename_samps: # add pop name to sample names in header
		headsampnames = [samppop[x] + '_' + x for x in sampnames]
	else:
		headsampnames =  sampnames
	if opt.depthsonly:
		header += headsampnames
	if opt.aims:
		header += ['***HEADER_INFO_TODO***']
	if not opt.no_calls:
		header += headsampnames
	if opt.segsites:
		header += ['SEGSITE']
	if opt.diversity: # needs --no_calls to line up with output
		for pop in poplist:
			header += [pop + '_N', pop + '_AC', pop + '_AF', pop + '_HET']
	if opt.Fst:
		header += ['Fst']
	fout.write('\t'.join(header) + '\n')


curchr = ''
for line in fin:
	if line.startswith('#'):
#		sys.stdout.write(line)
		continue

	tok = line.split()

	if tok[0] != curchr:
		sep = 0
		lastpos = 0
		curchr = tok[0]
		if opt.psmcfa:
			nextbin = PSMC_BINSIZE
			faout = FastaStream(sys.stdout)
			faout.newrec(curchr)
			inmask = False
			nN = 0
			outchar = 'T'

	if opt.replacecalls and tok[0]+tok[1] in reprecord:
		tok[VCF_FIXEDCOLS-1:] = reprecord[tok[0]+tok[1]]
	
	if opt.vars and tok[4] == '.':
		continue

#	sys.stdout.write(line)
	if not opt.indels and (len(tok[3]) > 1 or re.search('INDEL', tok[7])):
		continue

#	if len(tok[4]) > 1: # multiallelic sites
#		continue

	if opt.condsamp and not tok[condcol].startswith('0/1'): # condition on variant in cond_sample
		continue

	if opt.depths or opt.mediandep:
		formats = tok[VCF_FIXEDCOLS - 1].split(':')
		depfield = formats.index('DP')
		dep = [x.split(':')[depfield] for x in tok[VCF_FIXEDCOLS:(VCF_FIXEDCOLS+nsamp)]]

	if opt.depthsonly:
		fout.write('\t'.join(tok[0:2] + dep) + '\n')
		continue

	outvals = tok[0:3]

	if opt.haploid:
		if (tok[4] == "."):
			varvals = ['0' for x in tok[VCF_FIXEDCOLS:(VCF_FIXEDCOLS+nsamp)]]
			segsite = False
		else:
			varvals = [x[0] for x in tok[VCF_FIXEDCOLS:(VCF_FIXEDCOLS+nsamp)]]
		vargts = varvals
	else:
		if (tok[4] == "."):
			varvals = ['00' for x in tok[VCF_FIXEDCOLS:(VCF_FIXEDCOLS+nsamp)]]
			vargts = ['0/0' for x in tok[VCF_FIXEDCOLS:(VCF_FIXEDCOLS+nsamp)]]
			hets = [0 for x in tok[VCF_FIXEDCOLS:(VCF_FIXEDCOLS+nsamp)]]
			segsite = False
		else:
			varvals = [x[0] + x[2] for x in tok[VCF_FIXEDCOLS:(VCF_FIXEDCOLS+nsamp)]]
			vargts = [x[0:3] for x in tok[VCF_FIXEDCOLS:(VCF_FIXEDCOLS+nsamp)]]
			hets = [int(x[0] != x[2]) for x in tok[VCF_FIXEDCOLS:(VCF_FIXEDCOLS+nsamp)]]

		if opt.pseudodip:
			hetsite = False
			for gt in vargts:
				if gt[0] != gt[2]:
					hetsite = True
					break
			if hetsite:
				continue

	if opt.risk: # convert genotype to GRS
		sitealleles = [tok[3]] + tok[4].split(',')
		effectallele = str(sitealleles.index(SNPallele[tok[2]]))
		varvals = [SNPscore[tok[2]] * (int(x[0] == effectallele) + int(x[1] == effectallele)) for x in varvals]
		totalrisk = [x + y for (x, y) in zip(totalrisk, varvals)]
		varvals = [str(x) for x in varvals]

	if opt.pseudodip:
# combine consecutive varvals elements into new varvals (assumes none are het)
# TODO?: only check/exclude hets in consecutive pairs
		varvals = [varvals[i][0] + varvals[i+1][0] for i in range(0, len(varvals) - 1, 2)]
		vargts = [vargts[i][0] + '|' + vargts[i+1][0] for i in range(0, len(vargts) - 1, 2)]
#		debug(vargts)

# determine if segregating site
# ideally use FQ > 0 but vcf-subset doesn't currently adjust this
#		fq = re.search('FQ=(-?[0-9.]+)', tok[7])
#		if fq > 0: #het
#		sys.stdout.write(line)
	varvalstr = '-'.join(varvals)
#	debug(valstr)
	segsite = '0' in varvalstr and '1' in varvalstr

	if opt.alleles: # convert 0, 1 to alleles in output
		sitealleles = [tok[3]] + tok[4].split(',')
		trtab = maketrans(''.join([str(x) for x in range(len(sitealleles))]), ''.join(sitealleles))
		varalleles = [x.translate(trtab) for x in varvals]  
		varvals = varalleles  
		varvalstr = '-'.join(varvals)
		vargts = [x.translate(trtab) for x in vargts]  
#		debug([tok[1], sitealleles, vargts])


	if opt.aims:
		valpops = {}
		popvals = {}
		for isn, samp in enumerate(sampnames):
			spop = samppop[samp]
			for val in varvals[isn]:
				if val in valpops:
					valpops[val].add(spop)
				else:
					valpops[val] = set([spop])
				if spop in popvals:
					popvals[spop].add(val)
				else:
					popvals[spop] = set([val])
		for val in valpops:
			if len(valpops[val]) == 1: # require uniqueness of marker to pop
				spop = valpops[val].pop() 
				if len(popvals[spop]) == 1: # require fixation of marker in pop
					debug(varvals)
#					debug(valpops)
					debug(popvals)
					spvals = ['/'.join(list(popvals[x])) for x in poplist]
					debug(spvals)
					fout.write('\t'.join(outvals + spvals + [spop]) + '\n')
#					fout.write('\t'.join(outvals + [val] + spvals + [spop]) + '\n')
		continue

	if opt.segsep:
		pos = int(tok[1])

		while pos > lastpos:
			if opt.callmask and lastpos in callmask:
#				print('pos = %d; Masking %d to %d' % (pos, lastpos, callmask[lastpos]))
				lastpos = callmask[lastpos]
			else:
				sep += 1
			lastpos += 1
		if opt.callmask and lastpos in callmask:
#			print('pos = %d; Masking %d to %d' % (pos, lastpos, callmask[lastpos]))
			lastpos = callmask[lastpos] + 1
		elif lastpos == pos and segsite:
			if len(vargts) > 1:
				allstr = ','.join(phasecombs(vargts))
				fout.write('\t'.join(outvals + [str(sep)] + [allstr]) + '\n')
			else:
				fout.write('\t'.join(outvals + [str(sep)] + varvals) + '\n')
			sep = 0
		continue

	elif opt.psmcfa:
		if segsite:
			pos = int(tok[1])

#			print('site at %d' % pos)
			while lastpos < nextbin and pos > nextbin - PSMC_BINSIZE:
				lastpos += 1
				if opt.callmask and lastpos in callmask:
#					print('Mask %d to %d : %d' % (lastpos, callmask[lastpos], callmask[lastpos]-lastpos + 1))
					inmask = True
					endmask = callmask[lastpos]
				if inmask:
					nN += 1
					if lastpos == endmask:
						inmask = False
				if lastpos == pos:
					outchar = 'K'
				if lastpos == nextbin:
					if float(nN) / PSMC_BINSIZE > PSMC_N_RATIO:
						outchar = 'N'
					faout.write(outchar)
#					print(' %d-%d: %d Ns, pos = %d' % (nextbin - PSMC_BINSIZE + 1, nextbin, nN, pos))
					nextbin += PSMC_BINSIZE
					outchar = 'T'
					nN = 0
		continue

	if opt.rates:
		if opt.depths:
			depcount = [depcount[i] + int(dep[i]) for i in range(nsamp)]
		hetcount = [hetcount[i] + (varvals[i][0] != varvals[i][1]) for i in range(nsamp)]
		hnrcount = [hnrcount[i] + (varvals[i] == '11') for i in range(nsamp)]
		nsites += 1
		if not firstcoord:
			firstcoord = ':'.join(tok[0:2])
		lastcoord = ':'.join(tok[0:2])

	if not opt.nosites:
		if opt.mediandep:
			outvals.append('%d' % median(dep))

		if opt.segsites:
#			outvals.append(str(depfield - 1))
			outvals.append(str(int(segsite)))

		alcount_total = 0
		samps_total = 0
		Hs_mean = 0
		for pop in poplist:
			pvarvalstr = sampsep.join([varvals[i] for i in popsampnums[pop]])
			phets = [hets[i] for i in popsampnums[pop]]
			if not opt.no_calls:
				outvals.append(pvarvalstr)

			if opt.diversity or opt.Fst:
	#			alfreq = eval(re.search('AF1=([0-9.e\-]+)', tok[7]).group(1))# eval since vcf uses sci notation
	#			alcount = int(re.search('AC1=([0-9]+)', tok[7]).group(1))
	#			neidiv = 1 - alfreq**2 - (1 - alfreq)**2
	#			outvals += [str(nsamp), str(alfreq), str(neidiv), str(alcount), str(alcount/float(nsamp))]
				alcount = pvarvalstr.count('1')
				hetcount = sum(phets)
				pnsamp = len(popsampnums[pop])
				pfreq = alcount/float(2 * pnsamp)
				alcount_total += alcount
				samps_total += pnsamp
				Hs_mean += 2 * pfreq * (1 - pfreq)
				if opt.diversity:
					outvals += [str(pnsamp), str(alcount), str(pfreq), str(hetcount)]

		if opt.Fst:
			Hs_mean = Hs_mean / len(poplist)
			tfreq = alcount_total/float(2 * samps_total)
			Ht = 2 * tfreq * (1 - tfreq)
			Fst = (Ht - Hs_mean) / Ht
#			outvals += [str(Hs_mean), str(Ht), str(Fst)]
			outvals += [str(Fst)]

		fout.write('\t'.join(outvals) + '\n')

if opt.rates:
	fout.write('\t'.join(['#REGION'] + [firstcoord + '-' + lastcoord]) + '\n')
	fout.write('\t'.join(['#NSITES'] + [str(nsites)]) + '\n')
	if opt.depths:
		fout.write('\t'.join(['#MEANDEP'] + [str(float(x)/nsites) for x in depcount]) + '\n')
	fout.write('\t'.join(['#HETCOUNT'] + [str(x) for x in hetcount]) + '\n')
	fout.write('\t'.join(['#HETRATE'] + [str(float(x)/nsites) for x in hetcount]) + '\n')
	fout.write('\t'.join(['#HNRCOUNT'] + [str(x) for x in hnrcount]) + '\n')
	fout.write('\t'.join(['#HNRRATE'] + [str(float(x)/nsites) for x in hnrcount]) + '\n')
	fout.write('\t'.join(['#HNR/HET'] + [str(float(hnrcount[i])/hetcount[i]) for i in range(nsamp)]) + '\n')
#	fout.write('\t'.join(outvals) + '\n')

if opt.risk:
	fout.write('\t'.join(['#TOTALRISK', '.', '.'] + [str(x) for x in totalrisk]) + '\n')

try:
	sys.stdout.close()
except:
	pass
try:
	sys.stderr.close()
except:
	pass


772b0a538cc21031e1c6215be21ad0d2ec50b4f4
<mask token>


def phasecombs(gts, sep=''):
    if len(gts) > 1:
        for cc in phasecombs(gts[1:], sep):
            for ph in phases(gts[0]):
                yield sep.join(ph + cc)
    else:
        for ph in phases(gts[0]):
            yield ph


class FastaStream(object):

    def __init__(self, fh):
        self.fh = fh
        self.ccount = 0

    def write(self, c):
        self.fh.write(c)
        self.ccount += len(c)
        if self.ccount >= 60:
            self.fh.write('\n')
            self.ccount = 0

    def newrec(self, name):
        self.fh.write('\n>%s\n' % str(name))
        self.ccount = 0


<mask token>

<mask token>


def genotype(al1, al2):
    return ''.join(sorted([al1, al2]))


<mask token>


def dipsitetypes(n):
    for gt in ['00', '01', '11']:
        if n > 1:
            for cc in dipsitetypes(n - 1):
                yield '-'.join((gt, cc))
        else:
            yield gt


<mask token>


def phasecombs(gts, sep=''):
    if len(gts) > 1:
        for cc in phasecombs(gts[1:], sep):
            for ph in phases(gts[0]):
                yield sep.join(ph + cc)
    else:
        for ph in phases(gts[0]):
            yield ph


class FastaStream(object):

    def __init__(self, fh):
        self.fh = fh
        self.ccount = 0

    def write(self, c):
        self.fh.write(c)
        self.ccount += len(c)
        if self.ccount >= 60:
            self.fh.write('\n')
            self.ccount = 0

    def newrec(self, name):
        self.fh.write('\n>%s\n' % str(name))
        self.ccount = 0


<mask token>

<mask token>


def genotype(al1, al2):
    return ''.join(sorted([al1, al2]))


def hapsitetypes(n):
    for gt in ['0', '1']:
        if n > 1:
            for cc in hapsitetypes(n - 1):
                yield ''.join((gt, cc))
        else:
            yield gt


def dipsitetypes(n):
    for gt in ['00', '01', '11']:
        if n > 1:
            for cc in dipsitetypes(n - 1):
                yield '-'.join((gt, cc))
        else:
            yield gt


def phases(gt):
    yield gt[0] + gt[2]
    if gt[1] != '|' and gt[0] != gt[2]:
        yield gt[2] + gt[0]


def phasecombs(gts, sep=''):
    if len(gts) > 1:
        for cc in phasecombs(gts[1:], sep):
            for ph in phases(gts[0]):
                yield sep.join(ph + cc)
    else:
        for ph in phases(gts[0]):
            yield ph


class FastaStream(object):

    def __init__(self, fh):
        self.fh = fh
        self.ccount = 0

    def write(self, c):
        self.fh.write(c)
        self.ccount += len(c)
        if self.ccount >= 60:
            self.fh.write('\n')
            self.ccount = 0

    def newrec(self, name):
        self.fh.write('\n>%s\n' % str(name))
        self.ccount = 0


<mask token>

import sys
import optparse
import re
import os.path
from numpy import median
import logging
from logging import error, warning, info, debug, critical
import gzip
from string import maketrans
loglevel = logging.INFO
VCF_FIXEDCOLS = 9
PSMC_BINSIZE = 100
PSMC_N_RATIO = 0.9
sampsep = '-'


def genotype(al1, al2):
    return ''.join(sorted([al1, al2]))


def hapsitetypes(n):
    for gt in ['0', '1']:
        if n > 1:
            for cc in hapsitetypes(n - 1):
                yield ''.join((gt, cc))
        else:
            yield gt


def dipsitetypes(n):
    for gt in ['00', '01', '11']:
        if n > 1:
            for cc in dipsitetypes(n - 1):
                yield '-'.join((gt, cc))
        else:
            yield gt


def phases(gt):
    yield gt[0] + gt[2]
    if gt[1] != '|' and gt[0] != gt[2]:
        yield gt[2] + gt[0]


def phasecombs(gts, sep=''):
    if len(gts) > 1:
        for cc in phasecombs(gts[1:], sep):
            for ph in phases(gts[0]):
                yield sep.join(ph + cc)
    else:
        for ph in phases(gts[0]):
            yield ph


class FastaStream(object):

    def __init__(self, fh):
        self.fh = fh
        self.ccount = 0

    def write(self, c):
        self.fh.write(c)
        self.ccount += len(c)
        if self.ccount >= 60:
            self.fh.write('\n')
            self.ccount = 0

    def newrec(self, name):
        self.fh.write('\n>%s\n' % str(name))
        self.ccount = 0


tgt = ['A/B', 'C|D', 'E/F']
assert ','.join(phasecombs(tgt)) == 'ABCDEF,BACDEF,ABCDFE,BACDFE'
tgt = ['A/B', 'C/D', 'E/F']
assert ','.join(phasecombs(tgt)
    ) == 'ABCDEF,BACDEF,ABDCEF,BADCEF,ABCDFE,BACDFE,ABDCFE,BADCFE'
p = optparse.OptionParser()
p.add_option('-H', '--header', action='store_true', default=False)
p.add_option('--debug', action='store_true', default=False)
p.add_option('--depths', action='store_true', default=False, help=
    'include read depths in rates')
p.add_option('--depthsonly', action='store_true', default=False, help=
    'output per-individual read depths only')
p.add_option('--no_calls', action='store_true', default=False, help=
    'suppress output of calls')
p.add_option('--indels', action='store_true', default=False, help=
    'include indels')
p.add_option('--segsites', action='store_true', default=False, help=
    'output segregating site flag')
p.add_option('--mediandep', action='store_true', default=False, help=
    'output median read depth')
p.add_option('--diversity', action='store_true', default=False, help=
    'output diversity statistics')
p.add_option('--Fst', action='store_true', default=False, help=
    'output Fst (Weir & Cockerham)')
p.add_option('-c', '--condsamp', default='', help=
    'condition on het in sample CONDSAMP')
p.add_option('--rates', action='store_true', default=False, help=
    'output per-individual het and hom-non-ref rates')
p.add_option('--nosites', action='store_true', default=False, help=
    'suppress per-site output')
p.add_option('--vars', action='store_true', default=False, help=
    'output variant sites only')
p.add_option('--segsep', action='store_true', default=False, help=
    'output seg site separations (e.g. for input to msmc)')
p.add_option('--psmcfa', action='store_true', default=False, help=
    'output psmcfa (for input to psmc; NOTE psmc only works for two chrs)')
p.add_option('--replacecalls', default='', help=
    'replace calls at sites listed in REPLACECALLS, a vcf.gz file of replacement records (e.g. phased SNP calls). NOTE: no checking is done to ensure samples match'
    )
p.add_option('--risk', default='', help=
    'calculate risk scores based on allele risk scores in RISK; format: chr, pos, rsID, allele, score; assumes additive model'
    )
p.add_option('--effect_dir', action='store_true', default=False, help=
    'use only the risk score sign (+/- 1), not its magnitude')
p.add_option('--callmask', default='', help='bed.gz file of uncallable regions'
    )
p.add_option('--pops', default='', help=
    'file of sample population assignments (line format: sample pop)')
p.add_option('--aims', action='store_true', default=False, help=
    'output AIM sites (requires --pops)')
p.add_option('--alleles', action='store_true', default=False, help=
    'output alleles')
p.add_option('--pseudodip', action='store_true', default=False, help=
    'create pseudodiploids from consecutive pairs of input samples (assumed haploid so exclude hets)'
    )
p.add_option('--haploid', action='store_true', default=False, help=
    'haploid input calls')
opt, args = p.parse_args()
if opt.depthsonly:
    opt.depths = True
if opt.debug:
    loglevel = logging.DEBUG
logging.basicConfig(format=
    '%(module)s:%(lineno)d:%(levelname)s: %(message)s', level=loglevel)
if args:
    fin = open(args[0])
else:
    fin = sys.stdin
fout = sys.stdout
if opt.replacecalls:
    info('Reading replacement calls from %s' % opt.replacecalls)
    reprecord = {}
    for line in gzip.open(opt.replacecalls):
        if line.startswith('#'):
            continue
        tok = line.split()
        reprecord[tok[0] + tok[1]] = tok[VCF_FIXEDCOLS - 1:]
if opt.risk:
    info('Reading risk rsIDs, alleles and scores from %s' % opt.risk)
    SNPscore = {}
    SNPallele = {}
    for line in open(opt.risk):
        if line.startswith('#'):
            continue
        tok = line.split()
        SNPscore[tok[2]] = float(tok[4])
        if opt.effect_dir:
            SNPscore[tok[2]] /= abs(float(tok[4]))
        SNPallele[tok[2]] = tok[3]
    sampsep = '\t'
if opt.callmask:
    info('Reading uncallable regions from %s' % opt.callmask)
    callmask = {}
    if opt.callmask.endswith('.gz'):
        cmf = gzip.open(opt.callmask)
    else:
        cmf = open(opt.callmask)
    for line in cmf:
        if line.startswith('#'):
            continue
        tok = line.split()
        callmask[int(tok[1]) + 1] = int(tok[2])
if opt.pops:
    info('Reading sample population assignments from %s' % opt.pops)
    samppop = {}
    for line in open(opt.pops):
        if line.startswith('#'):
            continue
        tok = line.split()
        samppop[tok[0]] = tok[1]
    if opt.aims:
        opt.vars = True
        fout.write('\t'.join(['#chr', 'pos'] + poplist + ['AIM_pop']) + '\n')
for line in fin:
    if line.startswith('#CHROM'):
        tok = line.split()
        nsamp = len(tok) - VCF_FIXEDCOLS
        sampnames = tok[VCF_FIXEDCOLS:]
        sampn = dict([(x[1], int(x[0])) for x in enumerate(sampnames)])
        popsampnums = {}
        if opt.pops:
            poplist = []
            for samp in sampnames:
                if not samp in samppop:
                    warning('%s not found in %s' % (samp, opt.pops))
                    prefixes = [os.path.commonprefix([s, samp]) for s in
                        samppop.keys()]
                    longpref = max(prefixes, key=len)
                    if len(longpref) > 6:
                        match = samppop.keys()[prefixes.index(longpref)]
                        samppop[samp] = samppop[match]
                        warning('Using %s\t%s' % (match, samppop[match]))
                if not samppop[samp] in poplist:
                    poplist.append(samppop[samp])
                    popsampnums[samppop[samp]] = []
                popsampnums[samppop[samp]].append(sampn[samp])
        else:
            poplist = ['all_samples']
            popsampnums['all_samples'] = range(nsamp)
        if opt.condsamp:
            for i, name in enumerate(sampnames):
                if name.find(condsamp) >= 0:
                    condcol = i + VCF_FIXEDCOLS
                    break
        break
if opt.risk:
    totalrisk = [(0) for x in range(nsamp)]
if opt.rates:
    if opt.depths:
        depcount = [(0) for i in range(nsamp)]
    hetcount = [(0) for i in range(nsamp)]
    hnrcount = [(0) for i in range(nsamp)]
    nsites = 0
    firstcoord = ''
if opt.psmcfa:
    faout = FastaStream(sys.stdout)
if opt.header:
    header = []
    if opt.rates:
        header += ['#SAMPLES']
    else:
        header += ['CHR', 'POS', 'ID']
    if opt.pops:
        headsampnames = [(samppop[x] + '_' + x) for x in sampnames]
    else:
        headsampnames = sampnames
    if opt.depthsonly:
        header += headsampnames
    if opt.aims:
        header += ['***HEADER_INFO_TODO***']
    if not opt.no_calls:
        header += headsampnames
    if opt.segsites:
        header += ['SEGSITE']
    if opt.diversity:
        for pop in poplist:
            header += [pop + '_N', pop + '_AC', pop + '_AF', pop + '_HET']
    if opt.Fst:
        header += ['Fst']
    fout.write('\t'.join(header) + '\n')
curchr = ''
for line in fin:
    if line.startswith('#'):
        continue
    tok = line.split()
    if tok[0] != curchr:
        sep = 0
        lastpos = 0
        curchr = tok[0]
        if opt.psmcfa:
            nextbin = PSMC_BINSIZE
            faout = FastaStream(sys.stdout)
            faout.newrec(curchr)
            inmask = False
            nN = 0
            outchar = 'T'
    if opt.replacecalls and tok[0] + tok[1] in reprecord:
        tok[VCF_FIXEDCOLS - 1:] = reprecord[tok[0] + tok[1]]
    if opt.vars and tok[4] == '.':
        continue
    if not opt.indels and (len(tok[3]) > 1 or re.search('INDEL', tok[7])):
        continue
    if opt.condsamp and not tok[condcol].startswith('0/1'):
        continue
    if opt.depths or opt.mediandep:
        formats = tok[VCF_FIXEDCOLS - 1].split(':')
        depfield = formats.index('DP')
        dep = [x.split(':')[depfield] for x in tok[VCF_FIXEDCOLS:
            VCF_FIXEDCOLS + nsamp]]
    if opt.depthsonly:
        fout.write('\t'.join(tok[0:2] + dep) + '\n')
        continue
    outvals = tok[0:3]
    if opt.haploid:
        if tok[4] == '.':
            varvals = ['0' for x in tok[VCF_FIXEDCOLS:VCF_FIXEDCOLS + nsamp]]
            segsite = False
        else:
            varvals = [x[0] for x in tok[VCF_FIXEDCOLS:VCF_FIXEDCOLS + nsamp]]
        vargts = varvals
    else:
        if tok[4] == '.':
            varvals = ['00' for x in tok[VCF_FIXEDCOLS:VCF_FIXEDCOLS + nsamp]]
            vargts = ['0/0' for x in tok[VCF_FIXEDCOLS:VCF_FIXEDCOLS + nsamp]]
            hets = [(0) for x in tok[VCF_FIXEDCOLS:VCF_FIXEDCOLS + nsamp]]
            segsite = False
        else:
            varvals = [(x[0] + x[2]) for x in tok[VCF_FIXEDCOLS:
                VCF_FIXEDCOLS + nsamp]]
            vargts = [x[0:3] for x in tok[VCF_FIXEDCOLS:VCF_FIXEDCOLS + nsamp]]
            hets = [int(x[0] != x[2]) for x in tok[VCF_FIXEDCOLS:
                VCF_FIXEDCOLS + nsamp]]
        if opt.pseudodip:
            hetsite = False
            for gt in vargts:
                if gt[0] != gt[2]:
                    hetsite = True
                    break
            if hetsite:
                continue
    if opt.risk:
        sitealleles = [tok[3]] + tok[4].split(',')
        effectallele = str(sitealleles.index(SNPallele[tok[2]]))
        varvals = [(SNPscore[tok[2]] * (int(x[0] == effectallele) + int(x[1
            ] == effectallele))) for x in varvals]
        totalrisk = [(x + y) for x, y in zip(totalrisk, varvals)]
        varvals = [str(x) for x in varvals]
    if opt.pseudodip:
        varvals = [(varvals[i][0] + varvals[i + 1][0]) for i in range(0, 
            len(varvals) - 1, 2)]
        vargts = [(vargts[i][0] + '|' + vargts[i + 1][0]) for i in range(0,
            len(vargts) - 1, 2)]
    varvalstr = '-'.join(varvals)
    segsite = '0' in varvalstr and '1' in varvalstr
    if opt.alleles:
        sitealleles = [tok[3]] + tok[4].split(',')
        trtab = maketrans(''.join([str(x) for x in range(len(sitealleles))]
            ), ''.join(sitealleles))
        varalleles = [x.translate(trtab) for x in varvals]
        varvals = varalleles
        varvalstr = '-'.join(varvals)
        vargts = [x.translate(trtab) for x in vargts]
    if opt.aims:
        valpops = {}
        popvals = {}
        for isn, samp in enumerate(sampnames):
            spop = samppop[samp]
            for val in varvals[isn]:
                if val in valpops:
                    valpops[val].add(spop)
                else:
                    valpops[val] = set([spop])
                if spop in popvals:
                    popvals[spop].add(val)
                else:
                    popvals[spop] = set([val])
        for val in valpops:
            if len(valpops[val]) == 1:
                spop = valpops[val].pop()
                if len(popvals[spop]) == 1:
                    debug(varvals)
                    debug(popvals)
                    spvals = ['/'.join(list(popvals[x])) for x in poplist]
                    debug(spvals)
                    fout.write('\t'.join(outvals + spvals + [spop]) + '\n')
        continue
    if opt.segsep:
        pos = int(tok[1])
        while pos > lastpos:
            if opt.callmask and lastpos in callmask:
                lastpos = callmask[lastpos]
            else:
                sep += 1
            lastpos += 1
        if opt.callmask and lastpos in callmask:
            lastpos = callmask[lastpos] + 1
        elif lastpos == pos and segsite:
            if len(vargts) > 1:
                allstr = ','.join(phasecombs(vargts))
                fout.write('\t'.join(outvals + [str(sep)] + [allstr]) + '\n')
            else:
                fout.write('\t'.join(outvals + [str(sep)] + varvals) + '\n')
            sep = 0
        continue
    elif opt.psmcfa:
        if segsite:
            pos = int(tok[1])
            while lastpos < nextbin and pos > nextbin - PSMC_BINSIZE:
                lastpos += 1
                if opt.callmask and lastpos in callmask:
                    inmask = True
                    endmask = callmask[lastpos]
                if inmask:
                    nN += 1
                    if lastpos == endmask:
                        inmask = False
                if lastpos == pos:
                    outchar = 'K'
                if lastpos == nextbin:
                    if float(nN) / PSMC_BINSIZE > PSMC_N_RATIO:
                        outchar = 'N'
                    faout.write(outchar)
                    nextbin += PSMC_BINSIZE
                    outchar = 'T'
                    nN = 0
        continue
    if opt.rates:
        if opt.depths:
            depcount = [(depcount[i] + int(dep[i])) for i in range(nsamp)]
        hetcount = [(hetcount[i] + (varvals[i][0] != varvals[i][1])) for i in
            range(nsamp)]
        hnrcount = [(hnrcount[i] + (varvals[i] == '11')) for i in range(nsamp)]
        nsites += 1
        if not firstcoord:
            firstcoord = ':'.join(tok[0:2])
        lastcoord = ':'.join(tok[0:2])
    if not opt.nosites:
        if opt.mediandep:
            outvals.append('%d' % median(dep))
        if opt.segsites:
            outvals.append(str(int(segsite)))
        alcount_total = 0
        samps_total = 0
        Hs_mean = 0
        for pop in poplist:
            pvarvalstr = sampsep.join([varvals[i] for i in popsampnums[pop]])
            phets = [hets[i] for i in popsampnums[pop]]
            if not opt.no_calls:
                outvals.append(pvarvalstr)
            if opt.diversity or opt.Fst:
                alcount = pvarvalstr.count('1')
                hetcount = sum(phets)
                pnsamp = len(popsampnums[pop])
                pfreq = alcount / float(2 * pnsamp)
                alcount_total += alcount
                samps_total += pnsamp
                Hs_mean += 2 * pfreq * (1 - pfreq)
                if opt.diversity:
                    outvals += [str(pnsamp), str(alcount), str(pfreq), str(
                        hetcount)]
        if opt.Fst:
            Hs_mean = Hs_mean / len(poplist)
            tfreq = alcount_total / float(2 * samps_total)
            Ht = 2 * tfreq * (1 - tfreq)
            Fst = (Ht - Hs_mean) / Ht
            outvals += [str(Fst)]
        fout.write('\t'.join(outvals) + '\n')
if opt.rates:
    fout.write('\t'.join(['#REGION'] + [firstcoord + '-' + lastcoord]) + '\n')
    fout.write('\t'.join(['#NSITES'] + [str(nsites)]) + '\n')
    if opt.depths:
        fout.write('\t'.join(['#MEANDEP'] + [str(float(x) / nsites) for x in
            depcount]) + '\n')
    fout.write('\t'.join(['#HETCOUNT'] + [str(x) for x in hetcount]) + '\n')
    fout.write('\t'.join(['#HETRATE'] + [str(float(x) / nsites) for x in
        hetcount]) + '\n')
    fout.write('\t'.join(['#HNRCOUNT'] + [str(x) for x in hnrcount]) + '\n')
    fout.write('\t'.join(['#HNRRATE'] + [str(float(x) / nsites) for x in
        hnrcount]) + '\n')
    fout.write('\t'.join(['#HNR/HET'] + [str(float(hnrcount[i]) / hetcount[
        i]) for i in range(nsamp)]) + '\n')
if opt.risk:
    fout.write('\t'.join(['#TOTALRISK', '.', '.'] + [str(x) for x in
        totalrisk]) + '\n')
try:
    sys.stdout.close()
except:
    pass
try:
    sys.stderr.close()
except:
    pass

#!/usr/bin/env python
# Aylwyn Scally 2012-2018
#TODO: handle ./. output e.g. from vcf-merge

import sys
import optparse
#import os
import re
import os.path
from numpy import median
import logging
from logging import error, warning, info, debug, critical
import gzip
from string import maketrans
#import difflib

loglevel = logging.INFO

# defaults
VCF_FIXEDCOLS = 9
PSMC_BINSIZE = 100
PSMC_N_RATIO = 0.9
sampsep = '-'

#	fout.write('\t'.join([str(x) for x in args]) + '\n')

def genotype(al1, al2):
	return(''.join(sorted([al1, al2])))

def hapsitetypes(n):
# haploid allele combinations
	for gt in ['0', '1']:
		if n > 1:
			for cc in hapsitetypes(n-1):
				yield(''.join((gt, cc)))
		else:
			yield(gt)


def dipsitetypes(n):
# diploid genotype combinations
# to get segregating sites drop first and last elements of returned vals
	for gt in ['00', '01', '11']:
		if n > 1:
			for cc in dipsitetypes(n-1):
				yield('-'.join((gt, cc)))
		else:
			yield(gt)

def phases(gt):
# diploid allele combinations
	yield(gt[0] + gt[2])
	if gt[1] != '|' and gt[0] != gt[2]:
#		yield(gt[::-1])
		yield(gt[2] + gt[0])

def phasecombs(gts, sep = ''):
# different phase combinations of diploid genotypes in gts 
	if len(gts) > 1:
		for cc in phasecombs(gts[1:], sep):
			for ph in phases(gts[0]):
				yield(sep.join((ph + cc)))
	else:
		for ph in phases(gts[0]):
			yield(ph)

class FastaStream(object):
	def __init__(self, fh):
		self.fh = fh
		self.ccount = 0

	def write(self, c):
		self.fh.write(c)
		self.ccount += len(c)
		if self.ccount >= 60:
			self.fh.write('\n')
			self.ccount = 0

	def newrec(self, name):
		self.fh.write('\n>%s\n' % str(name))
		self.ccount = 0


#test
tgt = ['A/B', 'C|D', 'E/F']
assert ','.join(phasecombs(tgt)) == 'ABCDEF,BACDEF,ABCDFE,BACDFE'
tgt = ['A/B', 'C/D', 'E/F']
assert ','.join(phasecombs(tgt)) == 'ABCDEF,BACDEF,ABDCEF,BADCEF,ABCDFE,BACDFE,ABDCFE,BADCFE'

#p = optparse.OptionParser(usage = '%prog [file] [--segsites] [--diversity] [-c cond_sample]')
p = optparse.OptionParser()
p.add_option('-H', '--header', action='store_true', default = False)
p.add_option('--debug', action='store_true', default = False)
p.add_option('--depths', action='store_true', default = False, help = 'include read depths in rates')
p.add_option('--depthsonly', action='store_true', default = False, help = 'output per-individual read depths only')
p.add_option('--no_calls', action='store_true', default = False, help = 'suppress output of calls')
p.add_option('--indels', action='store_true', default = False, help = 'include indels')
p.add_option('--segsites', action='store_true', default = False, help = 'output segregating site flag')
p.add_option('--mediandep', action='store_true', default = False, help = 'output median read depth')
p.add_option('--diversity', action='store_true', default = False, help = 'output diversity statistics')
p.add_option('--Fst', action='store_true', default = False, help = 'output Fst (Weir & Cockerham)')
p.add_option('-c', '--condsamp', default = '', help = 'condition on het in sample CONDSAMP')
p.add_option('--rates', action='store_true', default = False, help = 'output per-individual het and hom-non-ref rates')
p.add_option('--nosites', action='store_true', default = False, help = 'suppress per-site output')
p.add_option('--vars', action='store_true', default = False, help = 'output variant sites only')
p.add_option('--segsep', action='store_true', default = False, help = 'output seg site separations (e.g. for input to msmc)')
p.add_option('--psmcfa', action='store_true', default = False, help = 'output psmcfa (for input to psmc; NOTE psmc only works for two chrs)')
p.add_option('--replacecalls', default='', help = 'replace calls at sites listed in REPLACECALLS, a vcf.gz file of replacement records (e.g. phased SNP calls). NOTE: no checking is done to ensure samples match')
p.add_option('--risk', default='', help = 'calculate risk scores based on allele risk scores in RISK; format: chr, pos, rsID, allele, score; assumes additive model')
p.add_option('--effect_dir', action='store_true', default= False, help = 'use only the risk score sign (+/- 1), not its magnitude')
p.add_option('--callmask', default='', help = 'bed.gz file of uncallable regions')
p.add_option('--pops', default='', help = 'file of sample population assignments (line format: sample pop)')
p.add_option('--aims', action='store_true', default = False, help = 'output AIM sites (requires --pops)')
p.add_option('--alleles', action='store_true', default = False, help = 'output alleles')
p.add_option('--pseudodip', action='store_true', default = False, help = 'create pseudodiploids from consecutive pairs of input samples (assumed haploid so exclude hets)')
p.add_option('--haploid', action='store_true', default = False, help = 'haploid input calls')

opt, args = p.parse_args()

if opt.depthsonly:
	opt.depths = True

#if opt.ratesonly:
#	opt.rates = True

if opt.debug:
	loglevel = logging.DEBUG
logging.basicConfig(format = '%(module)s:%(lineno)d:%(levelname)s: %(message)s', level = loglevel)

if args:
	fin = open(args[0])
else:
	fin = sys.stdin
fout = sys.stdout

if opt.replacecalls:
	info('Reading replacement calls from %s' % opt.replacecalls)
	reprecord = {}
	for line in gzip.open(opt.replacecalls):
		if line.startswith('#'):
			continue
		tok = line.split()
		reprecord[tok[0]+tok[1]] = tok[VCF_FIXEDCOLS-1:]

if opt.risk:
	info('Reading risk rsIDs, alleles and scores from %s' % opt.risk)
	SNPscore = {}
	SNPallele = {}
	for line in open(opt.risk):
		if line.startswith('#'):
			continue
		tok = line.split()
		SNPscore[tok[2]] = float(tok[4])
		if opt.effect_dir:
			SNPscore[tok[2]] /= abs(float(tok[4]))
		SNPallele[tok[2]] = tok[3]
	sampsep = '\t'

if opt.callmask:
	info('Reading uncallable regions from %s' % opt.callmask)
	callmask = {}
	if opt.callmask.endswith('.gz'):
		cmf = gzip.open(opt.callmask)
	else:
		cmf = open(opt.callmask)
	for line in cmf:
		if line.startswith('#'):
			continue
		tok = line.split()
		callmask[int(tok[1]) + 1] = int(tok[2])

if opt.pops:
	info('Reading sample population assignments from %s' % opt.pops)
	samppop = {}
	for line in open(opt.pops):
		if line.startswith('#'):
			continue
		tok = line.split()
		samppop[tok[0]] = tok[1]


	if opt.aims:
		opt.vars = True
		fout.write('\t'.join(['#chr', 'pos'] + poplist + ['AIM_pop']) + '\n')

#	for line in fin:
#		if line.startswith('#'):
#			sys.stdout.write(line)
#			continue
#		tok = line.split()
##		print(tok[0]+tok[1])
#		if tok[0]+tok[1] in reprecord:
##			print(tok[0]+tok[1])
#			tok[VCF_FIXEDCOLS-1:] = reprecord[tok[0]+tok[1]]
#		sys.stdout.write('\t'.join(tok) + '\n')
#	sys.exit(0)

for line in fin:
#	if not opt.noheader:
#		sys.stdout.write(line)
	if line.startswith('#CHROM'):
		tok = line.split()
		nsamp = len(tok) - VCF_FIXEDCOLS
		sampnames = tok[VCF_FIXEDCOLS:]
		sampn = dict([(x[1], int(x[0])) for x in enumerate(sampnames)])

		popsampnums = {}
		if opt.pops: # fill poplist with pops actually present
			poplist = []
			for samp in sampnames:
				if not samp in samppop: # look for matching prefix
					warning('%s not found in %s' % (samp, opt.pops))
					prefixes = [os.path.commonprefix([s, samp]) for s in samppop.keys()]
					longpref = max(prefixes, key=len)
					if len(longpref) > 6: #ARBITRARY HACK TO FIND BEST MATCHING SAMPLE
						match = samppop.keys()[prefixes.index(longpref)]
						samppop[samp] = samppop[match]
						warning('Using %s\t%s' % (match, samppop[match]))

				if not samppop[samp] in poplist: # error if samp not in opt.pops
					poplist.append(samppop[samp])
					popsampnums[samppop[samp]] = []
				popsampnums[samppop[samp]].append(sampn[samp])

		else:
			poplist = ['all_samples']
			popsampnums['all_samples'] = range(nsamp)

		if opt.condsamp:
			for i, name in enumerate(sampnames):
				if name.find(condsamp) >= 0:
					condcol = i + VCF_FIXEDCOLS
#				   print(i, condsamp, tok[condcol])
					break
		break

if opt.risk:
	totalrisk = [0 for x in range(nsamp)]

if opt.rates:
	if opt.depths:
		depcount = [0 for i in range(nsamp)]
	hetcount = [0 for i in range(nsamp)]
	hnrcount = [0 for i in range(nsamp)]
	nsites = 0
	firstcoord = ''

if opt.psmcfa:
	faout = FastaStream(sys.stdout)

#if opt.header:
#	if opt.pops:
#		fout.write('#VCF_POPULATIONS:\t%s\n' % ('\t'.join(poplist)))
#	if not opt.no_calls:
#		fout.write('\t'.join(['#SAMPLES'] + sampnames) + '\n')
if opt.header: #column headers
	header = []
	if opt.rates:
		header += ['#SAMPLES']
	else:
		header += ['CHR', 'POS', 'ID']
	if opt.pops:# and opt.rename_samps: # add pop name to sample names in header
		headsampnames = [samppop[x] + '_' + x for x in sampnames]
	else:
		headsampnames =  sampnames
	if opt.depthsonly:
		header += headsampnames
	if opt.aims:
		header += ['***HEADER_INFO_TODO***']
	if not opt.no_calls:
		header += headsampnames
	if opt.segsites:
		header += ['SEGSITE']
	if opt.diversity: # needs --no_calls to line up with output
		for pop in poplist:
			header += [pop + '_N', pop + '_AC', pop + '_AF', pop + '_HET']
	if opt.Fst:
		header += ['Fst']
	fout.write('\t'.join(header) + '\n')


curchr = ''
for line in fin:
	if line.startswith('#'):
#		sys.stdout.write(line)
		continue

	tok = line.split()

	if tok[0] != curchr:
		sep = 0
		lastpos = 0
		curchr = tok[0]
		if opt.psmcfa:
			nextbin = PSMC_BINSIZE
			faout = FastaStream(sys.stdout)
			faout.newrec(curchr)
			inmask = False
			nN = 0
			outchar = 'T'

	if opt.replacecalls and tok[0]+tok[1] in reprecord:
		tok[VCF_FIXEDCOLS-1:] = reprecord[tok[0]+tok[1]]
	
	if opt.vars and tok[4] == '.':
		continue

#	sys.stdout.write(line)
	if not opt.indels and (len(tok[3]) > 1 or re.search('INDEL', tok[7])):
		continue

#	if len(tok[4]) > 1: # multiallelic sites
#		continue

	if opt.condsamp and not tok[condcol].startswith('0/1'): # condition on variant in cond_sample
		continue

	if opt.depths or opt.mediandep:
		formats = tok[VCF_FIXEDCOLS - 1].split(':')
		depfield = formats.index('DP')
		dep = [x.split(':')[depfield] for x in tok[VCF_FIXEDCOLS:(VCF_FIXEDCOLS+nsamp)]]

	if opt.depthsonly:
		fout.write('\t'.join(tok[0:2] + dep) + '\n')
		continue

	outvals = tok[0:3]

	if opt.haploid:
		if (tok[4] == "."):
			varvals = ['0' for x in tok[VCF_FIXEDCOLS:(VCF_FIXEDCOLS+nsamp)]]
			segsite = False
		else:
			varvals = [x[0] for x in tok[VCF_FIXEDCOLS:(VCF_FIXEDCOLS+nsamp)]]
		vargts = varvals
	else:
		if (tok[4] == "."):
			varvals = ['00' for x in tok[VCF_FIXEDCOLS:(VCF_FIXEDCOLS+nsamp)]]
			vargts = ['0/0' for x in tok[VCF_FIXEDCOLS:(VCF_FIXEDCOLS+nsamp)]]
			hets = [0 for x in tok[VCF_FIXEDCOLS:(VCF_FIXEDCOLS+nsamp)]]
			segsite = False
		else:
			varvals = [x[0] + x[2] for x in tok[VCF_FIXEDCOLS:(VCF_FIXEDCOLS+nsamp)]]
			vargts = [x[0:3] for x in tok[VCF_FIXEDCOLS:(VCF_FIXEDCOLS+nsamp)]]
			hets = [int(x[0] != x[2]) for x in tok[VCF_FIXEDCOLS:(VCF_FIXEDCOLS+nsamp)]]

		if opt.pseudodip:
			hetsite = False
			for gt in vargts:
				if gt[0] != gt[2]:
					hetsite = True
					break
			if hetsite:
				continue

	if opt.risk: # convert genotype to GRS
		sitealleles = [tok[3]] + tok[4].split(',')
		effectallele = str(sitealleles.index(SNPallele[tok[2]]))
		varvals = [SNPscore[tok[2]] * (int(x[0] == effectallele) + int(x[1] == effectallele)) for x in varvals]
		totalrisk = [x + y for (x, y) in zip(totalrisk, varvals)]
		varvals = [str(x) for x in varvals]

	if opt.pseudodip:
# combine consecutive varvals elements into new varvals (assumes none are het)
# TODO?: only check/exclude hets in consecutive pairs
		varvals = [varvals[i][0] + varvals[i+1][0] for i in range(0, len(varvals) - 1, 2)]
		vargts = [vargts[i][0] + '|' + vargts[i+1][0] for i in range(0, len(vargts) - 1, 2)]
#		debug(vargts)

# determine if segregating site
# ideally use FQ > 0 but vcf-subset doesn't currently adjust this
#		fq = re.search('FQ=(-?[0-9.]+)', tok[7])
#		if fq > 0: #het
#		sys.stdout.write(line)
	varvalstr = '-'.join(varvals)
#	debug(valstr)
	segsite = '0' in varvalstr and '1' in varvalstr

	if opt.alleles: # convert 0, 1 to alleles in output
		sitealleles = [tok[3]] + tok[4].split(',')
		trtab = maketrans(''.join([str(x) for x in range(len(sitealleles))]), ''.join(sitealleles))
		varalleles = [x.translate(trtab) for x in varvals]  
		varvals = varalleles  
		varvalstr = '-'.join(varvals)
		vargts = [x.translate(trtab) for x in vargts]  
#		debug([tok[1], sitealleles, vargts])


	if opt.aims:
		valpops = {}
		popvals = {}
		for isn, samp in enumerate(sampnames):
			spop = samppop[samp]
			for val in varvals[isn]:
				if val in valpops:
					valpops[val].add(spop)
				else:
					valpops[val] = set([spop])
				if spop in popvals:
					popvals[spop].add(val)
				else:
					popvals[spop] = set([val])
		for val in valpops:
			if len(valpops[val]) == 1: # require uniqueness of marker to pop
				spop = valpops[val].pop() 
				if len(popvals[spop]) == 1: # require fixation of marker in pop
					debug(varvals)
#					debug(valpops)
					debug(popvals)
					spvals = ['/'.join(list(popvals[x])) for x in poplist]
					debug(spvals)
					fout.write('\t'.join(outvals + spvals + [spop]) + '\n')
#					fout.write('\t'.join(outvals + [val] + spvals + [spop]) + '\n')
		continue

	if opt.segsep:
		pos = int(tok[1])

		while pos > lastpos:
			if opt.callmask and lastpos in callmask:
#				print('pos = %d; Masking %d to %d' % (pos, lastpos, callmask[lastpos]))
				lastpos = callmask[lastpos]
			else:
				sep += 1
			lastpos += 1
		if opt.callmask and lastpos in callmask:
#			print('pos = %d; Masking %d to %d' % (pos, lastpos, callmask[lastpos]))
			lastpos = callmask[lastpos] + 1
		elif lastpos == pos and segsite:
			if len(vargts) > 1:
				allstr = ','.join(phasecombs(vargts))
				fout.write('\t'.join(outvals + [str(sep)] + [allstr]) + '\n')
			else:
				fout.write('\t'.join(outvals + [str(sep)] + varvals) + '\n')
			sep = 0
		continue

	elif opt.psmcfa:
		if segsite:
			pos = int(tok[1])

#			print('site at %d' % pos)
			while lastpos < nextbin and pos > nextbin - PSMC_BINSIZE:
				lastpos += 1
				if opt.callmask and lastpos in callmask:
#					print('Mask %d to %d : %d' % (lastpos, callmask[lastpos], callmask[lastpos]-lastpos + 1))
					inmask = True
					endmask = callmask[lastpos]
				if inmask:
					nN += 1
					if lastpos == endmask:
						inmask = False
				if lastpos == pos:
					outchar = 'K'
				if lastpos == nextbin:
					if float(nN) / PSMC_BINSIZE > PSMC_N_RATIO:
						outchar = 'N'
					faout.write(outchar)
#					print(' %d-%d: %d Ns, pos = %d' % (nextbin - PSMC_BINSIZE + 1, nextbin, nN, pos))
					nextbin += PSMC_BINSIZE
					outchar = 'T'
					nN = 0
		continue

	if opt.rates:
		if opt.depths:
			depcount = [depcount[i] + int(dep[i]) for i in range(nsamp)]
		hetcount = [hetcount[i] + (varvals[i][0] != varvals[i][1]) for i in range(nsamp)]
		hnrcount = [hnrcount[i] + (varvals[i] == '11') for i in range(nsamp)]
		nsites += 1
		if not firstcoord:
			firstcoord = ':'.join(tok[0:2])
		lastcoord = ':'.join(tok[0:2])

	if not opt.nosites:
		if opt.mediandep:
			outvals.append('%d' % median(dep))

		if opt.segsites:
#			outvals.append(str(depfield - 1))
			outvals.append(str(int(segsite)))

		alcount_total = 0
		samps_total = 0
		Hs_mean = 0
		for pop in poplist:
			pvarvalstr = sampsep.join([varvals[i] for i in popsampnums[pop]])
			phets = [hets[i] for i in popsampnums[pop]]
			if not opt.no_calls:
				outvals.append(pvarvalstr)

			if opt.diversity or opt.Fst:
	#			alfreq = eval(re.search('AF1=([0-9.e\-]+)', tok[7]).group(1))# eval since vcf uses sci notation
	#			alcount = int(re.search('AC1=([0-9]+)', tok[7]).group(1))
	#			neidiv = 1 - alfreq**2 - (1 - alfreq)**2
	#			outvals += [str(nsamp), str(alfreq), str(neidiv), str(alcount), str(alcount/float(nsamp))]
				alcount = pvarvalstr.count('1')
				hetcount = sum(phets)
				pnsamp = len(popsampnums[pop])
				pfreq = alcount/float(2 * pnsamp)
				alcount_total += alcount
				samps_total += pnsamp
				Hs_mean += 2 * pfreq * (1 - pfreq)
				if opt.diversity:
					outvals += [str(pnsamp), str(alcount), str(pfreq), str(hetcount)]

		if opt.Fst:
			Hs_mean = Hs_mean / len(poplist)
			tfreq = alcount_total/float(2 * samps_total)
			Ht = 2 * tfreq * (1 - tfreq)
			Fst = (Ht - Hs_mean) / Ht
#			outvals += [str(Hs_mean), str(Ht), str(Fst)]
			outvals += [str(Fst)]

		fout.write('\t'.join(outvals) + '\n')

if opt.rates:
	fout.write('\t'.join(['#REGION'] + [firstcoord + '-' + lastcoord]) + '\n')
	fout.write('\t'.join(['#NSITES'] + [str(nsites)]) + '\n')
	if opt.depths:
		fout.write('\t'.join(['#MEANDEP'] + [str(float(x)/nsites) for x in depcount]) + '\n')
	fout.write('\t'.join(['#HETCOUNT'] + [str(x) for x in hetcount]) + '\n')
	fout.write('\t'.join(['#HETRATE'] + [str(float(x)/nsites) for x in hetcount]) + '\n')
	fout.write('\t'.join(['#HNRCOUNT'] + [str(x) for x in hnrcount]) + '\n')
	fout.write('\t'.join(['#HNRRATE'] + [str(float(x)/nsites) for x in hnrcount]) + '\n')
	fout.write('\t'.join(['#HNR/HET'] + [str(float(hnrcount[i])/hetcount[i]) for i in range(nsamp)]) + '\n')
#	fout.write('\t'.join(outvals) + '\n')

if opt.risk:
	fout.write('\t'.join(['#TOTALRISK', '.', '.'] + [str(x) for x in totalrisk]) + '\n')

try:
	sys.stdout.close()
except:
	pass
try:
	sys.stderr.close()
except:
	pass

