def makeStringList(stringListFileName):
	# Make a list of strings from a file
	stringListFile = open(stringListFileName)
	stringList = []
	for line in stringListFile:
		# Iterate through the strings and add each to stringList
		stringList.append(line.strip())
	stringListFile.close()
	print len(stringList)
	return stringList

def getkMerMatches(kMer, featureList, numMismatches):
	# Get the indexes of all of the features in featureList that match the k-mer with <= numMismatches mismatches
	kMerMatches = []
	for i in range(len(featureList)):
		# Iterate through the features to see if each is a match
		mismatchCount = 0
		feature = featureList[i]
		for j in range(len(kMer)):
			# Check if the current base in the k-mer matches the current base in the feature
			if kMer[j] != feature[j]:
				# The current base is a mismatch
				mismatchCount = mismatchCount + 1
				if mismatchCount > numMismatches:
					# There are too many mismatches for the current feature
					break
		if mismatchCount <= numMismatches:
			# The current feature is sufficiently similar to the current k-mer
			kMerMatches.append(i)
	return kMerMatches

def getFeatureVecNoPairsClassification(refLineElements, altLineElements, featureList, direction, numMismatches):
	# Make a dictionary of features for the variants
	# Each feature value will be 1 if it is the ref. k-mer, -1 if it is the alt. k-mer (unless the ref. and alt. k-mers are the same), and 0 otherwise
	# If a k-mer is changed multiple times, the feature value will be (# of times k-mer is ref. allele) - (# of times k-mer is alt. allele)
	featureVec = []
	for i in range(len(featureList)):
		# Add a 0 for each feature
		featureVec.append(0)
	for i in range(2, len(refLineElements)):
		# Iterate through the k-mer changes and adjust each feature appropriately
		refkMer = refLineElements[i]
		altkMer = altLineElements[i]
		if refkMer == altkMer:
			# The k-mer does not change (should be due to indel)
			continue
		refkMerMatches = getkMerMatches(refkMer, featureList, numMismatches)
		altkMerMatches = getkMerMatches(altkMer, featureList, numMismatches)
		if direction > 0:
			# The correlation is positive so increment the reference k-mer feature and decrement the alternate k-mer feature
			for matchIndex in refkMerMatches:
				# Add 1 for each match
				featureVec[matchIndex] = featureVec[matchIndex] + 1
			for matchIndex in altkMerMatches:
				# Subtract 1 for each match
				featureVec[matchIndex] = featureVec[matchIndex] - 1
		else:
			# The correlation is positive so increment the reference k-mer feature and decrement the alternate k-mer feature
			for matchIndex in refkMerMatches:
				# Subtract 1 for each match
				featureVec[matchIndex] = featureVec[matchIndex] - 1
			for matchIndex in altkMerMatches:
				# Add 1 for each match
				featureVec[matchIndex] = featureVec[matchIndex] + 1
	return featureVec

def makeRandomForestClassificationInputsNoPairs(refFileName, altFileName, featureList, effectSizesFileName, numMismatches):
	# Run a random forest regression using k-mers
	# ASSUMES THAT refFile AND altFile ARE SORTED BY CHROMOSOME, POSITION
	# ASSUMES THAT effectSizesFile IS SORTED BY PEAK CHROMOSOME, PEAK START, PEAK END AND THAT ALL PEAKS' DIST. CUTOFFS ARE THE SAME
	# ASSUMES THAT ALL VARIANTS IN effectSizesFile ARE IN refFile and altFile
	# ASSUMES THAT refFile and altFile HAVE THE SAME VARIANTS
	effectSizesFile = open(effectSizesFileName)
	featureMat = []
	labelsVec = []
	lastVariants = []
	lastVariantsFeatureVecList = []
	currentVariants = []
	currentVariantsFeatureVecList = []
	lastPeakLocation = ("", 0, 0)
	refFile = open(refFileName)
	altFile = open(altFileName)
	refLine = refFile.readline()
	altLine = altFile.readline()
	for line in effectSizesFile:
		# Iterate though the variants and their effect sizes, and add the features for each to the feature matrix and the effect sizes to the effect sizes vector
		lineElements = line.split("\t")
		currentPeakLocation = (lineElements[0], int(lineElements[1]), int(lineElements[2]))
		if currentPeakLocation != lastPeakLocation:
			# At a new peak
			lastVariants = currentVariants
			lastVariantsFeatureVecList = currentVariantsFeatureVecList
			lastPeakLocation = currentPeakLocation
			currentVariants = []
			currentVariantsFeatureVecList = []
		variant = (lineElements[3], int(lineElements[4]))
		currentVariants.append(variant)
		if variant in lastVariants:
			# Store the feature vector for the variant by finding it in the feature vector for the variant in the previous peak
			variantIndex = lastVariants.index(variant)
			featureMat.append(lastVariantsFeatureVecList[variantIndex])
		else:
			# Make the feature vector for the current variant
			refLineElements = refLine.strip().split("\t")
			altLineElements = altLine.strip().split("\t")
			if (refLineElements[0], int(refLineElements[1])) != variant:
				print "Problem!"
				print variant
				print (refLineElements[0], int(refLineElements[1]))
			direction = 1
			if float(lineElements[7]) < 0:
				# The correlation is negative
				direction = -1
			featureVec = getFeatureVecNoPairsClassification(refLineElements, altLineElements, featureList, direction, numMismatches)
			featureMat.append(featureVec)
			currentVariants.append(variant)
			currentVariantsFeatureVecList.append(featureVec)
			refLine = refFile.readline()
			altLine = altFile.readline()
		label = int(lineElements[8])
		labelsVec.append(label)
	refFile.close()
	altFile.close()
	effectSizesFile.close()
	return [featureMat, labelsVec]

def runKMerRandomForestClassification(featureMat, labelsVec):
	# Run the random forest regression and record the results
	print len(featureMat)
	model = RandomForestClassifier(n_estimators=5000, max_features="sqrt", max_depth=100, n_jobs=-1)
	model.fit(featureMat, labelsVec)
	trainingAccuracy = model.score(featureMat, labelsVec)
	print trainingAccuracy
	return model

def outputFeatureImportances(model, featureList, outputFileName):
	# Record the feature importances
	outputFile = open(outputFileName, 'w+')
	for i in range(len(model.feature_importances_)):
		# Iterate through the feature importances and record each
		outputFile.write(featureList[i] + "\t" + str(model.feature_importances_[i]) + "\n")
	outputFile.close()
	
if __name__=="__main__":
	import sys
	from sklearn.ensemble import RandomForestClassifier
	refFileName = sys.argv[1]
	altFileName = sys.argv[2]
	featureListFileName = sys.argv[3]
	effectSizesFileName = sys.argv[4]
	outputFileName = sys.argv[5]
	numMismatches = int(sys.argv[6])
	featureList = makeStringList(featureListFileName)
	[featureMat, labelsVec] = makeRandomForestClassificationInputsNoPairs(refFileName, altFileName, featureList, effectSizesFileName, numMismatches)
	model = runKMerRandomForestClassification(featureMat, labelsVec)
	outputFeatureImportances(model, featureList, outputFileName)
	