/*
   Project: UL

   Copyright (C) 2005 Michael Johnston & Jordi Villa-Freixa

   Author: Michael Johnston

   Created: 2005-11-04 16:44:41 +0100 by michael johnston

   This application is free software; you can redistribute it and/or
   modify it under the terms of the GNU General Public
   License as published by the Free Software Foundation; either
   version 2 of the License, or (at your option) any later version.

   This application is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
   Library General Public License for more details.

   You should have received a copy of the GNU General Public
   License along with this library; if not, write to the 
   Free Software Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111 USA.
*/

#include "ULFramework/ULSimulation.h"

@implementation ULSimulation

- (void) _convertEnergies: (NSMutableArray*) energies
{
	int i;
	double value;
	
	for(i=0; i<[energies count]; i++)
	{
		value = [[energies objectAtIndex: i] doubleValue];	
		value *= CC;
		[energies replaceObjectAtIndex: i withObject: [NSNumber numberWithDouble: value]];
	}
}

- (id) _setOptions
{
	id subsystemDict, subsystems, obj;	
	id subsystem;
	id options;
	NSEnumerator* subsystemEnum;
	NSString* optionsFile;

	//FIXME. We have a problem with the resultsOptions.plist file
	//at the moment it is located in the resource dir of UL hence we have to access
	//it from there. We really should access it from the library resource bundle.
	//Also we assume UL is installed into $User/Gnustep
	
	optionsFile = [[[NSBundle mainBundle] 
			bundlePath] stringByAppendingPathComponent: @"resultsOptions.plist"];
	
	options = [NSMutableDictionary dictionaryWithContentsOfFile: optionsFile];
	if(options == nil)
	{
		optionsFile = [NSHomeDirectory()  
				stringByAppendingPathComponent: 
				@"GNUstep/Applications/UL.app/Resources/resultsOptions.plist"];
		options = [NSMutableDictionary 
				dictionaryWithContentsOfFile: optionsFile];
	}			

	NSDebugLLog(@"ULSimulation", @"Read options from resources dir at %@.\n%@", optionsFile, options);

	//add the available subsystems and terms
	
	subsystemDict = [options valueForKey:@"Subsystems"];
	subsystemEnum = [[self availableSubsystems] objectEnumerator];

	/*
	 * Again shallow copying the subsystem part of the options file for each subsystem
	 * causes problems later. Instead we'll set the subsystem part directly.
	 */

	while(subsystem = [subsystemEnum nextObject])
	{
		obj = [NSMutableDictionary dictionary];
		[obj setObject: [self availableTermsForSubsystem: subsystem] forKey: @"AvailableTerms"]; 
		[obj setObject: [NSMutableArray array] forKey: @"Selection"];
		[obj setObject: @"Multiple" forKey: @"Type"];
		[subsystemDict setObject: obj forKey: subsystem];
	}
	[subsystemDict removeObjectForKey: @"All"];

	return options;
}

- (void) _processEnergiesUsingDecoder: (NSCoder*) unarchiver
{
	int i, j, energyFrames; 
	NSMutableArray* frameStates;
	NSMutableDictionary* energyDict;
	NSString* stateKey;
	id obj, dict;

	subsystemsEnergies = [[NSMutableDictionary dictionaryWithCapacity: 1] retain];
	energyFrames =  [unarchiver decodeIntForKey: @"NumberOfStateFrames"];
	energyDict = [unarchiver decodeObjectForKey: @"StateFrames"];

	for(i=0; i<numberOfSubsystems; i++)
	{
		stateKey = subsystemInfoArray[i].stateKey;
		frameStates = [NSMutableArray arrayWithCapacity: 1];
		for(j=1; j<energyFrames+1; j++)
		{
			obj = [energyDict objectForKey: 
					[NSString stringWithFormat: @"%@.%d", stateKey, j]];
			[frameStates addObject: obj];
		}

		[subsystemsEnergies setObject: frameStates forKey: subsystemInfoArray[i].subsystemName];
	}
}

- (void) _convertTrajectoryToArchives: (NSData*) trajectoryData
{
	int i, j, energyFrames;
	int bytesLength, location;
	unsigned char* bytes;
	id names, name;
	NSData* archive;
	NSString* lastTag;
	NSKeyedUnarchiver* unarchiver;
	SubsystemResultsInfo* results_s;

	//each archive begins with <?
	//search through the data and break it into chunks
	//based on this

	bytesLength = 0;
	location = 0;
	bytes = (char*)[trajectoryData bytes];
	trajectoryArchives = [[NSMutableArray arrayWithCapacity: 1] retain];

	NSDebugLLog(@"ULSimulation", @"Pre-Processing Trajectory...");

	for(i = 0; i< [trajectoryData length] - 1; i++)
		if(bytes[i] == '<' && bytes[i+1] == '?')
			if(i != 0)
			{
				bytesLength = i - location;
				archive = [NSData dataWithBytes: (bytes + location) length: bytesLength];
				[trajectoryArchives addObject: archive];
				NSDebugLLog(@"ULSimulation", @"Archive at %d. Size %d", location, bytesLength);
				location = i;
			}

	//check the last archive is complete
	//i.e. it has a closing </plist> tag

	bytesLength = i - location + 1;

	lastTag = [[NSString alloc] initWithBytes: &bytes[i - 7] 
			length: 8 
			encoding: NSUTF8StringEncoding];
	[lastTag autorelease];

	if([lastTag isEqual: @"</plist>"])
	{
		archive = [NSData dataWithBytes: (bytes + location) length: bytesLength];
		[trajectoryArchives addObject: archive];
		NSDebugLLog(@"ULSimulation", @"Archive at %d. Size %d", location, bytesLength);
	}
	else
		NSWarnLog(@"The last configuration frame is not complete.");
	
	initialArchive = [trajectoryArchives objectAtIndex: 0];

	unarchiver = [NSKeyedUnarchiver alloc];

	[unarchiver initForReadingWithData: initialArchive];
	names = [unarchiver decodeObjectForKey: @"SubsystemNames"];
	[[names retain] autorelease];
	[unarchiver finishDecoding];
	[unarchiver release];
	numberOfSubsystems = [names count];
	NSDebugLLog(@"ULSimulation", @"Complete (%@)", names);
	subsystemInfoArray = malloc([names count]*sizeof(SubsystemResultsInfo));
	subsystemsInfo = [NSMutableDictionary dictionaryWithCapacity: 1];
	[subsystemsInfo retain];
	for(i=0; i<[names count]; i++)
	{
		name = [names objectAtIndex: i];
		results_s = &subsystemInfoArray[i];
		results_s->subsystemName = [name retain];
		results_s->stateKey = [[NSString stringWithFormat: @"%@.State", name] retain];
		results_s->dynamicsKey = [[NSString stringWithFormat: @"%@.Dynamics", name] retain];
		results_s->numberOfDynamicsFrames = [trajectoryArchives count];
		results_s->numberOfStateFrames = 0;
		[subsystemsInfo setObject: [NSValue valueWithPointer: results_s]
			forKey: name];
	}
}

- (void) _cacheEnergies
{	
	int i, j, energyFrames;
	int bytesLength, location;
	unsigned char* bytes;
	id names, name;
	NSData* archive;
	NSKeyedUnarchiver* unarchiver;
	SubsystemResultsInfo* results_s;

	NSDebugLLog(@"ULSimulation", @"Caching Energies...");	
	if([energyArchive length] > 0)
	{
		unarchiver = [NSKeyedUnarchiver alloc];
		[unarchiver initForReadingWithData: energyArchive];

		energyFrames = [unarchiver decodeIntForKey: @"NumberOfStateFrames"];
		for(i=0; i<[subsystemsInfo count]; i++)
		{
			results_s = &subsystemInfoArray[i];
			results_s->numberOfStateFrames = energyFrames;
		}
		[self _processEnergiesUsingDecoder: unarchiver];
		[unarchiver finishDecoding];
		[unarchiver release];
	}
	else
		subsystemsEnergies = nil; 
	
	[energyArchive release];
	cachedEnergies = YES;	

	NSDebugLLog(@"ULSimulation", @"Caching Complete");
}

/**************

Creation

***************/

- (id) init
{
	return [self initWithName: nil];	
}

- (id) initWithName: (NSString*) aName 
{
	NSArray* constants;
	NSArray* keys;

	if(self = [super init])
	{
		keys = [NSArray arrayWithObjects:
				@"KCalMol", 
				@"JouleMol", 
				@"Simulation", 
				nil];
		constants = [NSArray arrayWithObjects:
				[NSNumber numberWithDouble: 2390.05735688],
				[NSNumber numberWithDouble: 1E7],
				[NSNumber numberWithDouble: 1.0],
				nil];

		conversionConstants = [NSDictionary dictionaryWithObjects: constants 
					forKeys: keys];	
		[conversionConstants retain];
		CC = 1.0;
	
		if(aName != nil)
			[metadata setObject: aName forKey: @"Name"];
	}		

	return self;
}

//FIXME: loadData can only be called once
- (void) loadData
{
	NSAutoreleasePool *pool;
	NSData* trajectoryData, *energyData;
	NSError* error;
	
	//check we have access to a data store

	if(dataStorage == nil)
		[NSException raise: NSInternalInconsistencyException
			format: @"No data storage has been set."];
	
	//check the data store is accesible
	if(![dataStorage isAccessible])
	{
		NSWarnLog(@"Simulation data is not accesible raising exception");
		error = [dataStorage accessError];
		[NSException raise: NSInternalInconsistencyException
			format: [[error userInfo] 
					objectForKey: NSLocalizedDescriptionKey]];
	}
	
	pool = [[NSAutoreleasePool alloc] init];
	
	trajectoryData  = [dataStorage trajectoryData];
	energyData = [dataStorage energyData];

	NSDebugLLog(@"ULSimulation", @"Read in %lf KB of data",
		 ((double)[trajectoryData length])/1024);
	NSDebugLLog(@"ULSimulation", @"Read in %lf KB of energy data", 
		((double)[energyData length])/1024);

	if(trajectoryData != nil)
		[self _convertTrajectoryToArchives: trajectoryData];
	else
		NSWarnLog(@"Simulation %@ (%@) contains no trajectory data",
				[self name],
				[self identification]);

	energyArchive = energyData;
	[energyArchive retain];
	
	cachedEnergies = NO;
	
	NSDebugLLog(@"ULSimulation", @"Initialised ULSimulation");

	[pool release];
}

- (void) dealloc
{	
	int i;
	
	[dataStorage release];
	[trajectoryArchives release];
	[subsystemsEnergies release];
	
	for(i=0; i<numberOfSubsystems; i++)
	{
		[subsystemInfoArray[i].subsystemName release];
		[subsystemInfoArray[i].stateKey release];
		[subsystemInfoArray[i].dynamicsKey release];
	}

	free(subsystemInfoArray);
	[subsystemsInfo release];
	[conversionConstants release];
	[super dealloc];
}

- (NSMutableDictionary*) optionsDict
{
	/*
	Return a fresh instance of the options
	for each request.
	This is necessary since
	it is likely plugins will manipulate this
	dictionary to their own ends
	If we just return a copy is will only be
	a shallow copy and thus will affect the
	original object.
	*/
	return [self _setOptions];
}

- (void) printAvailableInfo
{
	int i;
	id term, terms;
	NSEnumerator* termEnum;
	SubsystemResultsInfo* results_s;
	
	if(!cachedEnergies)
		[self _cacheEnergies];

	GSPrintf(stderr, @"\nNumber of Subsystems %d\n\n", numberOfSubsystems);
	for(i=0; i<numberOfSubsystems; i++)
	{
		results_s = &subsystemInfoArray[i];
		GSPrintf(stderr, @"Subsystem name: %@\n", results_s->subsystemName); 
		GSPrintf(stderr, @"Number of state frames: %d\n", results_s->numberOfStateFrames); 
		GSPrintf(stderr, @"Available Terms:\n");
		terms = [[self availableTermsForSubsystem: results_s->subsystemName] 
				sortedArrayUsingSelector: @selector(caseInsensitiveCompare:)]; 
		termEnum = [terms objectEnumerator]; 
		while(term = [termEnum nextObject])
			GSPrintf(stderr, @"\t%@\n", term);
		GSPrintf(stderr, @"Number of dynamics frames: %d\n\n", results_s->numberOfDynamicsFrames); 
	}	
}

- (NSString*) availableInfo
{
	int i;
	id term, terms;
	NSEnumerator* termEnum;
	SubsystemResultsInfo* results_s;
	NSMutableString* info;
	
	if(!cachedEnergies)
		[self _cacheEnergies];

	info = [NSMutableString stringWithFormat: @"\nNumber of Subsystems %d\n\n", numberOfSubsystems];
	for(i=0; i<numberOfSubsystems; i++)
	{
		results_s = &subsystemInfoArray[i];
		[info appendString:
		 	[NSString stringWithFormat: @"Subsystem name: %@\n",
			 results_s->subsystemName]];
		
		[info appendString: 
			[NSString stringWithFormat: @"Number of state frames: %d\n",
			 results_s->numberOfStateFrames]]; 

		[info appendString: @"Available Terms:\n"];
		terms = [[self availableTermsForSubsystem: results_s->subsystemName] 
				sortedArrayUsingSelector: @selector(caseInsensitiveCompare:)]; 
		termEnum = [terms objectEnumerator]; 
		while(term = [termEnum nextObject])
			[info appendString: [NSString stringWithFormat: @"\t%@\n", term]];

		[info appendString:
		 	[NSString stringWithFormat: @"Number of dynamics frames: %d\n\n",
			 results_s->numberOfDynamicsFrames]]; 
	}	

	return [[info copy] autorelease];
}

/************

Public Methods

*************/

- (id) dataStorage
{
	return [[dataStorage retain] autorelease];
}

//FIXME: There is currently no way to be sure
//that the data accesible through \e object is 
//the data corresponding to this object. This validation
//had to be carried out be the database backend.
- (void) setDataStorage: (id) object
{
	if(dataStorage != object)
	{
		[dataStorage release];
		dataStorage = [object retain];
	}	
}

- (void) setEnergyUnit: (NSString*) name
{
	CC = [[conversionConstants valueForKey: name] doubleValue];
}

- (NSArray*) availableSubsystems
{
	return [subsystemsInfo allKeys];
}

- (NSArray*) availableTermsForSubsystem: (NSString*) name
{
	id array;
	
	if(!cachedEnergies)
		[self _cacheEnergies];

	array = [[[[subsystemsEnergies valueForKey: name] 
			objectAtIndex: 0] allKeys] mutableCopy];
	
	if(array == nil)
		array = [NSMutableArray arrayWithObject: @"None"];
	else
		[array removeObjectsInArray: [NSArray arrayWithObjects: @"Time", @"Temperature", nil]]; 	

	return array;
}

- (ULMatrix*) energiesForSubsystem: (NSString*) name
		terms: (NSArray*) terms 
		inFrames: (NSRange*) range 
		step: (int) stepsize
{
	int end, i;
	id energies, temp, dict;
	NSMutableArray *termsCopy;
	ULMatrix* table;
	NSEnumerator* termEnum;

	if(subsystemsEnergies == nil)
		return nil;

	if(!cachedEnergies)
		[self _cacheEnergies];

	if(range == NULL)
		*range = NSMakeRange(0, [[subsystemsEnergies valueForKey: name] count]);

	if(range->length <= 0)
		range->length = [[subsystemsEnergies valueForKey: name] count] - range->location;

	if(range->length > [[subsystemsEnergies valueForKey: name] count])
		[NSException raise: NSInvalidArgumentException 
			format: @"Length of requested range (%d) greater than number of frames(%d)", 
				range->length,[[subsystemsEnergies valueForKey: name] count]];

	if(stepsize <=0)
		[NSException raise: NSInvalidArgumentException 
			format: @"Stepsize must be greater than 0 (%d)", stepsize];

	energies = [subsystemsEnergies valueForKey: name];
	end = range->location + range->length;
	termsCopy = [terms mutableCopy];
	[termsCopy insertObject: @"Time" atIndex: 0];

	if(terms == nil)
		terms = [self availableTermsForSubsystem: name];

	table = [[[ULMatrix alloc] initWithRows: 0 columns: 0] autorelease];

	for(i=range->location; i < end; i += stepsize)
	{
		temp = [[[energies objectAtIndex: i] objectsForKeys: terms 
				notFoundMarker: @"NotAvailable"] mutableCopy];
		if(CC != 1.0)
			[self _convertEnergies: temp];

		[temp insertObject: [[energies objectAtIndex: i] objectForKey: @"Time"]
			atIndex: 0];
		[table extendMatrixWithRow: temp];
	}
	
	[table setColumnHeaders: termsCopy];	
	[table setName: name];

	return table;
}

- (ULMatrix*) energiesForSubsystem: (NSString*) name terms: (NSArray*) terms inFrames: (NSRange*) range
{
	return [self energiesForSubsystem: name 
			terms: terms 
			inFrames: range 
			step: 1];
}

- (ULMatrix*) energiesForSubsystem: (NSString*) name terms: (NSArray*) terms
{
	NSRange range;

	if(subsystemsEnergies == nil)
		return nil;

	if(!cachedEnergies)
		[self _cacheEnergies];

	range.location = 0;
	range.length = [[subsystemsEnergies valueForKey: name] count];

	return [self energiesForSubsystem: name 
			terms: terms
			inFrames: &range
			step: 1];
}

- (ULMatrix*) potentialEnergyForSubsystem: (NSString*) name frames: (NSRange*) range
{
	return [self energiesForSubsystem: name 
			terms: [NSArray arrayWithObject: @"PotentialEnergy"]
			inFrames: range
			step: 1];
}

- (ULMatrix*) kineticEnergyForSubsystem: (NSString*) name frames: (NSRange*) range
{
	return [self energiesForSubsystem: name 
			terms: [NSArray arrayWithObject: @"KineticEnergy"]
			inFrames: range
			step: 1];
}

- (ULMatrix*) totalEnergyForSubsystem: (NSString*) name frames: (NSRange*) range
{
	return [self energiesForSubsystem: name 
			terms: [NSArray arrayWithObject: @"TotalEnergy"]
			inFrames: range
			step: 1];
}

- (ULMatrix*) temperatureForSubsystem: (NSString*) name frames: (NSRange*) range
{
	return [self energiesForSubsystem: name 
			terms: [NSArray arrayWithObject: @"Temperature"]
			inFrames: range
			step: 1];
}

- (int) initialStateTimeForSubsystem: (NSString*) name
{
	if(!cachedEnergies)
		[self _cacheEnergies];

	return [[[[subsystemsEnergies valueForKey: name] 
			objectAtIndex: 0]
			valueForKey: @"Time"] intValue];
}

- (int) lastStateFrameTimeForSubsystem: (NSString*) name
{
	if(!cachedEnergies)
		[self _cacheEnergies];

	return [[[[subsystemsEnergies valueForKey: name] 
			lastObject]
			valueForKey: @"Time"] intValue];
}

- (int) timeForStateFrame: (int) frameNumber subsystem: (NSString*) name
{
	if(!cachedEnergies)
		[self _cacheEnergies];

	return [[[[subsystemsEnergies valueForKey: name] 
		objectAtIndex: frameNumber]
		valueForKey: @"Time"] intValue];
}

- (int) numberOfStateFramesForSubsystem: (NSString*) name
{
	if(!cachedEnergies)
		[self _cacheEnergies];

	return [[subsystemsEnergies valueForKey: name] count];
}

- (AdMatrix*) coordinatesForFrame: (int) frame subsystem: (NSString*) name;
{
	NSData* archive;
	id dynamics, source;
	NSKeyedUnarchiver* unarchiver;
	AdMatrix* matrix;

	if(frame > [trajectoryArchives count])
		[NSException raise: NSInvalidArgumentException
			format: @"Requested frame (%d) is greater than the number of available frames (%d)",
				frame, [trajectoryArchives count]];

	//FIXME: AdDynamics doesnt retain its data source
	//However we only store mementos of AdDynamics!
	//Have to get a ref to the data source a retain/autorelease it
	//However this indicates something is wrong with the memento scheme.
	archive = [trajectoryArchives objectAtIndex: frame];
	unarchiver = [[NSKeyedUnarchiver alloc] initForReadingWithData: archive];
	dynamics = [unarchiver decodeObjectForKey: 
			[NSString stringWithFormat: @"%@.dynamics", name]];
	[[dynamics retain] autorelease];
	source = [[[dynamics dataSource] retain] autorelease];		
	[unarchiver finishDecoding];

	[unarchiver release];

	return [[dynamics coordinates] pointerValue];
}

- (int) timeForCoordinateFrame: (int) frameNumber subsystem: (NSString*) name;
{
	
}

- (int) numberOfCoordinateFramesForSubsystem: (NSString*) name;
{
	NSWarnLog(@"Not fully implemented (%@, %@)", NSStringFromClass([self class]), 
			NSStringFromSelector(_cmd));
	return [trajectoryArchives count]; 
}

- (NSArray*) atomTypesForSubsystem: (NSString*) name
{
	id source;
	NSKeyedUnarchiver* unarchiver;

	unarchiver = [[NSKeyedUnarchiver alloc] initForReadingWithData: initialArchive];
	source = [unarchiver decodeObjectForKey:
			 [NSString stringWithFormat: @"%@.Source", name]];
	[[source retain] autorelease];
	[unarchiver finishDecoding];
	[unarchiver release];

	return [[[source objectValueForAtomTypes: self] retain] autorelease];
}

- (NSArray*) atomMassesForSubsystem: (NSString*) name
{
	id source, dynamics;
	NSKeyedUnarchiver* unarchiver;

	unarchiver = [[NSKeyedUnarchiver alloc] initForReadingWithData: initialArchive];
	source = [unarchiver decodeObjectForKey:
			 [NSString stringWithFormat: @"%@.Source", name]];
	dynamics = [unarchiver decodeObjectForKey:
		 	[NSString stringWithFormat: @"%@.dynamics", name]];
	
	[[dynamics retain] autorelease];
	[[source retain] autorelease];
	[unarchiver finishDecoding];
	[unarchiver release];

	return [dynamics atomMasses];
}

/*******
NSCoding Methods
Since the location of the data this
object uses is volatile the unarchiver must create
and set a dataStorage object providing access to 
the data when this object is decoded.
********/

- (id) initWithCoder: (NSCoder*) decoder
{
	if(self = [super initWithCoder: decoder])
	{
		if([decoder allowsKeyedCoding])
		{
			conversionConstants = [decoder decodeObjectForKey: @"ConversionConstants"];
		}	
		else
		{
			conversionConstants = [decoder decodeObject];
		}	
	}	

	[conversionConstants retain];

	CC = 1.0;

	return self;
}

- (void) encodeWithCoder: (NSCoder*) encoder
{
	[super encodeWithCoder: encoder];
	if([encoder allowsKeyedCoding])
	{
		[encoder encodeObject: conversionConstants forKey: @"ConversionConstants"];
	}
	else
	{
		[encoder encodeObject: conversionConstants];
	}
}

@end


