# To run:

#   python3 hq.py

#

# To profile:

#   python3 -m cProfile -s cumulative hq.py > hq-profile.txt

from hq_gen import (

	edges, expandQuadrant, genExpr2, genExpr3, genExpr4,

	permuteCase, simplifyWeights

	)

from collections import defaultdict

from itertools import count

from math import sqrt

import sys

# Variant classes:

class Variant(object):

	def __init__(self, pixelExpr, lite, table, narrow = None):

		self.lite = lite

		self.table = table

		if lite:

			pixelExpr = makeLite(pixelExpr, table)

		if narrow is not None:

			pixelExpr = narrow(pixelExpr)

                #   0 | 1 | 2

                #  ---+---+---

                #   3 | 4 | 5

                #  ---+---+---

                #   6 | 7 | 8

                # For the openGL shaders we want the following order:

                #   (4,6), (3,7), (3,4), (1,3), (0,4), (4,7), (1,4), (5,7), (4,8), (4,5), (2,4), (1,5)

                #   see comments in src/video/scalers/HQCommon.hh for why this order was chosen.

                # Compare this to the original order in 'edges' in hq_gen.py.

                # To switch between these different orderings we use the following permutations.

		pixelExpr = permuteCases(

			(9, 2, 7, 3, 4, 10, 5, 1, 11, 8, 6, 0)  # openGL

			if table else

			(2, 9, 7, 4, 3, 10, 11, 1, 8, 0, 6, 5),

			pixelExpr

			)

		self.pixelExpr = pixelExpr

	def writeSwitch(self, fileName):

		writeTextFile(fileName, genSwitch(self.pixelExpr))

# I/O:

def printText(contents):

	for text in contents:

		sys.stdout.write(text)

def writeTextFile(fileName, contents):

	with open(fileName, 'w') as out:

		for content in contents:

			out.write(content)

def writeBinaryFile(fileName, data):

	with open(fileName, 'wb') as out:

		out.write(bytes(data))

# Output as switch statement:

def genSwitch(pixelExpr):

	# Note: In practice, only the center subpixel of HQ3x is independent of

	#       the edge bits. So maybe this is a bit overengineered, but it does

	#       express why this optimization is correct.

	subExprForSubPixel = tuple(

		{expr[subPixel] for expr in pixelExpr if expr[subPixel] is not None}

		for subPixel in range(len(pixelExpr[0]))

		)

	isIndependentSubPixel = tuple(

		len(subExprs) == 1

		for subExprs in subExprForSubPixel

		)

	exprToCases = defaultdict(list)

	for case, expr in enumerate(pixelExpr):

		if expr[0] is None:

			assert all(subExpr is None for subExpr in expr)

			key = None

		else:

			key = tuple(tuple(subExpr) for subExpr in expr)

		exprToCases[key].append(case)

	yield 'switch (pattern) {\n'

	for cases, expr in sorted(

		( sorted(cases), expr )

		for expr, cases in exprToCases.items()

		):

		if expr is not None:

			for case in cases:

				yield 'case %d:\n' % case

			for subPixel, subExpr in enumerate(expr):

				if not isIndependentSubPixel[subPixel]:

					yield '\tpixel%s = %s;\n' % (

						hex(subPixel)[2 : ], getBlendCode(subExpr)

						)

			yield '\tbreak;\n'

	for case in sorted(exprToCases[None]):

		yield 'case %d:\n' % case

	yield 'default:\n'

	yield '\tUNREACHABLE;\n'

	yield '\t%s = 0; // avoid warning\n' % ' = '.join(

		'pixel%d' % subPixel

		for subPixel, independent_ in enumerate(isIndependentSubPixel)

		)

	yield '}\n'

	for subPixel, independent in enumerate(isIndependentSubPixel):

		if independent:

			subExpr, = subExprForSubPixel[subPixel]

			yield 'pixel%s = %s;\n' % (

				hex(subPixel)[2 : ], getBlendCode(subExpr)

				)

def getBlendCode(weights):

	wSum = sum(weights)

	assert isPow2(wSum)

	if wSum == 1:

		index, = (index for index, weight in enumerate(weights) if weight != 0)

		return 'c%d' % (index + 1)

	elif wSum <= 8:

		# Because the lower 3 bits of each color component (R,G,B)

		# are zeroed out, we can operate on a single integer as if it

		# is a vector.

		return ' + '.join(

			'(c%d / %d) * %d' % (index + 1, wSum, weight)

			for index, weight in enumerate(weights)

			if weight != 0

			)

	else:

		return '((%s) & 0xFF00FF00) | (((%s) / %d) & 0x00FF00FF)' % (

			' + '.join(

				'((c%d & 0xFF00FF00) / %d) * %d' % (index + 1, wSum, weight)

				for index, weight in enumerate(weights)

				if weight != 0

				),

			' + '.join(

				'(c%d & 0x00FF00FF) * %d' % (index + 1, weight)

				for index, weight in enumerate(weights)

				if weight != 0

				),

			wSum

			)

# Table output as text:

def formatOffsetsTable(pixelExpr):

	for case, expr in enumerate(pixelExpr):

		yield '// %d\n' % case

		for weights in expr:

			for x, y in transformOffsets(weights):

				yield ' %3d, %3d,' % (x, y)

			yield '\n'

def formatWeightsTable(pixelExpr, cellFunc):

	for case, expr in enumerate(pixelExpr):

		yield '// %d\n' % case

		for weights in expr:

			for weight in transformWeights(weights, cellFunc):

				yield ' %3d,' % weight

			yield '\n'

# The rest:

def getZoom(pixelExpr):

	zoom = int(sqrt(len(pixelExpr[0])))

	assert zoom * zoom == len(pixelExpr[0])

	return zoom

def isPow2(num):

	while num & 1 == 0:

		num >>= 1

	return num == 1

def permuteCases(permutation, pixelExpr):

	return [

		pixelExpr[permuteCase(case, permutation)]

		for case in range(len(pixelExpr))

		]

def computeNeighbours(weights):

	neighbours = [ i for i in range(9) if i != 4 and weights[i] != 0 ]

	assert len(neighbours) <= 2

	neighbours += [ None ] * (2 - len(neighbours))

	return neighbours

def transformOffsets(weights):

	for neighbour in computeNeighbours(weights):

		yield (

			min(255, (1 if neighbour is None else neighbour % 3) * 128),

			min(255, (1 if neighbour is None else neighbour // 3) * 128)

			)

def transformWeights(weights, cellFunc):

	scaledWeights = scaleWeights(weights, 256)

	for cell in cellFunc(scaledWeights):

		yield min(255, 0 if cell is None else scaledWeights[cell])

def computeOffsets(pixelExpr):

	'''Computes offsets for the fragment shader.

	Output is a 64N * 64N texture, where N is the zoom factor.

	'''

	zoom = getZoom(pixelExpr)

	for caseMajor in range(0, len(pixelExpr), 64):

		for subY in range(0, zoom ** 2, zoom):

			for caseMinor in range(64):

				for subX in range(zoom):

					weights = pixelExpr[caseMajor + caseMinor][subY + subX]

					for x, y in transformOffsets(weights):

						yield x

						yield y

def computeWeights(pixelExpr, cellFunc):

	'''Computes weights for the fragment shader.

	Output is a 64N * 64N texture, where N is the zoom factor.

	'''

	zoom = getZoom(pixelExpr)

	for caseMajor in range(0, len(pixelExpr), 64):

		for subY in range(0, zoom ** 2, zoom):

			for caseMinor in range(64):

				for subX in range(zoom):

					weights = pixelExpr[caseMajor + caseMinor][subY + subX]

					for tw in transformWeights(weights, cellFunc):

						yield tw

def computeWeightCells(weights):

	neighbours = computeNeighbours(weights)

	return (neighbours[0], neighbours[1], 4)

def computeLiteWeightCells(weights_):

	return (3, 4, 5)

def isContradiction(case):

	inv = case ^ 0xFFF

	return (

		(inv & 0x121) in [0x120, 0x101, 0x021] or

		(inv & 0x502) in [0x500, 0x402, 0x102] or

		(inv & 0x484) in [0x480, 0x404, 0x084] or

		(inv & 0x0A8) in [0x0A0, 0x088, 0x028] or

		(inv & 0x00F) in [0x00E, 0x00D, 0x00B, 0x007]

		)

def scaleWeights(weights, newSum):

	'''Returns the given weights, scaled such that their sum is the given one.

	'''

	oldSum = sum(weights)

	factor, remainder = divmod(newSum, oldSum)

	assert remainder == 0

	return tuple(weight * factor for weight in weights)

def blendWeights(weights1, weights2, factor1 = 1, factor2 = 1):

	factor1 *= sum(weights2)

	factor2 *= sum(weights1)

	return simplifyWeights(

		factor1 * w1 + factor2 * w2

		for w1, w2 in zip(weights1, weights2)

		)

def calcNeighbourToSet():

	# Compute equivalence classes.

	ret = []

	for case in range(1 << len(edges)):

		if isContradiction(case):

			ret.append(None)

		else:

			neighbourToSet = [{neighbour} for neighbour in range(9)]

			for edge, (neighbour1, neighbour2) in enumerate(edges):

				if (case & (1 << edge)) == 0:

					# No edge, so the two pixels are equal.

					# Merge the sets of equal pixels.

					set1 = neighbourToSet[neighbour1]

					set2 = neighbourToSet[neighbour2]

					set1 |= set2

					for n in set2:

						neighbourToSet[n] = set1

			ret.append(neighbourToSet)

	return ret

neighbourToSet = calcNeighbourToSet()

def lighten(case, weights, neighbourPreference):

	equalNeighboursOf = neighbourToSet[case]

	if equalNeighboursOf is None:

		return None

	else:

		done = set()

		newWeights = [ 0 ] * 9

		for neighbour in neighbourPreference:

			if neighbour not in done:

				equalNeighbours = equalNeighboursOf[neighbour]

				newWeights[neighbour] = sum(

					weights[n] for n in equalNeighbours

					)

				done |= equalNeighbours

		for c in (0, 1, 2, 6, 7, 8):

			# Only c4, c5, c6 have non-zero weight.

			assert newWeights[c] == 0

		return simplifyWeights(newWeights)

def makeLite(pixelExpr, biased):

	zoom = getZoom(pixelExpr)

	biasLeft  = (4, 3, 5, 1, 7, 0, 2, 6, 8)

	biasRight = (4, 5, 3, 1, 7, 2, 0, 8, 6)

	if biased:

		center = (zoom - 1) / 2

		neighbourPreferences = [

			biasRight if subPixel % zoom > center else biasLeft

			for subPixel in range(zoom ** 2)

			]

	else:

		neighbourPreferences = [ biasLeft ] * (zoom ** 2)

	return [

		[ lighten(case, weights, pref)

		  for weights, pref in zip(expr, neighbourPreferences) ]

		for case, expr in enumerate(pixelExpr)

		]

def makeNarrow2to1(pixelExpr):

	return tuple(

		(None, None) if a is None else (blendWeights(a, b), blendWeights(c, d))

		for a, b, c, d in pixelExpr

		)

def makeNarrow3to2(pixelExpr):

	return [

		[	blendWeights(expr[0], expr[1], 2, 1),

			blendWeights(expr[2], expr[1], 2, 1),

			blendWeights(expr[3], expr[4], 2, 1),

			blendWeights(expr[5], expr[4], 2, 1),

			blendWeights(expr[6], expr[7], 2, 1),

			blendWeights(expr[8], expr[7], 2, 1),

			]

		for expr in pixelExpr

		]

def genHQLiteOffsetsTable(pixelExpr):

	'''In the hqlite case, the result color depends on at most one neighbour

	color. Therefore, an offset into an interpolated texture is used instead

	of explicit weights.

	Output is a 64N * 64N texture, where N is the zoom factor.

	'''

	zoom = getZoom(pixelExpr)

	for caseMajor in range(0, len(pixelExpr), 64):

		for subY in range(zoom):

			for caseMinor in range(64):

				for subX in range(zoom):

					subPixel = zoom * subY + subX

					weights = pixelExpr[caseMajor + caseMinor][subPixel]

					if weights is None:

						neighbour = None

					else:

						neighbours = computeNeighbours(weights)

						assert neighbours[1] is None, neighbours

						neighbour = neighbours[0]

						factor = sum(weights)

					x = int(192.5 - 128 * (0.5 + subX) / zoom)

					y = int(192.5 - 128 * (0.5 + subY) / zoom)

					if neighbour == 3:

						x -= 128 * weights[3] // factor

					elif neighbour == 5:

						x += 128 * weights[5] // factor

					else:

						assert neighbour is None, neighbour

					assert 0 <= x < 256, x

					assert 0 <= y < 256, y

					yield x

					yield y

# Main:

def process2x():

	pixelExpr = expandQuadrant(genExpr2(), 2)

	fullTableVariant = Variant(pixelExpr, lite = False, table = True)

	liteTableVariant = Variant(pixelExpr, lite = True,  table = True)

	#printText(formatOffsetsTable(fullTableVariant.pixelExpr))

	#printText(formatOffsetsTable(liteTableVariant.pixelExpr))

	#printText(formatWeightsTable(

	#	fullTableVariant.pixelExpr, computeWeightCells

	#	))

	#printText(formatWeightsTable(

	#	liteTableVariant.pixelExpr, computeLiteWeightCells

	#	))

	writeBinaryFile(

		'HQ2xOffsets.dat',

		computeOffsets(fullTableVariant.pixelExpr)

		)

	writeBinaryFile(

		'HQ2xWeights.dat',

		computeWeights(fullTableVariant.pixelExpr, computeWeightCells)

		)

	writeBinaryFile(

		'HQ2xLiteOffsets.dat',

		genHQLiteOffsetsTable(liteTableVariant.pixelExpr)

		)

	# Note: HQ2xLiteWeights.dat is not needed, since interpolated texture

	#       offsets can perform all the blending we need.

	Variant(pixelExpr, lite = False, table = False).writeSwitch(

		'HQ2xScaler-1x1to2x2.nn'

		)

	Variant(pixelExpr, lite = False, table = False, narrow = makeNarrow2to1).writeSwitch(

		'HQ2xScaler-1x1to1x2.nn'

		)

	Variant(pixelExpr, lite = True,  table = False).writeSwitch(

		'HQ2xLiteScaler-1x1to2x2.nn'

		)

	Variant(pixelExpr, lite = True,  table = False, narrow = makeNarrow2to1).writeSwitch(

		'HQ2xLiteScaler-1x1to1x2.nn'

		)

def process3x():

	pixelExpr = expandQuadrant(genExpr3(), 3)

	fullTableVariant = Variant(pixelExpr, lite = False, table = True)

	liteTableVariant = Variant(pixelExpr, lite = True,  table = True)

	#printText(formatOffsetsTable(fullTableVariant.pixelExpr))

	#printText(formatOffsetsTable(liteTableVariant.pixelExpr))

	#printText(formatWeightsTable(

		#fullTableVariant.pixelExpr, computeWeightCells

		#))

	#printText(formatWeightsTable(

		#liteTableVariant.pixelExpr, computeLiteWeightCells

		#))

	writeBinaryFile(

		'HQ3xOffsets.dat',

		computeOffsets(fullTableVariant.pixelExpr)

		)

	writeBinaryFile(

		'HQ3xWeights.dat',

		computeWeights(fullTableVariant.pixelExpr, computeWeightCells)

		)

	writeBinaryFile(

		'HQ3xLiteOffsets.dat',

		genHQLiteOffsetsTable(liteTableVariant.pixelExpr)

		)

	# Note: HQ3xLiteWeights.dat is not needed, since interpolated texture

	#       offsets can perform all the blending we need.

	Variant(pixelExpr, lite = False, table = False).writeSwitch(

		'HQ3xScaler-1x1to3x3.nn'

		)

	Variant(pixelExpr, lite = True,  table = False).writeSwitch(

		'HQ3xLiteScaler-1x1to3x3.nn'

		)

def process4x():

	pixelExpr = expandQuadrant(genExpr4(), 4)

	fullTableVariant = Variant(pixelExpr, lite = False, table = True)

	liteTableVariant = Variant(pixelExpr, lite = True,  table = True)

	#printText(formatOffsetsTable(fullTableVariant.pixelExpr))

	#printText(formatWeightsTable(fullTableVariant.pixelExpr, computeWeightCells))

	#printText(formatWeightsTable(liteTableVariant.pixelExpr, computeLiteWeightCells))

	writeBinaryFile(

		'HQ4xOffsets.dat',

		computeOffsets(fullTableVariant.pixelExpr)

		)

	writeBinaryFile(

		'HQ4xWeights.dat',

		computeWeights(fullTableVariant.pixelExpr, computeWeightCells)

		)

	writeBinaryFile(

		'HQ4xLiteOffsets.dat',

		genHQLiteOffsetsTable(liteTableVariant.pixelExpr)

		)

	# Note: HQ4xLiteWeights.dat is not needed, since interpolated texture

	#       offsets can perform all the blending we need.

	#printText(genSwitch(Variant(

		#pixelExpr, lite = False, table = False

		#).pixelExpr))

if __name__ == '__main__':

	process2x()

	process3x()

	process4x()