#!/usr/bin/env python3

from collections import defaultdict
from gzip import GzipFile
from math import modf
from sys import stderr
from xml.sax import make_parser
from xml.sax.handler import feature_external_ges
from xml.dom import pulldom

cpuClock = 3579545
vdpClock = cpuClock * 6
masterClock = cpuClock * 960

vdpTicksPerLine = 1368

# The keys we're interested in: the letter to be used in the text file,
# followed by the keyboard matrix location (row, column).
# This should probably be made configurable at some point, but for now you
# can edit them here.
inputMap = {
	'r': (8, 7), 'd': (8, 6), 'u': (8, 5), 'l': (8, 4), 's': (8, 0)
	}
# Coleco:
#inputMap = {
#	'r': (0, 1), 'd': (0, 2), 'u': (0, 0), 'l': (0, 3), 's': (0, 6), 'o': (2, 1)
#	}
inputMapReverse = dict((pos, name) for name, pos in inputMap.items())

def readEvents(filename):
	'''Reads input events from an openMSX replay file (OMR).
	Returns a list of tuples consisting of a timestamp, keyboard row, press
	mask and release mask.
	'''
	print('reading replay:', filename, file=stderr)

	inputEvents = []

	def getText(node):
		return ''.join(
			child.data
			for child in node.childNodes
			if child.nodeType == child.TEXT_NODE
			)

	def parseItem():
		time = None
		row = None
		press = None
		release = None
		for event, node in xmlEvents:
			if event == pulldom.START_ELEMENT:
				doc.expandNode(node)
				if node.tagName == 'StateChange':
					assert time is None, time
					timeNode1, timeNode2 = node.getElementsByTagName('time')
					time = int(getText(timeNode2))
				elif node.tagName == 'row':
					assert row is None, row
					row = int(getText(node))
				elif node.tagName == 'press':
					assert press is None, press
					press = int(getText(node))
				elif node.tagName == 'release':
					assert release is None, release
					release = int(getText(node))
				else:
					pass #print(node.toxml())
			elif event == pulldom.END_ELEMENT:
				if node.tagName == 'item':
					break

		assert time is not None
		assert row is not None
		assert press is not None
		assert release is not None
		inputEvents.append((time, row, press, release))

	def parseEvents():
		for event, node in xmlEvents:
			if event == pulldom.START_ELEMENT:
				if node.tagName == 'item' \
						and node.getAttribute('type') == 'KeyMatrixState':
					parseItem()
			elif event == pulldom.END_ELEMENT:
				if node.tagName == 'events':
					break

	def parseTopLevel():
		for event, node in xmlEvents:
			if event == pulldom.START_ELEMENT:
				if node.tagName == 'events':
					parseEvents()

	parser = make_parser()
	parser.setFeature(feature_external_ges, False)

	with GzipFile(filename) as inp:
		doc = pulldom.parse(inp, parser)
		xmlEvents = iter(doc)
		parseTopLevel()

	print('read %d input events' % len(inputEvents), file=stderr)
	return inputEvents

def combineEvents(inputEvents):
	'''Combines multiple events on the same keyboard row at the same timestamp
	into a single event.
	'''
	pendingTime = None

	def outputPending():
		if pendingTime is not None:
			for row in sorted(pressForRow.keys() | releaseForRow.keys()):
				press = pressForRow[row]
				release = releaseForRow[row]
				if press != 0 or release != 0:
					yield pendingTime, row, press, release

	for time, row, press, release in inputEvents:
		if time != pendingTime:
			assert pendingTime is None or time > pendingTime, time

			# Output previous event.
			yield from outputPending()

			# Start new event.
			pendingTime = time
			pressForRow = defaultdict(int)
			releaseForRow = defaultdict(int)

		pressForRow[row] = (pressForRow[row] | press) & ~release
		releaseForRow[row] = (releaseForRow[row] | release) & ~press
	else:
		yield from outputPending()

def removeRedundantEvents(inputEvents):
	'''Remove (parts of) events that don't change the keyboard matrix state.
	'''
	matrix = [0] * 12
	for time, row, press, release in inputEvents:
		old = matrix[row]
		press &= ~old
		release &= old
		if press != 0 or release != 0:
			matrix[row] = (old | press) & ~release
			yield time, row, press, release

def filterEvents(inputEvents, wantedKeys):
	'''Filter the given input events to contain only presses and releases of
	the given keys.
	The wantedKeys argument should be mapping from row to a column bitmask,
	where bits are set if we want to preserve changes in the corresponding
	position in the keyboard matrix.
	'''
	for time, row, press, release in inputEvents:
		mask = wantedKeys.get(row, 0)
		press &= mask
		release &= mask
		if press != 0 or release != 0:
			yield time, row, press, release

def scaleTime(inputEvents, tickScale):
	'''Yields the given input events, with the time stamps divided by the
	given scale and rounded to the nearest integer.
	'''
	for time, row, press, release in inputEvents:
		yield int(time / tickScale + 0.5), row, press, release

def checkAlignment(timestamps, alignment):
	'''Returns the number of timestamps from the given sequence that are
	close to multiples of the given alignment.
	'''
	aligned = 0
	for time in timestamps:
		offset = modf(time / alignment)[0]
		if offset < 0.001 or offset > 0.999:
			aligned += 1
	return aligned

def detectTicksPerFrame(timestamps):
	'''Determine the most likely number of master clock ticks per frame,
	by looking at how well the timestamps align with frame boundaries when
	using the two frame timings that openMSX supports.
	This assumes that the frame timing is the same throughout the replay,
	which is not guaranteed in general, but will hopefully be the case for
	tool-assisted speedruns.
	'''
	ticksPerFrame50 = (masterClock * vdpTicksPerLine * 313) // vdpClock
	ticksPerFrame60 = (masterClock * vdpTicksPerLine * 262) // vdpClock

	timestamps = tuple(timestamps)
	aligned50 = checkAlignment(timestamps, ticksPerFrame50)
	aligned60 = checkAlignment(timestamps, ticksPerFrame60)

	if aligned50 > aligned60:
		print('event timestamps align with 50 fps timing: %.2f%%'
				% (100 * aligned50 / len(timestamps)), file=stderr)
		return ticksPerFrame50
	else:
		print('event timestamps align with 60 fps timing: %.2f%%'
				% (100 * aligned60 / len(timestamps)), file=stderr)
		return ticksPerFrame60

def eventsToState(inputEvents):
	'''Yields pairs of a timestamp and a set of the active keys at that time.
	'''
	active = set()
	prevTime = None
	for time, row, press, release in inputEvents:
		assert (press & release) == 0, (press, release)

		if prevTime != time:
			if prevTime is not None:
				yield prevTime, frozenset(active)
			prevTime = time

		col = 0
		while press:
			if press & 1:
				active.add((row, col))
			col += 1
			press >>= 1

		col = 0
		while release:
			if release & 1:
				active.remove((row, col))
			col += 1
			release >>= 1

	if prevTime is not None:
		yield prevTime, frozenset(active)

def formatState(active):
	return ' '.join(
		inputMapReverse[pos]
		for pos in sorted(active, key=lambda pos: (pos[0], -pos[1]))
		)

def convert(inFilename, outFilename):
	wantedKeys = {}
	for name, (row, col) in inputMap.items():
		wantedKeys[row] = wantedKeys.get(row, 0) | (1 << col)

	inputEvents = list(removeRedundantEvents(
			filterEvents(combineEvents(readEvents(inFilename)), wantedKeys)
			))
	print('after cleanup %d events remain' % len(inputEvents), file=stderr)
	if not inputEvents:
		print(
			"no events match; you should probably customize 'inputMap' "
			"at the top of this script", file=stderr
			)
		return

	ticksPerFrame = detectTicksPerFrame(evt[0] for evt in inputEvents)
	scaledEvents = list(removeRedundantEvents(
			combineEvents(scaleTime(inputEvents, ticksPerFrame))
			))

	with open(outFilename, 'w') as out:
		print('writing output:', outFilename)

		print('= base', inFilename, file=out)
		print('= out', 'replay.omr', file=out)
		print('= scale', ticksPerFrame, file=out)
		for name, (row, col) in sorted(
				inputMap.items(), key=lambda item: (item[1][0], -item[1][1])
				):
			print('= input', name, 'key', row, col, file=out)
		print(file=out)

		def printMilestone(frame, seconds):
			print('# frame %d (%d:%02d)' % ((frame,) + divmod(seconds, 60)),
					file=out)
		milestone = 10 # seconds

		prevFrame = 0
		prevSeconds = 0
		printMilestone(prevFrame, prevSeconds)
		prevActive = set()
		for frame, active in eventsToState(scaledEvents):
			delta = frame - prevFrame
			print(('%4d  %s' % (delta, formatState(prevActive))).rstrip(),
					file=out)
			seconds = (frame * ticksPerFrame) // masterClock
			if seconds // milestone != prevSeconds // milestone:
				printMilestone(frame, seconds)
			prevFrame = frame
			prevSeconds = seconds
			prevActive = active
		printMilestone(prevFrame, prevSeconds)

if __name__ == '__main__':
	from sys import argv
	if len(argv) != 2:
		print('Usage: omr2txt.py file1.omr', file=stderr)
		print('Converts an openMSX replay to a text file.', file=stderr)
		exit(2)
	else:
		inFilename = argv[1]
		if inFilename.endswith('.omr'):
			outFilename = inFilename[:-4] + '.txt'
			convert(inFilename, outFilename)
		else:
			print('File name does not end in ".omr":', inFilename)
			exit(2)