#!/usr/bin/env python


#############################################
#  Sept 29, 2011
#  Authors: Vlad Makarov, Chris Yoon
#  Language: Python
#  OS: UNIX/Linux, MAC OSX

#  Purpose:  Extracts GT field (genotype) from the VCF file format format
#  Input:  Text file in VCF format (http://www.1000genomes.org/wiki/Analysis/Variant%20Call%20Format/vcf-variant-call-format-version-41)
#  Output: Same file as above with 0, 1 or 2 value for each variant/each sample (no mutation, heterozygous, monozygous)


#  Arguments:
#  filename - name of VCF file

#  Note:
#  Current utility does not perform validity of VCF file format per ce.
#  Use VCFTools (http://vcftools.sourceforge.net/) for format verification if in doubt



#  Example to run:
#  python genotype.py myvars.vcf > mygenotype.txt
#

##############################################



import os.path
import csv
import os
import sys


def parse_field(text, key, sep1, sep2):
    fields = text.strip().split(sep1)
    for f in fields:
        pairs = f.split(sep2)
        if str(pairs[0]).find(str(key)) > -1:
            return str(pairs[1])
    return '.'


def togenotype(filename, header_line_start='#CHROM', comment_chars='##'):
    fh = open(filename, "r")
    colnum = 0
    i=0
    for line in fh:
        line = line.strip()
        if line.startswith(comment_chars):
            continue
        elif line.startswith(header_line_start):
            line = line.strip()
            fields = line.split('\t')
            samples = fields[9:len(fields)]
            newline = '\t'.join(samples)
            print newline.strip()
        else:

            line = line.strip()
            fields = line.split('\t')
            samples = fields[9:len(fields)]
            genotype = '0'
            genotypes = []
            for s in samples:
                if s.startswith('0/0') or s.startswith('./.') :
                    genotype = '0'
                elif s.startswith('0/1') or s.startswith('1/0'):
                    genotype = '1'
                elif s.startswith('1/1'):
                    genotype = '2'
                else:
                    genotype = '0'
                genotypes.append(genotype)

            newline =  '\t'.join(genotypes)
            print(newline.strip())


filename=sys.argv[1]

togenotype(filename=filename, header_line_start='#CHROM', comment_chars='##')
