#!/usr/bin/env python
# -*- coding: utf-8 vi:noet
# What if we made crazy QR codes with extra data in chroma...

"""
To reduce flicker, the extra data can be gray-encoded
"""

import sys, struct, ctypes
import cv2
import numpy as np


def gray_enc(value):
	return (value^(value>>1))

def gray_dec(value, bits):
	out = 0
	for idx_bit in range(bits-1, -1, -1):
		out |= (((value >> idx_bit) & 1) ^ ((out >> (idx_bit+1)) & 1)) << idx_bit
	return out

def sbits(value, bits):
	return bin(value | (1<<bits))[3:]

bits = 8
for i in range(1<<bits):
	e = gray_enc(i)
	d = gray_dec(e, bits)
	assert d == i
	continue
	print("%s %s %s" % (
	 sbits(i, bits),
	 sbits(e, bits),
	 sbits(d, bits),
	))


def enc_flt(x):
	b = struct.pack("<f", float(x))
	i = struct.unpack("<I", b)[0]
	g = gray_enc(i)
	return struct.pack("<I", g)

class QREncoder(object):
	def __init__(self):
		self._lib = ctypes.CDLL("libqrencode.so")
		self._lib.QRcode_encodeData.argtypes = ctypes.c_int, ctypes.c_char_p, ctypes.c_int, ctypes.c_int
		class QRCode(ctypes.Structure):
			_fields_ = [
			 ("version", ctypes.c_int),
			 ("width", ctypes.c_int),
			 ("data", ctypes.POINTER(ctypes.c_uint8)),
			]
		self._lib.QRcode_encodeData.restype = ctypes.POINTER(QRCode)
	def encode(self, value):
		assert isinstance(value, bytes)
		version = 2
		level = 3
		qrcode = self._lib.QRcode_encodeData(len(value), value, version, level)
		res = qrcode.contents
		version = int(res.version)
		size = int(res.width)
		data = bytearray(size*size)
		for idx_data in range(size*size):
			data[idx_data] = (1 - (res.data[idx_data] & 1)) * 255
		data = bytes(data)
		self._lib.QRcode_free(qrcode)
		return version, size, data

def make_samesize(datas):
	# make the QR codes the same size (somehow)
	sizes = [None] * len(datas)
	payloads = [None] * len(datas)
	encoder = QREncoder()
	while sizes[0] is None or not reduce(lambda x, y: x and y, map(lambda x: x == sizes[0], sizes)):

		if sizes[0] is not None:
			for idx_data, data in enumerate(datas):
				if sizes[idx_data] < max(sizes):
					datas[idx_data] += b" "


		for idx_data, data in enumerate(datas):
			version, size, data = encoder.encode(data)
			payloads[idx_data] = data
			sizes[idx_data] = size

	return size, payloads

def example_encode_luma():

	datas = []
	datas.append("http://zougloub.eu")
	datas.append("hello world!")
	datas.append("hello again!")
	datas.append("hello ter!")

	size, payloads = make_samesize(datas)

	scale = 16
	img = np.zeros((size+2, size+2), dtype=np.uint8) + 255

	layer_weights = 32, 16, 8, 4, 2

	assert len(datas) <= len(layer_weights)

	img_upd = img[1:-1,1:-1]
	img_qr0 = np.fromstring(payloads[0], dtype=np.uint8).reshape((size, size))
	img_upd[:] = img_qr0

	for idx_data, data in enumerate(payloads[1:]):
		img_qr = np.int32(np.fromstring(data, dtype=np.uint8).reshape((size, size)))
		weight = layer_weights[idx_data]
		comp = img_qr0 != img_qr
		compimg = np.ones_like(img_upd)
		compimg[comp] = weight
		cv2.imwrite("comp-%d.png" % idx_data, compimg)
		img_upd[img_upd >= 128] -= compimg[img_upd >= 128]
		img_upd[img_upd <  128] += compimg[img_upd <  128]
		cv2.imwrite("upd-%d.png" % idx_data, img)

	img = cv2.resize(img, ((size+2)*scale, (size+2)*scale), interpolation=cv2.INTER_NEAREST)
	cv2.imwrite("qrcode-luma.png", img)


def example_encode_chroma():

	datas = []

	datas.append("http://zougloub.eu")

	# some structure...
	extradata = {"x": 1.33}
	# more data with some string encoding

	datas.append("f" + enc_flt(1.33))

	# more data as a string
	datas.append("hello world!")

	assert len(datas) == 3, "Unimplemented"

	size, payloads = make_samesize(datas)

	scale = 16
	img = np.zeros((size+2, size+2), dtype=np.uint8) + 255

	img_upd = img[1:-1,1:-1]
	img_qr0 = np.fromstring(payloads[0], dtype=np.uint8).reshape((size, size))
	img_upd[:] = img_qr0
	img = cv2.merge((img, img, img))

	img_yuv = cv2.cvtColor(img, cv2.COLOR_RGB2YUV)

	for idx_plane in range(1, 3):
		info = payloads[idx_plane]
		img_qr = np.fromstring(info, dtype=np.dtype('B')).reshape((size, size))
		img_upd = img_yuv[1:-1,1:-1,idx_plane]
		img_upd[img_qr0 != img_qr] -= 80

	img = cv2.cvtColor(img_yuv, cv2.COLOR_YUV2RGB)
	img = cv2.resize(img, ((size+2)*scale, (size+2)*scale), interpolation=cv2.INTER_NEAREST)
	cv2.imwrite("qrcode-chroma.png", img)

if __name__ == '__main__':
	example_encode_chroma()
	example_encode_luma()