def GetAlleles(filename):
    """ Read the alleles in filename """ 
    allele = dict()
    with open(filename, 'rb') as tsvfile:
        for row in tsvfile:
            line = row[:row.find('#')].strip(' \r\n')   # strip comments and line endings
            if (len(line) == 0):                        # skip empty lines
                continue
            field = line.split('\t')
            if (len(field) != 4):                       # print error for lines missing fields
                print "ERROR: %s" % str(field)
                continue
            allele[field[0]] = field[3]
    return allele

def GetCoordinates(filename):
    """ Read the chromosome and positions of alleles in filename """ 
    chrPos = dict()
    with open(filename, 'rb') as tsvfile:
        for row in tsvfile:
            line = row[:row.find('#')].strip(' \r\n')   # strip comments and line endings
            if (len(line) == 0):                        # skip empty lines
                continue
            field = line.split('\t')
            if (len(field) != 4):                       # print error for lines missing fields
                print "ERROR: %s" % str(field)
                continue
            chrPos[field[0]] = (field[1], int(field[2]))
    return chrPos

phaseRule = {
    ('--', '--', 'AA'): 'AA', 
    ('--', '--', 'AC'): 'ac', 
    ('--', '--', 'AG'): 'ag', 
    ('--', '--', 'AT'): 'at', 
    ('--', '--', 'CC'): 'CC', 
    ('--', '--', 'CG'): 'cg', 
    ('--', '--', 'CT'): 'ct', 
    ('--', '--', 'DD'): 'DD',
    ('--', '--', 'GG'): 'GG',
    ('--', '--', 'GT'): 'gt', 
    ('--', '--', 'II'): 'II',
    ('--', '--', 'TT'): 'TT', 
    ('--', 'AA', '--'): '-A',
    ('--', 'AA', 'AA'): 'AA', 
    ('--', 'AA', 'AC'): 'CA', 
    ('--', 'AA', 'AG'): 'GA',
    ('--', 'AA', 'AT'): 'TA', 
    ('--', 'AC', 'AA'): 'AA', 
    ('--', 'AC', 'AC'): 'ac',
    ('--', 'AC', 'CC'): 'CC', 
    ('--', 'AG', 'AA'): 'AA', 
    ('--', 'AG', 'AG'): 'ag',
    ('--', 'AG', 'GG'): 'GG',
    ('--', 'AT', 'AA'): 'AA',
    ('--', 'AT', 'AT'): 'at',
    ('--', 'AT', 'TT'): 'TT', 
    ('--', 'CC', '--'): '-C',
    ('--', 'CC', 'AC'): 'AC', 
    ('--', 'CC', 'CC'): 'CC', 
    ('--', 'CC', 'CG'): 'GC', 
    ('--', 'CC', 'CT'): 'TC',
    ('--', 'CG', 'CC'): 'CC',
    ('--', 'CG', 'CG'): 'cg',
    ('--', 'CG', 'GG'): 'GG',
    ('--', 'CT', 'CC'): 'CC', 
    ('--', 'CT', 'CT'): 'ct',
    ('--', 'CT', 'TT'): 'TT',
    ('--', 'DD', 'DD'): 'DD',
    ('--', 'GG', '--'): '-G', 
    ('--', 'GG', 'AG'): 'AG',
    ('--', 'GG', 'CG'): 'CG',
    ('--', 'GG', 'GG'): 'GG', 
    ('--', 'GG', 'GT'): 'TG', 
    ('--', 'GT', 'GG'): 'GG',
    ('--', 'GT', 'GT'): 'gt', 
    ('--', 'GT', 'TT'): 'TT', 
    ('--', 'II', '--'): '-I',
    ('--', 'II', 'II'): 'II', 
    ('--', 'TT', '--'): '-T', 
    ('--', 'TT', 'AT'): 'AT',
    ('--', 'TT', 'CT'): 'CT',
    ('--', 'TT', 'GT'): 'GT', 
    ('--', 'TT', 'TT'): 'TT', 
    ('AA', '--', '--'): 'A-',
    ('AA', '--', 'AA'): 'AA', 
    ('AA', '--', 'AC'): 'AC', 
    ('AA', '--', 'AG'): 'AG',
    ('AA', '--', 'AT'): 'AT', 
    ('AA', 'AA', '--'): 'AA', 
    ('AA', 'AA', 'AA'): 'AA',
    ('AA', 'AC', '--'): 'A-',
    ('AA', 'AC', 'AA'): 'AA',
    ('AA', 'AC', 'AC'): 'AC',
    ('AA', 'AG', '--'): 'A-',
    ('AA', 'AG', 'AA'): 'AA',
    ('AA', 'AG', 'AG'): 'AG',
    ('AA', 'AT', '--'): 'A-',
    ('AA', 'AT', 'AA'): 'AA',
    ('AA', 'AT', 'AT'): 'AT',
    ('AA', 'CC', '--'): 'AC',
    ('AA', 'CC', 'AC'): 'AC',
    ('AA', 'GG', '--'): 'AG',
    ('AA', 'GG', 'AG'): 'AG',
    ('AA', 'TT', '--'): 'AT',
    ('AA', 'TT', 'AT'): 'AT',
    ('AC', '--', 'AA'): 'AA',
    ('AC', '--', 'AC'): 'ac',
    ('AC', '--', 'CC'): 'CC',
    ('AC', 'AA', '--'): '-A',
    ('AC', 'AA', 'AA'): 'AA',
    ('AC', 'AA', 'AC'): 'CA',
    ('AC', 'AC', 'AA'): 'AA',
    ('AC', 'AC', 'AC'): 'ac',
    ('AC', 'AC', 'CC'): 'CC',
    ('AC', 'CC', '--'): '-C',
    ('AC', 'CC', 'AC'): 'AC',
    ('AC', 'CC', 'CC'): 'CC',
    ('AG', '--', 'AA'): 'AA',
    ('AG', '--', 'AG'): 'ag',
    ('AG', '--', 'GG'): 'GG',
    ('AG', 'AA', '--'): '-A',
    ('AG', 'AA', 'AA'): 'AA',
    ('AG', 'AA', 'AG'): 'GA',
    ('AG', 'AG', 'AA'): 'AA',
    ('AG', 'AG', 'AG'): 'ag',
    ('AG', 'AG', 'GG'): 'GG',
    ('AG', 'GG', '--'): '-G',
    ('AG', 'GG', 'AG'): 'AG',
    ('AG', 'GG', 'GG'): 'GG',
    ('AT', '--', 'AA'): 'AA',
    ('AT', '--', 'AT'): 'at',
    ('AT', '--', 'TT'): 'TT',
    ('AT', 'AA', '--'): '-A',
    ('AT', 'AA', 'AA'): 'AA',
    ('AT', 'AA', 'AT'): 'TA',
    ('AT', 'AT', 'AA'): 'AA',
    ('AT', 'AT', 'AT'): 'at',
    ('AT', 'AT', 'TT'): 'TT',
    ('AT', 'TT', '--'): '-T',
    ('AT', 'TT', 'AT'): 'AT',
    ('AT', 'TT', 'TT'): 'TT',
    ('CC', '--', '--'): 'C-',
    ('CC', '--', 'AC'): 'CA',
    ('CC', '--', 'CC'): 'CC',
    ('CC', '--', 'CG'): 'CG',
    ('CC', '--', 'CT'): 'CT',
    ('CC', 'AA', '--'): 'CA',
    ('CC', 'AA', 'AC'): 'CA',
    ('CC', 'AC', '--'): 'C-',
    ('CC', 'AC', 'AC'): 'CA',
    ('CC', 'AC', 'CC'): 'CC',
    ('CC', 'CC', '--'): 'CC',
    ('CC', 'CC', 'CC'): 'CC',
    ('CC', 'CG', '--'): 'C-',
    ('CC', 'CG', 'CC'): 'CC',
    ('CC', 'CG', 'CG'): 'CG',
    ('CC', 'CT', '--'): 'C-',
    ('CC', 'CT', 'CC'): 'CC',
    ('CC', 'CT', 'CT'): 'CT',
    ('CC', 'GG', '--'): 'CG',
    ('CC', 'GG', 'CG'): 'CG',
    ('CC', 'TT', '--'): 'CT',
    ('CC', 'TT', 'CT'): 'CT',
    ('CG', '--', 'CC'): 'CC',
    ('CG', '--', 'CG'): 'cg',
    ('CG', '--', 'GG'): 'GG',
    ('CG', 'CC', 'CC'): 'CC',
    ('CG', 'CC', 'CG'): 'GC',
    ('CG', 'CG', 'CC'): 'CC',
    ('CG', 'CG', 'CG'): 'cg',
    ('CG', 'CG', 'GG'): 'GG',
    ('CG', 'GG', '--'): '-G',
    ('CG', 'GG', 'CG'): 'CG',
    ('CG', 'GG', 'GG'): 'GG',
    ('CT', '--', 'CC'): 'CC',
    ('CT', '--', 'CT'): 'ct',
    ('CT', '--', 'TT'): 'TT',
    ('CT', 'CC', '--'): '-C',
    ('CT', 'CC', 'CC'): 'CC',
    ('CT', 'CC', 'CT'): 'TC',
    ('CT', 'CT', 'CC'): 'CC',
    ('CT', 'CT', 'CT'): 'ct',
    ('CT', 'CT', 'TT'): 'TT',
    ('CT', 'TT', '--'): '-T',
    ('CT', 'TT', 'CT'): 'CT',
    ('CT', 'TT', 'TT'): 'TT',
    ('DD', '--', 'DD'): 'DD',
    ('DD', 'DD', '--'): 'DD',
    ('DD', 'DD', 'DD'): 'DD',
    ('DD', 'DI', 'DI'): 'DI',
    ('DD', 'DI', 'DD'): 'DD',
    ('DD', 'II', 'DI'): 'DI',
    ('DI', 'DD', 'DD'): 'DD',
    ('DI', 'DD', 'DI'): 'DI',
    ('DI', 'DI', 'DD'): 'DD',
    ('DI', 'DI', 'DI'): 'di',
    ('DI', 'DI', 'II'): 'II',
    ('DI', 'II', 'II'): 'II',
    ('DI', 'II', 'DI'): 'DI',
    ('GG', '--', '--'): 'G-',
    ('GG', '--', 'AG'): 'GA',
    ('GG', '--', 'CG'): 'GC',
    ('GG', '--', 'GG'): 'GG',
    ('GG', '--', 'GT'): 'GT',
    ('GG', 'AA', '--'): 'GA',
    ('GG', 'AA', 'AG'): 'GA',
    ('GG', 'AG', '--'): 'G-',
    ('GG', 'AG', 'AG'): 'GA',
    ('GG', 'AG', 'GG'): 'GG',
    ('GG', 'CC', 'CG'): 'GC',
    ('GG', 'CG', 'CG'): 'GC',
    ('GG', 'CG', 'GG'): 'GG',
    ('GG', 'GG', '--'): 'GG',
    ('GG', 'GG', 'GG'): 'GG',
    ('GG', 'GT', '--'): 'G-',
    ('GG', 'GT', 'GG'): 'GG',
    ('GG', 'GT', 'GT'): 'GT',
    ('GG', 'TT', '--'): 'GT',
    ('GG', 'TT', 'GT'): 'GT',
    ('GT', '--', 'GG'): 'GG',
    ('GT', '--', 'GT'): 'gt',
    ('GT', '--', 'TT'): 'TT',
    ('GT', 'GG', '--'): '-G',
    ('GT', 'GG', 'GG'): 'GG',
    ('GT', 'GG', 'GT'): 'TG',
    ('GT', 'GT', 'GG'): 'GG',
    ('GT', 'GT', 'GT'): 'gt',
    ('GT', 'GT', 'TT'): 'TT',
    ('GT', 'TT', '--'): '-T',
    ('GT', 'TT', 'GT'): 'GT',
    ('GT', 'TT', 'TT'): 'TT',
    ('II', '--', 'II'): 'II',
    ('II', 'DD', 'DI'): 'ID',
    ('II', 'DI', 'DI'): 'ID',
    ('II', 'DI', 'II'): 'II',
    ('II', 'II', '--'): 'II',
    ('II', 'II', 'II'): 'II',
    ('TT', '--', '--'): 'T-',
    ('TT', '--', 'AT'): 'TA',
    ('TT', '--', 'CT'): 'TC',
    ('TT', '--', 'GT'): 'TG',
    ('TT', '--', 'TT'): 'TT',
    ('TT', 'AA', '--'): 'TA',
    ('TT', 'AA', 'AT'): 'TA',
    ('TT', 'AT', '--'): 'T-',
    ('TT', 'AT', 'AT'): 'TA',
    ('TT', 'AT', 'TT'): 'TT',
    ('TT', 'CC', '--'): 'TC',
    ('TT', 'CC', 'CT'): 'TC',
    ('TT', 'CT', '--'): 'T-',
    ('TT', 'CT', 'CT'): 'TC',
    ('TT', 'CT', 'TT'): 'TT',
    ('TT', 'GG', '--'): 'TG',
    ('TT', 'GG', 'GT'): 'TG',
    ('TT', 'GT', '--'): 'T-',
    ('TT', 'GT', 'GT'): 'TG',
    ('TT', 'GT', 'TT'): 'TT',
    ('TT', 'TT', '--'): 'TT',
    ('TT', 'TT', 'TT'): 'TT',
    ('--', 'A', '--'): '-A',
    ('--', 'A', 'AA'): 'AA',
    ('--', 'A', 'AC'): 'CA',
    ('--', 'A', 'AG'): 'GA',
    ('--', 'A', 'AT'): 'TA',
    ('--', 'C', '--'): '-C',
    ('--', 'C', 'AC'): 'AC',
    ('--', 'C', 'CC'): 'CC',
    ('--', 'C', 'CG'): 'GC',
    ('--', 'C', 'CT'): 'TC',
    ('--', 'D', 'DD'): 'DD',
    ('--', 'G', '--'): '-G',
    ('--', 'G', 'AG'): 'AG',
    ('--', 'G', 'CG'): 'CG',
    ('--', 'G', 'GG'): 'GG',
    ('--', 'G', 'GT'): 'TG',
    ('--', 'T', '--'): '-T',
    ('--', 'T', 'AT'): 'AT',
    ('--', 'T', 'CT'): 'CT',
    ('--', 'T', 'GT'): 'GT',
    ('--', 'T', 'TT'): 'TT',
    ('AA', 'A', 'AA'): 'AA',
    ('AA', 'C', 'AC'): 'AC',
    ('AA', 'G', 'AG'): 'AG',
    ('AA', 'T', 'AT'): 'AT',
    ('AC', 'A', 'AA'): 'AA',
    ('AC', 'A', 'AC'): 'CA',
    ('AC', 'C', 'AC'): 'AC',
    ('AC', 'C', 'CC'): 'CC',
    ('AG', 'A', 'AA'): 'AA',
    ('AG', 'A', 'AG'): 'GA',
    ('AG', 'G', 'AG'): 'AG',
    ('AG', 'G', 'GG'): 'GG',
    ('AT', 'A', 'AA'): 'AA',
    ('AT', 'A', 'AT'): 'TA',
    ('AT', 'T', 'AT'): 'AT',
    ('AT', 'T', 'TT'): 'TT',
    ('CC', 'A', 'AC'): 'CA',
    ('CC', 'C', 'CC'): 'CC',
    ('CC', 'G', 'CG'): 'CG',
    ('CC', 'T', 'CT'): 'CT',
    ('CG', 'C', 'CC'): 'CC',
    ('CG', 'C', 'CG'): 'GC',
    ('CG', 'G', 'CG'): 'CG',
    ('CG', 'G', 'GG'): 'GG',
    ('CT', 'C', 'CC'): 'CC',
    ('CT', 'C', 'CT'): 'TC',
    ('CT', 'T', 'CT'): 'CT',
    ('CT', 'T', 'TT'): 'TT',
    ('DD', 'D', 'DD'): 'DD',
    ('GG', 'A', 'AG'): 'GA',
    ('GG', 'C', 'CG'): 'GC',
    ('GG', 'G', 'GG'): 'GG',
    ('GG', 'T', 'GT'): 'GT',
    ('GT', 'G', 'GG'): 'GG',
    ('GT', 'G', 'GT'): 'TG',
    ('GT', 'T', 'GT'): 'GT',
    ('GT', 'T', 'TT'): 'TT',
    ('II', 'I', 'II'): 'II',
    ('TT', 'A', 'AT'): 'TA',
    ('TT', 'C', 'CT'): 'TC',
    ('TT', 'G', 'GT'): 'TG',
    ('TT', 'T', 'TT'): 'TT',
    ('--', 'MT', 'A'): 'A',
    ('--', 'MT', 'C'): 'C',
    ('--', 'MT', 'D'): 'D',
    ('--', 'MT', 'G'): 'G',
    ('--', 'MT', 'I'): 'I',
    ('--', 'MT', 'T'): 'T',
    ('A', 'MT', '--'): 'A',
    ('C', 'MT', '--'): 'C',
    ('D', 'MT', '--'): 'D',
    ('G', 'MT', '--'): 'G',
    ('I', 'MT', '--'): 'I',
    ('T', 'MT', '--'): 'T',
    ('A', 'MT', 'A'): 'A',
    ('C', 'MT', 'C'): 'C',
    ('D', 'MT', 'D'): 'D',
    ('G', 'MT', 'G'): 'G',
    ('I', 'MT', 'I'): 'I',
    ('T', 'MT', 'T'): 'T',
    ('--', 'X', 'A'): 'A',
    ('--', 'X', 'C'): 'C',
    ('--', 'X', 'D'): 'D',
    ('--', 'X', 'G'): 'G',
    ('--', 'X', 'I'): 'I',
    ('--', 'X', 'T'): 'T',
    ('AA', 'X', 'A'): 'A',
    ('AC', 'X', 'A'): 'A',
    ('AC', 'X', 'C'): 'C',
    ('AG', 'X', 'A'): 'A',
    ('AG', 'X', 'G'): 'G',
    ('AT', 'X', 'A'): 'A',
    ('AT', 'X', 'T'): 'T',
    ('CC', 'X', 'C'): 'C',
    ('CG', 'X', 'C'): 'C',
    ('CG', 'X', 'G'): 'G',
    ('CT', 'X', 'C'): 'C',
    ('CT', 'X', 'T'): 'T',
    ('DD', 'X', 'D'): 'D',
    ('GG', 'X', 'G'): 'G',
    ('GT', 'X', 'G'): 'G',
    ('GT', 'X', 'T'): 'T',
    ('II', 'X', 'I'): 'I',
    ('TT', 'X', 'T'): 'T',
    ('Y', '--', 'A'): 'A',
    ('Y', '--', 'C'): 'C',
    ('Y', '--', 'D'): 'D',
    ('Y', '--', 'G'): 'G',
    ('Y', '--', 'I'): 'I',
    ('Y', '--', 'T'): 'T',
    ('Y', 'A', '--'): 'A',
    ('Y', 'C', '--'): 'C',
    ('Y', 'D', '--'): 'D',
    ('Y', 'G', '--'): 'G',
    ('Y', 'I', '--'): 'I',
    ('Y', 'T', '--'): 'T',
    ('Y', 'A', 'A'): 'A',
    ('Y', 'C', 'C'): 'C',
    ('Y', 'D', 'D'): 'D',
    ('Y', 'G', 'G'): 'G',
    ('Y', 'I', 'I'): 'I',
    ('Y', 'T', 'T'): 'T'
}

invalidTrio = {}

def phase(mom, pop, kid, ChrPos):
    """ Use trio alleles to phase kid. Returns a new kid allele dictionary """
    modified = 0
    unchanged = 0
    invalid = 0
    for gId, call in kid.iteritems():
        chromo = ChrPos[gId][0]
        try:
            momA = mom[gId] if (chromo != 'Y') else 'Y'
            popA = pop[gId] if (chromo != 'MT') else 'MT'
        except KeyError:
            # print gId
            continue
        if ((chromo == 'X') and (len(call) == 1)):
            popA = 'X'
        try: 
            newCall = phaseRule[momA, popA, call]
        except KeyError:
            if (call != '--'):
                invalidTrio[momA, popA, call] = invalidTrio.get((momA, popA, call), 0) + 1
                invalid += 1
            newCall = '--'
        if (newCall != call):
            kid[gId] = newCall
            modified += 1
        else:
            unchanged += 1
    print modified, unchanged, invalid

allelesM = GetAlleles("genomeF.txt")
allelesP = GetAlleles("genomeG.txt")
allelesC = GetAlleles("genomeI.txt")
coordinates = GetCoordinates("genomeI.txt")

phase(allelesM, allelesP, allelesC, coordinates)

print len(phaseRule), len(invalidTrio)
for key in sorted(invalidTrio.iterkeys()):
    print "%s: %d" % (key, invalidTrio[key])

%matplotlib inline
import matplotlib.pyplot as plot
import numpy

for sample in "ABCDEFGHI":
    alleles[sample] = GetAlleles("genome%s.txt" % sample)
coordinates = GetCoordinates("genomeI.txt")

cmap = { 'A':'r', 'C':'y', 'G':'b', 'T':'c', 'D':'g', 'I':'m' }
plot.figure(figsize=(10, 5))

for gid, chrpos in coordinates.iteritems():
    if (chrpos[0] != 'MT'):
        continue
    x = []
    y = []
    color = []
    for i, sample in enumerate("ABCDEFGHI"):
        try:
            call = alleles[sample][gid]
        except KeyError:
            continue
        if (call == '--'):
            continue
        x.append(chrpos[1])
        y.append(9-i)
        color.append(cmap[call])
    if (len(x) < 9):
        continue
    for c in color:
        if (c != color[0]):
            break
    else:
        continue
    plot.scatter(x, y, c=color, s=100, marker='|')

for i, sample in enumerate("ABCDEFGHI"):
    plot.text(-500, 8.8-i, sample, color='k') 
plot.show()

def kNearest(k, geneId, coordinates):
    ''' returns tuples of (dist, id) for the specified geneId's k-nearest neighbors '''
    neighbor = [(2**31, 'foo') for i in xrange(k)]
    chromo, position = coordinates[geneId]
    for gid, chrpos in coordinates.iteritems():
        c, p = chrpos
        if (c != chromo):
            continue
        if (gid == geneId):
            continue
        delta = abs(p - position)
        if (delta < neighbor[-1][0]):
            neighbor.pop()
            neighbor.append((delta, gid))
            neighbor.sort()
    return neighbor

coordinates = GetCoordinates("genomeI.txt")
for gId, call in allelesC.iteritems():
    if (call.islower()):
        break
        
print gId, call
knn = kNearest(10, gId, coordinates)
knn.append((0,gId))
print knn

plot.figure(figsize=(10, 5))
chromo, position = coordinates[gid]
for i, sample in enumerate("ABCDEFGHI"):
    x = []
    y = []
    color = []
    for delta, gId in knn:
        call = alleles[sample][gId]
        diff = coordinates[gId][1]-position
        offset = 0.4 if ((sample == 'I') and (delta > 0)) else 0.2
        x.append(diff)
        y.append(9-i+offset)
        color.append(cmap[call[0]])
        x.append(diff)
        y.append(9-i-offset)
        color.append(cmap[call[0]])
    plot.scatter(x, y, c=color, s=20, marker='o')

for i, sample in enumerate("ABCDEFGHI"):
    plot.text(-500, 8.8-i, sample, color='k') 
plot.show()