//
//  MPT_Scheduler.m
//  AUCoderDemo
//
//  Created by Andrew Salamon on 6/5/08.
//  Copyright 2008 Machine Perception Laboratory, University of California San Diego.. All rights reserved.
//

#import "MPT_Scheduler.h"

#import "MPT_CERT.h"
#import "MPT_Pair.h"

// For finding the number of available cpus
#include <mach/machine.h>
#include <mach/mach.h>
#include <mach/mach_host.h>
#include <mach/host_info.h>

#ifdef CERT_TIMING
#import "MPT_Timing.h"
#define STAMPTIME(ct_inputs,ct_id) [[(ct_inputs) objectForKey:@"CERT.timing"] stampWithID:(ct_id)]
#else
#define STAMPTIME
#endif

static int numThreads = -1;

typedef enum {
	PROC_AVAIL = 1,
	PROC_NOAVAIL,
	PROC_STARTING,
	PROC_FULL
} PROC_CONDITION;

@interface MPT_Scheduler(PrivateMethods)
- (void)makeAvailable:(id)proc;
- (id)getNextAvailable;
- (void)runProcessor:(MPT_Pair *)pair;
@end

@implementation MPT_Scheduler
+ (unsigned int)numThreads
{
    return numThreads;
}

+ (void)setNumThreads:(unsigned int)value
{
    numThreads = value;
}

/*
 On Mac OS X 10.5 or later we could use:
 [[NSProcessInfo processInfo] activeProcessorCount]
 */
+ (unsigned int)availableCPUs
{
	host_basic_info_data_t hostInfo;
	mach_msg_type_number_t infoCount;
	
	infoCount = HOST_BASIC_INFO_COUNT;
	host_info(mach_host_self(), HOST_BASIC_INFO, (host_info_t)&hostInfo, &infoCount);
	
	//		printf( "Max CPUs: %d\n", (unsigned int)hostInfo.max_cpus );
	//		printf( "Used CPUs: %d\n", (unsigned int)hostInfo.avail_cpus );
	return hostInfo.avail_cpus;
}

+ (void)initialize
{
	NSUserDefaults *_prefs = [NSUserDefaults standardUserDefaults];
	NSDictionary *defPrefs = [NSDictionary dictionaryWithObjectsAndKeys:
							  [NSNumber numberWithBool:YES], @"autoMaxThreads",
							  [NSNumber numberWithInt:[self availableCPUs]], @"maxThreads", nil];
	[_prefs registerDefaults:defPrefs];
}

- (id)initWithDelegate:(id)_delegate
{
	if( self = [super init] )
	{
		stopping = NO;
		delegate = _delegate;
		
		available = [[NSMutableArray alloc] init];
		processors = nil;
		
		NSNumber *autoMax = [[NSUserDefaults standardUserDefaults] objectForKey:@"autoMaxThreads"];
		if( !autoMax || [autoMax boolValue] )
		{
			maxThreads = [[self class] availableCPUs];
		}
		else
		{
			NSNumber *maxNum = [[NSUserDefaults standardUserDefaults] objectForKey:@"maxThreads"];
			if( !maxNum )
				maxThreads = [[self class] availableCPUs];
			else
				maxThreads = [maxNum intValue];
		}
		
		[[NSUserDefaults standardUserDefaults] addObserver:self forKeyPath:@"autoMaxThreads" options:NSKeyValueObservingOptionNew context:NULL];
		[[NSUserDefaults standardUserDefaults] addObserver:self forKeyPath:@"maxThreads" options:NSKeyValueObservingOptionNew context:NULL];
		
		condLock = [[NSConditionLock alloc] initWithCondition:PROC_STARTING];
	}
	
	return self;
}

- (void)dealloc
{
	[[NSUserDefaults standardUserDefaults] removeObserver:self forKeyPath:@"autoMaxThreads"];
	[[NSUserDefaults standardUserDefaults] removeObserver:self forKeyPath:@"maxThreads"];
	[processors release];
	processors = nil;
	[available release];
	available = nil;
	[condLock release];
	condLock = nil;
	[super dealloc];
}

- (void)observeValueForKeyPath:(NSString *)keyPath ofObject:(id)object change:(NSDictionary *)change context:(void *)context
{
	if( [keyPath isEqualToString:@"maxThreads"] )
	{
		NSNumber *autoMax = [[NSUserDefaults standardUserDefaults] objectForKey:@"autoMaxThreads"];		
		if( !autoMax || [autoMax boolValue] )
		{
			[self setMaxThreads:[[self class] availableCPUs]];
		}
		else
		{
			[self setMaxThreads:[[change objectForKey:NSKeyValueChangeNewKey] intValue]];
		}
	}
	else if( [keyPath isEqualToString:@"autoMaxThreads"] )
	{
		NSNumber *newAuto = [change objectForKey:NSKeyValueChangeNewKey];
		if( !newAuto || [newAuto boolValue] )
		{
			[self setMaxThreads:[[self class] availableCPUs]];
		}
		else
		{
			[self setMaxThreads:[[[NSUserDefaults standardUserDefaults] objectForKey:@"maxThreads"] intValue]];
		}
	}
	else
	{
#ifdef DEBUG
		NSLog( @"Received Unsupported key value change notification in MPT_Scheduler: %@", keyPath );
#endif
	}	
}

- (id)delegate
{
    return delegate;
}
- (void)setDelegate:(id)value
{
	delegate = value;
}

- (int)maxThreads {
    return maxThreads;
}

- (void)setMaxThreads:(int)value
{
	if( value < 1 ) value = [[self class] availableCPUs];
	
    if (maxThreads != value)
	{
        maxThreads = value;
		[self createProcessors];
    }
}

- (void)resetThreadsDefault
{
	[[NSUserDefaults standardUserDefaults] setInteger:maxThreads forKey:@"maxThreads"];
}

- (void)createProcessors
{
	if( [processors count] == maxThreads )
	{
#ifdef DEBUG
		NSLog( @"Processor Count: %d", [processors count] );
#endif
		return;
	}
	
	if( [processors count] > maxThreads )
	{
		// remove enough to equal maxThreads
		[condLock lock];
		NSMutableArray *tmp = [[processors mutableCopy] autorelease];
		[processors release];
		processors = nil;
		while( [tmp count] > maxThreads )
		{
			// just delete one from processors? That would work if we're not running any images
			// which we can check...
			id proc = [tmp lastObject];
			if( [available containsObject:proc] )
			{
				[available removeObject:proc];
				[tmp removeObject:proc];
			}
			else
			{
				processors = [tmp copy];
				maxThreads = [processors count];
				[condLock unlockWithCondition:(([available count] > 0)?PROC_AVAIL:PROC_NOAVAIL)];
				[self performSelector:@selector(resetThreadsDefault) withObject:nil afterDelay:0.1];
				NSLog( @"Unable to reduce processor count, some processors were in use." );
#ifdef DEBUG
				NSLog( @"Processor Count: %d", [processors count] );
#endif
				
				return;
			}
		}
		processors = [tmp copy];
		[condLock unlockWithCondition:(([available count] > 0)?PROC_AVAIL:PROC_NOAVAIL)];
#ifdef DEBUG
		NSLog( @"Processor Count: %d", [processors count] );
#endif
		
		return;
	}
	
	NSMutableArray *tmp = nil;
	if( [processors count] > 0 )
	{
		tmp = [[processors mutableCopy] autorelease];
		[processors release];
		processors = nil;
	}
	else
		tmp = [NSMutableArray array];
	
	[condLock lock];
	while( [tmp count] < maxThreads )
	{
		[delegate performSelectorOnMainThread:@selector(incrementWeightsProgressBy:) withObject:nil waitUntilDone:NO];
		if( stopping )
		{
			[condLock unlockWithCondition:PROC_STARTING];
			return;
		}
		MPT_CERT *cert = [[MPT_CERT alloc] init];
		[cert makeThreadsafe];
		[tmp addObject:cert];
		[available addObject:cert];
		[cert release];
	}
	processors = [tmp copy];
	
	[condLock unlockWithCondition:(([available count] > 0)?PROC_AVAIL:PROC_NOAVAIL)];
	
#ifdef DEBUG
	NSLog( @"Processor Count: %d", [processors count] );
#endif
}

- (BOOL)isEmpty
{
	BOOL scheduleIsEmpty = NO;
	[condLock lock];
	scheduleIsEmpty = [available count] == [processors count];
	[condLock unlockWithCondition:(([available count] > 0)?PROC_AVAIL:PROC_NOAVAIL)];
	return scheduleIsEmpty;
}

- (BOOL)isAvailable
{
	int cnt = 0;
	[condLock lock];
	cnt = [available count];
	[condLock unlockWithCondition:((cnt > 0)?PROC_AVAIL:PROC_NOAVAIL)];
	return cnt > 0;
}

- (void)makeAvailable:(id)proc
{
	BOOL scheduleIsEmpty = NO;
	[condLock lock];
	[available addObject:proc];
	scheduleIsEmpty = [available count] == [processors count];
	if( stopping )
	{
		if( scheduleIsEmpty )
			[condLock unlockWithCondition:PROC_FULL];
		else
			[condLock unlockWithCondition:PROC_AVAIL];
	}
	else
	{
		[condLock unlockWithCondition:PROC_AVAIL];
		if( scheduleIsEmpty )
			[delegate performSelectorOnMainThread:@selector(scheduleIsEmpty) withObject:nil waitUntilDone:NO];
	}
}

- (id)getNextAvailable
{
	[condLock lockWhenCondition:PROC_AVAIL];
	id proc = [available lastObject];
	[available removeLastObject];
	[condLock unlockWithCondition:(([available count] > 0)? PROC_AVAIL: PROC_NOAVAIL)];
	return proc;
}

- (BOOL)doesScheduleInput:(NSDictionary *)inputs
{
	[self scheduleInput:inputs];
	return YES;
}

- (void)scheduleInput:(NSDictionary *)inputs
{
	if( stopping ) return;
	
	if( oneProc )
	{
		[oneProc process:inputs];
		[delegate handleResults:[oneProc pluginResults]];
		[delegate performSelector:@selector(scheduleIsEmpty)];
	}
	else
	{
		STAMPTIME(inputs,@"Before getNextAvailable");
		id proc = [self getNextAvailable];
		STAMPTIME(inputs,@"After getNextAvailable");

		MPT_Pair *pair = [MPT_Pair pairWith:proc and:inputs];
		
//		NSLog( @"Detaching cert: %d", proc );
        STAMPTIME(inputs,@"Before runProcessor");
		[NSThread detachNewThreadSelector:@selector(runProcessor:) toTarget:self withObject:pair];
        STAMPTIME(inputs,@"After runProcessor");
	}
}

- (void)runProcessor:(MPT_Pair *)pair
{
	NSAutoreleasePool *pool = [[NSAutoreleasePool alloc] init];
	id proc = [pair first];
	id inputs = [pair second];
	
	//NSLog( @"Default thread priority: %lf", [NSThread threadPriority] );
	// Setting the priority to 1.0 should have sped CERT up, but all it seems to do is slow the GUI way down, without any increase in framerate
	//[NSThread setThreadPriority:1.0];
	
	STAMPTIME(inputs,@"Before process");
	[proc process:inputs];
	STAMPTIME(inputs,@"After process");

//	NSLog( @"Finishing cert: %d", proc );
	if( !stopping )
		[delegate performSelectorOnMainThread:@selector(handleUnsortedResults:) withObject:[proc pluginResults] waitUntilDone:NO];
	
	STAMPTIME(inputs,@"Before makeAvailable");
	[self makeAvailable:proc];
	STAMPTIME(inputs,@"After makeAvailable");

	[pool release];
}

/// Give all CERT processors a chance to stop otherwise we can get crashes when quitting.
- (void)cleanup
{
	stopping = YES;
	
	[condLock lock];
	if( [available count] != [processors count] )
	{
		NSLog( @"Not all CERT processors are finished." );
		[condLock unlockWithCondition:PROC_NOAVAIL];
		[condLock lockWhenCondition:PROC_FULL beforeDate:[NSDate dateWithTimeIntervalSinceNow:(2.5 * maxThreads)]];
		NSLog( @"All CERT processors are now finished." );
	}
	
	[available removeAllObjects];
	available = nil;
	[processors release];
	processors = nil;
	[condLock unlockWithCondition:PROC_NOAVAIL]; // probably a waste, since we're quitting, but it's safer this way.
}
@end
