//
//  MPT_Emotions2.m
//  Emotions detector for AFECT 1.0
//
//  Created by Andrew Salamon on Nov 16, 2010.
//  Copyright 2010 Machine Perception Technologies. All rights reserved.
//

#import "MPT_Emotions2.h"

//#import <DemoPlugin/NSImageToRimage.h>

@interface MPT_Emotions2(MPT_Emotions2_PrivateMethods)
-(void)makeNewEmotions;
- (void)setAuPluginID:(NSString *)value;
- (void)updateMatrix;
@end

@implementation MPT_Emotions2
- (id)initWithBundle:(NSBundle *)_bundle
{
	if( self = [super initWithBundle:_bundle] )
	{
		NSAssert( ([[self labels] count] == mpt::Emotions2::emotionCount),
				 @"Number of labels doesn't match number of emotions in MPT_Emotions2" );
		// This is the ID of the plugin we need to get the AU values from.
		[self setAuPluginID:@"com.mpt4u.afect.plugins.facs4_4"];
		[self makeNewEmotions];
		
		if( ![NSBundle loadNibNamed:@"EmotionsPref.nib" owner:self] )
		{
			NSDictionary *context = [NSDictionary dictionaryWithObject:self forKey:@"NSOwner"];
			if( ![[self bundle] loadNibFile:@"EmotionsPref.nib" externalNameTable:context withZone:[self zone]] )
				NSLog( @"Unable to load pref pane for Emotion Detector plugin (%@).", [self internalName] );			
		}
		
		[self updateMatrix];
	}
	
	return self;
}

- (void)dealloc
{
	delete emotions;
	emotions = nil;
	[super dealloc];
}

- (void)awakeFromNib
{
	[self updateMatrix];
}

- (id)copyWithZone:(NSZone *)zone
{
	id cp = [super copyWithZone:zone];
	return cp;
}

- (void)makeThreadsafe
{
	[super makeThreadsafe];
	emotions = nil;
	[self makeNewEmotions];
}

- (NSString *)localizedName
{
	return NSLocalizedStringFromTableInBundle(@"Basic Emotions 4.4", nil, [self bundle], @"");
}

- (NSString *)internalName
{
	return @"emotions2";
}

- (NSString *)step
{
	return MPT_PluginStep_Results; 
}

- (NSView *)prefPane
{
	return prefPane;
}

- (NSArray *)dependencies
{
	return [[[NSArray arrayWithObject:[self auPluginID]] retain] autorelease];
}

- (NSString *)auPluginID {
    return [[auPluginID retain] autorelease];
}

- (void)setAuPluginID:(NSString *)value {
    if (auPluginID != value) {
        [auPluginID release];
        auPluginID = [value copy];
    }
}

- (void)updateMatrix
{
	[linearButton setIntValue:NO];
	[sigmoidButton setIntValue:NO];
	[softmaxButton setIntValue:NO];
	switch( [self activation] )
	{
		case mpt::Emotions2::LINEAR: [linearButton setIntValue:YES]; break;
		case mpt::Emotions2::SIGMOID: [sigmoidButton setIntValue:YES]; break;
		case mpt::Emotions2::SOFTMAX: [softmaxButton setIntValue:YES]; break;
	}
}

- (IBAction)matrixChange:(id)sender
{
	int row = [sender selectedRow];
	[self setActivation:row];
}

- (int)activation
{
	return emotions->getActivation();
}
- (void)setActivation:(int)activ
{
	emotions->setActivation(((mpt::Emotions2::activationtype)activ));
}

-(void)makeNewEmotions
{
	NSString *iname = [self internalName];
	NSString *weights11Path = [[self bundle] pathForResource:[NSString stringWithFormat:@"%@_11", iname] ofType:@"txt"];
	NSString *weights8Path = [[self bundle] pathForResource:[NSString stringWithFormat:@"%@_8", iname] ofType:@"txt"];
	NSString *weightsMorePath = [[self bundle] pathForResource:[NSString stringWithFormat:@"%@_more", iname] ofType:@"txt"];
	NSString *thresholdPath = [[self bundle] pathForResource:[NSString stringWithFormat:@"%@_threshold", iname] ofType:@"txt"];

	if( emotions )
		delete emotions;
	
	if( 0 == [weights11Path length] )
		NSLog( @"Unable to read weights file 11 from %@.", weights11Path );
	if( 0 == [weights8Path length] )
		NSLog( @"Unable to read weights file 8 from %@.", weights8Path );
	if( 0 == [weightsMorePath length] )
		NSLog( @"Unable to read weights file more from %@.", weightsMorePath );
	if( 0 == [thresholdPath length] )
		NSLog( @"Unable to read threshold file from %@.", thresholdPath );

	if( (0 == [weights11Path length])
	   || (0 == [weights8Path length])
	   || (0 == [weightsMorePath length])
	   || (0 == [thresholdPath length]) )
	{
		NSLog( @"Unable to read one or more weights file in %@ plugin.", [self localizedName] );
		return;
	}

	emotions =  new mpt::Emotions2( [weights11Path UTF8String], [weights8Path UTF8String], [weightsMorePath UTF8String], [thresholdPath UTF8String] );
}

- (NSMutableDictionary *)processResults:(NSDictionary *)results usingController:(id<MPT_PluginControllerProtocol>)controller;
{
	NSMutableDictionary *res = [NSMutableDictionary dictionary];
	
	NSDictionary *aus = [results objectForKey:[self auPluginID]];
	
	if( !aus )
	{
		// Maintain support for the older plugins if the newer one isn't available
		aus = [results objectForKey:@"edu.ucsd.mplab.plugins.CERT4_4_Weights"];
		NSDictionary *addAUs = [results objectForKey:@"edu.ucsd.mplab.plugins.AdditionalAUs3"];
		
		if( !aus || !addAUs )
			return res;
		
		NSMutableDictionary *temp = [NSMutableDictionary dictionaryWithDictionary:aus];
		[temp addEntriesFromDictionary:addAUs];
		aus = temp;
	}
	
	NSNumber *val = nil;
	std::vector<float> inputs;
	std::vector<float> add;
	std::vector<float> more( mpt::Emotions2::moreCount, 0.0 );
	
	// Find the inputs we need and add them to a vector

	val = [aus objectForKey:@"(AU 1) Inner Brow Raise"];
	if( val ) inputs.push_back( [val floatValue] );
	
	val = [aus objectForKey:@"(AU 2) Outer Brow Raise"];
	if( val ) inputs.push_back( [val floatValue] );
	
	val = [aus objectForKey:@"(AU 4) Brow Lower"];
	if( val ) inputs.push_back( [val floatValue] );
	
	val = [aus objectForKey:@"(AU 5) Eye Widen"];
	if( val ) inputs.push_back( [val floatValue] );
	
	val = [aus objectForKey:@"(AU 9) Nose Wrinkle"];
	if( val ) inputs.push_back( [val floatValue] );
	
	val = [aus objectForKey:@"(AU 10) Lip Raise"];
	if( val ) inputs.push_back( [val floatValue] );
	
	val = [aus objectForKey:@"(AU 12) Lip Corner Pull"];
	if( val ) inputs.push_back( [val floatValue] );
	
	val = [aus objectForKey:@"(AU 14) Dimpler"];
	if( val ) inputs.push_back( [val floatValue] );
	
	val = [aus objectForKey:@"(AU 15) Lip Corner Depressor"];
	if( val ) inputs.push_back( [val floatValue] );
	
	val = [aus objectForKey:@"(AU 17) Chin Raise"];
	if( val ) inputs.push_back( [val floatValue] );
	
	val = [aus objectForKey:@"(AU 20) Lip stretch"];
	if( val ) inputs.push_back( [val floatValue] );
	
	
	val = [aus objectForKey:@"(AU 6) Cheek Raise"];
	if( val ) add.push_back( [val floatValue] );
	
	val = [aus objectForKey:@"(AU 7) Lids Tight"];
	if( val ) add.push_back( [val floatValue] );
	
	val = [aus objectForKey:@"(AU 18) Lip Pucker"];
	if( val ) add.push_back( [val floatValue] );
	
	val = [aus objectForKey:@"(AU 23) Lip Tightener"];
	if( val ) add.push_back( [val floatValue] );
	
	val = [aus objectForKey:@"(AU 24) Lip Presser"];
	if( val ) add.push_back( [val floatValue] );
	
	val = [aus objectForKey:@"(AU 25) Lips Part"];
	if( val ) add.push_back( [val floatValue] );
	
	val = [aus objectForKey:@"(AU 26) Jaw Drop"];
	if( val ) add.push_back( [val floatValue] );
	
	val = [aus objectForKey:@"(AU 28) Lips Suck"];
	if( val ) add.push_back( [val floatValue] );
	
	// now fill in whichever parts of the 'more' vector is available. Rest will be zero.
	int ind = 0;
	
	val = [aus objectForKey:@"(AU 45) Blink/Eye Closure"];
	if( val ) more[ind] = [val floatValue];
	++ind;
	
	val = [aus objectForKey:@"Fear Brow (1+2+4)"];
	if( val ) more[ind] = [val floatValue];
	++ind;
	
	val = [aus objectForKey:@"Distress Brow (1, 1+4)"];
	if( val ) more[ind] = [val floatValue];
	++ind;
	
	NSDictionary *unilats = [results objectForKey:@"edu.ucsd.mplab.plugins.Unilaterals"];
	if( unilats )
	{
		val = [unilats objectForKey:@"AU 10 Left"];
		if( val ) more[ind] = [val floatValue];
		++ind;
		
		val = [unilats objectForKey:@"AU 12 Left"];
		if( val ) more[ind] = [val floatValue];
		++ind;
		
		val = [unilats objectForKey:@"AU 14 Left"];
		if( val ) more[ind] = [val floatValue];
		++ind;
		
		val = [unilats objectForKey:@"AU 10 Right"];
		if( val ) more[ind] = [val floatValue];
		++ind;
		
		val = [unilats objectForKey:@"AU 12 Right"];
		if( val ) more[ind] = [val floatValue];
		++ind;
		
		val = [unilats objectForKey:@"AU 14 Right"];
		if( val ) more[ind] = [val floatValue];
		++ind;
		
	}
	else
	{
		ind += 6;
	}

	NSNumber *gender = [results objectForKey:@"edu.ucsd.mplab.plugins.GenderDetector"];
	if( gender ) more[ind] = [gender floatValue];
	++ind;
	
	NSNumber *glasses = [results objectForKey:@"edu.ucsd.mplab.plugins.GlassesDetector"];
	if( glasses ) more[ind] = [glasses floatValue];
	++ind;
	
	NSDictionary *pose = [results objectForKey:@"edu.ucsd.mplab.plugins.PoseDetector"];
	if( pose )
	{
		val = [pose objectForKey:@"Yaw"];
		if( val ) more[ind] = [val floatValue];
		++ind;

		val = [pose objectForKey:@"Pitch"];
		if( val ) more[ind] = [val floatValue];
		++ind;

		val = [pose objectForKey:@"Roll"];
		if( val ) more[ind] = [val floatValue];
		++ind;
	}
	else
	{
		ind += 3;
	}
	
	NSNumber *smile = [results objectForKey:@"edu.ucsd.mplab.plugins.SmileDetector2"];
	if( smile ) more[ind] = [smile floatValue];
	++ind;
		
	if( (inputs.size() == mpt::Emotions2::auCount) && (add.size() == mpt::Emotions2::addAuCount) )
	{
		NSAssert( emotions, @"No emotions detector." );
		std::vector<float> ans = emotions->calc( inputs, add, more );
		
		NSEnumerator *labelEnum = [[self labels] objectEnumerator];
		NSString *label;
		int i = 0;
		while( label = [labelEnum nextObject] )
		{
			[res setObject:[NSNumber numberWithFloat:ans[i++]] forKey:label];
		}
	}
	else
	{
		NSAssert( (inputs.size() == mpt::Emotions2::auCount), @"Number of input aus is incorrect in MPT_Emotions2" );
		NSAssert( (add.size() == mpt::Emotions2::addAuCount), @"Number of additional aus is incorrect in MPT_Emotions2" );
	}
	return res;
}

@end