//
//  EyeDeeThreeLearner.m
//  Pennyworth
//
//  Created by Chris Karr on 1/22/08.
//  Copyright 2008 Chris J. Karr. All rights reserved.
//

#import "DecisionTreeLearner.h"

#define EXAMPLES_KEY @"Training Examples"
#define FEATURE_KEY @"Feature"

#define LEARNER_PATH [NSString stringWithFormat:@"%@/Library/Application Support/Pennyworth/Decision Tree (KEY).learner", NSHomeDirectory ()]

NSInteger TreeSort (id one, id two, void * context)
{
	if (([one isEqual:@"true"] || [one isEqual:@"false"]) &&
		([two isEqual:@"true"] || [two isEqual:@"false"]))
		return [two compare:one];
	
	return [one compare:two];
}


@implementation DecisionTreeLearner

@synthesize examples;
@synthesize decisionTree;

- (NSString *) bestFeatureForExamples:(NSArray *) exampleArray
{
	NSMutableSet * features = [NSMutableSet set];
	
	for (NSDictionary * example in exampleArray)
		[features addObjectsFromArray:[example allKeys]];
		
	[features removeObject:LABEL_KEY];
	
	return [features anyObject];
}

- (NSMutableDictionary *) treeForExamples:(NSArray *) exampleArray
{
	NSMutableDictionary * node = [NSMutableDictionary dictionary];
	
	NSString * last = nil;
	BOOL match = YES;
	
	for (NSDictionary * example in exampleArray)
	{
		if (last == nil)
			last = [example valueForKey:LABEL_KEY];
		else if (![last isEqual:[example valueForKey:LABEL_KEY]])
			match = NO;
	}
	
	if (match)
		[node setValue:last forKey:LABEL_KEY];
	else
	{
		NSString * feature = [self bestFeatureForExamples:exampleArray];
		
		NSMutableDictionary * matchingExamples = [NSMutableDictionary dictionary];
		
		if (feature != nil)
		{
			[node setValue:feature forKey:FEATURE_SENSOR];

			for (NSMutableDictionary * example in exampleArray)
			{
				NSMutableDictionary * thisExample = [NSMutableDictionary dictionaryWithDictionary:[example copy]];
			
				NSObject * value = [thisExample valueForKey:feature];
			
				if (value == nil)
				{
					[thisExample setValue:UNKNOWN_VALUE forKey:feature];
					value = UNKNOWN_VALUE;
				}
			
				NSMutableArray * values = [NSMutableArray array];
				
				if ([value isKindOfClass:[NSArray array]])
					[values addObjectsFromArray:(NSArray *) value];
				else
					[values addObject:value];

				for (NSString * observation in values)
				{
					NSMutableArray * matches = [matchingExamples valueForKey:observation];
			
					if (matches == nil)
					{
						matches = [NSMutableArray array];
						[matchingExamples setValue:matches forKey:observation];
					}
				
					[matches addObject:thisExample];
				}
				
				[thisExample removeObjectForKey:feature];
			}

			for (NSString * exampleKey in [matchingExamples allKeys])
				[node setValue:[self treeForExamples:[matchingExamples valueForKey:exampleKey]] forKey:exampleKey];
		}
	}
	
	return node;
}

- (CGFloat) prune:(NSMutableDictionary *) tree
{
	NSLog (@"default implementation doesn't prune");
	
	return 0;
}

- (DecisionTreeLearner *) init
{
	if (self = [super init])
	{
		[self load];
	}
	
	return self;
}

- (void) load
{
	BOOL isDir = NO;
		
	NSFileManager * fm = [NSFileManager defaultManager];
	
	NSString * path = [LEARNER_PATH stringByReplacingOccurrencesOfString:@"KEY" withString:key];
		
	[fm fileExistsAtPath:[path stringByDeletingLastPathComponent] isDirectory:&isDir];

	if (!isDir)
		[fm createDirectoryAtPath:[path stringByDeletingLastPathComponent] attributes:nil];

	if ([fm fileExistsAtPath:path isDirectory:&isDir])
	{
		if (self.examples == nil)
		{
			NSData * data = [NSData dataWithContentsOfFile:path];
			NSKeyedUnarchiver * unarchiver = [[NSKeyedUnarchiver alloc] initForReadingWithData:data];
			
			self.examples = [NSMutableArray arrayWithArray:[unarchiver decodeObjectForKey:EXAMPLES_KEY]];
			
			[unarchiver release];
		}
	}
	else
		self.examples = [NSMutableArray array];
		
	self.decisionTree = [self treeForExamples:self.examples];
	
	[self prune:self.decisionTree];
}

/* - (void) awakeFromNib
{
	NSMutableArray * testExamples = [NSMutableArray array];
	
	[testExamples addObject:[NSMutableDictionary dictionaryWithObjectsAndKeys:
								@"No", LABEL_KEY,
								@"Rain", @"Outlook",
								@"Mild", @"Temp",
								@"High", @"Hum", 
								@"Strong", @"Wind", nil]];

	[testExamples addObject:[NSMutableDictionary dictionaryWithObjectsAndKeys:
								@"Yes", LABEL_KEY,
								@"Overcast", @"Outlook",
								@"Hot", @"Temp",
								@"Normal", @"Hum", 
								@"Weak", @"Wind", nil]];

	[testExamples addObject:[NSMutableDictionary dictionaryWithObjectsAndKeys:
								@"Yes", LABEL_KEY,
								@"Overcast", @"Outlook",
								@"Mild", @"Temp",
								@"High", @"Hum", 
								@"Strong", @"Wind", nil]];

	[testExamples addObject:[NSMutableDictionary dictionaryWithObjectsAndKeys:
								@"Yes", LABEL_KEY,
								@"Sunny", @"Outlook",
								@"Mild", @"Temp",
								@"Normal", @"Hum", 
								@"Strong", @"Wind", nil]];

	[testExamples addObject:[NSMutableDictionary dictionaryWithObjectsAndKeys:
								@"Yes", LABEL_KEY,
								@"Rain", @"Outlook",
								@"Mild", @"Temp",
								@"Normal", @"Hum", 
								@"Weak", @"Wind", nil]];

	[testExamples addObject:[NSMutableDictionary dictionaryWithObjectsAndKeys:
								@"Yes", LABEL_KEY,
								@"Sunny", @"Outlook",
								@"Cool", @"Temp",
								@"Normal", @"Hum", 
								@"Weak", @"Wind", nil]];

	[testExamples addObject:[NSMutableDictionary dictionaryWithObjectsAndKeys:
								@"No", LABEL_KEY,
								@"Sunny", @"Outlook",
								@"Mild", @"Temp",
								@"High", @"Hum", 
								@"Weak", @"Wind", nil]];

	[testExamples addObject:[NSMutableDictionary dictionaryWithObjectsAndKeys:
								@"Yes", LABEL_KEY,
								@"Overcast", @"Outlook",
								@"Cool", @"Temp",
								@"Normal", @"Hum", 
								@"Strong", @"Wind", nil]];

	[testExamples addObject:[NSMutableDictionary dictionaryWithObjectsAndKeys:
								@"No", LABEL_KEY,
								@"Rain", @"Outlook",
								@"Cool", @"Temp",
								@"Normal", @"Hum", 
								@"Strong", @"Wind", nil]];

	[testExamples addObject:[NSMutableDictionary dictionaryWithObjectsAndKeys:
								@"Yes", LABEL_KEY,
								@"Rain", @"Outlook",
								@"Cool", @"Temp",
								@"Normal", @"Hum", 
								@"Weak", @"Wind", nil]];

	[testExamples addObject:[NSMutableDictionary dictionaryWithObjectsAndKeys:
								@"Yes", LABEL_KEY,
								@"Rain", @"Outlook",
								@"Mild", @"Temp",
								@"High", @"Hum", 
								@"Weak", @"Wind", nil]];

	[testExamples addObject:[NSMutableDictionary dictionaryWithObjectsAndKeys:
								@"Yes", LABEL_KEY,
								@"Overcast", @"Outlook",
								@"Hot", @"Temp",
								@"High", @"Hum", 
								@"Weak", @"Wind", nil]];

	[testExamples addObject:[NSMutableDictionary dictionaryWithObjectsAndKeys:
								@"No", LABEL_KEY,
								@"Sunny", @"Outlook",
								@"Hot", @"Temp",
								@"High", @"Hum", 
								@"Strong", @"Wind", nil]];

	[testExamples addObject:[NSMutableDictionary dictionaryWithObjectsAndKeys:
								@"No", LABEL_KEY,
								@"Sunny", @"Outlook",
								@"Hot", @"Temp",
								@"High", @"Hum", 
								@"Weak", @"Wind", nil]];

	NSLog (@"tree = %@", [self treeForExamples:testExamples]);
} */

- (NSXMLElement *) getXML:(NSDictionary *) tree
{
	NSXMLElement * element = [[NSXMLElement alloc] initWithName:@"node"];

	NSString * label = [tree valueForKey:LABEL_KEY];

	// [element addAttribute:[NSXMLNode attributeWithName:@"count" stringValue:[[tree valueForKey:EXAMPLE_COUNT] description]]];
	 
	if (label != nil)
		[element addAttribute:[NSXMLNode attributeWithName:@"label" stringValue:label]];
	else
	{
		NSDictionary * feature = [tree valueForKey:FEATURE_SENSOR];
		
		[element addAttribute:[NSXMLNode attributeWithName:@"feature" stringValue:[feature valueForKey:FEATURE_NAME]]];
		
		NSNumber * split = [feature valueForKey:SPLIT_VALUE];
		if (split != nil)
			[element addAttribute:[NSXMLNode attributeWithName:@"split_value" stringValue:[split description]]];
		
		for (NSString * k in [[tree allKeys] sortedArrayUsingFunction:TreeSort context:NULL])
		{
			NSObject * value = [tree valueForKey:k];
			
			if ([value isKindOfClass:[NSDictionary class]] && ![k isEqualToString:FEATURE_SENSOR])
			{
				NSXMLElement * valueElement = [[NSXMLElement alloc] initWithName:@"value"];
				[valueElement addAttribute:[NSXMLNode attributeWithName:@"observed" stringValue:k]];
				
				[valueElement addChild:[self getXML:((NSDictionary *) value)]];
				
				[element addChild:valueElement];
				[valueElement release];
			}
		}
	}

	return element;
}

- (void) saveXMLTree:(NSString *) path
{
	NSXMLDocument * doc = [[NSXMLDocument alloc] init];
	
	[doc setRootElement:[self getXML:self.decisionTree]];
	[[doc rootElement] release];
	
	[[doc XMLData] writeToFile:path atomically:YES];

	[doc release];
}

- (void) save
{
	NSMutableData * data = [NSMutableData data];
	
	NSKeyedArchiver * archiver = [[NSKeyedArchiver alloc] initForWritingWithMutableData:data];
	
	[archiver encodeObject:self.examples forKey:EXAMPLES_KEY];
	[archiver finishEncoding];
	
	[archiver release];

	NSString * path = [LEARNER_PATH stringByReplacingOccurrencesOfString:@"KEY" withString:key];
	
	[data writeToFile:path atomically:YES];
	
	// [self saveXMLTree:[[NSString stringWithFormat:@"~/Desktop/%@.xml", key, nil] stringByExpandingTildeInPath]];
}

- (void) setKey:(NSString *) newKey
{
	self.examples = nil;
	
	[super setKey:newKey];
	[self load];
}

- (void) reset
{
	[super reset];
	
	[self.examples removeAllObjects];

	[self save];
	[self load];
}

- (NSNumber *) exampleCount
{
	return [NSNumber numberWithInt:[examples count]];
}

- (NSString *) labelForTree:(NSDictionary *) treeDict features:(NSArray *) features
{
	NSString * label = [treeDict valueForKey:LABEL_KEY];
	
	if (label != nil)
		return label;
	else
	{
		NSString * featureKey = [treeDict valueForKey:FEATURE_SENSOR];

		NSMutableArray * labels = [NSMutableArray array];

		NSMutableArray * matchingFeatures = [NSMutableArray array];

		for (NSDictionary * feature in features)
		{
			if ([featureKey isEqual:[feature valueForKey:FEATURE_SENSOR]])
				[matchingFeatures addObject:feature];
		}
		
		NSMutableArray * newFeatures = [NSMutableArray arrayWithArray:features];
		[newFeatures removeObjectsInArray:matchingFeatures];
		
		for (NSDictionary * feature in matchingFeatures)
		{
			NSDictionary * child = [treeDict valueForKey:[feature valueForKey:FEATURE_OBSERVATION]];
			
			if (child != nil)
			{
				NSString * featureLabel = [self labelForTree:child features:newFeatures];

				if (featureLabel != nil)
					[labels addObject:featureLabel];
			}
		}
		
		[labels sortUsingSelector:@selector(compare:)];
		
		NSMutableDictionary * labelCounts = [NSMutableDictionary dictionary];
		
		for (NSString * newLabel in labels)
		{
			NSNumber * count = [labelCounts valueForKey:newLabel];
			
			if (count == nil)
				count = [NSNumber numberWithInt:1];
			else
				count = [NSNumber numberWithInt:(1 + [count intValue])];
			
			[labelCounts setValue:count forKey:newLabel];
		}
		
		int currentCount = 0;
		
		for (NSString * featureKey in [labelCounts allKeys])
		{
			NSNumber * labelCount = [labelCounts valueForKey:featureKey];
			
			if ([labelCount intValue] > currentCount)
			{
				currentCount = [labelCount intValue];
				label = featureKey;
			}
		}
	}
	
	return label;
}

- (NSString *) getLabelForExample:(NSArray *) features
{
	return [self labelForTree:self.decisionTree features:features];
}

- (void) addExample:(NSArray *) features forClass:(NSString *) label
{
	NSMutableDictionary * example = [NSMutableDictionary dictionaryWithObject:label forKey:LABEL_KEY];
	
	for (NSDictionary * feature in features)
		[example setValue:[feature valueForKey:FEATURE_OBSERVATION] forKey:[feature valueForKey:FEATURE_SENSOR]];

	[self.examples addObject:example];

	[self save];
	[self load];
}

- (void) removeLastExample:(NSNotification *) note
{
	if ([self.examples count] > 0)
	{
		[self.examples removeLastObject];

		[self save];
		[self load];
	}
}


@end