import sys
import numpy as np
from itertools import izip

FIRSTFILENAME = sys.argv[1]
SECONDFILENAME = sys.argv[2]
OUTPUTFILENAME = sys.argv[3]
LABELSFILENAME = sys.argv[4]
CHROMLIST = []
for chrom in sys.argv[5:len(sys.argv)]:
	# Iterate through the desired chromosomes and add them to the list
	CHROMLIST.append(chrom)
	
def recordSequencesAndLabels(inputFileName, outputFile, labelsFile, label):
	# Record the sequences on the desired chromosomes and their corresponding label
	inputFile = open(inputFileName)
	line = inputFile.readline().strip()
	while line != "":
		# Iterate through the lines of the files and record the sequences on the proper chromosomes
		lineElements = line.strip().split(":")
		sequence = inputFile.readline().strip().upper()
		if lineElements[0][1:len(lineElements[0])] in CHROMLIST:
			# Record the current fasta entry
			outputFile.write(sequence + "\n")
			labelsFile.write(str(label) + "\n")
		line = inputFile.readline().strip()
	inputFile.close()

def selectChromosomesAndMakeLabelsFromFastaFilePair():
	# Select rows from pairs of files where the first (and the second if CHROMINBOTHFILES is 1) have a desired chromosome
	outputFile = open(OUTPUTFILENAME, 'w+')
	labelsFile = open(LABELSFILENAME, 'w+')
	recordSequencesAndLabels(FIRSTFILENAME, outputFile, labelsFile, 1)
	recordSequencesAndLabels(SECONDFILENAME, outputFile, labelsFile, 0)
	outputFile.close()
	labelsFile.close()

if __name__=="__main__":
	selectChromosomesAndMakeLabelsFromFastaFilePair()
