import sys
import argparse
from itertools import izip
import math

def parseArgument():
	# Parse the input
	parser = argparse.ArgumentParser(description = "Make regions with 0 signal the average of their surrounding regions")
	parser.add_argument("--signalsFileName", required=True, help='Signals file')
	parser.add_argument("--peakIndexesFileName", required=True, help='Peak indexes file')
	parser.add_argument("--outputFileName", required=True, help='Output file, where signals that were 0 will be the average of their surrounding signals')
	options = parser.parse_args();
	return options
	
def averageZeroSignalsWithinPeaks(options):
	signalsFile = open(options.signalsFileName)
	peakIndexesFile = open(options.peakIndexesFileName)
	outputFile = open(options.outputFileName, 'w+')
	lastSignal = None
	lastLastSignal = None
	lastPeakIndex = None
	lastLastPeakIndex = None
	for signalsLine, peakIndexesLine in izip(signalsFile, peakIndexesFile):
		# Iterate through the signals and set those that are zero to the average of those of the surrounding regions
		signal = float(signalsLine.strip())
		peakIndex = int(peakIndexesLine.strip())
		if lastSignal == 0:
			# The previous signal was a zero, so set it to the average of the surrounding signals
			

		if signal != 0:
			# The current signal is not 0, so record it
			outputFile.write(str(signal) + "\n")
		lastLastSignal = lastSignal
		lastLastPeakIndex = lastPeakIndex
		lastSignal = signal
		lastPeakIndex = peakIndex
	if lastSignal == 0:
		# The final signal was a zero, so set it to the signal before it
		if (lastPeakIndex == lastLastPeakIndex) and (not math.isnan(lastLastSignal)):
			# Set the output to the region before it
			outputFile.write(str(lastLastSignal) + "\n")
		else:
			outputFile.write(str(lastSignal) + "\n")
	signalsFile.close()
	peakIndexesFile.close()
	outputFile.close()

if __name__=="__main__":
	options = parseArgument()
	averageZeroSignalsWithinPeaks(options)

if (peakIndex == lastPeakIndex) and (not math.isnan(lastSignal)):
				# Include the current region in the average
				if (lastPeakIndex == lastLastPeakIndex) and (not math.isnan(lastLastSignal)):
					# Include the region before the previous region in the average
					if not math.isnan(signal):
						# The current signal is not a nan, so include it in the average
						lastSignalCorrected = (signal + lastLastSignal)/2.0
						outputFile.write(str(lastSignalCorrected) + "\n")
					else:
						# The current signal is a nan, so use only the previous signal
						outputFile.write(str(lastLastSignal) + "\n")
				elif not math.isnan(signal):
					outputFile.write(str(signal) + "\n")
				else:
					outputFile.write(str(lastSignal) + "\n")
			elif (lastPeakIndex == lastLastPeakIndex) and (not math.isnan(lastLastSignal)):
				# Set the output to the region before it
				outputFile.write(str(lastLastSignal) + "\n")
			else:
				outputFile.write(str(lastSignal) + "\n")