##################################
#                                #
# Last modified 2019/06/20       # 
#                                #
# Georgi Marinov                 #
#                                # 
##################################

import sys
import gzip

def run():

    if len(sys.argv) < 3:
        print 'usage: python %s fastq wanted fieldID' % sys.argv[0]
        print '\tuse - instead of a wanted filename to indicate standard input' 
        print '\tthe input file can also be in .gz format' 
        print '\tthe fastq file can be in .gz format' 
        print '\tthe script will print to stodut by default' 
        sys.exit(1)

    fastq = sys.argv[1]
    inputfilename = sys.argv[2]
    fieldID = int(sys.argv[3])

    WantedDict = {}

    doStdInput = False
    if inputfilename == '-':
        doStdInput = True

    if doStdInput:
        input_stream = sys.stdin
    else:
        if inputfilename.endswith('.gz'):
            input_stream = gzip.open(inputfilename)
        else:
            input_stream = open(inputfilename)
    for line in input_stream:
        if line.startswith('#'):
            continue
        fields = line.strip().split('\t')
        WantedDict[fields[fieldID]] = 1

    i=0
    pos=1
    if fastq.endswith('.gz'):
        input_stream = gzip.open(fastq)
    else:
        input_stream = open(fastq)
    for line in input_stream:
        if pos==1:
            if line.startswith('@'):
                ID = line.strip()[1:].split(' ')[0]
                if WantedDict.has_key(ID):
                    doPrint = True
                else:
                    doPrint = False
                if doPrint:
                    print line.strip()
                pos=2
                continue
            else:
                print 'invalid read', line
                break
        if pos==2:
            if doPrint:
                print line.strip()
            pos=3
            continue
        if pos==3 and line.startswith('+'):
            if doPrint:
                print line.strip()
            pos=4
            continue
        if pos==4:
            if doPrint:
                print line.strip()
            pos=1
            continue

run()

