//
//  MPT_CERT.mm
//  AUCoderDemo
//
//  Created by Andrew Salamon on 5/17/06.
//  Copyright (c) 2006 Machine Perception Laboratory 
//  University of California San Diego.
//

#import "MPT_CERT.h"

#import <map>
#import <vector>
#import <exception>

#import <CERTFramework/rimage.h>
#import <CERTFramework/rimage_ser.h>
#import <CERTFramework/visualobject.h>
#import <CERTFramework/mp_eyefinderWrapper.h>
#import <CERTFramework/mp_rotatingEyefinder.h>
#import <CERTFramework/mp_FDEyes.h>

#import <CERTFramework/mp_rotateFace.h>
#import <CERTFramework/mp_SVMWeights.h>
#import <CERTFramework/mp_gabor.h>
#import <CERTFramework/mp_threadedGabor.h>
#import <CERTFramework/ExpressionFilterWrapper.h>

#import "NSImageToRimage.h"
#import "MPT_Face.h"
#import "MPT_RimageAdaptor.h"
#import "MPT_PluginController.h"
#import "MPT_PluginProtocolPlusPlus.h"
#import "MPT_PluginProtocolSVM.h"
#import "MPT_PluginProtocolFeatures.h"
#import "MPT_PluginProtocolFrequencies.h"

#ifdef __APPLE_CC__
#include <Accelerate/Accelerate.h>
#endif

//typedef MP_Gabor GABOR_CLASS;
typedef mp_threadedGabor GABOR_CLASS;
//#define USING_FEATURE_DETECTOR
#ifdef USING_FEATURE_DETECTOR
typedef mp_FDEyes EYEFINDER_CLASS;
#else
typedef MP_eyefinderWrapper EYEFINDER_CLASS;
//typedef mp_rotatingEyefinder EYEFINDER_CLASS;
#endif
typedef std::map<int,ExpressionFilter*> filterMapType;

static const int auValues[] = { 1, 2, 4, 5, 10, 12, 14, 20, 101 };

@interface MPT_CERT(PrivateMethods)
- (void)processRImageImp:(RImage<float> *)rimage;
- (double)filterValue:(double)value forAU:(int)AU withFace:(bool)faceFound;
- (void)rotateImage:(RImage<float> *)rimage into:(RImage<float> *)destImage byRadians:(float)rad;
- (BOOL)loadFeaturePlugins;
@end

@implementation MPT_CERT
- (id)init
{
	if( self = [super init] )
	{
		results = [[NSMutableArray alloc] init];
		pluginResults = [[NSMutableDictionary alloc] init];
		faceArray = [[NSMutableArray alloc] init];
		featuresByName = [[NSMutableDictionary alloc] init];
			
		NSBundle *CERTFramework = [NSBundle bundleWithIdentifier:@"edu.ucsd.mplab.CERTFramework"];

		findFaces = YES;
		linearCombo = NO;
		
		float _angles[] = { 15.0, -15.0, 30.0, -30.0, 45.0, -45.0 };
		std::vector<float> angles( _angles, _angles+6 );
		rotatingEyefinder = new mp_rotatingEyefinder( angles, false );
		((mp_rotatingEyefinder *)rotatingEyefinder)->setUseSequence( true );
		
		mpeyefinder = new EYEFINDER_CLASS();
//		((EYEFINDER_CLASS *)mpeyefinder)->setWinshift( 0.0 );
		((EYEFINDER_CLASS *)mpeyefinder)->setUseSequence( true );
//		((EYEFINDER_CLASS *)mpeyefinder)->setMinScale( 1 );
		
		[self loadFeaturePlugins];		 
		
		patch = new RImage<float>( 96, 96 );
		
		gabor = NULL;
		RImage<float> &lpatch = *((RImage<float> *)patch);
		gabor = new GABOR_CLASS( lpatch, lpatch.width, lpatch.height );
					
		rotator = new MP_RotateFace( 96 ); // face patches are 96x96 pixels
		((MP_RotateFace *)rotator)->setEyeLenRate( 0.48 );
		//((MP_RotateFace *)rotator)->setEyeTopRate( 0.20 );

		{
			NSString *tablesDir = [[CERTFramework resourcePath] stringByAppendingPathComponent:@"likelihood_tables"];
			std::string dir( ([tablesDir length]?[tablesDir UTF8String]:"") );

			std::vector<int> auLabels;
			auLabels.push_back(1);
			auLabels.push_back(2);
			auLabels.push_back(4);
			auLabels.push_back(5);
			auLabels.push_back(10);
			auLabels.push_back(12);
			auLabels.push_back(14);
			auLabels.push_back(20);
			auLabels.push_back(101);
			filters = new ExpressionFilterWrapper( auLabels, dir );
		}
		
		[self setEyefinderOffset:1.0];
		[self setEyefinderDropoff:1.0];
	}
	
	return self;
}

- (void)dealloc
{
	[grayImage release];
	[results release];
	[faceArray release];
	[featuresByName release];
	if( mpeyefinder ) delete (EYEFINDER_CLASS *)mpeyefinder;
	if( rotatingEyefinder ) delete (mp_rotatingEyefinder *)rotatingEyefinder;
	if( gabor ) delete (GABOR_CLASS *)gabor;
	if( rotator ) delete (MP_RotateFace *)rotator;
	if( patch ) delete (RImage<float> *)patch;
	if( filters ) delete (ExpressionFilterWrapper *)filters;

	[super dealloc];
}

//#define GRAY_IS_PATCH
- (void)processImage:(NSImage *)image
{
	
	RImage<float> *rimage = [image rimageRep];
//#ifdef USE_IMAGEMAGICK
//	RImage<float> *rimage = [image rimageFromImageMagick];
//#endif
	
	if( rimage )
	{
		[self processRImageImp:rimage];
		[grayImage release];
		if( [self displayGray] )
		{
#ifdef GRAY_IS_PATCH
			grayImage = [[NSImage imageFromRImage:(RImage<float> *)patch] retain];
#else
			grayImage = [[NSImage imageFromRImage:rimage] retain];
#endif
		}
		else
			grayImage = nil;

		delete rimage;
	}
}

- (void)processRImageFile:(NSString *)rimagePath
{
	RImage<float> rimage;
	std::string filename( [rimagePath UTF8String] );
	
	if( restoreRImageFromFile( rimage, filename ) )
	{
		[self processRImageImp:&rimage];
		[grayImage release];
		if( [self displayGray] )
		{
#ifdef GRAY_IS_PATCH
			grayImage = [[NSImage imageFromRImage:(RImage<float> *)patch] retain];
#else
			grayImage = [[NSImage imageFromRImage:&rimage] retain];
#endif
		}
		else
			grayImage = nil;
	}
}

- (void)processRImageAdaptor:(MPT_RimageAdaptor *)rimageAdaptor
{
	RImage<float> *rimage = (RImage<float> *)[rimageAdaptor rimage];
	
	if( rimage )
	{
		[self processRImageImp:rimage];
		[grayImage release];
		if( [self displayGray] )
		{
#ifdef GRAY_IS_PATCH
			grayImage = [[NSImage imageFromRImage:(RImage<float> *)patch] retain];
#else
			grayImage = [[NSImage imageFromRImage:rimage] retain];
#endif			
		}
		else
			grayImage = nil;
	}
}

- (void)processRImageImp:(RImage<float> *)rimage
{
	EyeFaceList faces;
	FaceObject *largestFace = nil;
	MPT_PluginController *pluginController = [MPT_PluginController defaultController];
	bool foundRotatedFace = false;
	
	[self clearResults];

	// rimage->print( 5 );

	if( findFaces )
	{
		try {
			//((EYEFINDER_CLASS *)mpeyefinder)->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." );
			return;
		}
	}
	
	if( [self multipleRotations] && (0 == faces.size()) )
	{
		((mp_rotatingEyefinder *)rotatingEyefinder)->findRotatedEyes( *rimage, faces );
		
		if( faces.size() > 0 )
		{	
			*rimage = ((mp_rotatingEyefinder *)rotatingEyefinder)->getBestImage();
			foundRotatedFace = true;
//			NSLog( @"Best Angle: %f", ((EYEFINDER_CLASS *)mpeyefinder)->getBestAngle() );
		}
	}
	
	double largestArea = 0.0;
	
	for( EyeFaceList::iterator faceIter = faces.begin(), last_face = faces.end(); faceIter != last_face; ++faceIter )
	{
		NSRect oneFace;
		NSPoint oneLeftEye;
		NSPoint oneRightEye;
		
		oneFace.origin.x = faceIter->x;
		oneFace.origin.y = faceIter->y;
		oneFace.size.width = faceIter->xSize;
		oneFace.size.height = faceIter->ySize;
		
		oneLeftEye.x = faceIter->both_eyes.xLeft;
		oneLeftEye.y = faceIter->both_eyes.yLeft;
		oneRightEye.x = faceIter->both_eyes.xRight;
		oneRightEye.y = faceIter->both_eyes.yRight;
		
		double faceArea = faceIter->xSize * faceIter->ySize;
		
		if( (faceIter->x > 0.0)	&& (faceIter->y > 0.0) && (faceArea > largestArea) )
		{
			face = oneFace;
			leftEye = oneLeftEye;
			rightEye = oneRightEye;
			largestArea = faceArea;
			largestFace = &(*faceIter);
		}
		
#ifndef USING_FEATURE_DETECTOR
		MPT_Face *mptFace = [[MPT_Face alloc] initWithFace:oneFace];
		
		MPT_Feature *feature = [MPT_Feature featureOfType:@"left_eye" AtLocation:oneLeftEye];
		[feature setColor:[NSColor blueColor]];
		[mptFace addFeature:feature];
		
		feature = [MPT_Feature featureOfType:@"right_eye" AtLocation:oneRightEye];
		[feature setColor:[NSColor blueColor]];
		[mptFace addFeature:feature];
		
		[faceArray addObject:mptFace];
		[mptFace release];
#endif
	}
	
#ifdef USING_FEATURE_DETECTOR
	ObjectWithFeatures *largestObject = ((EYEFINDER_CLASS *)mpeyefinder)->largestFace();
	if( largestObject )
	{
		face.origin.x = largestObject->x;
		face.origin.y = largestObject->y;
		face.size.width = largestObject->xSize;
		face.size.height = largestObject->ySize;
		
		MPT_Face *mptFace = [[MPT_Face alloc] initWithFace:face];
		[faceArray addObject:mptFace];
		[mptFace release];
		
//		for( std::vector<VisualObject>::iterator featureIter = largestObject->finalBestFeature.begin(), lastFeature = largestObject->finalBestFeature.end();
//			 featureIter != lastFeature;
//			 ++featureIter )
//		{
//			NSString *type = [NSString stringWithCString:featureIter->type.c_str()];
//			NSPoint loc = NSMakePoint( featureIter->x, featureIter->y);
////			NSLog( @"Type %@ at loc %f,%f", type, loc.x, loc.y );
//		}
		
		typedef std::vector< std::vector< VisualObject > > bestFeaturesVector;
		
		for( std::vector<VisualObject>::iterator featureIter = largestObject->finalBestFeature.begin(), lastFeature = largestObject->finalBestFeature.end();
//		for( bestFeaturesVector::iterator featureIter = largestObject->featureDistribution.begin(), lastFeature = largestObject->featureDistribution.end();
			 featureIter != lastFeature;
			 ++featureIter )
		{
//			std::vector< VisualObject > &oneFeature = *featureIter;
//			if( oneFeature.size() > 0 )
			{
				//VisualObject &obj = *featureIter; // oneFeature[ oneFeature.size()-1 ];
				NSString *type = [NSString stringWithCString:featureIter->type.c_str()];
				NSPoint loc = NSMakePoint( featureIter->x, featureIter->y);
//				NSLog( @"Type %@ at loc %f,%f", type, loc.x, loc.y );
				MPT_Feature *feature = [MPT_Feature featureOfType:type AtLocation:loc];
				[mptFace addFeature:feature];
				
				id plugin = [featuresByName objectForKey:type];
				if( plugin )
				{
					[pluginResults setObject:NSStringFromPoint(loc) forKey:[plugin pluginID]];
					[feature setColor:[plugin displayColor]];
				}
				
				if( [type isEqualToString:@"left_eye"] )
				{
					leftEye = loc;
				}
				else if( [type isEqualToString:@"right_eye"] )
				{
					rightEye = loc;
				}
			}
		}	
	}
#endif USING_FEATURE_DETECTOR
	
	// 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( (leftEye.x > 0) && (rightEye.x > 0) && ( (leftEye.x-rightEye.x) > 0) )
	{
		NSArray *plugins = [pluginController enabledPluginsForStep:MPT_PluginStep_SVM];
		if( [plugins count] > 0 )
		{
			// Before doing any CERT processing we should see if there are any enabled SVM plugins.
			RImage<float> &lpatch = *((RImage<float> *)patch);
			RIntegral<float> *pfIntegral = nil;
			// The rotating eyefinder probably won't return the correct integral image, so force the face rotator to generate a new integral image.
			// Otherwise, use the integral image from the eyefinder.
			if( !foundRotatedFace )
				pfIntegral = static_cast<RIntegral<float>*>(((EYEFINDER_CLASS *)mpeyefinder)->getIntegralPtr().get());
			BOOL didRotate = ((MP_RotateFace *)rotator)->rotateImageIntoPatch( rimage, leftEye.x, rightEye.x, leftEye.y, rightEye.y, lpatch, pfIntegral );
			
			if( didRotate )
			{
				if( !gabor )
					gabor = new GABOR_CLASS( lpatch, lpatch.width, lpatch.height );
				else
					((GABOR_CLASS *)gabor)->setPixels( lpatch );
				
				((GABOR_CLASS *)gabor)->CERT_Gabor();
				
				// Collect results for each enabled plugin
				{
					NSEnumerator *pluginEnum = [plugins objectEnumerator];
					
					while( id plugin = [pluginEnum nextObject] )
					{
						std::vector<double> resultsForCategory = ((GABOR_CLASS *)gabor)->resultsForCategory( [[plugin internalName] UTF8String] );
						NSMutableDictionary *tmpResults = [NSMutableDictionary dictionary];
						NSEnumerator *labelEnum = [[plugin labels] objectEnumerator];
						NSString *label;
						unsigned int ind = 0;
						
						//NSAssert( (resultsForCategory.size() == [[plugin labels] count]), @"Mismatch between results and labels." );
						
						while( label = [labelEnum nextObject] )
						{
							[tmpResults setObject:[NSNumber numberWithDouble:resultsForCategory[ind++]] forKey:label];
						}
						[pluginResults setObject:tmpResults forKey:[plugin pluginID]];
					}
				}
				
				// Get results for any plugins that need the gabor frequency data
				{
					NSArray *mlrPlugins = [pluginController enabledPluginsForStep:MPT_PluginStep_Frequencies];
					std::map< std::string, mp_AUResults > freqs = ((GABOR_CLASS *)gabor)->getFrequencies();
					NSEnumerator *pluginEnum = [mlrPlugins objectEnumerator];
					
					while( id<MPT_PluginProtocolFrequencies> plugin = [pluginEnum nextObject] )
					{
						double mlrRes = 0.0;
						try {
							mlrRes = [plugin processFrequencyResults:freqs];
							[pluginResults setObject:[NSNumber numberWithDouble:mlrRes] forKey:[plugin pluginID]];
						} catch( std::exception &exc ) {
							NSLog( @"Unable to process frequency plugin '%@': %s", [plugin localizedName], exc.what() );
						}
					}
				}
			}
		}

		// We need to check this here because sometimes the feature detector returns valid eyes, but an invalid face (y==0).
		if( largestFace )
		{
			plugins = [pluginController enabledPluginsForStep:MPT_PluginStep_AfterFaces];
			EyeFaceList largestFaceList;
			NSEnumerator *pluginEnum = [plugins objectEnumerator];
			
			largestFaceList.push_back( *largestFace );
			while( id plugin = [pluginEnum nextObject] )
			{
				double value = 0.0;
				largestFaceList.begin()->activation = 0.0;
				[plugin processFaces:largestFaceList fromEyefinder:(EYEFINDER_CLASS *)mpeyefinder withImage:*rimage];
				value = largestFaceList.begin()->activation;
				[pluginResults setObject:[NSNumber numberWithDouble:value] forKey:[plugin pluginID]];
			}
		}
	}
	else
	{
		// No face found
		for( int i=0; i < [self AUCnt]; ++i )
		{
			[self filterValue:0.0 forAU:[self AUAtIndex:i] withFace:false];
		}
	}
}

- (unsigned int)AUCnt
{
	static const unsigned int auCnt = 9;
	return auCnt;
}
- (int)AUAtIndex:(unsigned int)ind
{
	if( ind < [self AUCnt] )
		return auValues[ind];
	return 0;
}
- (unsigned int)indexForAU:(int)au
{
	unsigned int ind=-1;
	switch( au )
	{
		case 1: ind=0;
		case 2: ind=1;
		case 4: ind=2;
		case 5: ind=3;
		case 10: ind=4;
		case 12: ind=5;
		case 14: ind=6;
		case 20: ind=7;
		case 101: ind=8;
	}
	return ind;
}
- (NSString *)nameForAU:(int)au
{
	static const NSString *auNames[] = { @"AU 1", @"AU 2", @"AU 4", @"AU 5", @"AU 10", @"AU 12", @"AU 14", @"AU 20", @"Smile" };
	unsigned int ind = [self indexForAU:au];
	if( ind != -1 )
		return auNames[ind];
	
	return @"Unkown";
}
- (double)filterValue:(double)value forAU:(int)AU withFace:(bool)faceFound
{
	if( [self doFilters] )
	{
		ExpressionFilterWrapper *locFilter = static_cast<ExpressionFilterWrapper *>(filters);
		return locFilter->filterValue( AU, value, faceFound );
	}
	else
	{
		return value;
	}
}

- (NSArray *)results
{
	return results;
}

- (NSMutableDictionary *)pluginResults
{
    return pluginResults;
}

- (NSImage *)grayImage
{
	return grayImage;
}
- (NSImage *)videoImage
{
	return videoImage;
}

- (NSPoint)left
{
	return leftEye;
}
- (NSPoint)right
{
	return rightEye;
}
- (NSRect)face
{
	return face;
}
- (NSArray *)faces
{
	if( [faceArray count] > 0 )
		return faceArray;
	else
		return nil;
}
- (MPT_Face *)largestFace
{
	MPT_Face *largest = nil;
	float area = 0.0;
	NSEnumerator *faceEnum = [faceArray objectEnumerator];
	MPT_Face *aFace;
	while( aFace = [faceEnum nextObject] )
	{
		NSRect faceRect = [aFace face];
		float faceArea = faceRect.size.width * faceRect.size.height;
		if( faceArea > area )
		{
			largest = aFace;
			area = faceArea;
		}
	}
	return largest;
}

- (void)clearResults
{
	leftEye = NSMakePoint( -1.0, -1.0 );
	rightEye = leftEye;
	face = NSMakeRect( -1.0, -1.0, -1.0, -1.0 );
	[results removeAllObjects];
	[faceArray removeAllObjects];
	[pluginResults removeAllObjects];
}

- (void)setDoingSequence:(BOOL)seq
{
	((EYEFINDER_CLASS *)mpeyefinder)->setUseSequence( seq );
}

- (BOOL)findFaces
{
	return findFaces;
}
- (void)setFindFaces:(BOOL)value
{
	if( findFaces != value )
	{
		findFaces = value;
	}
}

- (BOOL)displayGray
{
	NSUserDefaults *prefs = [NSUserDefaults standardUserDefaults];
	bool gray = [prefs boolForKey:@"displayGray"];
	
    return gray;
}

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

- (float)eyefinderOffset
{
    return eyefinderOffset;
}
- (float)defaultEyefinderOffset
{
    return 1.0f; // cascadeThreshAdjustor::default_cascade_adjust[0];
}
- (void)setEyefinderOffset:(float)value
{
    eyefinderOffset = value;
	((EYEFINDER_CLASS *)mpeyefinder)->setCascadeThresholds( eyefinderOffset, eyefinderDropoff );
}

- (float)eyefinderDropoff
{
    return eyefinderDropoff;
}
- (float)defaultEyefinderDropoff
{
    return 1.0f; // cascadeThreshAdjustor::default_cascade_adjust[1];
}
- (void)setEyefinderDropoff:(float)value
{
    eyefinderDropoff = value;
	((EYEFINDER_CLASS *)mpeyefinder)->setCascadeThresholds( eyefinderOffset, eyefinderDropoff );
}

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

- (BOOL)linearCombo
{
    return linearCombo;
}

- (void)setLinearCombo:(BOOL)value
{
    linearCombo = value;
}

/** One or more plugins has had their enabled state changed.
 * We need to let the gabor processor know which SVM plugins are still enabled.
 */
- (void)pluginStateHasChanged:(id)plugin
{
	if( plugin )
	{
		if( MPT_PluginStep_SVM == [plugin step] )
		{
			std::string name( [[plugin internalName] UTF8String] );
			
			if( [plugin enabled] )
			{
				((GABOR_CLASS *)gabor)->addCategory( name );
			}
			else
			{
				((GABOR_CLASS *)gabor)->removeCategory( name );
			}
		}
	}
	else
	{
		std::vector< std::string > categories;
		NSArray *plugins = [[MPT_PluginController defaultController] pluginsForStep:MPT_PluginStep_SVM];
		NSEnumerator *pluginEnum = [plugins objectEnumerator];
		id plugin;
		
		while( plugin = [pluginEnum nextObject] )
		{
			if( [plugin enabled] )
			{
				std::string name( [[plugin internalName] UTF8String] );
				categories.push_back(name);
			}
		}
		((GABOR_CLASS *)gabor)->setCategories( categories );
	}

	[self loadFeaturePlugins];
}

- (BOOL)loadFeaturePlugins
{
#ifdef USING_FEATURE_DETECTOR
	NSArray *plugins = [[MPT_PluginController defaultController] pluginsForStep:MPT_PluginStep_Features];
	NSEnumerator *pluginEnum = [plugins objectEnumerator];
	id plugin;
	
	((EYEFINDER_CLASS *)mpeyefinder)->clearFeatures();
	
	[featuresByName removeAllObjects];
	
	while( plugin = [pluginEnum nextObject] )
	{
		[featuresByName setObject:plugin forKey:[plugin internalName]];
		
		if( [plugin enabled] )
		{
			NSString *weights = [plugin featureDataFile];
			NSString *priors = [plugin gpriorFile];
			
			((EYEFINDER_CLASS *)mpeyefinder)->loadFeatureFromFiles( [weights UTF8String], [priors UTF8String] );
		}
	}
#endif USING_FEATURE_DETECTOR
	return YES;
}

// This method is no longer being used.
- (void)rotateImage:(RImage<float> *)rimage into:(RImage<float> *)destImage byRadians:(float)rad
{
	// source image
	vImage_Buffer vImgSrc = {(void*)rimage->array, rimage->height, rimage->width, rimage->width*sizeof(float)};
				
	// destination image
	vImage_Buffer vImgDst = {(void*)destImage->array, destImage->height, destImage->width, destImage->width*sizeof(float)};
				
	Pixel_F black = 0.0;
	vImageRotate_PlanarF( &vImgSrc, &vImgDst,  NULL,  rad,  black, 
						  (kvImageBackgroundColorFill | kvImageHighQualityResampling) );				
}
@end