#!/usr/bin/python
# -*- coding: utf-8 -*-

from nijmetrics import Application, Error, Database
import pg
import json
import urllib, urllib2
import re
import itertools
import cPickle as pickle
import zlib
import nijmetrics
import PIL

'''
Format of the extracted objects (saved as .pickle files):
<fooID> denotes a key that *is a* fooID (an ID of a foo) - there can be 0..n of them
[type] denotes a value of the given type

<exerciseID>
	exercise
		id [int]
		name [string]
	compressed [boolean]
	mark
		href [string]
		width [int]
		height [int]
		resolution [int]
	print
		href [string]
		width [int]
		height [int]
		resolution [int]
	consensus
		qmaps
			<quality_level> [qmap]
		masks
			<threshold> [qmap]
		masks_raw
			<threshold> [qmap]
		numcontributors [int]
	users
		<userID>
			username [string]
			status [string]
			qmaps
				<quality_level> [qmap]
			qmap_error [None or string describing the error]
			primary_quality [int]
			analysis [annotations]
			comparison [annotations]

<quality_level> is a constant in (1,2,3) - that is, 'high', 'medium' and 'low' quality
[qmap] type is an array of integer numbers of length mark['width']*mark['height']
		If 'compressed' is True, this array is stringified using repr() and compressed
		using zlib
		
'''

class Extractor(Application):
	
	def __init__(self, db_url, maskThresholds, exfilter=None, userfilter=None, outputprefix=None, compress=False, limit=None):
		Application.__init__(self)
		self.db_url = db_url
		self.exfilter = exfilter
		self.userfilter = userfilter
		self.outputprefix = outputprefix
		self.compress = compress
		self.maskThresholds = maskThresholds
		self.limit = limit
		
		if type(self.maskThresholds) not in (type([]), type(())):
			self.maskThresholds = (self.maskThresholds,)
			
		for th in self.maskThresholds:
			if th < 0.0 or th > 1.0:
				raise Error("Invalid quality consensus mask threshold: %f. Value must be in [0..1]" % self.maskThreshold)
			
		# The ANSI/NIST color scale (which can be found in pianos4_constants.js). Used to detect
		# ANSI/NIST quality maps, and convert if needed.
		self.ansiNistTable = {
			'#00ffff': {'id': 1, 'color': '#00ffff', 'label': 'Definitive Pores', 'conv': 1},
			'#3366ff': {'id': 2, 'color': '#3366ff', 'label': 'Definitive Ridge Edges', 'conv': 1},
			'#01ff02': {'id': 3, 'color': '#01ff02', 'label': 'Definitive Minutiae', 'conv': 2},
			'#ffff00': {'id': 4, 'color': '#ffff00', 'label': 'Debatable Minutiae', 'conv': 2},
			'#fe0000': {'id': 5, 'color': '#fe0000', 'label': 'Debatable Ridge Flow', 'conv': 3},
			'#010101': {'id': 6, 'color': '#010101', 'label': 'Background', 'conv': 3},
		}
		self.standardTable = {
			'#2ca02c': {'id': 1, 'color': '#2ca02c', 'label': 'High'},
			'#ff7f0e': {'id': 2, 'color': '#ff7f0e', 'label': 'Medium'},
			'#d62728': {'id': 3, 'color': '#d62728', 'label': 'Low'},
		}

		self.maskFilter = None


	def run(self):
		try:
				
			dbargs = (self.db_url)
			self.db = Database(self.db_url)
			self.db.connect()
			
			# The global accumulation of all quality maps. First-level keys are exercise ids, they
			# contain dicts. Second-level keys are the quality levels (1,2 or 3). The values are
			# arrays of size imageWidth*imageHeight with integers accumulating the values of each
			# user's maps.
			maps = {}
			consensusCompressed = False
			if self.compress:
				print "Using compression for output quality maps"
			
			self.assignments = self.collectAssignments(self.userfilter, self.exfilter, self.limit)

			
			progress = 0
			for a in self.assignments:
				print "%s - %s  (%.2f%%)" % (a['exname'],a['username'],(float(progress)/float(len(self.assignments))*100.0))
				progress += 1

				# Create the entry in the 'maps' global object if this is the first time we see
				# this exercise
				if not a['exid'] in maps:
					maps[a['exid']] = {'exercise':{'id':a['exid'],'name':a['exname']}, 'consensus':{'qmaps':{}, 'masks_raw':{}, 'masks':{}}, 'users':{},'compressed':self.compress, 'mark': self.getImage('mark', a['exid']), 'print': self.getImage('print', a['exid'])}

				# Handy shortcut variables for future use
				markDims = (maps[a['exid']]['mark']['width'], maps[a['exid']]['mark']['height'])

				# Create the entry in the 'users' portion of the results
				maps[a['exid']]['users'][a['userid']] = { 'qmaps':{}, 'qmap_error':None, 'primary_quality':None, 'username':a['username'], 'status':a['status'] }

				r = self.db.query("SELECT * FROM result WHERE assignment = "+str(a['id'])+" AND step = (SELECT id FROM resultstep WHERE short = 'ANALYSIS') ORDER BY updatedon DESC LIMIT 1;")
				if len(r) == 1:
					# Since users are dumb and did not use the "compute final quality map" feature of
					# PiAnoS, we'll have to do it for them. Unfortunately, it is not as straightforward
					# for 2 reasons:
					# 1) the feature relies on C++ and used through an API. The rather technical
					#    side of the problem is easily addressed with urllib, but the real difficulty
					#    lies in the need to re-code parts of the feature that are in Javascript.
					# 2) the real problem: the less-than-ideal state of the feature in PiAnoS makes
					#    it necessary to specify a so-called "primary quality" - the quality area that
					#    is mainly present on the trace, which is drawn first by the user, with higher/lower
					#    qualities that are drawn above. Thus the primary quality area is "cut" by other
					#    areas. We do not have this information, we'll have to guess it.
					#    The chosen heuristic, absolutely not proven and bullshittedly elected by my guts,
					#    is to rely on the global areas (in cm^2, that is): the larger the zone, the higher
					#    the probability it is the "primary quality". So we need a first step to compute the
					#    overall areas of all quality levels, then choose the largest and re-extract the
					#    areas using getQualityMaps(), this time specifying the primary quality.
					# If users were not dumb, i.e. they actually used the feature, we'll certainly(*) find
					# holes in their quality maps. So if there are already holes, assume the job is done
					# and do not guess again - just use the quality maps "as is".
					# (*) well, this is debatable, but my guts tell me we'll find holes if they did. Although
					# it is not necessary, it is highly likely. Or quite likely. Or simply probable. Maybe we
					# should consider calculating likelihood ratios for that, so we can take the right decision.
					# OOops, no. So we can know how wrong we are when we take a given decision.
					
					obj = json.loads(r[0]['data'])
					
					# Store the overall analysis result object for this user
					maps[a['exid']]['users'][a['userid']]['analysis'] = obj
						
					# All areas are in the obj['mark']['areas'] dictionary - we want to keep only those
					# that are quality maps (PiAnoS stores them as areas with a special flag)
					try:
						areas = filter(lambda x: 'type' in x and x['type'] == 'quality', obj['mark']['areas'].values())
						if len(areas) == 0:
							maps[a['exid']]['users'][a['userid']]['qmap_error'] = 'No quality maps'
							continue
					except AttributeError,err:
						print "WARNING:",err
						print "NOTE: the offending thing that should be a dict of quality maps and isn't is:"
						print repr(obj['mark']['areas'])
						maps[a['exid']]['users'][a['userid']]['qmap_error'] = 'No quality maps (invalid format)'
						continue
						
					# Check that the quality has been entered using Standard mode, and not ANSI/NIST. If
					# it was ANSI/NIST, it needs a conversion.
					# BONUS: there is no clear indication of the actual quality scale that was used, only
					#        2 hints may (seriously, hopefully) drive us in the good direction: the Standard
					#        scale has only levels 1,2 and 3, and the colors are all different. The level is
					#        not enough since we can have an ANSI/NIST map with only the three levels (if 
					#        nothing was drawn with levels 4,5 and 6). But the color is enough: there is no
					#        overlap, so detecting a single ANSI/NIST color is sufficient.
					# TODO: slap me
					hasAnsiNistColors = False
					hasStandardColors = False
					for x in areas:
						if x['color'] in self.ansiNistTable:
							hasAnsiNistColors = True
						if x['color'] in self.standardTable:
							hasStandardColors = True
						
					if not hasAnsiNistColors and not hasStandardColors:
						print "WARNING: unknown quality levels found in the annotations, skipping"
						print "NOTE: this is very weird and should never happen, please inspect"
						maps[a['exid']]['users'][a['userid']]['qmap_error'] = 'unknown quality levels found in the annotations'
						continue
						
					if hasAnsiNistColors and hasStandardColors:
						print "WARNING: use of mixed quality scales detected, skipping"
						print "NOTE: while not expected, the user should not be blamed since PiAnoS allows that (but shouldn't, indeed)"
						maps[a['exid']]['users'][a['userid']]['qmap_error'] = 'use of mixed quality scales detected'
						continue
						
					if hasAnsiNistColors:
						areas = self.convertQualityScale(areas)
						
					# Check if the final quality map feature has been used already.
					# Note that we could anyway re-calculate the final quality, but since the we are
					# detecting the primary quality using heuristics, we might do more harm than good.
					hasUsedFeature = False
					for area in areas:
						if 'holes' in area and area['holes']:
							hasUsedFeature = True
							maps[a['exid']]['users'][a['userid']]['primary_quality'] = -1
							break
					
					try:					
						if not hasUsedFeature:
							# Get the areas, without any primary quality specified for the moment.
							array, mapping = self.getQualityMaps(areas)
							
							# Make the first API call to get the global areas of all quality levels
							result = self.callCppJsonApi('calcAreas', {'areas':array})
							# Sort the list so we find the largest zone - our chosen primary quality level
							tuples = [(int(x[0]),x[1]) for x in result['areas'].items()]
							stuples = sorted(tuples, key=lambda x: x[1], reverse=True)
							primaryQuality = mapping[stuples[0][0]] # very logical, indeed
							#print "Primary quality found to be",primaryQuality
							maps[a['exid']]['users'][a['userid']]['primary_quality'] = primaryQuality
							
							# Now calculate the actual quality map
							array, mapping = self.getQualityMaps(areas, primaryQuality)
							result = self.callCppJsonApi('intersect', {'areas':array})
							
							# And update the areas in the result object
							areas = []
							id = 1
							for layer,polyset in result['areas'].iteritems():
								for poly in polyset:
									qual = mapping[int(layer)]
									area = {'id':id, 'type':'quality', 'color':'#000000', 'value':qual, 'points':poly['outer'], 'holes':poly['rings']}
									areas.append(area)
									id += 1
						
						# Send the quality maps to the C++ side to calculate the rasters
						array, mapping = self.getQualityMaps(areas)
						result = self.callCppJsonApi('rasterAreas', {'areas':array}, {'width':markDims[0],'height':markDims[1]})
					
					except Error,err:	
						print "WARNING:",err
						print "Skipping result"
						maps[a['exid']]['users'][a['userid']]['qmap_error'] = str(err)
						continue
						
					map = {}
					for x,y in result['areas'].iteritems():
						# Results are returned using a simple run-length encoding, so they need decoding
						stringres = self.decodeRLE(y)
						# Now we have a big string with spaces and dots. Convert that to a list of ints
						map[mapping[int(x)]] = [1 if k=='.' else 0 for k in stringres]
					
						
					# Maps for this exercise. Indices are the quality levels.
					#print "Summing the values"
					consensus = maps[a['exid']]['consensus']['qmaps']
					for quality,array in map.iteritems():
						if not quality in consensus:
							qexmap = [0]*(markDims[0]*markDims[1])
							consensus[quality] = qexmap
						else:
							qexmap = consensus[quality]
						
						assert len(qexmap) == (markDims[0]*markDims[1]) == len(array)
						
						# This sums the user's map with the consensus (for this quality level)
						# See http://stackoverflow.com/questions/12215766/python-element-wise-in-place-addition
						qexmap[:] = ((x+y) for x,y in itertools.izip(qexmap,array))
						
						# Store the user's map as well - we'll need it to calculate the correlation between
						# the user's map and the consensus
						if self.compress:
							maps[a['exid']]['users'][a['userid']]['qmaps'][quality] = zlib.compress(repr(array))
						else:
							maps[a['exid']]['users'][a['userid']]['qmaps'][quality] = array
				
				
				r = self.db.query("SELECT * FROM result WHERE assignment = "+str(a['id'])+" AND step = (SELECT id FROM resultstep WHERE short = 'COMPARISON') ORDER BY updatedon DESC LIMIT 1;")
				if len(r) == 1:
					assert a['exid'] in maps
					obj = json.loads(r[0]['data'])
					
					# Store the overall comparison result object for this user
					maps[a['exid']]['users'][a['userid']]['comparison'] = obj
			
			
			if len(self.assignments):
				print "Finished extraction (%.2f%%)" % ((float(progress)/float(len(self.assignments))*100.0))

			# Calculate the mask for this consensus. The mask will have 1s where at least 75% of users
			# have annotated a quality (whatever the quality) and 0s elsewhere
			# We first gather the number users that contributed to the consensus
			print "Generating consensus maps"
			for exid,exdata in maps.iteritems():
				print " ",a['exname'],"..."
				
				# Get the number of users that actually annotated quality (quality maps that this script 
				# could understand, that is)
				numcontributors = 0
				for uid,udata in exdata['users'].iteritems():
					if udata['qmap_error'] is None:
						numcontributors += 1
				exdata['consensus']['numcontributors'] = numcontributors
				assert 'mark' in exdata
				
				# Generate the accumulation map: for each pixel, the number of persons who covered it
				# using either quality zone (we would expect each pixel to be covered exactly 0 or 1 times
				# but results show that some user manage to cover the same pixel with 2 zones. This is
				# tolerated and counted 2 times because it does not occur that much and has a very low
				# impact)
				accum = [0]*(exdata['mark']['width']*exdata['mark']['height'])
				for quality in exdata['consensus']['qmaps']:
					qm = exdata['consensus']['qmaps'][quality]
					assert len(accum) == len(qm)
					for i in range(len(accum)):
						accum[i] += qm[i]
						if accum[i] > numcontributors:
							print "WARNING: incoherency in the mask calculation. (i, accum[i], numcontributors) =",(i,accum[i],numcontributors)
				fnumc = float(numcontributors)
				
				# Calculate the masks for all thresholds (0.25, 0.5, 0.75 + others)
				for threshold in self.maskThresholds:
					print "   ",threshold,"raw mask ...",
					mask = [0]*(exdata['mark']['width']*exdata['mark']['height'])
					for i in range(len(accum)):
						prop = float(accum[i]) / fnumc
						mask[i] = 1 if prop >= threshold else 0
					exdata['consensus']['masks_raw'][threshold] = mask
					
					# Apply a post-treatment to the mask to "blur" it - make edges look less jaggy and fill
					# small holes. The result is kept binarized.
					print "filtering ...",
					exdata['consensus']['masks'][threshold] = self.filterMask(mask, (exdata['mark']['width'], exdata['mark']['height']))
					print "ok"
					
				print " ",a['exname'],"done ("+str(numcontributors)+" contributors)"
				
			# Compress the consensus and mask if needed
			if self.compress:
				print "Compressing consensus maps and masks"
				for exid,exdata in maps.iteritems():
					for quality in exdata['consensus']['qmaps']:
						exdata['consensus']['qmaps'][quality] = zlib.compress(repr(exdata['consensus']['qmaps'][quality]))
					for threshold in exdata['consensus']['masks_raw']:
						exdata['consensus']['masks_raw'][threshold] = zlib.compress(repr(exdata['consensus']['masks_raw'][threshold]))
					for threshold in exdata['consensus']['masks']:
						exdata['consensus']['masks'][threshold] = zlib.compress(repr(exdata['consensus']['masks'][threshold]))
				consensusCompressed = True
			
				
			print "Saving data"
			self.saveData(maps)
			return True
				
		except (pg.InternalError,pg.ProgrammingError),err:
			print err
			return False
			
		except Error,err:
			print "ERROR:",err
			return False
		
		except KeyboardInterrupt:
			print "WARNING: Ctrl^C pressed, interrupting. Saving results to '*_interrupted.pickle'"
			if self.compress and not consensusCompressed:
				for exid,exdata in maps.iteritems():
					for quality in exdata['consensus']['qmaps']:
						exdata['consensus']['qmaps'][quality] = zlib.compress(repr(exdata['consensus']['qmaps'][quality]))
					for threshold in exdata['consensus']['masks_raw']:
						exdata['consensus']['masks_raw'][threshold] = zlib.compress(repr(exdata['consensus']['masks_raw'][threshold]))
					for threshold in exdata['consensus']['masks']:
						exdata['consensus']['masks'][threshold] = zlib.compress(repr(exdata['consensus']['masks'][threshold]))
			self.saveData(maps, 'interrupted')
			return True
			
		return True
		

	def saveData(self, data, suffix=None):
		tokens = []
		if self.userfilter:
			tokens.append(nijmetrics.sanitizefn(self.userfilter))
		if self.exfilter:
			tokens.append(nijmetrics.sanitizefn(self.exfilter))
		if suffix:
			tokens.append(suffix)
		fname = self.outputprefix+'_'.join(tokens)
		f = file(fname+'.pickle', 'w')
		pickle.dump(data, f)
		f.close()

		
	def decodeRLE(self, string):
		# Decodes the provided runlength encoded string. This is a special RLE in that
		# the input string must not contain any digits, and the output is packed into
		# a string with repeating tuples (count,char), e.g.
		# INPUT: xaaaaaaaaaaaabbbbbbbbbcYYYYYYYu
		# RLE: 1x12a9b1c7Y1u
		# This scheme makes it easier to share the data between the C++ side and us, because
		# the pythonic RLE with a list of actual tuples (i.e. [(count,char),(count,char),...])
		# is not easily represented in C++, where there are two lists of counts and chars 
		# (i.e. [[count,count,...],[char,char,...]])
		tokens = re.split('(\d+)', string)
		assert tokens[0] == ''
		chars = tokens[2::2]
		counts = [int(x) for x in tokens[1::2]]
		return ''.join(c*n for n,c in zip(counts,chars))
		
		
	def getQualityMaps(self, areas, primaryQuality=0):
		# Build an array of areas using the "primary quality" concept. More information can
		# be found in the Javascript function getQualityMaps() in 'pianos4_tool_ridge.js'
		# This function does exactly the same, except that the mapping is returned instead
		# of being passed by reference.
		# Reminder of the output (array,mapping = getQualityMaps([...]))
		# - array: unordered list of 4-item elements: the ID, the layer
		# - mapping: a dict that maps layers sent to the API and quality levels (remember
		#   that those are not the same because the primaryQuality concept re-orders them)
		ret = []
		mapping = {}
		for area in areas:
			a = primaryQuality-area['value']+1
			ispos = 1 if a > 0 else 0
			layer = 2*abs(a)-ispos
			mapping[layer] = area['value']
			adata = [area['id'], layer, area['points'], []]
			if 'holes' in area:
				adata[3] = area['holes']
			ret.append(adata)
		
		return ret, mapping

		
	def convertQualityScale(self, areas):
		print "NOTE: ANSI/NIST quality scale detected, converting to standard scale"
		for area in areas:
			assert area['color'] in self.ansiNistTable and area['value'] == self.ansiNistTable[area['color']]['id']
			area['value'] = self.ansiNistTable[area['color']]['conv']
		return areas
		
		
	def callCppJsonApi(self, method, data, moreParams=None):
		headers = {
			'Content-Type': 'application/json'
		}
		baseURL = 'http://127.0.0.1:8081/uflr'
		params = {'action':method}
		if moreParams:
			params = dict(params.items()+moreParams.items())
		url = baseURL + '?' + urllib.urlencode(params)
		#print "Calling API",url
		req = urllib2.Request(url, json.dumps(data), headers)
		response = urllib2.urlopen(req)
		rawresult = response.read()
		result = json.loads(rawresult)
		if not result['success']:
			raise Error(result['message'])
		return result['ret']
		
	
	def getImage(self, type, exid):
		r = self.db.query("SELECT imagev.* FROM imagev JOIN exercise ON exercise."+type+"=imagev.id WHERE exercise.id = "+str(int(exid))+";")
		assert len(r) == 1
		return r[0]
		
	
	def filterMask(self, mask, dims):
		if not self.maskFilter:
			kernel = [
				1, 1, 1, 1, 1,
				1, 1, 1, 1, 1,
				1, 1, 0, 1, 1,
				1, 1, 1, 1, 1,
				1, 1, 1, 1, 1,
			]
			self.maskFilter = PIL.ImageFilter.Kernel((5,5),kernel)

		surface = PIL.Image.new('L', dims, 0)
		dc = PIL.ImageDraw.Draw(surface)
		y = 0
		for line in itertools.izip_longest(*[iter(mask)]*dims[0]):
			for x in range(dims[0]):
				value = line[x]
				if value:
					dc.point((x,y), fill=255)
				else:
					dc.point((x,y), fill=0)
			y += 1			
			
		# blur + thresholding at 0.25 + blur + threshold 0.25
		# -> results in slight expansion of the borders - smoothes quite well but still artifacts
		# filtered = surface.filter(PIL.ImageFilter.BLUR).point(lambda x: 0 if x < 64 else 255)\
		# .filter(PIL.ImageFilter.BLUR).point(lambda x: 0 if x < 64 else 255)
		
		# Same blurs but with reduction in second threshold instead of expansion
		# filtered = surface.filter(PIL.ImageFilter.BLUR).point(lambda x: 0 if x < 64 else 255)\
		 # .filter(PIL.ImageFilter.BLUR).point(lambda x: 0 if x < 196 else 255)

		# custom kernel. Blurring Expansion followed by blurring reduction
		#filtered = surface.filter(PIL.ImageFilter.Kernel((5,5),[1]*25)).point(lambda x: 0 if x < 64 else 255)\
		#	.filter(PIL.ImageFilter.Kernel((5,5),[1]*25)).point(lambda x: 0 if x < 196 else 255)				

		# stronger custom kernel. Blurring Expansion followed by blurring reduction, applied twice
		filtered = surface.filter(self.maskFilter).point(lambda x: 0 if x < 64 else 255)\
			.filter(self.maskFilter).point(lambda x: 0 if x < 196 else 255)\
			.filter(self.maskFilter).point(lambda x: 0 if x < 64 else 255)\
			.filter(self.maskFilter).point(lambda x: 0 if x < 196 else 255)					

		return [1 if x==255 else 0 for x in filtered.getdata()]
	
	