/*
   Project: Adun

   Copyright (C) 2005 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/AdunDynamics.h"

@implementation AdDynamics

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

Coordinate Manipulation Methods

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

//make this method public

- (void) _calculateCentreOfMass
{
	int i, j;
	
	if(centreOfMass != NULL)
		[memoryManager freeArray: centreOfMass];	

	//malloc memory for the centre of mass
	
	centreOfMass = (double*)[memoryManager allocateArrayOfSize: 3*sizeof(double)];
	for(mass = 0, i=0; i<numberOfAtoms; i++)
	{
		mass += coordinates->matrix[i][3];
		for(j=0; j<3;j++)
			centreOfMass[j] += coordinates->matrix[i][j]*coordinates->matrix[i][3];
	}

	for(i=0; i<3; i++)
		centreOfMass[i] = centreOfMass[i]/mass;

	NSDebugLLog(@"AdDynamics", @"Mass %lf. Centre of Mass is: %-6.2lf%-6.2lf%-6.2lf", 
			mass, centreOfMass[0], centreOfMass[1], centreOfMass[2]);
}

- (void) moveCentreOfMassToOrigin
{	
	int i,j;

	NSDebugLLog(@"AdDynamics", @"Moving centreOfMass of mass to (0,0,0)");

	for(i=0; i<numberOfAtoms; i++)
		for(j=0; j<3; j++)
			coordinates->matrix[i][j] -= centreOfMass[j];
}

- (void) centerOnPoint: (double*) point
{
	int i,j;

	NSDebugLLog(@"AdDynamics", @"Moving %-6.2lf%-6.2lf%-6.2lf to origin", point[0], point[1], point[2]);

	for(i=0; i<numberOfAtoms; i++)
		for(j=0; j<3; j++)
			coordinates->matrix[i][j] -= point[j];
}

- (void) centerOnAtom: (int) atom
{
	int i,j;

	NSDebugLLog(@"AdDynamics", @"Moving atom %d (%-6.2lf%-6.2lf%-6.2lf) to origin", atom,
		coordinates->matrix[atom][0],
		coordinates->matrix[atom][1], 
		coordinates->matrix[atom][2]);

	for(i=0; i<numberOfAtoms; i++)
		for(j=0; j<3; j++)
			coordinates->matrix[i][j] -= coordinates->matrix[atom][j];
}

- (void) zeroAccelerations
{
	int i, j;

	for(i=0; i<numberOfAtoms; i++)
		for(j=0; j<3; j++)
			accelerations->matrix[i][j] = 0;
}

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

Accessing Data Source & Matrix Creation

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

- (void) _createVelocityMatrix
{
	int i, j;
	double var, hold;
	gsl_rng* twister;
	
	velocities = (AdMatrix*)[memoryManager allocateMatrixWithRows: numberOfAtoms 
				withColumns: 3];	
	
	NSDebugLLog(@"AdDynamics", @"Randomly initialising velocities.");
	twister = gsl_rng_alloc(gsl_rng_mt19937);
	gsl_rng_set(twister, seed);

	/*
   	 * Draw velocites from a Maxwell Distribution at the given temperature
	 * Maxwell Distribution is a gaussian with variance = sqrt(KB*T/mass)
	 * Catch when the temperature is set to zero
  	 * Not sure what the temperature limit should be..
	 */

	if(targetTemperature > 0.001)
	{
		var = targetTemperature*KB;
		for(i=0; i<numberOfAtoms; i++)
		{
			hold = sqrt(var*coordinates->matrix[i][4]);		
			for(j=0; j<3; j++)
				velocities->matrix[i][j] = gsl_ran_gaussian(twister, hold);
		}
	}
	else
		for(i=0; i<numberOfAtoms; i++)
			for(j=0; j<3; j++)
				velocities->matrix[i][j] = 0;

	ownsVelocities = YES;
	gsl_rng_free(twister);
}

- (void) _createAccelerationMatrix
{
	accelerations = (AdMatrix*)[memoryManager allocateMatrixWithRows: numberOfAtoms
					withColumns: 3];	
	ownsAccelerations = YES;
}

- (void) _retrieveCoordinates
{
	NSDebugLLog(@"AdDynamics", @"Accesing data source for coordinates");
	coordinates = [[dataSource objectValueForCoordinates: self] pointerValue];
	numberOfAtoms = coordinates->no_rows;
	NSDebugLLog(@"AdDynamics", @"Number of coordinates = %d\n", numberOfAtoms);
}

- (void) _retrieveAtomTypes
{
	int i;

	NSDebugLLog(@"AdDynamics", @"Accesing data source for atomTypes");
	atomTypes = [dataSource objectValueForAtomTypes: self];
	if((int)[atomTypes count] != numberOfAtoms)
	{
		NSDebugLLog(@"AdDynamics", @"Number of atoms %d. Number of atom types %d",
		 	numberOfAtoms,
		 	[atomTypes count]);
		[NSException raise: NSInternalInconsistencyException
			format: @"The dataSource doesnt contain the correct number of atomType or doesnt return an Array!"];
	}
}

- (void) _retrieveVelocities
{
	id retVal;

	NSDebugLLog(@"AdDynamics", @"Accesing data source for velocities\n");
	retVal = [dataSource objectValueForVelocities: self];
	if(retVal == nil)
	{
		NSDebugLLog(@"AdDynamics", @"Data source contains no velocities information.");
		[self _createVelocityMatrix];
	}
	else
	{		
		NSDebugLLog(@"AdDynamics", @"Retrieved velocities from data source");
		velocities = [retVal pointerValue];
		ownsVelocities = NO;
	}
}

- (void) _retrieveAccelerations
{
	id retVal;

	NSDebugLLog(@"AdDynamics", @"Accesing data source for accelerations\n");
	retVal = [dataSource objectValueForAccelerations: self];
	if(retVal == nil)
	{
		NSDebugLLog(@"AdDynamics", @"Data source contains no accelerations information.");
		[self _createAccelerationMatrix];
	}
	else
	{		
		NSDebugLLog(@"AdDynamics", @"Retrieved accelerations from data source");
		accelerations = [retVal pointerValue];
		ownsAccelerations = YES;
	}
}

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

Initialisation

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


- (void) _useEnvironmentDefaults
{
	seed = 500;
	targetTemperature = 300;
}

- (id) initWithEnvironment: (id) object
{
	return [self initWithEnvironment: object observe: YES];
}

- (id) initWithEnvironment: (id) object observe: (BOOL) value
{
	//default values for environment parameters
	//and data source protocol setup

	if(self = [super initWithEnvironment: object observe: value])
	{
		memoryManager = [AdMemoryManager appMemoryManager];
		dataSourceProtocolName = @"AdDynamicsDataSource";
		[dataSourceProtocolName retain];
		dataSourceProtocol = @protocol(AdDynamicsDataSource);
		currentCaptureMethod = @"Standard";
		[currentCaptureMethod retain];
		memento = NO;
		dependantsDict = [NSMutableDictionary dictionaryWithCapacity: 1];
		if(environment == nil)
			[self _useEnvironmentDefaults];
		else
		{
			[self synchroniseWithEnvironment];
			[self registerWithEnvironment];
		}
	}

	return self;
}

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

- (void) dealloc
{
	[currentCaptureMethod release];	
	[dataSourceProtocolName release];
	
	if(ownsVelocities)
		[memoryManager freeMatrix: velocities];

	if(ownsAccelerations)
		[memoryManager freeMatrix: accelerations];

	[super dealloc];
}

- (void) deregisterWithEnvironment
{
	[environment removeObserver: self forKey: @"TargetTemperature"];
	[environment removeObserver: self forKey: @"Seed"];
}

- (void) registerWithEnvironment
{
	//register for observation of environment keys

	if(observesEnvironment)
	{
		[environment addObserver: self forKey: @"TargetTemperature"];
		[environment addObserver: self forKey: @"Seed"];
	}
}

/*
 * Environment Observation
 */

- (void) updateForKey: (NSString*) key value: (id) value object: (id) object
{
	if([key isEqual: @"TargetTemperature"])
		targetTemperature = [value doubleValue];
	else if([key isEqual: @"Seed"])
		seed = [value intValue];
}

- (void) synchroniseWithEnvironment
{
	if(environment == nil)
		[NSException raise: NSInternalInconsistencyException
			format: @"Cannot synchronise with environment as there is none set"];

	targetTemperature = [[environment valueForKey: @"TargetTemperature"] intValue];

	//FIXME: This is in lieu of rigorous option checking by the environment or UL
	//which still has to be implemented

	if(targetTemperature < 0.0)
		[NSException raise: NSInvalidArgumentException
			format: @"Temperature cannot be negative"];
	
	seed = [[environment valueForKey: @"Seed"] intValue];
}

- (void) setEnvironment: (id) object
{
	[self deregisterWithEnvironment];	
	environment = object;
	[self registerWithEnvironment];
}

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

 Accessors 

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

- (int) numberOfAtoms
{
	return coordinates->no_rows;
}

- (NSArray*) atomTypes
{
	return atomTypes;
}

- (NSArray*) atomMasses
{
	int i;
	NSMutableArray* masses =  [NSMutableArray array];

	for(i=0; i < coordinates->no_rows; i++)
		[masses addObject: 
			[NSNumber numberWithDouble: coordinates->matrix[i][3]]];

	return [[masses copy] autorelease];		
}

- (NSValue*) coordinates
{
	return [NSValue valueWithPointer: coordinates];
}

- (NSValue*) velocities
{
	return [NSValue valueWithPointer: velocities];
}

- (NSValue*) accelerations	
{
	return [NSValue valueWithPointer: accelerations];
}

- (double) mass
{
	return mass;
}

- (double*) centreOfMass
{
	return centreOfMass;
}

- (double) targetTemperature
{
	return targetTemperature;
}

- (int) seed
{
	return seed;
}

- (void) setSeed: (int) aNumber
{
	seed = aNumber;
}

- (void) setTargetTemperature: (double) aNumber
{
	if(targetTemperature < 0.0)
		[NSException raise: NSInvalidArgumentException
			format: @"Temperature cannot be negative"];
		
	targetTemperature = aNumber;
}

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

Protocol Implementations

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

//AdDataSource delegation methods

- (void) reloadData
{
	id enumerator, obj;

	NSDebugLLog(@"AdDynamics", @"Loading data from data source");

	if(dataSource != nil)
	{
		[self _retrieveCoordinates];
		[self _retrieveAtomTypes];
	
		if(ownsVelocities)
			[memoryManager freeMatrix: velocities];

		[self _retrieveVelocities];

		if(ownsAccelerations)
			[memoryManager freeMatrix: accelerations];
		[self _retrieveAccelerations];
		[self _calculateCentreOfMass];
	}
	else
		[NSException raise: NSInternalInconsistencyException
			format: @"No data source has been set (%@)", [self description]];

	NSDebugLLog(@"AdDynamics", @"Loading of data complete");
}

//NSCoding Methods

- (id) initWithCoder: (NSCoder*) decoder
{	
	int i, j;
	int length;
	double* array, *buffer;
	NSString* archiveType;

	self  = [super initWithCoder: decoder];
	if([decoder allowsKeyedCoding])
	{
		memoryManager = [AdMemoryManager appMemoryManager];
		archiveType = [decoder decodeObjectForKey: @"ArchiveType"];
		seed = [decoder decodeIntForKey: @"Seed"];
		targetTemperature = [decoder decodeDoubleForKey: @"TargetTemperature"];
		dataSource = [decoder decodeObjectForKey: @"DataSource"];
		ownsVelocities = [decoder decodeBoolForKey: @"OwnsVelocities"];
		ownsAccelerations = [decoder decodeBoolForKey: @"OwnsAccelerations"];
		currentCaptureMethod = [[decoder decodeObjectForKey: @"CurrentCaptureMethod"] retain];

		if([decoder decodeIntForKey: @"Velocities.Rows"] != 0)
		{
			velocities = (AdMatrix*)malloc(sizeof(AdMatrix));
			buffer = (double*)[decoder decodeBytesForKey: @"VelocitiesMatrix"
						 returnedLength: &length];
			velocities->no_rows = [decoder decodeIntForKey: @"Velocities.Rows"];
			velocities->no_columns = 3;
			velocities->matrix = (double**)malloc(velocities->no_rows*sizeof(double*));
			array = malloc(length);
			for(i=0, j=0; i<velocities->no_rows; i ++, j = j + velocities->no_columns)
				velocities->matrix[i] = array + j;
			for(i=0; i< velocities->no_rows*velocities->no_columns; i++)
				array[i] = buffer[i];
		}

		if([decoder decodeIntForKey: @"Accelerations.Rows"] != 0)
		{
			accelerations = (AdMatrix*)malloc(sizeof(AdMatrix));
			buffer = (double*)[decoder decodeBytesForKey: @"AccelerationsMatrix"
						 returnedLength: &length];
			accelerations->no_rows = [decoder decodeIntForKey: @"Accelerations.Rows"];
			accelerations->no_columns = 3;
			accelerations->matrix = (double**)malloc(accelerations->no_rows*sizeof(double*));
			array = malloc(length);
			for(i=0, j=0; i<accelerations->no_rows; i ++, j = j + accelerations->no_columns)
				accelerations->matrix[i] = array + j;
			for(i=0; i< accelerations->no_rows*accelerations->no_columns; i++)
				array[i] = buffer[i];
		}

		[self _retrieveCoordinates];
		if([archiveType isEqual: @"Complete"])
		{
			[self _retrieveAtomTypes];
			[self _calculateCentreOfMass];
		}

		numberOfAtoms = coordinates->no_rows;
	}
	else
		[NSException raise: NSInvalidArgumentException 
			format: @"%@ does not support non keyed coding", [self classDescription]];

	//connect with the current environment (if it exists)

	environment = [AdEnvironment globalEnvironment];
	if(environment != nil)
	{
		[self synchroniseWithEnvironment];
		[self registerWithEnvironment];
	}
	else
		[self _useEnvironmentDefaults];	
	
	dataSourceProtocolName = @"AdDynamicsDataSource";
	[dataSourceProtocolName retain];
	dataSourceProtocol = @protocol(AdDynamicsDataSource);

	return self;
}

- (void) _fullEncodeWithCoder: (NSCoder*) encoder
{
	int bytesLength;

	NSDebugLLog(@"Encode", @"Encoding %@", [self description]);
	[encoder encodeInt: seed forKey: @"Seed"];
	[encoder encodeDouble: targetTemperature forKey: @"TargetTemperature"];
	[encoder encodeBool: ownsVelocities forKey: @"OwnsVelocities"];
	[encoder encodeBool: ownsAccelerations forKey: @"OwnsAccelerations"];
	[encoder encodeConditionalObject: dataSource forKey: @"DataSource"];
	[encoder encodeObject: currentCaptureMethod forKey: @"CurrentCaptureMethod"];

	//velocities

	if(ownsVelocities)
	{
		bytesLength = sizeof(AdMatrixSize)*velocities->no_rows*velocities->no_columns;	
		[encoder encodeBytes: (uint8_t*)velocities->matrix[0] 
			length: bytesLength 
			forKey: @"VelocitiesMatrix"];
		[encoder encodeInt: velocities->no_rows forKey: @"Velocities.Rows"];	
	}

	//acceleration

	if(ownsAccelerations)
	{
		bytesLength = sizeof(AdMatrixSize)*accelerations->no_rows*accelerations->no_columns;	
		[encoder encodeBytes: (uint8_t*)accelerations->matrix[0]
			length: bytesLength 
			forKey: @"AccelerationsMatrix"];
		[encoder encodeInt: accelerations->no_rows forKey: @"Accelerations.Rows"];		
		NSDebugLLog(@"Encode", @"Encoding %d bytes", bytesLength);
	}
	NSDebugLLog(@"Encode", @"Complete %@", [self description]);
}

- (void) _mementoEncodeWithCoder: (NSCoder*) encoder
{
	int bytesLength;

	[dataSource captureStateWithArchiver: encoder key: @"DataSource"];

	//velocities

	if([currentCaptureMethod isEqual: @"Full"])
	{
		if(ownsVelocities)
		{
			bytesLength = sizeof(AdMatrixSize)*velocities->no_rows*velocities->no_columns;	
			[encoder encodeBytes: (uint8_t*)velocities->matrix[0] 
				length: bytesLength 
				forKey: @"VelocitiesMatrix"];
			[encoder encodeInt: velocities->no_rows forKey: @"Velocities.Rows"];	
			NSDebugLLog(@"Encode", @"Encoding %d bytes", bytesLength);
		}

		//acceleration

		if(ownsAccelerations)
		{
			bytesLength = sizeof(AdMatrixSize)*accelerations->no_rows*accelerations->no_columns;	
			[encoder encodeBytes: (uint8_t*)accelerations->matrix[0] 
				length: bytesLength 
				forKey: @"AccelerationsMatrix"];
			[encoder encodeInt: accelerations->no_rows forKey: @"Accelerations.Rows"];		
			NSDebugLLog(@"Encode", @"Encoding %d bytes", bytesLength);
		}
		NSDebugLLog(@"Encode", @"Complete %@", [self description]);
	}
}

- (void) encodeWithCoder: (NSCoder*) encoder
{
	[super encodeWithCoder: encoder];
	if([encoder allowsKeyedCoding])
	{
		[encoder encodeObject: currentCaptureMethod forKey: @"CaptureMethod"];

		if(memento)
		{
			[encoder encodeObject: @"Memento" forKey: @"ArchiveType"];
			[self _mementoEncodeWithCoder: encoder];
		}
		else
		{
			[encoder encodeObject: @"Complete" forKey: @"ArchiveType"];
			[self _fullEncodeWithCoder: encoder];
		}
	}
	else
		[NSException raise: NSInvalidArgumentException
			format: @"%@ class does not support non keyed coding", [self class]];
}

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

AdMemento

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

- (id) setCaptureMethod: (NSString*) name
{
	if([name isEqual: @"Standard"] || [name isEqual: @"Full"])
	{
		[currentCaptureMethod release];
		currentCaptureMethod = [name retain];
	}
	else
		[NSException raise: NSInvalidArgumentException
			format: @"Capture Method %@ is not supported by this object"]; 	
}

- (void) captureStateWithArchiver: (NSCoder*) archiver key: (NSString*) key
{
	memento = YES;
	[archiver encodeObject: self forKey: key];
	memento = NO;
}

- (void) returnToState: (id) stateMemento
{	
	AdMatrix* mementoVelocities, *mementoAccelerations;
	int i, j;
	id dataSourceMemento;

	//check that this is actually a memento of this class

	if(![stateMemento isMemberOfClass: [self class]])	
		[NSException raise: NSInternalInconsistencyException
			format: @"StateMemento is not of the correct type (%@)", 
			[stateMemento class]];	

	//Also check that there is the correct number of particles in the decoded matrix
	//N:B: ownsVelocities and ownsAccelerations must be NO when this object is initialised

	if(ownsVelocities)
	{
		mementoVelocities = [[stateMemento valueForKey: @"Velocities"] pointerValue];

		if(mementoVelocities->no_rows != velocities->no_rows)
			[NSException raise: NSInternalInconsistencyException
				format: @"StateMemento system is a different size (%d v %d)", 
					mementoVelocities->no_rows, velocities->no_rows];

		for(i=0; i< velocities->no_rows; i++)
			for(j=0; j<velocities->no_columns; j++)
				velocities->matrix[i][j] = mementoVelocities->matrix[i][j];
	}

	if(ownsAccelerations)
	{
		mementoAccelerations = [[stateMemento valueForKey: @"Accelerations"] pointerValue];

		if(mementoAccelerations->no_rows != accelerations->no_rows)
			[NSException raise: NSInternalInconsistencyException
				format: @"StateMemento system is a different size (%d v %d)", 
					mementoAccelerations->no_rows, accelerations->no_rows];

		for(i=0; i< accelerations->no_rows; i++)
			for(j=0; j<accelerations->no_columns; j++)
				accelerations->matrix[i][j] = mementoAccelerations->matrix[i][j];
	}

	//get the datasource memento from the  stateMemento

	dataSourceMemento = [stateMemento valueForKey: @"DataSource"];
	[dataSource returnToState: dataSourceMemento];
}

- (id) captureState
{
	[NSException raise: NSInternalInconsistencyException
		format: [NSString stringWithFormat: 
			@"Warning method %@ not implemented yet", NSStringFromSelector(_cmd)]];
}

@end



