import os
from PIL import Image, ImageDraw,ImageTk
import argparse
import json
import time
import glob
import numpy as np
import datetime
import Tkinter
import moviepy.editor as mpy
import sys
import wave
	
FPS = 1

facialParameters = [
	"smile",
	"innerBrowRaise",
	"browRaise",
	"browFurrow",
	"noseWrinkle",
	"upperLipRaise",
	"lipCornerDepressor",
	"chinRaise",
	"lipPucker",
	"lipPress",
	"lipSuck",
	"mouthOpen",
	"smirk",
	#~ "attention",
	"eyeClosure",
	"eyeWiden",
	"cheekRaise",
	"lidTighten",
	"dimpler",
	"lipStretch",
	"jawDrop",
]

defaultFacialParameters = [
	"smile",
	"innerBrowRaise",
	"browRaise",
	"browFurrow",
	"noseWrinkle",
	"upperLipRaise",
	"lipCornerDepressor",
	"chinRaise",
	#~ "lipPucker",
	"lipPress",
	"lipSuck",
	"mouthOpen",
	#~ "smirk",
	#~ "attention",
	"eyeClosure",
	"eyeWiden",
	"cheekRaise", #
	"lidTighten",
	"dimpler", #
	"lipStretch",
	"jawDrop", #
]

parser = argparse.ArgumentParser(description='Parses opencv-webcam-demo json output files and collects statistics')
parser.add_argument('--frameOutput', '-o', required=True, help='directory to look for frames & json')
parser.add_argument('--targetDir', help='Directory in which files are stored')
parser.add_argument('--status', action='store_true', help='Keep status of last frame')
parser.add_argument('--stats', action='store_true', help='Show statistics for all frames')
parser.add_argument('--cutAllFaces', action='store_true', help='Cut out all faces from all frames')
parser.add_argument('--sum', action='store_true', help='Get total scores over all time')
parser.add_argument('--unique', action='store_true', help='Get most unique window')
parser.add_argument('--avg', action='store_true', help='Get most average window')
parser.add_argument('--json', action='store_true', help='Get output as JSON for graphs (graphs.php)')
parser.add_argument('--wav', '-w', required=False, default=False, help='split wave file into windows')
parser.add_argument('--disonant', action='store_true', help='Get most disonant faces over time')
parser.add_argument('--window-size', '-s', type=int, default=10, help='The nr of frames to group in one sliding window for analysis')
parser.add_argument("--params", "--metrics", "-p",  type=str, nargs='+', default=defaultFacialParameters, choices=facialParameters, help="The parameters used to calculate the statistics")

args = parser.parse_args()

faces = []
class Face:
	def __init__(self, frame, data):
		self.id = data['id']
		self.frame = frame # Frame class
		self.data = data # json data
		self.disonanceScore = None # a first attempt, can be deprecated?
		self.anomalyScore = None
	
	def getFaceImg(self):
		r = self.data['rect']
		return self.frame.getImg().crop((int(r['x']), int(r['y']), int(r['x']+r['w']), int(r['y']+r['h'])))
	
	def getCharacteristicVector(self, params):
		self.vector = [self.data[p] for p in params]
		return self.vector
	
	def setAnomalyScore(self, score):
		self.anomalyScore = score

class Frame:
	"""
	Everything for an analysed frame
	"""
	def __init__(self, outputPath, nr):
		self.outputPath = outputPath
		self.nr = nr
		self.name = "frame%06d" % nr
		self.jsonPath = os.path.join(outputPath, ("frame%06d" % (nr)) + ".json")
		self.imgPath = os.path.join(outputPath, self.name + ".jpg")
		self.faces = None # init with getFaces
	
	def getTime(self):
		"""
		@return datetime
		"""
		with Image.open(self.imgPath) as i:
			#~ print(i._getexif()[36867])
			return datetime.datetime.strptime(i._getexif()[36867], "%Y:%m:%d %H:%M:%S")
		#~ return datetime.datetime.fromtimestamp(os.path.getmtime(self.imgPath))
		
	def getJson(self):
		#~ try:
			with open(self.jsonPath) as fp:
				return json.load(fp)
		#~ except Exception as e:
			#~ # no json file yet?
			#~ return None
			
	def getImg(self, markFaces = False):
		img = Image.open(self.imgPath)
		if not markFaces:
			return img
		
		draw = ImageDraw.Draw(img)
		for f in self.faces:
			xy1 = (int(f.data['rect']['x']), int(f.data['rect']['y']))
			xy2 = (int(f.data['rect']['x'] + f.data['rect']['w']), int(f.data['rect']['y'] + f.data['rect']['h']))
			draw.rectangle([xy1, xy2], outline="#ff0000")
		return img
		
	def getFaces(self):
		if self.faces is None:
			j = self.getJson()
			
			self.faces = [Face(self, f) for f in j['faces']]
			faces.extend(self.faces)
		
		return self.faces
	
	def updateDisonanceScores(self):
		totalValence = 0.0
		totalFaces = 0
		for face in self.getFaces():
			totalValence += face.data['valence']
			totalFaces += 1

		if totalFaces == 0:
			return
		
		avgValence = totalValence / totalFaces
		
		for face in self.getFaces():
			face.disonanceScore = abs(face.data['valence'] - avgValence)
	
	def getAverageV(self, params):
		vectors = [face.getCharacteristicVector(params) for face in self.getFaces()]
		vAvg = np.mean(vectors, axis=0)
		return vAvg
	
	def getStdDev(self, params):
		"""
		Get standard deviation of the faces within the frame for given params
		"""
		vectors = [f.getCharacteristicVector(params) for f in self.getFaces()]
		return np.std(vectors)
	
	def updateAnomalyScores(self, params):
		vAvg = self.getAverageV(params)
		for face in self.getFaces():
			face.setAnomalyScore(np.linalg.norm(face.getCharacteristicVector() - vAvg))
		
	
	def exists(self):
		return os.path.exists(self.jsonPath) and os.path.exists(self.imgPath)

frames = {}

class Window:
	def __init__(self, frameSubset):
		"""
		Init a sliding window for given Frame-s
		"""
		self.frames = frameSubset
		self.deviation = None
		self.standardDeviation = None
		
	def getStartTime(self):
		"""get time of first frame in window. Returns datetime"""
		return self.frames[0].getTime()
		
	def getEndTime(self):
		"""get time of last frame in window. Returns datetime"""
		return self.frames[-1].getTime()
	
	def getFaces(self):
		faces = []
		for frame in self.frames:
			faces.extend(frame.getFaces())
		return faces
	
	def getStdDev(self, params):
		"""
		Get standard deviation of the faces within the window for given params
		"""
		vectors = [f.getCharacteristicVector(params) for f in self.getFaces()]
		return np.std(vectors)
	
	def getAverageV(self, params):
		vectors = [f.getCharacteristicVector(params) for f in self.getFaces()]
		vAvg = np.mean(vectors, axis=0)
		return vAvg
	
	@staticmethod
	def createWindows(windowSize, frames):
		"""
		Give a full list of frames and tunrn it into a collection of sliding windows
		"""
		frames = sorted(frames.items(), key=lambda f: f[0])
		frames = [f[1] for f in frames]
		windows = []
		windowCount = len(frames) / windowSize
		#~ windowCount = len(frames) - windowSize + 1
		if windowCount < 1:
			raise Exception("Not enough frames ({}) for a window of size {}".format(len(frames), windowSize))
		
		for offset in range(0, windowCount):
			frameSubset = frames[offset*windowSize:(offset+1)*windowSize]
			#~ frameSubset = frames[offset:offset+windowSize]
			windows.append(Window(frameSubset))
		
		return windows

class WindowCollection:
	def __init__(self, windowSize, frames):
		self.windows = Window.createWindows(windowSize, frames)
		self.frames = frames
		#~ self.faces = [face for face in frame.getFaces() for frame in frames]
	
	#~ def getMostWindowsClosestToMedian(self, nr = 5):
		#~ """
		#~ Get windows with the faces closest to the median
		#~ """
		#~ self.faces
	
	def getWindowVectors(self, params):
		return [window.getAverageV(params) for window in self.windows]
	
	def getWindowsByDeviation(self, params):
		vectors = self.getWindowVectors(params)
		vAvg = np.mean(vectors, axis=0)
		#~ diffs = [numpy.linalg.norm(v-vAvg) for v in vectors]
		#~ min_index, min_value = min(enumerate(diffs), key=lambda p: p[1])
		#~ max_index, max_value = max(enumerate(diffs), key=lambda p: p[1])
		
		return sorted(self.windows, key=lambda w: np.linalg.norm(w.getAverageV(params)-vAvg))
	
	def getDeviationForWindow(self, window, params):
		vectors = self.getWindowVectors(params)
		vAvg = np.mean(vectors, axis=0)
		return np.linalg.norm(w.getAverageV(params)-vAvg)
		
	def getUniqueWindows(self, params, nr=5):
		windows = self.getWindowsByDeviation(params)
		return windows[0: nr]
	
	def getMostAvgWindows(self, params, nr=5):
		windows = self.getWindowsByDeviation(params)
		windows.reverse()
		return windows[0:nr]
	
	def getMostContrastingWindows(self, params, nr=5):
		sortedWindows = sorted(self.windows, key=lambda w: w.getStdDev(params), reverse=True)
		return sortedWindows[0:nr]
		
def loadFrames(frameDir):
	global frames
	nr = 2
	nextFrame = Frame(frameDir, nr)
	# TODO; make threaded and infinite loop that updates global frames
	while nextFrame.exists():
		frames[nr] = nextFrame
		nr+=1
		nextFrame = Frame(frameDir, nr)
	return frames

def getLastFrame(frameDir):
	jsons = sorted(glob.glob(os.path.join(frameDir, "*.json")))
	if len(jsons):
		lastJson = jsons[-1]
		lastNr = int(lastJson[-11:-5])
		frame = Frame(frameDir, lastNr)
		return frame
	return None

def cutOutFaces(frame, targetDir):
	for faceNr, face in enumerate(frame.getFaces()):
		print(faceNr, face)
		img = face.getFaceImg()
		faceImgPath = os.path.join(targetDir, frame.name + "-%s.jpg" % face.id)
		print(faceImgPath)
		img.save(faceImgPath)
	pass

def validateJsonTimes():
	lastTime = None
	for frameNr, frame in loadFrames(args.frameOutput).items():
		thisTime = frame.getJson()['t']
		#print(frameNr, thisTime)
		if not (lastTime is None) and lastTime > thisTime:
			sys.stderr.write("ERRROR!! Time error at %s. Restarted scanner there?\n" % frameNr)
		lastTime = thisTime

def sumEmotions():
	total = 0.
	summed = 0.
	items = 0
	for frameNr, frame in loadFrames(args.frameOutput).items():
		for face in frame.getFaces():
			total += abs(face.data['valence'])
			summed += face.data['valence']
			items += 1
	
	average = summed / items
	print ("Total emotion %d, positivity score %d (average: %s)" % (total, summed, average))
	
def getMostDisonant(nr = 5):
	for frameNr, frame in loadFrames(args.frameOutput).items():
		frame.updateDisonanceScores()
	faces.sort(key=lambda x: x.disonanceScore, reverse=True)
	
	mostDisonantFaces = faces[:nr]
	for face in mostDisonantFaces:
		print("Frame %d, face %d, score %d, valence %d" % (face.frame.nr, face.id, face.disonanceScore, face.data['valence']))
		face.getFaceImg().show()

def getAnomalies(params, nr = 5):
	for frameNr, frame in loadFrames(args.frameOutput).items():
		frame.updateAnomalyScores(params)
	faces.sort(key=lambda x: x.anomalyScore, reverse=True)
	
	anomalies = faces[:nr]
	for face in anomalies:
		print("Frame %d, face %d, score %d" % (face.frame.nr, face.id, face.anomalyScore))
		#~ getCharacteristicVector
		face.getFaceImg().show()

def printFrameStats(frame, params):
	os.system('clear')
	print(time.time())
	print( ("Nr: %d" % frame.nr).ljust(40) + ("t: {}".format(frame.getJson()['t'])) )
	#~ print
	faces = frame.getFaces()
	print("Faces: %d" % len(faces))
	
	if len(faces) < 1:
		return
	
	print " ".ljust(20), "0%".rjust(13), "q1".rjust(13), "median".rjust(13), "q3".rjust(13), "100%".rjust(13)
	
	for p in params:
		q0 = np.percentile(np.array([f.data[p] for f in faces]),0)
		q1 = np.percentile(np.array([f.data[p] for f in faces]),25)
		q2 = np.percentile(np.array([f.data[p] for f in faces]),50)
		q3 = np.percentile(np.array([f.data[p] for f in faces]),75)
		q4 = np.percentile(np.array([f.data[p] for f in faces]),100)
		print p.ljust(20), ("%f%%" % q0).rjust(13), ("%f%%" % q1).rjust(13),("%f%%" % q2).rjust(13),("%f%%" % q3).rjust(13),("%f%%" % q4).rjust(13)
	
	
	#~ TODO: speaker stats
	
	frame.updateDisonanceScores()
	dissonantFace = max(faces,key=lambda f: f.disonanceScore)
	#~ dissonantFace.getFaceImg()
	

def monitorStatus(frameDir, params):
	while True:
		frame = getLastFrame(frameDir)
		if not frame is None:
			printFrameStats(frame, params)
			
		# don't check too often	
		time.sleep(.5)

def printStats(frameDir, params):
	frames = loadFrames(frameDir)
	vAvg = np.mean([f.getAverageV(params) for i,f in frames.items()], axis=0)
	for nr, frame in frames.items():
		max_index, max_value = max(enumerate(frame.getAverageV(params)), key=lambda a: a[1])
		distance = np.linalg.norm(vAvg - frame.getAverageV(params))
		print("{:06d}: {: >4} faces (d: {: 6.2f}, internal s {: 6.2f}, max: {})".format(frame.nr, len(frame.getFaces()),distance, frame.getStdDev(params), params[max_index]))
	
	print("{} frames".format(len(frames)))

def playWindowStopmotion(window, params):
	"""
	Play a set of sliding window frames as stop motion video
	"""
	root = Tkinter.Tk()
	root.geometry('%dx%d+%d+%d' % (1000,1000,0,0))	
	canvas = Tkinter.Canvas(root,width=1000,height=1000)
	canvas.pack()
	
	old_label_image = None
	for frame in window.frames:
		image = frame.getImg(markFaces=True)
		basewidth = 1000
		wpercent = (basewidth / float(image.size[0]))
		hsize = int((float(image.size[1]) * float(wpercent)))
		image = image.resize((basewidth, hsize), Image.ANTIALIAS)
		
		tkpi = ImageTk.PhotoImage(image)
		canvas.delete("IMG")
		imagesprite = canvas.create_image(500,500,image=tkpi, tags="IMG")
		root.update()
		time.sleep(1)

def createWindowVideo(window, params, targetDir):
	global FPS
	basewidth = 1920
	i=0
	images = []
	for frame in window.frames:
		image = frame.getImg(markFaces=True)
		wpercent = (basewidth / float(image.size[0]))
		hsize = int((float(image.size[1]) * float(wpercent)))
		image = image.resize((basewidth, hsize), Image.ANTIALIAS)
		imgName = os.path.join("/tmp/","renderframe{:06d}.jpg".format(i))
		print("(VIDEO) Create frame: {}".format(imgName))
		image.save(imgName)
		images.append(imgName)
		i+=1
	
	clip = mpy.ImageSequenceClip(images, fps=FPS)
	#~ audio = mpy.AudioFileClip(filename="")
	#~ audio.set_start(...)
	#~ audio.set_duration(len(images)/FPS)
	#~ clip = clip.set_audio(audio)
	
	targetName = "video_" + "-".join(params) + ".mp4"
	print("(VIDEO) Write video: {}".format(targetName))
	clip.write_videofile(os.path.join(targetDir,targetName), fps=FPS)

def getGraphData(frames, params, window_size, outputNr):
	"""
	Get data to generate graph with
	"""
	collection = WindowCollection(window_size, frames)
	windows = collection.windows
	jsonOutput = {'windows':[], 'outputNr': outputNr}
	for window in windows:
		w = {
			'start': window.getStartTime().strftime('%Y-%m-%d %H:%M:%S'),
			'end': window.getEndTime().strftime('%Y-%m-%d %H:%M:%S'),
			'startFrame': window.frames[0].name,
			'params': {},
			'param_values': {},
		}
		
		pVs = []
		for param in params:
			paramV = window.getAverageV([param])
			value = paramV[0]
			pVs.append(value)
			w['params'][param] = value
			#~ values = []
			#~ for frame in window.frames:
				#~ for face in frame.getFaces():
					#~ values.append(face.data[param])
			#~ w['param_values'][param] = values
		
		w['avg'] =  np.mean(pVs)
		w['stddev'] = window.getStdDev(params)
		jsonOutput['windows'].append(w)
	
	jsonOutput['range'] = [
		jsonOutput['windows'][0]['start'],
		jsonOutput['windows'][-1]['end'],
	]
	return jsonOutput
	
		
def splitAudio(frames, window_size, srcWav, targetDir):
	"""
	Get audio cllips
	"""
	global FPS
	if not os.path.exists(srcWav):
		raise Exception("Wav does not exist")
	if not os.path.exists(targetDir):
		raise Exception("Target dir does not exist")
		
	collection = WindowCollection(window_size, frames)
	origAudio = wave.open(srcWav,'r')
	
	frameRate = origAudio.getframerate()
	nChannels = origAudio.getnchannels()
	sampWidth = origAudio.getsampwidth()
	
	windows = collection.windows
	for i, window in enumerate(windows):
		print("Window {}, {}".format(i, window.frames[0].name))
		start = i * window_size * FPS
		end = i+1 * window_size * FPS
		origAudio.setpos(start*frameRate)
		print(origAudio.tell())
		chunkData = origAudio.readframes(int((end-start)*frameRate))
		print(origAudio.tell())
		targetFile = os.path.join(targetDir, window.frames[0].name + '.wav')
		print("\t{}".format(targetFile))
		chunkAudio = wave.open(targetFile,'w')
		chunkAudio.setnchannels(nChannels)
		chunkAudio.setsampwidth(sampWidth)
		chunkAudio.setframerate(frameRate)
		chunkAudio.writeframes(chunkData)
		chunkAudio.close()

#~ def createGraphPage(windows, params, targetFilename):
	#~ jsonOutput = {'windows':[]}
	#~ for window in windows:
		#~ window
		

validateJsonTimes()
	
if args.sum:
	sumEmotions()

if args.disonant:
	getMostDisonant()
	
if args.cutAllFaces:
	faceDir = os.path.join(args.frameOutput, 'faces')

	if not os.path.exists(faceDir):
		os.mkdir(faceDir)
	for frameNr, frame in loadFrames(args.frameOutput).items():
		cutOutFaces(faceDir)

if args.unique:
	collection = WindowCollection(args.window_size, frames)
	windows = collection.getUniqueWindows(args.params)
	#~ print(windows)
	if args.targetDir:
		createWindowVideo(windows[0], args.params, args.targetDir)
	else:
		playWindowStopmotion(windows[0], args.params)

if args.avg:
	collection = WindowCollection(args.window_size, frames)
	windows = collection.getMostAvgWindows(args.params)
	#~ print(windows)
	if args.targetDir:
		createWindowVideo(windows[0], args.params, args.targetDir)
	else:
		playWindowStopmotion(windows[0], args.params)

if args.status:
	monitorStatus(args.frameOutput, args.params)
	
if args.stats:
	printStats(args.frameOutput, args.params)

if args.json:
	nr = int(args.frameOutput[-1])
	print json.dumps(getGraphData(frames, args.params, args.window_size, nr))
	
if args.wav:
	print splitAudio(frames, args.window_size, args.wav, args.targetDir)