//
//  FaceCropController.m
//  FaceCrop
//
//  Created by Andrew Salamon on 10/9/06.
//  Copyright 2006 Machine Perception Laboratory, University of California San Diego.. All rights reserved.
//

#import "FaceCropController.h"
#import "MPT_ResolveAlias.h"
#import "NSImageToRimage.h"
#import "NSImageFromPPM.h"

#ifdef USE_PPMLIB
#include "RImageFromPPM.h"
#endif

#import <CERTFramework/rimage_ser.h>
#import <CERTFramework/mp_eyefinderWrapper.h>
#import <CERTFramework/mp_rotateFace.h>

@interface FaceCropController(PrivateMethods)
- (void)processFolderStep;
- (void)processImage:(NSImage *)image fromPath:(NSString *)path;
- (void)processRimage:(RImage<float> *)rimage fromPath:(NSString *)path display:(BOOL)display;
- (void)findFacesIn:(RImage<float> *)rimage fromPath:(NSString *)path;
- (NSString *)newPathForPath:(NSString *)path withType:(NSBitmapImageFileType)type;
- (BOOL)setupDataFileForDir:(NSString *)baseDir;
- (void)writeToDataFile:(NSString *)message;
- (NSString *)fullVersion;
- (NSBitmapImageFileType)outputFileTypeForPath:(NSString *)path;
- (BOOL)writePatchToPath:(NSString *)path;
- (NSString *)frameNumberForPath:(NSString *)path;

- (NSDictionary *)eyeDataFromFile:(NSString *)file;
@end

@implementation FaceCropController
- (id)init
{
	if( self = [super init] )
	{
		filesToProcess = [[NSMutableArray alloc] init];
		
		NSArray *rimageTypes = [NSArray arrayWithObjects:@"rimage", @"rimage_bin", @"rimage_xml", @"ppm", nil];
		imageTypes = [[[NSImage imageFileTypes] arrayByAddingObjectsFromArray:rimageTypes] retain];
		
		eyefinder = new MP_eyefinderWrapper();
		eyefinder->setUseSequence( true );	
		
		patch = new RImage<float>( 96, 96 );
		
		rotator = new MP_RotateFace( 96 ); // face patches are 96x96 pixels
		rotator->setEyeLenRate( 0.48 );
		//rotator->setEyeTopRate( 0.20 );
		
		hasGUI = YES;
		
		eyeData = nil;
	}
	return self;
}

- (void)awakeFromNib
{
}

- (void)dealloc
{
	[filesToProcess release];
	[imageTypes release];
	[eyeDataFile release];
	[eyeData release];
	
	delete eyefinder;
	delete patch;
	delete rotator;
	[super dealloc];
}

- (NSWindow *)window
{
	return window;
}


- (NSString *)outputDir
{
    return outputDir; // [[NSUserDefaults standardUserDefaults] stringForKey:@"outputDir"];
}

- (void)setOutputDir:(NSString *)value
{
    if( [self outputDir] != value )
    {
		outputDir = [value copy]; // [[NSUserDefaults standardUserDefaults] setObject:value forKey:@"outputDir"];
    }
}

- (IBAction)cancel:(id)sender
{
	interrupt = YES;
}

- (BOOL)outputFace
{
	return [[NSUserDefaults standardUserDefaults] boolForKey:@"outputFace"];
}
- (BOOL)outputData
{
	return [[NSUserDefaults standardUserDefaults] boolForKey:@"outputData"];
}

- (int)imageCount
{
    return imageCount;
}

- (void)setImageCount:(int)value
{
    imageCount = value;
}

- (int)imagesDone
{
    return imagesDone;
}

- (void)setImagesDone:(int)value
{
    imagesDone = value;
}

- (int)minImages
{
	return 0;
}


- (NSString *)eyeDataFile
{
    return eyeDataFile;
}

- (void)setEyeDataFile:(NSString *)value
{
    if( eyeDataFile != value )
    {
        [eyeDataFile release];
        eyeDataFile = [value retain];
    }
}

- (void)reset
{
	[filesToProcess removeAllObjects];
	[fileEnumerator release];
	fileEnumerator = nil;
	[dataFile closeFile];
	[dataFile release];
	dataFile = nil;
	[loopPool release];
	loopPool = nil;
	[self setImageCount:0];
	[self setImagesDone:0];	
}

- (float)cropLength
{
    return rotator->getEyeLenRate();
}
- (void)setCropLength:(float)value
{
	rotator->setEyeLenRate( value );
}

- (float)cropTop
{
    return rotator->getEyeTopRate();
}
- (void)setCropTop:(float)value
{
	rotator->setEyeTopRate( value );
}


- (IBAction)chooseOutputDirectory:(id)sender
{
	NSOpenPanel *panel = [NSOpenPanel openPanel];
	[panel setCanChooseDirectories:YES];
	[panel setAllowsMultipleSelection:NO];
	[panel setCanChooseFiles:NO];
	[panel setCanCreateDirectories:YES];
	[panel setMessage:@"Choose the folder where the new face image and data file will be saved."];
	
	[panel beginSheetForDirectory:nil
							 file:nil
							types:nil
				   modalForWindow:[self window]
					modalDelegate:self
				   didEndSelector:@selector(outputPanelDidEnd:returnCode:contextInfo:)
					  contextInfo:nil];
}

- (void)outputPanelDidEnd:(NSOpenPanel *)panel returnCode:(int)returnCode  contextInfo:(void  *)contextInfo
{
	if( returnCode == NSOKButton )
	{
		if( 1 == [[panel filenames] count] )
		{
			[self setOutputDir:[[panel filenames] objectAtIndex:0]];
		}
		else
		{
			[self setOutputDir:nil];
		}
	}
}

- (IBAction)chooseDataFile:(id)sender
{
	NSOpenPanel *panel = [NSOpenPanel openPanel];
	[panel setCanChooseDirectories:NO];
	[panel setAllowsMultipleSelection:NO];
	[panel setCanChooseFiles:YES];
	[panel setCanCreateDirectories:NO];
	[panel setMessage:@"Choose the data file that contains eye position data."];
	
	[panel beginSheetForDirectory:nil
							 file:nil
							types:[NSArray arrayWithObjects:@"text",@"txt",@"TEXT",@"TXT",nil]
				   modalForWindow:[self window]
					modalDelegate:self
				   didEndSelector:@selector(dataFilePanelDidEnd:returnCode:contextInfo:)
					  contextInfo:nil];
}

- (void)dataFilePanelDidEnd:(NSOpenPanel *)panel returnCode:(int)returnCode  contextInfo:(void  *)contextInfo
{
	if( returnCode == NSOKButton )
	{
		if( 1 == [[panel filenames] count] )
		{
			[self setEyeDataFile:[[[panel filenames] objectAtIndex:0] retain]];
			[eyeData release];
			eyeData = [[self eyeDataFromFile:eyeDataFile] retain];
		}
	}
}

- (IBAction)open:(id)sender
{
	if( [self outputData] || [self outputFace] )
	{
		if( 0 == [[self outputDir] length] )
		{
			NSRunAlertPanel( @"No output folder",@"Choose an output folder before processing images",@"OK",nil,nil);
			return;
		}
	}
	
	NSOpenPanel *panel = [NSOpenPanel openPanel];
	[panel setCanChooseDirectories:YES];
	[panel setAllowsMultipleSelection:YES];
	[panel setCanChooseFiles:YES];
	[panel setMessage:@"Choose the folders and/or files to process."];

	[panel beginSheetForDirectory:nil
							 file:nil
							types:imageTypes
				   modalForWindow:[self window]
					modalDelegate:self
				   didEndSelector:@selector(openPanelDidEnd:returnCode:contextInfo:)
					  contextInfo:nil];
}

- (void)openPanelDidEnd:(NSOpenPanel *)panel returnCode:(int)returnCode  contextInfo:(void  *)contextInfo
{
	if( returnCode == NSOKButton )
	{
		NSString *baseDir = nil;
		NSFileManager *fileManager = [NSFileManager defaultManager];
		NSArray *filenames = [panel filenames];
		NSEnumerator *fileEnum = [filenames objectEnumerator];
		NSString *path;
		while( path = [fileEnum nextObject] )
		{
			BOOL isDir;
			if( [fileManager fileExistsAtPath:path isDirectory:&isDir] )
			{
				if( isDir )
				{
					NSArray *dirFiles = [fileManager directoryContentsAtPath:path];
					NSEnumerator *dirEnum = [dirFiles objectEnumerator];
					NSString *file;
					while( file = [dirEnum nextObject] )
					{
						if( [imageTypes containsObject:[file pathExtension]] )
							[filesToProcess addObject:[path stringByAppendingPathComponent:file]];
					}
					if( !baseDir )
						baseDir = [path copy];
				}
				else
				{
					[filesToProcess addObject:path];
					if( !baseDir )
						baseDir = [[path stringByDeletingLastPathComponent] copy];
				}
			}
		}

		if( [self setupDataFileForDir:baseDir] )
		{
			interrupt = NO;
			fileEnumerator = [[filesToProcess objectEnumerator] retain];
			
			[self setImageCount:[filesToProcess count]];
			[self setImagesDone:0];
			
			[self performSelector:@selector(processFolderStep) withObject:nil afterDelay:0.01];
		}
		[baseDir release];
	}
}

// For use by a command line version of the app
- (BOOL)processFolder:(NSString *)source into:(NSString *)dest
{
	NSFileManager *fileManager = [NSFileManager defaultManager];
	BOOL isDir = NO;
	if( ![fileManager fileExistsAtPath:source isDirectory:&isDir] || !isDir )
	{
		return NO;
	}
	if( ![fileManager fileExistsAtPath:dest isDirectory:&isDir] || !isDir )
	{
		return NO;
	}
	
	rotator->setEyeLenRate( 0.48 );

	hasGUI = NO;
	[self reset];
	[self setOutputDir:dest];

	NSArray *dirFiles = [fileManager directoryContentsAtPath:source];
	NSEnumerator *dirEnum = [dirFiles objectEnumerator];
	NSString *file;
	while( file = [dirEnum nextObject] )
	{
		if( [imageTypes containsObject:[file pathExtension]] )
			[filesToProcess addObject:[source stringByAppendingPathComponent:file]];
	}
	
	if( [self setupDataFileForDir:source] )
	{
		interrupt = NO;
		fileEnumerator = [[filesToProcess objectEnumerator] retain];
		
		[self setImageCount:[filesToProcess count]];
		[self setImagesDone:0];
		
		[self processFolderStep];
	}
	
	return YES;
}

- (void)processFolderStep
{
	NSString *imagePath;
	
	if( 0 == ([self imagesDone] % 1000) )
	{
		[loopPool release];
		loopPool = [[NSAutoreleasePool alloc] init];
	}
	
	if( !interrupt && (imagePath = [fileEnumerator nextObject]) )
	{
		NSFileManager *fileManager = [NSFileManager defaultManager];
		imagePath = [[[imagePath stringByStandardizingPath] stringByResolvingSymlinksInPath] stringByResolvingAliases];
		NSDictionary *fileAttributes = [fileManager fileAttributesAtPath:imagePath traverseLink:YES];
		
		[self setImagesDone:[self imagesDone] + 1];

		if( [[fileAttributes objectForKey:NSFileType] isEqualToString:NSFileTypeRegular]
			|| [[fileAttributes objectForKey:NSFileType] isEqualToString:NSFileTypeSymbolicLink] )
		{
			NSImage *image = [[NSImage alloc] initWithContentsOfFile:imagePath];
			if( image )
			{
				[self processImage:image fromPath:imagePath];
				[image release];
			}
			else
			{
				if( [[imagePath pathExtension] isEqualToString:@"rimage"] )
				{
					// read the rimage, find faces, crop and rotate, save face, save file info
					RImage<float> rimage;
					std::string filename( [imagePath UTF8String] );
					
					if( restoreRImageFromFile( rimage, filename ) )
					{
						[self processRimage:&rimage fromPath:imagePath display:YES];
					}
					else
					{
						[self writeToDataFile:[NSString stringWithFormat:@"Unable to read image from %@", imagePath]];
					}
				}
				else if( [[imagePath pathExtension] isEqualToString:@"ppm"] )
				{
					image = [[NSImage alloc] initWithContentsOfPPMFile:imagePath];
					if( image )
					{
						[self processImage:image fromPath:imagePath];
						[image release];
					}
#ifdef USE_PPMLIB
					RImage<float> rimage;
					std::string filename( [imagePath UTF8String] );
					if( restoreRimageFromPPM( rimage, filename ) )
					{
						[self processRimage:&rimage fromPath:imagePath display:YES];
					}
					else
					{
						[self writeToDataFile:[NSString stringWithFormat:@"Unable to read image from %@", imagePath]];
					}
#endif
				}
				else
				{
					[self writeToDataFile:[NSString stringWithFormat:@"Unknown image type: %@", imagePath]];
				}
			}			
		}
		if( hasGUI )
			[self performSelector:@selector(processFolderStep) withObject:nil afterDelay:0.01];
		else
			[self processFolderStep];
	}
	else
	{
		[self reset];
	}
}

- (void)processImage:(NSImage *)image fromPath:(NSString *)path
{
	[imageView setImage:image];
	RImage<float> *rimage = [image rimageRep];
	[self processRimage:rimage fromPath:path display:NO];
	delete rimage;
}

- (void)processRimage:(RImage<float> *)rimage fromPath:(NSString *)path display:(BOOL)display
{
	if( display )
	{
		NSImage *image = [NSImage imageFromRImage:rimage];
		[imageView setImage:image];
	}
	[self findFacesIn:rimage fromPath:path];
}

- (void)findFacesIn:(RImage<float> *)rimage fromPath:(NSString *)path
{
	NSPoint oneLeftEye;
	NSPoint oneRightEye;
	int frameNum = [[self frameNumberForPath:[path lastPathComponent]] intValue];
	NSArray *eyes = [eyeData objectForKey:[NSNumber numberWithInt:frameNum]];
	
	if( 4 == [eyes count] )
	{
		oneLeftEye.x = [[eyes objectAtIndex:0] floatValue];
		oneLeftEye.y = [[eyes objectAtIndex:1] floatValue];
		oneRightEye.x = [[eyes objectAtIndex:2] floatValue];
		oneRightEye.y = [[eyes objectAtIndex:3] floatValue];
	}
	else
	{
		EyeFaceList faces;
		
		try {
			eyefinder->findEyesSpeed( *rimage, faces );
			//		((EYEFINDER_CLASS *)mpeyefinder)->findEyes( *rimage, faces );		
		} catch( ... ) {
			// Probably not much we can do, except possibly keep track of how many exceptions are thrown in what amount of time
			// and possibly warn the user if it happens too often
			NSLog( @"Exception thrown while trying to find eyes in image: %@", path );
			return;
		}
		
		double largestArea = 0.0;
		
		for( EyeFaceList::iterator faceIter = faces.begin(), last_face = faces.end(); faceIter != last_face; ++faceIter )
		{
			double faceArea = faceIter->xSize * faceIter->ySize;
			
			if( (faceIter->x > 0.0)	&& (faceIter->y > 0.0) && (faceArea > largestArea) )
			{
				oneLeftEye.x = faceIter->both_eyes.xLeft;
				oneLeftEye.y = faceIter->both_eyes.yLeft;
				oneRightEye.x = faceIter->both_eyes.xRight;
				oneRightEye.y = faceIter->both_eyes.yRight;
				
				largestArea = faceArea;
			}
		}
	}
	
	BOOL noFace = YES;
	// Criteria are eye x-positions are nonzero and at least 50 pix apart.
	// That criteria is really only for data sets, but we might want it to be a preference.
	if( (oneLeftEye.x > 0) && (oneRightEye.x > 0) && ( (oneLeftEye.x-oneRightEye.x) > 30.0) )
	{		
		if( rotator->rotateImageIntoPatch( rimage, oneLeftEye.x, oneRightEye.x, oneLeftEye.y, oneRightEye.y, *patch ) )
		{
			NSImage *face = [NSImage imageFromRImage:patch];
			[faceView setImage:face];
			
			if( [self outputFace] )
			{
				NSBitmapImageFileType newType = [self outputFileTypeForPath:path];
				NSString *facePath = [self newPathForPath:path withType:newType];
				
				if( NSJPEGFileType == newType )
				{
					NSBitmapImageRep *rep = [face bitmapImageRep];
					NSData *data = [rep representationUsingType:newType properties:nil];
					[data writeToFile:facePath atomically:NO];
				}
				else
				{
					[self writePatchToPath:[[facePath stringByDeletingPathExtension] stringByAppendingPathExtension:@"pgm"]];
				}
			}
			
			if( [self outputData] )
			{
				[self writeToDataFile:[NSString stringWithFormat: @"%@\t%f\t%f\t%f\t%f\n", [self frameNumberForPath:[path lastPathComponent]], oneLeftEye.x, oneLeftEye.y, oneRightEye.x, oneRightEye.y]];
			}
			noFace = NO;
		}
	}		
	if( noFace )
	{
		[self writeToDataFile:[NSString stringWithFormat: @"%@\tNAN\n", [self frameNumberForPath:[path lastPathComponent]]]];
		[faceView setImage:nil];
	}
}

- (NSString *)newPathForPath:(NSString *)path withType:(NSBitmapImageFileType)type
{
	NSString *name = [path lastPathComponent];
	NSString *ext = (NSTIFFFileType == type)?@"tiff":@"jpeg";
	
	name = [name stringByDeletingPathExtension];
	name = [name stringByAppendingString:@"_face"];
	name = [name stringByAppendingPathExtension:ext];
	
	// If an output directory has been set, use it, otherwise use the same directory as the given path
	if( [[self outputDir] length] > 0 )
	{
		name = [[self outputDir] stringByAppendingPathComponent:name];
	}
	else
	{
		NSString *oldPath = [path stringByDeletingLastPathComponent];
		name = [oldPath stringByAppendingPathComponent:name];
	}
	
	return name;
}

- (BOOL)setupDataFileForDir:(NSString *)baseDir
{
	if( [self outputData] )
	{
		NSAssert( (nil == dataFile), @"Data file handle should be nil" );
		
		if( [[self outputDir] length] > 0 )
		{
			NSString *dataFilePath = [[self outputDir] stringByAppendingPathComponent:@"FaceData.txt"];
			NSFileManager *fileManager = [NSFileManager defaultManager];
			if( ![fileManager isWritableFileAtPath:dataFilePath] )
			{
				if( ![fileManager createFileAtPath:dataFilePath contents:nil attributes:nil] )
				{
					NSLog( @"Unable to create data file at %@", dataFilePath );
					return NO;
				}
			}
			
			dataFile = [[NSFileHandle fileHandleForUpdatingAtPath:dataFilePath] retain];
			unsigned long long len = [dataFile seekToEndOfFile];
			if( len > 0 )
				[self writeToDataFile:@"\nResuming output"];
		}
		
		if( !baseDir )
			baseDir = @"Unknown source";
		
		[self writeToDataFile:[NSString stringWithFormat:@"Faces and data generated with %@ at %@\n%@\n\n", [self fullVersion], [[NSCalendarDate calendarDate] description], baseDir]];
		
		return (dataFile != nil);	
	}
	else
	{
		return YES;
	}
}

- (void)writeToDataFile:(NSString *)message
{
	if( dataFile )
	{
		NSData *data = [message dataUsingEncoding:NSUTF8StringEncoding];
		[dataFile writeData:data];
	}
	else
	{
		NSLog( @"Attempted to write to closed file: %@", message );
	}
}

- (NSString *)fullVersion
{
	NSBundle *certBundle = [NSBundle bundleWithIdentifier:@"edu.ucsd.mplab.CERTFramework"];
	NSString *mptVersion = [certBundle objectForInfoDictionaryKey:@"MPTVersion"];
	NSString *certVersion = [certBundle objectForInfoDictionaryKey:@"CERTVersion"];
	//	int rev = getCertRevision();
	////	char *revDate = getCertRevisionDate();
	//	const char *mptRev = getMPTRevisionString();
	return [NSString stringWithFormat:@"build %@::%@", mptVersion, certVersion];
}

- (IBAction)displayAboutPanel:(id)sender
{
	NSDictionary *options = [NSDictionary dictionaryWithObjectsAndKeys:[self fullVersion], @"Version", nil]; // , @"1.0", @"ApplicationVersion", nil
	[NSApp orderFrontStandardAboutPanelWithOptions:options];
}

- (NSBitmapImageFileType)outputFileTypeForPath:(NSString *)path
{
	NSString *ext = [path pathExtension];
	if( (NSOrderedSame == [ext compare:@"jpeg" options:NSCaseInsensitiveSearch])
		|| (NSOrderedSame == [ext compare:@"jpg" options:NSCaseInsensitiveSearch]) )
	{
		return NSJPEGFileType;
	}
	
	return NSTIFFFileType;
}

- (float)maxPatchPixel
{
	float *vals = patch->array;
	int cnt = patch->numpixels;
	float max = 0.0;
	for( int idx = 0; idx < cnt; ++idx )
	{
		if( vals[idx] > max )
			max = vals[idx];
	}
	
	return max;
}

- (BOOL)writePatchToPath:(NSString *)path
{
	unsigned short maxPixel = 255;
	float adjust = maxPixel; // ((float)maxPixel) / [self maxPatchPixel];
	NSString *header = [NSString stringWithFormat:@"P5\n%d %d\n%d\n", patch->width, patch->height, maxPixel];

	NSFileManager *fileManager = [NSFileManager defaultManager];
	if( ![fileManager isWritableFileAtPath:path] )
	{
		if( ![fileManager createFileAtPath:path contents:nil attributes:nil] )
		{
			NSLog( @"Unable to create output PGM file at %@", path );
			return NO;
		}
	}
	
	NSFileHandle *pgmFile = [NSFileHandle fileHandleForWritingAtPath:path];
	if( !pgmFile )
		return NO;
	
	[pgmFile writeData:[header dataUsingEncoding:NSASCIIStringEncoding allowLossyConversion:YES]];
	
	float *vals = patch->array;
	int cnt = patch->numpixels;
	unsigned char row[patch->width];
	unsigned int rowIdx = 0;
	for( int idx = 0; idx < cnt; ++idx )
	{
		unsigned char val = roundf( vals[idx] * adjust );
		if( val > maxPixel ) val = maxPixel;
		row[rowIdx++] = val;
		if( rowIdx == patch->width )
		{
			NSData *rowData = [NSData dataWithBytesNoCopy:row length:patch->width freeWhenDone:NO];
			[pgmFile writeData:rowData];
			rowIdx = 0;
		}
	}
	
	[pgmFile closeFile];
	return YES;
}

- (NSString *)frameNumberForPath:(NSString *)path
{
	static NSRange frameRange = { 5, 5 };
	return [path substringWithRange:frameRange];
}

- (NSDictionary *)eyeDataFromFile:(NSString *)file
{
	NSMutableDictionary *data = [NSMutableDictionary dictionary];
	NSString *dataString = [[NSString alloc] initWithContentsOfFile:file];
	NSArray *lines = [dataString componentsSeparatedByString:@"\n"];
	NSEnumerator *lineEnum = [lines objectEnumerator];
	NSString *line;
	
	// Skip the first three (non-data) lines
	[lineEnum nextObject]; [lineEnum nextObject]; [lineEnum nextObject];
	
	while( line = [lineEnum nextObject] )
	{
		// 00036	NAN
		// 00037	375.296173	140.509933	283.162048	140.773544
		
		NSArray *items = [line componentsSeparatedByString:@"\t"];
		if( 2 == [items count] )
		{
			// No face/eyes found last time, skip it.
		}
		else if( 5 == [items count] )
		{
			// get the frame number
			// get the eye position numbers
			// store the eye positions (as an NSRect?) in a dictionary with the frame number as the key
			int frame = [[items objectAtIndex:0] intValue];
			NSArray *eyes = [NSArray arrayWithObjects:[NSNumber numberWithFloat:[[items objectAtIndex:1] floatValue]],
				[NSNumber numberWithFloat:[[items objectAtIndex:2] floatValue]],
				[NSNumber numberWithFloat:[[items objectAtIndex:3] floatValue]],
				[NSNumber numberWithFloat:[[items objectAtIndex:4] floatValue]], nil];
			[data setObject:eyes forKey:[NSNumber numberWithInt:frame]];
		}
		else
		{
			NSLog( @"Bad data in FaceData.txt file: %@", line );
		}
	}
	
	return [[data copy] autorelease];
}
@end