/*
   Project: AdunCore

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

   Author: 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 "AdunKernel/AdunSimulationData.h"

@implementation AdSimulationData

- (void) _loadStateData: (NSData*) data;
{	
	NSDebugLLog(@"AdSimulationData", @"Load Energies...");	
	stateData = [NSKeyedUnarchiver unarchiveObjectWithData: data];
	[stateData retain];
	NSDebugLLog(@"AdSimulationData", @"Finished");
	
	if(stateData == nil || [[stateData dataMatrices] count] == 0)
		NSWarnLog(@"Simulation %@ (%@) contains no energy data",
				[self name],
				[self identification]);
	else		
		//All systems have the same number of energy checkpoints
		numberEnergyCheckpoints = [[[stateData dataMatrices] 
						objectAtIndex: 0] numberOfRows];
}

- (void) _loadSystemCollection
{
	NSKeyedUnarchiver* unarchiver;
	NSData* data;

	data = [dataStorage systemData];
	unarchiver = [NSKeyedUnarchiver alloc];
	[unarchiver initForReadingWithData: data];
	systemCollection = [unarchiver decodeObjectForKey: @"SystemCollection"];
	[systemCollection retain];
	[unarchiver finishDecoding];
	[unarchiver release];
}

- (void) _loadFrames
{
	NSKeyedUnarchiver* unarchiver;
	NSData* data;

	data = [dataStorage frameData];
	if(data != nil && [data length] > 0)
	{
		unarchiver = [NSKeyedUnarchiver alloc];
		[unarchiver initForReadingWithData: data];
		frames = [unarchiver decodeObjectForKey: @"root"];
		[frames retain];
		[unarchiver finishDecoding];
		[unarchiver release];
	}	
	else
		NSWarnLog(@"No frame data present");
}

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

Creation

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

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

- (id) initWithName: (NSString*) aName 
{
	if((self = [super init]))
	{
		stateData = nil;
		systemCollection = nil;
		frames = nil;
		numberEnergyCheckpoints = 0;
		if(aName != nil)
			[self setValue: aName forMetadataKey: @"Name"];
	}		

	return self;
}

- (void) dealloc
{	
	[frames release];
	[dataStorage release];
	[stateData release];
	[systemCollection release];
	[super dealloc];
}

- (void) loadData
{
	NSData *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]];
	}
	
	//Get rid of anything we loaded previously
	if(stateData != nil)
	{
		[stateData release];
		[systemCollection release];
		[frames release];
		stateData = nil;
		systemCollection = nil;
	}

	//Call update on the data store in case anything
	//has been written to it since it was initialised
	[dataStorage update];

	energyData = [dataStorage energyData];
	NSDebugLLog(@"AdSimulationData", 
			@"There are %d trajectory frames available",
			[dataStorage numberTrajectoryCheckpoints]);
	NSDebugLLog(@"AdSimulationData", 
		@"Read in %lf KB of energy data", 
		((double)[energyData length])/1024);

	[self _loadStateData: energyData];
	NSDebugLLog(@"AdSimulationData", @"Unarchiving system manager");
	[self _loadSystemCollection];
	NSDebugLLog(@"AdSimulationData", @"Unarchiving frame data");
	[self _loadFrames];
	NSDebugLLog(@"AdSimulationData", @"Complete");
		
	
	NSDebugLLog(@"AdSimulationData", @"Initialised AdSimulationData");
}

- (BOOL) checkDataStorageIdentification
{
	return checkDataStorageIdentification;
}

- (void) setCheckDataStorageIdentification: (BOOL) value
{
	checkDataStorageIdentification = value;
}

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

- (void) setDataStorage: (id) object
{
	NSString* dataId;

	if(dataStorage != object)
	{
		[dataStorage release];
		dataStorage = [object retain];
	}	

	//Check if the data in the store corresponds
	//to the simulation data for this object
	//by checking its id
	if(checkDataStorageIdentification)
		if((dataId = [dataStorage identification]) != nil)
			if(dataId != identification)
			{
				NSWarnLog(@"Supplied data storage is not related to this object");
				return;
			}
}

- (AdDataSet*) energies
{
	return [[stateData retain] autorelease];
}

- (AdSystemCollection*) systemCollection
{
	return [[systemCollection retain] autorelease];
}

- (unsigned int) numberTrajectoryCheckpoints
{
	return [dataStorage numberTrajectoryCheckpoints];
}

- (unsigned int) numberTopologyCheckpoints
{
	return [dataStorage numberTopologyCheckpoints];
}

- (id) dataSourceForSystem: (id) system inTopologyCheckpoint: (unsigned int) number
{
	NSData* archive;
	NSKeyedUnarchiver* unarchiver;
	NSString* name;
	id dataSource;

	if(number >= [self numberTopologyCheckpoints])
		[NSException raise: NSInvalidArgumentException
			format: @"Requested checkpoint (%d) is greater than the number of available checkpoints (%d)",
				number, [self numberTopologyCheckpoints]];

	name = [system systemName];
	archive = [dataStorage topologyCheckpoint: number];
	unarchiver = [[NSKeyedUnarchiver alloc] initForReadingWithData: archive];
	dataSource = [unarchiver decodeObjectForKey: name];
	[[dataSource retain] autorelease];
	[unarchiver finishDecoding];
	[unarchiver release];

	return dataSource;
}

- (id) mementoForSystem: (id) system inTrajectoryCheckpoint: (unsigned int) number
{
	NSData* archive;
	NSKeyedUnarchiver* unarchiver;
	NSString* name;
	id memento;

	if(number >= [self numberTrajectoryCheckpoints])
		[NSException raise: NSInvalidArgumentException
			format: @"Requested checkpoint (%d) is greater than the number of available checkpoints (%d)",
				number, [self numberTrajectoryCheckpoints]];

	name = [system systemName];
	archive = [dataStorage trajectoryCheckpoint: number];
	unarchiver = [[NSKeyedUnarchiver alloc] initForReadingWithData: archive];
	memento = [unarchiver decodeObjectForKey: name];
	[[memento retain] autorelease];
	[unarchiver finishDecoding];
	[unarchiver release];

	return memento;
}

- (NSString*) description
{
	NSMutableString *string = [NSMutableString string];
	NSEnumerator* systemEnum;
	id system, energies;

	[string appendFormat: @"Name: %@ - ", [self name]];
	if(systemCollection == nil)
	{
		[string appendString: @"No data loaded\n"];
		return string;
	}	
	
	[string appendFormat: @"%d System(s):\n\n",
		[[systemCollection allSystems] count]];
	systemEnum = [[systemCollection allSystems] objectEnumerator];
	while((system = [systemEnum nextObject]))
	{
		[string appendFormat: @"%@\n", [system systemName]];
		[string appendFormat: @"\t%d Trajectory Frames\n",
			[self numberTrajectoryCheckpoints]];
		[string appendFormat: @"\t%d Topology Frames\n",
			[self numberTopologyCheckpoints]];
		energies = [stateData dataMatrixWithName: [system systemName]];
		if(energies != nil)
		{
			[string appendFormat: @"\t%d Energy Frames\n\n",
				[energies numberOfRows]];	
		}	
		else
			[string appendString: @"\tNo available energies\n"];
	}

	return string;
}

/**
Frame related methods - Wont work with pre 0.71 stores 
where the frame data wasn't encoded.
*/

- (unsigned int) numberOfFrames
{
	if(frames == nil)
		return 0;
	else
		return [frames numberOfRows];
}

/**
The first frame is 0.
Checkpoints are also numbered from 0.
*/
- (unsigned int) _frameForCheckpoint: (unsigned int) number inColumn: (NSString*) aString
{
	int i;
	unsigned int counter = -1; 	//Accounting for checkpoint offset
	NSMutableArray* array = [NSMutableArray new];

	[frames addColumnWithHeader: aString toArray: array];
	for(i=0; i<(int)[array count]; i++)
	{
		counter += [[array objectAtIndex: i] boolValue];
		if(counter == number)
			break;
	}		
	[array release];

	return i;
}

- (unsigned int) frameForTrajectoryCheckpoint: (unsigned int) number
{
	if(number >= [self numberTrajectoryCheckpoints])
		[NSException raise: NSInvalidArgumentException
			format: 
			@"Requested checkpoint index (%d) is greater than the number of available checkpoints (%d)",
			number, [self numberTrajectoryCheckpoints]];

	return [self _frameForCheckpoint: number inColumn: @"Trajectory"];
}

- (unsigned int) frameForEnergyCheckpoint: (unsigned int) number
{
	if(number >= numberEnergyCheckpoints)
		[NSException raise: NSInvalidArgumentException
			format:
			@"Requested checkpoint index (%d) is greater than the number of available checkpoints (%d)",
			number, numberEnergyCheckpoints];

	return [self _frameForCheckpoint: number inColumn: @"Energy"];
}

- (unsigned int) frameForTopologyCheckpoint: (unsigned int) number
{
	if(number >= [self numberTopologyCheckpoints])
		[NSException raise: NSInvalidArgumentException
			format:
			@"Requested checkpoint index (%d) is greater than the number of available checkpoints (%d)",
			number, [self numberTopologyCheckpoints]];

	return [self _frameForCheckpoint: number inColumn: @"Topology"];
}

- (NSArray*) dataRecordedInFrame: (unsigned int) number
{
	NSMutableArray* array = [NSMutableArray array];

	if(number >= [frames numberOfRows])
		[NSException raise: NSInvalidArgumentException
			format: 
			@"Requested frame index (%d) is greater than the number of available frames (%d)",
			number, [self numberOfFrames]];

	if([[frames elementAtRow: number ofColumnWithHeader: @"Trajectory"] boolValue])
		[array addObject: @"Trajectory"];
	if([[frames elementAtRow: number ofColumnWithHeader: @"Topology"] boolValue])
		[array addObject: @"Topology"];
	if([[frames elementAtRow: number ofColumnWithHeader: @"Energy"] boolValue])
		[array addObject: @"Energy"];

	return [[array copy] autorelease];
}

- (id) lastRecordedDataSourceForSystem: (id) system inRange: (NSRange) range
{
	BOOL value;
	int end, i;
	id dataSource;

	end = NSMaxRange(range);
	if(end > (int)[frames numberOfRows])
		[NSException raise: NSInvalidArgumentException
			format: @"Specfied range (%d,%d) falls outside the available frames %d",
			range.location, range.length, [frames numberOfRows]];
	
	NSDebugLLog(@"AdSimulationData", 
		@"Searching range %d-%d for data source", 
		range.location, end);
	dataSource = nil;
	for(i = ([self numberTopologyCheckpoints] - 1); i>=0; i++ )
	{
		value = [self frameForTopologyCheckpoint: i];
		if(NSLocationInRange(value, range))
		{
			NSDebugLLog(@"AdSimulationData", 
				@"Topology checkpoint %d is in the specfied range", i);
			//This checkpoint is in the range.
			//Now we must check it contains the required data.
			dataSource = [self dataSourceForSystem: system
					inTopologyCheckpoint: i];
			if(dataSource != nil)
			{
				NSDebugLLog(@"AdSimulationData", @"Found data source");
				break;
			}	

			NSDebugLLog(@"AdSimulationData", 
				@"Checkpoint does not contain data for system");	
		}
	}

	return dataSource;
}

/**
Deprecated
*/
- (unsigned int) numberOfFramesForSystem: (id) system
{
	return [self numberTrajectoryCheckpoints];
}

/**
Deprecated
*/
- (id) mementoForFrame: (unsigned int) frame ofSystem: (id) system
{
	return [self mementoForSystem: system 
			inTrajectoryCheckpoint: frame];
}

@end


@implementation AdSimulationDataWriter

- (void) _createStateMatrices
{
	NSArray* headers, *forceFieldArray;
	NSMutableArray* stateHeaders;
	NSEnumerator* systemEnum;
	AdForceField* systemForceField;
	id system, stateMatrix;

	if(iterationHeader == nil)
		iterationHeader = [@"Iteration" retain];

	headers = [NSArray arrayWithObjects:
			iterationHeader,
			@"PotentialEnergy",
			@"KineticEnergy", 
			@"Temperature",
			@"TotalEnergy",
			nil];

	systemEnum = [[systemCollection allSystems] objectEnumerator];
	//Create an entry in stateData for each system
	while((system = [systemEnum nextObject]))
	{
		//FIXME: Supports only one force field per system
		forceFieldArray = [forceFieldCollection forceFieldsForSystem: system];
		if([forceFieldArray count] == 0)
		{
			NSWarnLog(@"No force fields operating on system %@", [system systemName]);
			NSWarnLog(@"Skipping");
			continue;
		}	
			
		systemForceField = [forceFieldArray objectAtIndex: 0];
		stateHeaders = [NSMutableArray arrayWithArray: headers];
		if([system isKindOfClass: [AdInteractionSystem class]])
		{
			[stateHeaders removeObject: @"KineticEnergy"];
			[stateHeaders removeObject: @"Temperature"];
		}	
		//FIXME: Does not handle custom terms.
		[stateHeaders addObjectsFromArray: [systemForceField coreTerms]];		
		stateMatrix = [[AdMutableDataMatrix alloc]
					initWithNumberOfColumns: [stateHeaders count]
					columnHeaders: stateHeaders
					columnDataTypes: nil];
		[stateMatrix autorelease];			
		[(AdMutableDataMatrix*)stateMatrix setName: [system systemName]];	
		[stateData addDataMatrix: stateMatrix];
	}
}

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

- (id) initWithDataStorage: (id) aDataStore
{
	return [self initWithDataStorage: aDataStore
		systems: nil
		forceFields: nil
		iterationHeader: nil];
}

- (id) initWithDataStorage: (id) aDataStore
	systems: (AdSystemCollection*) aSystemCollection
	forceFields: (AdForceFieldCollection*) aForceFieldCollection
	iterationHeader: (NSString*) aString
{
	NSArray* headers;

	if((self = [super init]))
	{
		if(aDataStore == nil)
			[NSException raise: NSInvalidArgumentException
				format: @"aDataStore cannot be nil"];
	
		if([aDataStore storageMode] != AdSimulationStorageWriteMode)
			[NSException raise: NSInternalInconsistencyException
				format: @"You can only write data to a store in AdSimulationStorageWriteMode"];
	
		lastFrame = -1;
		systemCollection = nil;
		forceFieldCollection = nil;
		iterationHeader = [aString retain];
		iterationValue = nil;
		stateData = [[AdDataSet alloc] 
				initWithName: @"StateData"
				inputReferences: nil
				dataGenerator: [NSBundle mainBundle]];
		topologyData = [NSMutableDictionary new];
		energyData = [NSMutableDictionary new];
		//This is allocated on each trajectory checkpoint request
		trajectoryData = nil;
		dataStorage = [aDataStore retain];	
		headers = [NSArray arrayWithObjects: 
				@"Energy",
				@"Trajectory",
				@"Topology",
				nil];
		frames = [[AdMutableDataMatrix alloc]
				initWithNumberOfColumns: 3
				columnHeaders: headers
				columnDataTypes: nil];		
		
		[self setSystems: aSystemCollection];
		[self setForceFields: aForceFieldCollection];

	}

	return self;
}

- (void) dealloc
{
	[self synchToStore];
	[iterationValue release];
	[iterationHeader release];
	[topologyData release];
	[energyData release];
	[trajectoryData release];
	[stateData release];
	[dataStorage release];
	[frames release];
	[systemCollection release];
	[forceFieldCollection release];
	[super dealloc];
}

- (void) setSystems: (AdSystemCollection*) aSystemCollection
{
	NSMutableData* data;
	NSKeyedArchiver* archiver;

	if(systemCollection != nil)
		[NSException raise: NSInternalInconsistencyException
			format: @"The system collection cannot be changed"];

	systemCollection = [aSystemCollection retain];		

	data = [NSMutableData new];
	archiver = [[NSKeyedArchiver alloc] initForWritingWithMutableData: data];
	[archiver setOutputFormat: NSPropertyListXMLFormat_v1_0];
	[archiver encodeObject: systemCollection forKey: @"SystemCollection"];
	[archiver finishEncoding];
	[archiver release];
	[dataStorage setSystemData: data];
	[data release];
	if(forceFieldCollection != nil)
		[self _createStateMatrices];
}

- (void) updateSystems
{
	NSWarnLog(@"Method %@ not implemented", NSStringFromSelector(_cmd));
}

- (void) setForceFields: (AdForceFieldCollection*) aForceFieldCollection
{
	[forceFieldCollection release];
	forceFieldCollection = [aForceFieldCollection retain];
	if(systemCollection != nil)
		[self _createStateMatrices];
}

- (void) setIterationHeader: (NSString*) aString
{
	if(iterationHeader != nil)
	{
		NSWarnLog(@"State matrices have already been created.");
		return;
	}

	iterationHeader = [aString retain];
}

- (void) openFrame: (NSNumber*) value
{
	NSDebugLLog(@"AdSimulationDataWriter",
		@"Recieved open frame request");
	NSDebugLLog(@"AdSimulationDataWriter",
		@"Frame open? - %d", frameOpen);
	NSDebugLLog(@"AdSimulationDataWriter",
		@"Current iteration value %@", iterationValue);
	NSDebugLLog(@"AdSimulationDataWriter",
		@"Requested iteration value %@", value);

	if(frameOpen)
		return;

	if([value isEqual: iterationValue])
		return;

	NSDebugLLog(@"AdSimulationDataWriter", 
		@"Opening frame");

	frameOpen =YES;
	[iterationValue release];
	iterationValue = [value retain];
}

- (BOOL) isOpenFrame
{
	return frameOpen;
}

- (void) _addEnergyCheckpoint
{
	NSEnumerator* dataEnum;
	NSString* systemName;
	id state, stateMatrix;

	dataEnum = [energyData keyEnumerator];
	while((systemName = [dataEnum nextObject]))
	{
		stateMatrix = [stateData dataMatrixWithName: 
				systemName];
		state = [energyData objectForKey: systemName];		
		[stateMatrix extendMatrixWithRow: state];
	 }		
}

- (void) closeFrame
{
	NSMutableData* data;
	NSMutableArray* checkpointed = [NSMutableArray new];
	NSKeyedArchiver* archiver;
	NSEnumerator* systemEnum;
	id systemName, topology;

	if(!frameOpen)
		return;

	//Write all the data
	frameOpen = NO;
	lastFrame++;
	
	NSDebugLLog(@"AdSimulationDataWriter",
		@"Closing frame %d", lastFrame);
	if(trajectoryCheckpoint)
	{
		[dataStorage addTrajectoryCheckpoint: trajectoryData];
		NSDebugLLog(@"AdSimulationDataWriter",
			@"Checkpointed %d bytes of trajectory data",
			[trajectoryData length]);
		[checkpointed addObject: [NSNumber numberWithInt: 1]];
	}
	else
		[checkpointed addObject: [NSNumber numberWithInt: 0]];

	if(energyCheckpoint)
	{
		if(forceFieldCollection == nil)
			[checkpointed addObject: [NSNumber numberWithBool: NO]];
		else
		{
			[self _addEnergyCheckpoint];
			[checkpointed addObject: [NSNumber numberWithInt: 1]];
			[energyData removeAllObjects];
		}
	}
	else
		[checkpointed addObject: [NSNumber numberWithInt: 0]];

	if(topologyCheckpoint)
	{
		data = [NSMutableData new];
		archiver = [[NSKeyedArchiver alloc] initForWritingWithMutableData: data];
		[archiver setOutputFormat: NSPropertyListXMLFormat_v1_0];
		systemEnum = [topologyData keyEnumerator];
		while((systemName = [systemEnum nextObject]))
		{
			//We assume the data source has been synched
			//to the current configuration
			NSDebugLLog(@"AdSimulationDataWriter",
				@"Checkpointing topology of %@", systemName);
			topology = [topologyData objectForKey: systemName];
			[archiver encodeObject: topology
				forKey: systemName];
		}
		
		[archiver finishEncoding];
		[archiver release];
		[dataStorage addTopologyCheckpoint: data];
		NSDebugLLog(@"AdSimulationDataWriter",
			@"Checkpointed %d bytes of topology data", [data length]);
		[data release];
		[checkpointed addObject: [NSNumber numberWithInt: 1]];
		[topologyData removeAllObjects];
	}
	else
		[checkpointed addObject: [NSNumber numberWithInt: 0]];

	[frames extendMatrixWithRow: checkpointed];

	//Set for next frame
	[checkpointed release];
	energyCheckpoint = NO;
	trajectoryCheckpoint = NO;
	topologyCheckpoint = NO;
}

- (int)  lastFrame
{
	return lastFrame;
}

/**
The first frame should indicate a checkpoint for everything
*/
- (unsigned int) _lastCheckpointInColumn: (NSString*) column
{	
	int i, numberOfRows;

	if(lastFrame == -1)
		return lastFrame;

	numberOfRows = [frames numberOfRows];
	for(i=numberOfRows-1; i>=0; i--)
		if([[frames elementAtRow: i ofColumnWithHeader: column] boolValue])
			break;
	
	return i;
}

- (unsigned int) lastEnergyCheckpoint
{
	return [self _lastCheckpointInColumn: @"Energy"];
}

- (unsigned int) lastTopologyCheckpoint
{
	return [self _lastCheckpointInColumn: @"Topology"];
}

- (unsigned int) lastTrajectoryCheckpoint
{
	return [self _lastCheckpointInColumn: @"Trajectory"];
}

- (void) addTopologyCheckpoint
{
	NSEnumerator* systemEnum;
	id system;

	//Clear previously collected data
	[topologyData removeAllObjects];

	//Add new checkpoints
	systemEnum = [[systemCollection fullSystems] objectEnumerator];
	while((system = [systemEnum nextObject]))
		[self addTopologyCheckpointForSystem: system];
}

- (void) addTopologyCheckpointForSystem: (id) aSystem
{
	id topology;

	topologyCheckpoint = YES;
	topology = [[[aSystem dataSource] copy] autorelease];
	[topologyData setObject: topology
		forKey: [aSystem systemName]];
}

- (void) addTrajectoryCheckpoint
{
	NSKeyedArchiver* archiver;
	NSEnumerator* systemEnum;
	id system, state;
	
	trajectoryCheckpoint = YES;

	//Release previous trajectory checkpoint
	[trajectoryData release];
	
	//Checkpoint current trajectory
	
	trajectoryData = [NSMutableData new];
	archiver = [[NSKeyedArchiver alloc] 
			initForWritingWithMutableData: trajectoryData];
	[archiver setOutputFormat: NSPropertyListXMLFormat_v1_0];
	systemEnum = [[systemCollection fullSystems] objectEnumerator];
	while((system = [systemEnum nextObject]))
	{
		state = [system captureState];
		[archiver encodeObject: state 
			forKey: [system systemName]];
	}	

	[archiver finishEncoding];
	[archiver release];
}

- (void) addEnergyCheckpoint
{
	double potentialEnergy, totalEnergy, kineticEnergy;
	NSEnumerator* systemEnum;
	AdForceField* systemForceField;
	id system, state;

	energyCheckpoint = YES;
	
	//Check that a forceFieldCollection is available before continuing
	if(forceFieldCollection == nil)
	{
		NSWarnLog(@"No force fields are provided - No energy data will be checkpointed");
		return;
	}	

	//Remove last energy checkpoint
	[energyData removeAllObjects];
	
	//Checkpoint new energies
	//Go through each system and add the relevant data
	//to an array. Each array is added to energyData dictionary.
	//When the frame is close the arrays are added to the
	//state matrix for each system (in closeFrame:)
	systemEnum = [[systemCollection allSystems] objectEnumerator];
	while((system = [systemEnum nextObject]))
	{
		state = [NSMutableArray array];
		//FIXME: Only one force field per system supported
		systemForceField = [[forceFieldCollection forceFieldsForSystem: system]
					objectAtIndex: 0];
		[state addObject: iterationValue];
		potentialEnergy = [systemForceField totalEnergy];
		[state addObject:
			[NSNumber numberWithDouble: potentialEnergy]];

		totalEnergy = potentialEnergy;
		if([system respondsToSelector: @selector(kineticEnergy)])
		{
			kineticEnergy =  [system kineticEnergy];
			totalEnergy += kineticEnergy;
			[state addObject:
				[NSNumber numberWithDouble: kineticEnergy]];
			[state addObject:
				[NSNumber numberWithDouble: [system temperature]]];
		}
			
		[state addObject:
			[NSNumber numberWithDouble: totalEnergy]];
		[state addObjectsFromArray: 
			[systemForceField arrayOfCoreTermEnergies]];
		//Add the array containing the checkpointed energies
		//to the energyData dictionary
		[energyData setObject: state
			forKey: [system systemName]];
	 }		
}

- (void) rollBackToFrame: (unsigned int) value
{
	BOOL wasCheckpointed;
	int start, i;
	int trajectory, energy, topology;
	NSRange energyRange;
	NSEnumerator* matrixEnum;
	NSIndexSet* indexSet;
	AdMutableDataMatrix* matrix;

	if(value > [frames numberOfRows])
		[NSException raise: NSRangeException
			format: @"(%@) %d is out of range %d",
			NSStringFromSelector(_cmd),
			value,
			[frames numberOfRows]];

	trajectory = energy = topology = 0;
	//From the frame after the requested frame to the
	//last recorded frame
	start = value + 1;
	for(i=start; i<lastFrame+1;i++)
	{
		wasCheckpointed = [[frames elementAtRow: i 
					ofColumnWithHeader: @"Trajectory"] 
			 		boolValue];

		if(wasCheckpointed)
			trajectory++;
			
		wasCheckpointed = [[frames elementAtRow: i
					ofColumnWithHeader: @"Energy"] 
			 		boolValue];

		if(wasCheckpointed)
			energy++;

		wasCheckpointed = [[frames elementAtRow: i 
					ofColumnWithHeader: @"Topology"] 
			 		boolValue];

		if(wasCheckpointed)
			topology++;
	}

	NSDebugLLog(@"AdSimulationDataWriter", @"Rolling back to frame %d", value);
	NSDebugLLog(@"AdSimulationDataWriter", @"Deleting %d trajectory checkpoints", trajectory);
	NSDebugLLog(@"AdSimulationDataWriter", @"Deleting %d energy checkpoints", energy);
	NSDebugLLog(@"AdSimulationDataWriter", @"Deleting %d topology checkpoints", topology);

	[dataStorage removeTrajectoryCheckpoints: trajectory];
	[dataStorage removeTopologyCheckpoints: topology];

	//Delete the energy frames
	matrixEnum = [[stateData dataMatrices] objectEnumerator];
	energyRange.length = energy;
	while((matrix = [matrixEnum nextObject]))
	{
		energyRange.location = [matrix numberOfRows] - energyRange.length;
		indexSet = [NSIndexSet indexSetWithIndexesInRange: energyRange];
		NSDebugLLog(@"AdSimulationDataWriter", 
			@"Removing %@ rows from matrix %@", indexSet, [matrix name]);
		[matrix removeRowsWithIndexes: indexSet];
	}	

	//Delete the frame information
	[frames removeRowsWithIndexes: 
		[NSIndexSet indexSetWithIndexesInRange: 
			NSMakeRange(start, lastFrame - start + 1)]];

	//Remove any data marked for checkpointing so far
	if(frameOpen)
	{
		energyCheckpoint = NO;
		trajectoryCheckpoint = NO;
		topologyCheckpoint = NO;
		[topologyData removeAllObjects];
		[energyData removeAllObjects];
		[trajectoryData release];
		trajectoryData = nil;
		frameOpen = NO;
	}

	lastFrame = value;
	[self synchToStore];
}

- (void) synchToStore
{
	NSKeyedArchiver* archiver;
	NSMutableData* data;
	
	NSDebugLLog(@"AdSimulationDataWriter", @"Synching to store");
	
	data = [NSMutableData new];
	archiver = [[NSKeyedArchiver alloc] initForWritingWithMutableData: data];
	[archiver setOutputFormat: NSPropertyListXMLFormat_v1_0];
	[archiver encodeObject: stateData forKey: @"root"];
	[archiver finishEncoding];
	[dataStorage setEnergyData: data];
	NSDebugLLog(@"AdSimulationDataWriter", @"Wrote %d bytes of energy data", [data length]);
	[archiver release];
	[data release];

	data = [NSMutableData new];
	archiver = [[NSKeyedArchiver alloc] initForWritingWithMutableData: data];
	[archiver setOutputFormat: NSPropertyListXMLFormat_v1_0];
	[archiver encodeObject: frames forKey: @"root"];
	[archiver finishEncoding];
	[dataStorage setFrameData: data];
	NSDebugLLog(@"AdSimulationDataWriter", @"Wrote %d bytes of frame data", [data length]);
	[archiver release];

	[dataStorage synchronizeStore];
	NSDebugLLog(@"AdSimulationDataWriter", @"Complete");
}

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

@end