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

import pianos
import argparse
import sys
import os
# You can find a prebuilt download for PythonMagick on Windows on the 
# magnificent Unofficial Windows Binaries for Python Extension Packages 
# site, which is found at: http://www.lfd.uci.edu/~gohlke/pythonlibs/
import PythonMagick
import PIL, PIL.Image
import cStringIO as StringIO


		
class SD27:
	
	images_dir = '1000ppi'
	marks_dir = 'Latents-JPEG2000-Lossless'
	prints_dir = 'Exemplars-JPEG2000-Lossless'
	
	Q_GOOD = 1
	Q_BAD = 2
	Q_UGLY = 3
	
	mark_quality_by_prefix = {
		'G': Q_GOOD,
		'B': Q_BAD,
		'U': Q_UGLY,
	}
	mark_quality_labels = {
		Q_GOOD: 'good',
		Q_BAD : 'bad',
		Q_UGLY: 'ugly',
	}
	
	resolution = 1000
	
	def __init__( self, path ):
		self.path = path
		self.prints = []
		self.prints_by_mark = {}
		self.marks = []
		if not os.path.isdir( self.path ):
			raise OSError( "'%s' is not a valid directory" % self.path )
		print "Checking input paths"
		self.check_path( SD27.images_dir, SD27.marks_dir )
		self.check_path( SD27.images_dir, SD27.prints_dir )
		
	def init( self ):
		self.index_images()
		
	def check_path( self, *args ):
		path = os.path.join( self.path, *args )
		print "  %s" % path,
		if not os.path.isdir( path ):
			raise OSError( "Could not find the required data directory '%s'" % path )
		print "(OK)"
		
	def index_images( self ):
		print "Scanning input directories"
		# As per the NIST SD27A Readme.txt file:
		# The format of the prints is NNN-ZZ[-XXXX].jp2 where
		# - NNN is the subject (person) number
		# - ZZ is the finger number (1..10)
		# - XXXX is the corresponding latent image file
		prints_dir = os.path.join( self.path, SD27.images_dir, SD27.prints_dir )
		count = 0
		for filename in os.listdir( prints_dir ):
			name, ext = os.path.splitext( filename )
			if ext != '.jp2':
				# We are only looking for images - not that there should be something 
				# else in this directory but who knows
				continue
			count += 1
			tokens = name.split( '-' )
			if len( tokens ) == 3:
				_print = {
					'filename': filename,
					'name': name,
					'person': int( tokens[0] ),
					'finger': int( tokens[1] ),
					'mark': tokens[2],
				}
				self.prints.append( _print )
				if _print['mark'] not in self.prints_by_mark:
					self.prints_by_mark[_print['mark']] = []
				self.prints_by_mark[_print['mark']].append( _print )
		print "  Found %d prints with a corresponding latent (%d scanned)" % ( len( self.prints ), count )
		
		marks_dir = os.path.join( self.path, SD27.images_dir, SD27.marks_dir )
		count, useless = 0, 0
		for filename in os.listdir( marks_dir ):
			name, ext = os.path.splitext( filename )
			if ext != '.jp2':
				continue
			count += 1
			if name not in self.prints_by_mark:
				useless += 1
				continue
			prefix = filename[0]
			if prefix not in SD27.mark_quality_by_prefix:
				raise ValueError( "Mark '%s' has an unknown prefix: '%s' (must be one of %s)" % ( filename, prefix, ', '.join( SD27.mark_quality_by_prefix.keys() ) ) )
			mark = {
				'filename': filename,
				'name': name,
				'quality': SD27.mark_quality_by_prefix[prefix],
			}
			self.marks.append( mark )
		print "  Found %d marks (%d scanned, %d with no associated print)" % ( len( self.marks ), count, useless )
			
		
	def read_jp2_to_pil( self, filename ):
		image = PythonMagick.Image( filename )
		blob = PythonMagick.Blob()
		image.write( blob, "png" )
		buf = StringIO.StringIO( blob.data )
		image = PIL.Image.open( buf )
		image.filename = os.path.basename( filename )
		return image

		
	def generate_exercises_by_mark( self ):
		for mark in self.marks:
			_print = self.prints_by_mark[mark['name']][0]
			yield self.get_exercise( mark, _print )

			
	def generate_exercises_by_mark_all( self ):
		for mark in self.marks:
			for _print in self.prints_by_mark[mark['name']]:
				yield self.get_exercise( mark, _print )
			
	def get_exercise( self, mark, _print ):
		q = SD27.mark_quality_labels[mark['quality']]
		qlabel = q[0].upper() + q[1:]
		return {
			'mark': mark,
			'print': _print,
			'name': mark['name'] + ' vs. ' + _print['name'],
			'description': '%s mark against finger %d of individual %d\nSource: NIST Special Database 27A' % ( qlabel, _print['finger'], _print['person'] )
		}

		
	def load_exercise( self, ex ):
		print "Loading exercise '%s'" % ex['name']
		mark_filename = os.path.join( self.path, SD27.images_dir, SD27.marks_dir, ex['mark']['filename'] )
		print "  Loading mark '%s'" % mark_filename
		ex['mark']['image'] = self.read_jp2_to_pil( mark_filename )
		print_filename = os.path.join( self.path, SD27.images_dir, SD27.prints_dir, ex['print']['filename'] )
		print "  Loading print '%s'" % print_filename
		ex['print']['image'] = self.read_jp2_to_pil( print_filename )


		
def main():
	parser = argparse.ArgumentParser( description = 'Imports the pairs from NIST Special Database 27A into PiAnoS.', epilog = 'To find NIST Special Database 27A, visit http://www.nist.gov/itl/iad/ig/sd27a.cfm' )
	
	default_mode = 1
	default_folder_name = 'NIST Special Database 27A'

	mode_table = {
		1: SD27.generate_exercises_by_mark,
		2: SD27.generate_exercises_by_mark_all,
	}

	parser.add_argument( 'src_path', metavar = 'PATH_TO_SD27A', help = "Path to the unpacked NIST SD 27A. This must point to the root of the archive (this is the directory containing 'Readme.txt')" )
	parser.add_argument( 'db_url', metavar = 'PIANOS_PGDB_URL', help = pianos.argparse_help )
	parser.add_argument( '-m', '--mode', metavar = 'NUM', type = int, choices = mode_table.keys(), default = default_mode, help = "Mode to operate on. Defaults to %d. In mode 1, one exercise per mark is created (selecting the first matching print). In mode 2, one exercise per mark-print is created (superset of mode 1)." % default_mode )
	parser.add_argument( '-d', '--dry-run', action = 'store_true', help = "Perform a dry run, leaving the database unchanged. Helpful to check the input and the number of exercises that will be created." )
	parser.add_argument( '-f', '--folder', metavar = 'NAME', default = default_folder_name, help = "Name of the PiAnoS folder to create. Defaults to '%s'" % default_folder_name )
	args = parser.parse_args()
	
	sd27 = SD27( args.src_path )
	db = pianos.Database( args.db_url )
	
	sd27.init()
	
	# Get bound method from unbound descriptor
	# http://stackoverflow.com/questions/1015307/python-bind-an-unbound-method
	generator = mode_table[args.mode].__get__( sd27 )
	
	
	if args.dry_run:
		count = 0
		for ex in generator():
			count += 1
		print ""
		print "Dry run results:"
		print "  would create folder '%s'" % args.folder
		print "  would create %d exercises" % count
	
	else:

		folder_id = db.create_folder( args.folder )
		print "Created folder ID=%d" % folder_id
		
		for ex in generator():
			sd27.load_exercise( ex )
			print "  Saving to the database ...",
			ex_id = db.create_exercise( folder_id, ex['name'], ex['description'], ex['mark']['image'], SD27.resolution, ex['print']['image'], SD27.resolution )
			print "Done (ID=%d)" % ex_id
		db.conn.commit()
		
	return 0
	
if __name__ == '__main__':
	sys.exit( main() )
