from hq import getZoom, scaleWeights
from hq_gen import (
	edges, expandQuadrant, genExpr2, genExpr3, genExpr4, simplifyWeights
	)

from collections import defaultdict
from itertools import count

def normalizeWeights(pixelExpr):
	maxSum = max(
		sum(weights)
		for expr in pixelExpr
		for weights in expr
		)
	return [
		[ scaleWeights(weights, maxSum) for weights in expr ]
		for expr in pixelExpr
		]

def extractTopLeftWeights(weights):
	assert all(weights[n] == 0 for n in (2, 5, 6, 7, 8)), weights
	return weights[0 : 2] + weights[3 : 5]

def extractTopLeftQuadrant(pixelExpr):
	zoom = getZoom(pixelExpr)
	quadrantWidth = (zoom + 1) // 2
	quadrantMap = [
		qy * zoom + qx
		for qy in range(quadrantWidth)
		for qx in range(quadrantWidth)
		]
	for expr in [
		[ expr[subPixel] for subPixel in quadrantMap ]
		for expr in pixelExpr
		]:
		for weights in expr:
			for neighbour in (2, 5, 6, 7, 8):
				assert weights[neighbour] == 0, weights
	return [
		[ extractTopLeftWeights(expr[subPixel]) for subPixel in quadrantMap ]
		for expr in pixelExpr
		]

def convertExpr4to2(case, expr4):
	weights2 = [0] * 4
	for weights4 in expr4:
		for neighbour, weight in enumerate(scaleWeights(weights4, 256)):
			weights2[neighbour] += weight
	weights2 = simplifyWeights(weights2)
	if ((case >> 4) & 15) in (2, 6, 8, 12):
		assert sorted(weights2) == [0, 2, 7, 23]
		weightMap = { 0: 0, 2: 1, 7: 1, 23: 2 }
	elif ((case >> 4) & 15) in (0, 1, 4, 5):
		assert sorted(weights2) == [0, 3, 3, 10]
		weightMap = { 0: 0, 3: 1, 10: 2 }
	else:
		weightMap = None
	if weightMap:
		weights2 = tuple(weightMap[weight] for weight in weights2)
	return [weights2]

def convert4to2(topLeftQuadrant4):
	return [
		convertExpr4to2(case, expr4)
		for case, expr4 in enumerate(topLeftQuadrant4)
		]

def analyzeCaseFunction(caseToWeights):
	weightsToCase = defaultdict(set)
	for case, weights in enumerate(caseToWeights):
		weightsToCase[weights].add(case)
	for weights in sorted(weightsToCase.keys()):
		cases = weightsToCase[weights]
		partitions = {
			tuple((case >> edgeNum) & 1 for edgeNum in range(11, -1, -1))
			for case in cases
			}
		# Repeatedly merge partitions until we have a minimal set.
		for edgeNum in range(12):
			changed = True
			while changed:
				changed = False
				for part in list(partitions):
					if part not in partitions:
						continue
					if part[edgeNum] < 2:
						pre = part[ : edgeNum]
						post = part[edgeNum + 1 : ]
						dual = pre + (part[edgeNum] ^ 1,) + post
						if dual in partitions:
							partitions.remove(part)
							partitions.remove(dual)
							partitions.add(pre + (2,) + post)
							changed = True
		yield weights, [
			''.join('01x'[bit] for bit in partition)
			for partition in sorted(partitions)
			]

# Various analysis:

def findRelevantEdges():
	for zoom, genExpr in zip(count(2), (genExpr2, genExpr3, genExpr4)):
		quadrant = genExpr()
		quadrantWidth = (zoom + 1) // 2
		assert quadrantWidth ** 2 == len(quadrant[0])
		subPixelOutput = [[] for _ in range(quadrantWidth * 10)]
		for subPixel in range(quadrantWidth ** 2):
			neighbourOutput = [[] for _ in range(8)]
			for neighbour in range(4):
				relevant = [
					edgeNum
					for edgeNum in range(12)
					if any(
						quadrant[case][subPixel][neighbour] !=
							quadrant[case ^ (1 << edgeNum)][subPixel][neighbour]
						for case in range(len(quadrant))
						)
					]
				zero = len(relevant) == 0 and all(
					quadrant[case][subPixel][neighbour] == 0
					for case in range(len(quadrant))
					)
				center = ('.' if zero else str(neighbour))
				for rowNum, row in enumerate(formatEdges(relevant)):
					if rowNum == 1:
						assert row[1] == 'o'
						row = row[0] + center + row[2]
					neighbourOutput[(neighbour // 2) * 4 + rowNum].append(row)
				neighbourOutput[(neighbour // 2) * 4 + 3].append('   ')
			for lineNum, line in enumerate(neighbourOutput):
				lineOutput = '  %s  |' % '  '.join(line)
				subPixelOutput[
					(subPixel // quadrantWidth) * 10 + lineNum + 1
					].append(lineOutput)
				if lineNum == 7:
					subPixelOutput[(subPixel // quadrantWidth) * 10].append(
						lineOutput
						)
					subPixelOutput[(subPixel // quadrantWidth) * 10 + 9].append(
						'%ss%d' % ('-' * (len(lineOutput) - 2), subPixel)
						)
		print('Relevant edges for zoom %d:' % zoom)
		print()
		for line in subPixelOutput:
			print('  %s' % ''.join(line))
		print()

# Visualization:

def formatEdges(edgeNums):
	cells = ['.' for _ in range(9)]
	cells[4] = 'o'
	def combine(index, ch):
		old = cells[index]
		if old == '.':
			cells[index] = ch
		elif (old == '/' and ch == '\\') or (old == '\\' and ch == '/'):
			cells[index] = 'X'
		else:
			assert False, (index, old, ch)
	for edgeNum in edgeNums:
		edge = edges[edgeNum]
		if 4 in edge:
			other = sum(edge) - 4
			combine(other, '\\|/-'[other if other < 4 else 8 - other])
		else:
			assert edge in ((1, 5), (5, 7), (3, 7), (1, 3))
			x = 0 if 3 in edge else 2
			y = 0 if 1 in edge else 2
			combine(y * 3 + x, '/' if x == y else '\\')
	for y in range(3):
		yield ''.join(cells[y * 3 : (y + 1) * 3])

def formatWeights(weights):
	return ' '.join('%3d' % weight for weight in weights)

def comparePixelExpr(pixelExpr1, pixelExpr2):
	zoom = getZoom(pixelExpr1)
	assert zoom == getZoom(pixelExpr2)
	mismatchCount = 0
	for case, (expr1, expr2) in enumerate(zip(pixelExpr1, pixelExpr2)):
		if expr1 != expr2:
			binStr = bin(case)[2 : ].zfill(12)
			print('case: %d (%s)' % (
				case,
				' '.join(binStr[i : i + 4] for i in range(0, 12, 4))
				))
			for sy in range(zoom):
				matrices = []
				for sx in range(zoom):
					subPixel = sy * zoom + sx
					subExpr1 = expr1[subPixel]
					subExpr2 = expr2[subPixel]
					rows = []
					for ny in range(3):
						rows.append('%s   %s %s' % (
							formatWeights(subExpr1[ny * 3 : (ny + 1) * 3]),
							'.' if subExpr1 == subExpr2 else '!',
							formatWeights(subExpr2[ny * 3 : (ny + 1) * 3])
							))
					matrices.append(rows)
				for ny in range(3):
					print('  %s' % '       '.join(
						matrices[sx][ny]
						for sx in range(zoom)
						))
				print()
			mismatchCount += 1
	print('Number of mismatches: %d' % mismatchCount)

# Sanity checks:

def checkConvert4to2():
	topLeftQuadrant4 = genExpr4()
	topLeftQuadrant2 = convert4to2(topLeftQuadrant4)
	pixelExpr2 = expandQuadrant(topLeftQuadrant2, 2)

	comparePixelExpr(pixelExpr2, expandQuadrant(genExpr2(), 2))

# Main:

if __name__ == '__main__':
	findRelevantEdges()
	checkConvert4to2()