import picamera
import io, os
import datetime
import csv
from subprocess import Popen, PIPE
from PIL import Image
import numpy as np
import cPickle as pickle
import requests
import time
import thread
from websocket import create_connection
import logging
import json
from io import BytesIO

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger('portraits')

camera = picamera.PiCamera()
camera.rotation = 180
camera.resolution = (1920, 1080)
# camera.resolution = (1280, 720)

outputResolution = (1000, 1000)

# the binary genders as outputted by Affectiva
genders = ['male', 'female', 'unknown']
perspectives = ['side', 'front']
gender_perspectives = [g+"_"+p for p in perspectives for g in genders]
# we also want dedicated composites for the perspectives
gender_perspectives.extend(perspectives)

curdir = os.path.dirname(os.path.abspath(__file__))
tmpimage = '/tmp/piimage.jpg'
tmpimageResults = '/tmp/piimage.csv'
cmd = [
	os.path.join(curdir, 'build/video-demo/video-demo'),
	'--input', tmpimage,
	'--data', os.path.join(curdir, 'affdex-sdk/data'),
	'--draw', '0',
	'--numFaces', '20',
]

# without these vars video-demo yields a segmentation fault
environment = {
	'LC_LANG': 'en_GB.UTF-8',
	'LD_PRELOAD': '/usr/lib/arm-linux-gnueabihf/libopencv_core.so.2.4',
}

updateGoesOk = True;

with open(os.path.join(curdir, 'uploadkey.json')) as fp:
	uploadkey = json.load(fp)

def updateStats(type, name, count, image_filename):
	params = {
		'type': type,
		'name': name,
		'time': int(time.time()),
		'case_count': int(count),
		'key': uploadkey,
	}
	# try:
	# 	ws = create_connection("ws://localhost:8888/ws")
	# 	js = json.dumps({
	# 		'type': type,
	# 		'name': name,
	# 		'img_src': os.path.basename(image_filename),
	# 		'case_count': int(count),
	# 	})
	# 	ws.send(js)
	# except Exception as e:
	# 	raise

	url = 'https://artstats.rubenvandeven.com/composites/views.php'
	try:
		if count % 10 == 0:
			# only send every one in x image, so that the server never can
			# retrace _exact_ faces by comparing the sent frames.
			with open(image_filename) as fp:
				im = Image.open(fp)
				bytesImg = BytesIO()
				im.save(bytesImg, format='jpeg')
				bytesImg.seek(0)

				logger.info('send request including image')
				r = requests.post(
					url ,
					files={'image': bytesImg},
					params=params
				)
		else:
			logger.info('send request')
			r = requests.post(
				url,
				params=params
			)
		updateGoesOk = True
	except Exception as e:
		updateGoesOk = False
		logger.critical("Error when updating statistics")
		logger.exception(e)

class CompositeImage:
	def __init__(self, name, resolution):
		self.name = name
		self.count = 0
		self.resolution = resolution
		self.image = np.zeros((resolution[0],resolution[1],3))
		# use state to determine whether a save is necessary
		self.state_dirty = True

	def addFace(self, img):
		img_a = np.array(img.resize(self.resolution))
		self.count += 1

		self.image = (self.image * (self.count - 1)/float(self.count) + img_a / float(self.count))
		self.state_dirty = True

	def restore(self, i, dir):
		'''
		Restore from pickle nr
		'''
		self.count = i
		name = self.get_frame_filename(self.count)
		img_file = os.path.join(dir, name)
		logger.info("\trestore {}".format(img_file))
		self.image = np.array(Image.open(img_file)).astype('float64')
		self.state_dirty = False

	def get_image(self):
		return Image.fromarray(self.image.astype('uint8'),'RGB')

	def get_frame_filename(self, i):
		return "{}-{}x{}-{}.png".format(self.name, self.resolution[0], self.resolution[1], i)

	def get_current_filename(self):
		return self.get_frame_filename(self.count)

	def save_image(self, dir):
		if self.state_dirty is False:
			# don't save if no changes are made since last save
			return

		name = self.get_current_filename()
		filename = os.path.join(dir, name)
		self.get_image().save(filename)

		thread.start_new_thread( updateStats, ('gender', self.name, self.count, filename) )

		self.state_dirty = False

class CompositeCollection:
	"""
	Store/save the composite images
	"""
	def __init__(self, names, size, target_dir = None):
		self.id = "{}-{}x{}".format("-".join(names), size[0], size[1])
		self.names = names
		self.size = size
		self.target_dir = os.path.dirname(os.path.abspath(__file__)) if target_dir is None else target_dir

		self.load()

	def get_pickle_filename(self):
		return os.path.join(self.target_dir, self.id + ".p")

	def get_json_filename(self):
		return os.path.join(self.target_dir, "composites.json")

	def load(self):
		pickle_file_name = self.get_pickle_filename()
		# if os.path.exists(pickle_file_name):
		composites = {}
		try:
			with open( pickle_file_name, "rb" ) as fp:
				data = pickle.load( fp )
			for name in data['c']:
				composites[name] = CompositeImage(name, self.size)
				composites[name].restore( data['c'][name], self.target_dir)
		except Exception as e:
			logger.info("Create new composite", e)
			for name in self.names:
				composites[name] = CompositeImage(name, self.size)

		self.composites = composites

	def save(self):
		data = { 'size' : self.size, 'c': {} }
		for name in self.composites:
			data['c'][name] = self.composites[name].count

		with open( self.get_pickle_filename(), "wb" ) as fp:
			logger.info("Save", data)
			pickle.dump( data, fp )

		self.save_json()

	def save_json(self):
		"""
		Save statistics as json
		"""
		data = {}
		for name in self.composites:
			data[name] = {
				'count': self.composites[name].count,
				'img': self.composites[name].get_current_filename()
			}

		with open( self.get_json_filename(), "w" ) as fp:
			logger.debug("Json to {}".format(self.get_json_filename()))
			json.dump(data, fp)

		ws = create_connection("ws://localhost:8888/ws")
		ws.send("update")



	def save_img(self, name):
		self.get(name).save_image(self.target_dir)

	def get_as_percentages(self, precision = 3):
		total = sum([c.count for c in self.composites])
		percentages = {}
		if total < 1:
			# assert: in the beginning, we were all made equal
			for c in self.composites:
				percentages[c.name] = round(100 / len(self.composites), precision)
		else:
			for c in self.composites:
				percentages[c.name] = round(100 *  (c.count / total), precision)
		return percentages

	def get(self, name):
		return self.composites[name]

	def clean(self):
		for name in self.names:
			c = self.get(name)
			start = max(0, c.count - 10)
			end = max(0, c.count - 5)
			for i in range(start, end):
				filename = os.path.join(self.target_dir, c.get_frame_filename(i))
				if os.path.exists(filename):
					logger.info("Clean {}".format(filename))
					os.unlink(filename)



def append_face(row, image, composites):
	# degrees to distinguish side (as we will never be able to use 90 :-( )
	# Plus, we now have a wide angle lens.
	suffix = 'side' if abs(float(row['yaw'])) > 12 else 'front'
	names = [suffix, "{}_{}".format(row['gender'], suffix)]

	compositesToUse = []
	for name in names:
		if name not in composites.names:
			return
		compositesToUse.append(composites.get(name))

	# TODO: matrix transform the image, to skew the face into being a flat-ish surface
	# This might yield less blurry composites

	# crop image, bt keep it bigger than the found face
	grow_x = .2 # in every direction, so .2 becomes 1.4 * width
	grow_y = grow_x

	face_w = int(row['width'])
	face_h = int(row['height'])
	face_x = int(row['x'])
	face_y = int(row['y'])

	# we go square:
	size_x = max(face_w, face_h) * (1 + grow_x * 2)
	size_y = size_x

	dx = (face_w - size_x) / 2
	dy = (face_h - size_y) / 2

	# PIL.Image handles cropping outside the canvas by filling with black/transparent
	x = face_x + dx
	y = face_y + dy
	logger.debug('crop')
	i = image.crop((x,y, x + size_x, y + size_y))

	if suffix == 'side' and float(row['yaw']) < 0:
		logger.debug('\tflip')
		i = i.transpose(Image.FLIP_LEFT_RIGHT)

	for composite in compositesToUse:
		logger.debug('add')
		composite.addFace(i)
		logger.debug('added')

composites = CompositeCollection(gender_perspectives, outputResolution, os.path.join(curdir, 'output'))

while True:
	start = datetime.datetime.utcnow()
	# stream = io.BytesIO()
	camera.capture(tmpimage, format='jpeg')
	process = Popen(cmd, env=environment)
	process.wait()
	img = Image.open(tmpimage)
	os.unlink(tmpimage)
	with open(tmpimageResults) as csvfile:
		logger.debug("open csv")
		data = csv.DictReader(csvfile)
		faces = 0
		for row in data:
			if row['faceId'] == 'nan':
				# not a valid face
				continue
			faces += 1
			logger.info("append face")
			append_face(row, img, composites)

		if faces > 0:
			logger.info("save :-)")
			for name in composites.names:
				logger.info("\tsave img '{}'".format(name))
				c = composites.save_img(name)
			# save pickle after images, so they can be restored
			composites.save()
			composites.clean()