# Utils for dealing with VCF generated by nextstrain.py, with a particular style of ID
# and clades appearing in genotypes...

from collections import defaultdict

def readVcfSampleClades(vcfFile):
    """Read VCF sample IDs from the #CHROM line, and parse out clades from the first row GT cols"""
    samples = []
    sampleClades = defaultdict()
    with open(vcfFile, 'r') as vcfF:
        line = vcfF.readline().strip()
        while (line):
            if (line.startswith('#CHROM')):
                samples = line.split('\t')[9:]
            elif (not line.startswith('#')):
                gts = line.split('\t')[9:]
                if (len(gts) != len(samples)):
                    die("VCF file '%s' has %d samples but %d genotypes in first row" %
                        (vcfFile, len(samples), len(gts)));
                for sample, gt in zip(samples, gts):
                    gtVal, clade = gt.split(':')
                    sampleClades[sample] = clade
                break
            line = vcfF.readline().strip()
        vcfF.close()
    return samples, sampleClades