//
//  MPT_GKRPlugin.m
//  MPT_GKR
//
//  Created by Andrew Salamon on 8/27/08.
//  Copyright 2008 Machine Perception Laboratory, University of California San Diego.. All rights reserved.
//

#import "MPT_GKRPlugin.h"

#import <DemoPlugin/MPT_AUView.h>

@interface MPT_GKRPlugin(PrivateMethods)
- (IBAction)saveTrainingData:(id)sender;
@end

@implementation MPT_GKRPlugin
+ (void)initialize
{
	// Make sure the correct alignment method is used by default unless the user overrides it.
	NSUserDefaults *_prefs = [NSUserDefaults standardUserDefaults];
	NSDictionary *defPrefs = [NSDictionary dictionaryWithObjectsAndKeys:
							  @"edu.ucsd.mplab.plugins.CERT4_0_Weights", @"targetID",
							  @"edu.ucsd.mplab.plugins.PoseDetector", @"poseID",
							  [NSNumber numberWithBool:NO], @"GKR",
							  [NSNumber numberWithBool:YES], @"logistic",
							  [NSNumber numberWithBool:YES], @"CKF",
							  nil];
	NSString *ident = [[NSBundle bundleForClass:[self class]] objectForInfoDictionaryKey:(NSString *)kCFBundleIdentifierKey];
	if( [ident length] > 0 )
		[_prefs registerDefaults:[NSDictionary dictionaryWithObject:defPrefs forKey:ident]];
}

- (id)initWithBundle:(NSBundle *)_bundle
{
	if( self = [super initWithBundle:_bundle] )
	{
		saving = NO;
		
		if( 0 == [[self targetID] length] )
			[self setTargetID:@"edu.ucsd.mplab.plugins.CERT4_0_Weights"];
		if( 0 == [[self poseID] length] )
			[self setPoseID:@"edu.ucsd.mplab.plugins.PoseDetector"];

		[self setTraining:NO];
		
		models = new GKR::Models;

		if( ![NSBundle loadNibNamed:@"GKRPrefPane" owner:self] )
			NSLog( @"Unable to load pref pane for MPT_GKR plugin." );
		
		models->setFlags( [self training], [self GKR], [self logistic], [self CKF] );
		
		if( [[self modelPath] length] > 0 )
			models->loadModels( [[self modelPath] UTF8String] );
		
		NSString *ckfPath = [_bundle pathForResource:@"CKFmap" ofType:@"xml"];
		if( [ckfPath length] > 0 )
		{
			std::string ckfString( [ckfPath UTF8String] );
			if( !models->loadCKFValues( [ckfPath UTF8String] ) )
			{
				NSLog( @"Unable to load the CKF data for the GKR plugin." );
			}
		}
	}
	return self;
}

- (id)copyWithZone:(NSZone *)zone
{
	return nil;
}

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

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

- (NSString *)step
{
	return MPT_PluginStep_Results;
}

- (NSString *)localizedName
{
	return NSLocalizedStringFromTableInBundle(@"GKR", nil, [NSBundle bundleForClass:[self class]], @"");
}

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

- (NSView *)prefPane
{
	return prefView;
}

- (void)setView:(MPT_AUView *)newView forLabel:(NSString *)label
{
	[super setView:newView forLabel:label];
	
	[newView setDefaultMin:0.0];
	[newView setDefaultMax:1.0];
	[newView setMin:0.0];
	[newView setMax:1.0];
	[newView resetDynamicRange];
}

- (void)pluginsFinishedLoading:(NSArray *)plugins
{
	id<MPT_PluginProtocol> targetPlugin = nil; // [[self delegate] pluginForPluginID:targetID];
	
	NSEnumerator *pluginEnum = [plugins objectEnumerator];
	id<MPT_PluginProtocol> plugin;
	while( !targetPlugin && (plugin = [pluginEnum nextObject]) )
	{
		if( [[plugin pluginID] isEqualToString:[self targetID]] && [[plugin step] isEqualToString:MPT_PluginStep_SVM] && (plugin != self) )
		{
			targetPlugin = plugin;
			
			targetLabels = [[plugin labels] copy];
			
			NSMutableArray *myLabels = [NSMutableArray array];
			NSEnumerator *labelEnum = [targetLabels objectEnumerator];
			NSString *label;
			std::vector<std::string> labelVec;
			while( label = [labelEnum nextObject] )
			{
				[myLabels addObject:[@"GKR: " stringByAppendingString:label]];
				labelVec.push_back( std::string( [label UTF8String] ) );
			}
			[labels release];
			labels = [myLabels copy];
			models->setLabels( labelVec );
		}
	}
	[self initViews];
}

- (NSMutableDictionary *)processResults:(NSDictionary *)results usingController:(id<MPT_PluginControllerProtocol>)controller
{
	if( saving ) return nil;
	
	BOOL faceFound = ([[results objectForKey:@"CERT.faces"] count] > 0);
	std::vector<double> pose;
	std::vector<double> aus;
	
	if( faceFound )
	{
		// Fill the two vectors with results from the pose and target plugins.
		id poseRes = [results objectForKey:[self poseID]];
		id auRes = [results objectForKey:[self targetID]];

		if( poseRes && auRes )
		{		
			if( ![poseRes isKindOfClass:[NSDictionary class]] || ![auRes isKindOfClass:[NSDictionary class]] )
			{
				NSLog( @"Unexpected result types in GKR Plugin." );
				return nil;
			}
			
			if( ([poseRes count] != 3) || ([auRes count] != [targetLabels count]) )
			{
				NSLog( @"Unexpected result counts in GKR Plugin." );
				return nil;
			}
			
			pose.push_back( [[poseRes objectForKey:@"Yaw"] doubleValue] );
			pose.push_back( [[poseRes objectForKey:@"Pitch"] doubleValue] );
			pose.push_back( [[poseRes objectForKey:@"Roll"] doubleValue] );
			
			NSEnumerator *labelEnum = [targetLabels objectEnumerator];
			NSString *label;
			while( label = [labelEnum nextObject] )
			{
				aus.push_back( [[auRes objectForKey:label] doubleValue] );
			}
		}
		else
		{
			// Fall through and treat it just like a face not found.
			faceFound = NO;
		}
	}

    std::vector<double> gkrRes = models->addFrame( pose, aus, faceFound );
	
	NSMutableDictionary *resDict = nil;
	if( !training )
	{
		if( gkrRes.size() != [[self labels] count] )
		{
			NSLog( @"Length mis-match between GKR results and labels in GKR Plugin." );
			return nil;
		}
		
		resDict = [NSMutableDictionary dictionary];
		
		NSEnumerator *labelEnum = [[self labels] objectEnumerator];
		NSString *label;
		int cnt = 0;
		while( label = [labelEnum nextObject] )
		{
			[resDict setObject:[NSNumber numberWithDouble:gkrRes[cnt++]] forKey:label];
		}
	}
	
	return resDict;
}

- (IBAction)startTraining:(id)sender
{
	[self setTraining:YES];
}
- (IBAction)stopTraining:(id)sender
{
	saving = YES;
	[self setTraining:NO];
	[self saveTrainingData:self];
}

- (IBAction)saveTrainingData:(id)sender
{
	if( [delegate isEvaluation] )
	{
		models->train();
		saving = NO;
	}
	else
	{
		NSSavePanel *panel = [NSSavePanel savePanel];
		[panel setPrompt:@"Save Training Data"];
		[panel setTitle:@"Save Training Data"];
		
		NSWindow *window = [[self prefPane] window];
		[window makeKeyAndOrderFront:self];
		[panel beginSheetForDirectory:nil
								 file:nil
					   modalForWindow:window
						modalDelegate:self
					   didEndSelector:@selector(savePanelDidEnd:returnCode:contextInfo:)
						  contextInfo:nil];
	}
}
- (void)savePanelDidEnd:(NSOpenPanel *)panel returnCode:(int)returnCode  contextInfo:(void  *)contextInfo
{
	if( returnCode == NSOKButton )
	{
		NSArray *filenames = [panel filenames];
		if( [filenames count] > 0 )
		{
			NSString *path = [filenames objectAtIndex:0];
			if( [path length] > 0 )
			{
				models->train();
				models->saveModels( [path UTF8String] );
				// Automatically set the modelPath to the new model (without reloading it)?
			}
		}
	}
	else
	{
		models->clearTraining();
//		if( [[self modelPath] length] > 0 )
//			models->loadModels( [[self modelPath] UTF8String] );
	}
	saving = NO;
}

- (BOOL)GKR
{
    return [[self getPrefForKey:@"GKR"] boolValue];
}
- (void)setGKR:(BOOL)value
{
    if( [self GKR] != value )
	{
		[self setPrefForKey:@"GKR" to:[NSNumber numberWithBool:value]];
		models->setGKR( value );
    }
}

- (BOOL)logistic
{
    return [[self getPrefForKey:@"logistic"] boolValue];
}
- (void)setLogistic:(BOOL)value
{
    if( [self logistic] != value )
	{
		[self setPrefForKey:@"logistic" to:[NSNumber numberWithBool:value]];
		models->setLogistic( value );
    }
}

- (BOOL)CKF
{
    return [[self getPrefForKey:@"CKF"] boolValue];
}
- (void)setCKF:(BOOL)value
{
    if( [self CKF] != value )
	{
		[self setPrefForKey:@"CKF" to:[NSNumber numberWithBool:value]];
		models->setTemporal( value );
    }
}

- (BOOL)training {
    return training;
}
- (void)setTraining:(BOOL)value {
    if (training != value) {
        training = value;
		models->setTraining( value );
    }
}

- (NSString *)targetID
{
    return [self getPrefForKey:@"targetID"];
}
- (void)setTargetID:(NSString *)value
{
    if( ![[self targetID] isEqualToString:value] )
	{
		[self setPrefForKey:@"targetID" to:value];
    }
}

- (NSString *)poseID
{
    return [self getPrefForKey:@"poseID"];
}
- (void)setPoseID:(NSString *)value
{
    if( ![[self poseID] isEqualToString:value] )
	{
		[self setPrefForKey:@"poseID" to:value];
    }
}

- (NSString *)modelPath
{
    return [self getPrefForKey:@"modelPath"];
}
- (void)setModelPath:(NSString *)value
{
    if( ![[self modelPath] isEqualToString:value] )
	{
		[self setPrefForKey:@"modelPath" to:value];
        if( ![self training] )
            models->loadModels( [value UTF8String] );
    }
}

- (IBAction)chooseModelPath:(id)sender
{
	NSWindow *window = [[self prefPane] window];
	
	[window makeKeyAndOrderFront:self];
	
	NSOpenPanel *panel = [NSOpenPanel openPanel];
	[panel setCanChooseDirectories:NO];
	[panel setAllowsMultipleSelection:NO];
	[panel setCanChooseFiles:YES];
	[panel setMessage:@"Choose a GKR Model file."];
		
	[panel beginSheetForDirectory:nil
							 file:nil
							types:[NSArray arrayWithObjects:@"xml", @"XML", @"txt", @"TXT", @"text", @"TEXT", @"GKR", @"gkr", nil]
				   modalForWindow:window
					modalDelegate:self
				   didEndSelector:@selector(openModelPanelDidEnd:returnCode:contextInfo:)
					  contextInfo:nil];
}
- (void)openModelPanelDidEnd:(NSOpenPanel *)panel returnCode:(int)returnCode  contextInfo:(void  *)contextInfo
{
	if( returnCode == NSOKButton )
	{
		NSArray *filenames = [panel filenames];
		if( [filenames count] > 0 )
		{
			NSString *path = [filenames objectAtIndex:0];
			if( [path length] > 0 )
			{
				[self setModelPath:path];
			}
		}
	}
}

@end