/*
   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/AdunState.h"

static inline double CalcKineticEnergy(double**, double**, int);

@implementation AdState

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

Object Creation and Maintainence

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

- (void) clearTimer
{
	[scheduler removeMessageWithName: @"ShortRangeNonbonded"];
	[scheduler removeMessageWithName: @"LogOutput"];
	//[scheduler removeMessageWithName: @"LongRangeNonbonded"];
}

- (void) _initialiseTimer
{
	[scheduler sendMessage: @selector(update) 
		toObject: nonbonded
		interval: update_interval
		name: @"ShortRangeNonbonded"];
/*	[scheduler sendMessage: @selector(update) 
		toObject: longRangeNonbonded
		interval: update_interval
		name: @"LongRangeNonbonded"];*/
	[scheduler sendMessage: @selector(_writeEnergies) 
		toObject: self
		interval: log_interval
		name: @"LogOutput"];
}

- (void) _setupDynamicState
{
	NSError* error;
	NSException* exception;

	//FIXME: We assume here that the translational degrees of 
	//freedom have been moved. However this assumption is invalid.
	//It is possible that more or less DOFs could have been removed.
	DOF =  3*no_of_atoms - 3;
	if(DOF <= 0)
	{
		error = AdKnownExceptionError(10,
			@"There are 0 degress of freedom",
			@"This is most likely because you have only 1 atom in the simulation.\n \
Since translational degrees of freedom have been removed the atom is essentially frozen.\n",
			nil);
		exception = [NSException exceptionWithName: NSInternalInconsistencyException
				reason: @"Inconsistency detected in system state."
				userInfo: [NSDictionary dictionaryWithObject: error
						forKey: @"AdKnownExceptionError"]];
	}
	ke_2_temp =  (2*KB_1)/(DOF);

	//calculate initial temperature and kinetic energy

	kinetic = CalcKineticEnergy(coordinates->matrix, velocities->matrix, no_of_atoms);
	temperature = ke_2_temp*kinetic;
	total = 0;
}

- (void) _loadSystemData
{
	dynamics = [system valueForKey:@"dynamics"];
	if(dynamics != nil)
	{
		coordinates = [[dynamics valueForKey: @"coordinates"] pointerValue];
		atomTypes = [dynamics valueForKey: @"atomTypes"];
		velocities = [[dynamics valueForKey: @"velocities"] pointerValue];
		no_of_atoms = coordinates->no_rows;
		[self _setupDynamicState];
	}

	bonded = [system valueForKey:@"bondedTopology"];
	nonbonded = [system valueForKey:@"shortRangeNonbondedTopology"];
	longRangeNonbonded = [system valueForKey:@"longRangeNonbondedTopology"];
	
	if(systemName != nil)
		[systemName release];
	systemName = [system valueForKey: @"systemName"];	
	[systemName retain];
	
	[self _initialiseTimer];
}

- (void) _useEnvironmentDefaults
{
	timestep = 1.0;
	update_interval = 20;
	log_interval = 100;
}

- (id) initWithEnvironment: (id) object system: (id) aSystem
{
	return [self initWithEnvironment: object system: aSystem observe: YES];
}
 
- (id) initWithEnvironment: (id) object system: (id) aSystem observe: (BOOL) value
{
	if(self = [super initWithEnvironment: object observe: value])
	{
		if(aSystem == nil)
			[NSException raise: NSInvalidArgumentException
				format: @"System cannot be nil"];

		dependantsDict = [NSMutableDictionary dictionaryWithCapacity: 1];
		[dependantsDict setObject: [NSArray arrayWithObject: @"_DOF"]
				forKey: @"NumberOfAtoms"];
		[dependantsDict setObject: [NSArray arrayWithObject: @"_KineticToTemperature"]
				forKey: @"_DOF"];
		[dependantsDict retain];
		
		potential = 0;
		kinetic = 0;
		total = 0;
		DOF = 0;
		kineticUpdate = YES;
		scheduler = [AdTimer new];

		if(environment == nil)
			[self _useEnvironmentDefaults];
		else
		{
			[self synchroniseWithEnvironment];
			[self registerWithEnvironment];
		}

		system = aSystem;
		[self _loadSystemData];
	}

	return self;		
}

- (id) initWithEnvironment: (id) object
{
	[NSException raise: NSInvalidArgumentException
		format: @"You must use initWithEnvironment:system to initialise this object"];

	return nil;
}

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

- (void) dealloc
{
	[systemName release];
	[dependantsDict release];
	[scheduler release];
	[super dealloc];
}

/*
 * Environment observation
 */

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

	if(observesEnvironment)
	{
		timestep = [[environment valueForKey: @"TimeStep"] doubleValue];
		update_interval = [[environment valueForKey: @"UpdateInterval"] intValue];
		log_interval = [[environment valueForKey: @"LogInterval"] intValue];
	}
}

- (void) deregisterWithEnvironment
{
	[environment removeObserver: self forKey: @"TimeStep"];
	[environment removeObserver: self forKey: @"UpdateInterval"];
	[environment removeObserver: self forKey: @"LogInterval"];
}

- (void) registerWithEnvironment
{
	if(observesEnvironment)
	{
		[environment addObserver: self forKey: @"TimeStep"];
		[environment addObserver: self forKey: @"UpdateInterval"];
		[environment addObserver: self forKey: @"LogInterval"];
	}
}


- (void) updateForKey: (NSString*) key value: (id) value object: (id) object
{
	if([key isEqual: @"TimeStep"])
		timestep = [value doubleValue];
	else if([key isEqual: @"UpdateInterval"])
		update_interval = [value intValue];
	else if([key isEqual: @"LogInterval"])
		log_interval = [value intValue];
}

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

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

Scheduler related methods

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

- (void) _writeEnergies
{
	NSEnumerator* potentialEnum;
	NSMutableDictionary* termPotentials;
	id energy;
	
	GSPrintf(stderr, @"%@\n", [NSString stringWithFormat: @"\nSystem: %@\n", systemName]);
	GSPrintf(stderr, @"%@\n", 
		[NSString stringWithFormat:
		 	@"Time:%-8.2lf Total:%-12.3lf Kinetic:%-12.3lf Potential:%-12.3lf Temperature: %-12.3lf\n", 
			time, STCAL*total, STCAL*kinetic, 
			STCAL*potential, temperature]);

	termPotentials = [currentState valueForKey:@"TermPotentials"];
	potentialEnum = [termPotentials keyEnumerator];
	while(energy = [potentialEnum nextObject])
		 GSPrintf(stderr, @"%@\n", 
			[NSString stringWithFormat: @"%@:%8.2lf ", 
				energy,	
				STCAL*(*(double*)[[termPotentials objectForKey: energy] pointerValue])]);

	GSPrintf(stderr, @"\n");	
}

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

Public Methods

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

//care required when archiving the potential values since 
//they are owned by the ForceField - probably archive a copy

- (void) setCurrentForceFieldState: (id) state
{
	currentState = state;
	potential = *(double*)[[currentState valueForKey: @"TotalPotential"] pointerValue];
}

- (id) currentForceFieldState
{
	return currentState;
}

- (void) update
{
	kinetic = CalcKineticEnergy(coordinates->matrix, velocities->matrix, no_of_atoms);
	total = kinetic + potential;
	temperature = ke_2_temp*kinetic;
}

- (void) frameUpdate
{
	time += timestep;

	if(kineticUpdate)
		[self update];

	[scheduler increment];
}

/**
Reaquires system variables and recalculates various quantities
e.g. DOF, Temperature etc. Does not clear the timer or reset 
anything
*/
- (void) updateSystemData
{
	NSError* error;
	NSException* exception;

	//Require references to various quantities

	dynamics = [system valueForKey:@"dynamics"];
	if(dynamics != nil)
	{
		coordinates = [[dynamics valueForKey: @"coordinates"] pointerValue];
		atomTypes = [dynamics valueForKey: @"atomTypes"];
		velocities = [[dynamics valueForKey: @"velocities"] pointerValue];
		no_of_atoms = coordinates->no_rows;

		//update dynamics related quanitites
		
		DOF =  3*no_of_atoms - 3;
		if(DOF <= 0)
		{
			error = AdKnownExceptionError(10,
				@"There are 0 degress of freedom", 
				@"This is most likely because you have only 1 atom in the simulation.\n \
Since translational degrees of freedom have been removed the atom is essentially frozen.\n",
				nil);
			exception = [NSException exceptionWithName: NSInternalInconsistencyException
					reason: @"Inconsistency detected in system state."
					userInfo: [NSDictionary dictionaryWithObject: error
							forKey: @"AdKnownExceptionError"]];
		}
		ke_2_temp =  (2*KB_1)/(DOF);
		kinetic = CalcKineticEnergy(coordinates->matrix, velocities->matrix, no_of_atoms);
		temperature = ke_2_temp*kinetic;
	}
}

- (void) setSystem: (id) object
{	
	system = object;
	potential = kinetic = total = DOF = 0;
	[self clearTimer];
	[self _loadSystemData];
}

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

Protocols

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

- (id) initWithCoder: (NSCoder*) decoder
{
	self  = [super initWithCoder: decoder];
	if([decoder allowsKeyedCoding])
	{
		time = [decoder decodeDoubleForKey: @"Time"];
		total = [decoder decodeDoubleForKey: @"TotalEnergy"];
		kinetic = [decoder decodeDoubleForKey: @"KineticEnergy"];
		temperature = [decoder decodeDoubleForKey: @"Temperature"];
		system = [decoder decodeObjectForKey: @"System"];
		scheduler = [[decoder decodeObjectForKey: @"Scheduler"] retain];		

		//reconstruct the force field state

		currentState = [[NSMutableDictionary dictionaryWithCapacity: 1] retain];
		[currentState setObject: [NSNumber numberWithDouble:
			[decoder decodeDoubleForKey: @"TotalPotential"]]
			forKey: @"TotalPotential"];
		[currentState setObject: [decoder decodeObjectForKey: @"ForceField"]  
			forKey: @"ForceField"];
		[currentState setObject: [decoder decodeObjectForKey: @"NonbondedCalculationMethod"]
			forKey: @"NonbondedCalculationMethod"];
		[currentState setObject: [decoder decodeObjectForKey: @"InactiveTerms"]
			forKey: @"InactiveTerms"];
		[currentState setObject: [decoder decodeObjectForKey: @"CustomTerms"]
			forKey: @"CustomTerms"];
		[currentState setObject: [decoder decodeObjectForKey: @"TermPotentials"]
			forKey: @"TermPotentials"];

		kineticUpdate = [decoder decodeBoolForKey: @"KineticUpdate"];
	
		dynamics = [system valueForKey:@"dynamics"];
		if(dynamics != nil)
		{
			coordinates = [[dynamics valueForKey: @"coordinates"] pointerValue];
			atomTypes = [dynamics valueForKey: @"atomTypes"];
			velocities = [[dynamics valueForKey: @"velocities"] pointerValue];
			no_of_atoms = coordinates->no_rows;
		}

		bonded = [system valueForKey:@"bondedTopology"];
		nonbonded = [system valueForKey:@"shortRangeNonbondedTopology"];
		longRangeNonbonded = [system valueForKey:@"longRangeNonbondedTopology"];
		systemName = [system systemName];	
		DOF =  3*no_of_atoms - 3;
		ke_2_temp =  (2*KB_1)/(DOF);
	}
	else
		[NSException raise: NSInvalidArgumentException
			format: @"%@ class does not support non keyed coding", [self class]];
	
	environment = [AdEnvironment globalEnvironment];
	if(environment != nil)
	{
		[self synchroniseWithEnvironment];
		[self registerWithEnvironment];
	}
	else
		[self _useEnvironmentDefaults];	

	return self;
}

- (void) encodeWithCoder: (NSCoder*) encoder
{
	double temp;
	id key;
	NSMutableDictionary* encodedState, *termPotentials;
	NSEnumerator* potentialEnum;

	encodedState = [NSMutableDictionary dictionaryWithCapacity: 1];
	termPotentials = [currentState objectForKey: @"TermPotentials"];
	potentialEnum = [termPotentials keyEnumerator];
	while(key = [potentialEnum nextObject])
	{	
		temp = *(double*)[[termPotentials objectForKey: key] pointerValue];
		[encodedState setObject: [NSNumber numberWithDouble: temp]
			forKey: key];
	}

	[super encodeWithCoder: encoder];
	if([encoder allowsKeyedCoding])
	{
		[encoder encodeDouble: time forKey: @"Time"];
		[encoder encodeDouble: total forKey: @"TotalEnergy"];
		[encoder encodeDouble: kinetic forKey: @"KineticEnergy"];
		[encoder encodeDouble: potential forKey: @"TotalPotential"];
		[encoder encodeDouble: temperature forKey: @"Temperature"];
		[encoder encodeConditionalObject: system forKey: @"System"];
		[encoder encodeObject: [currentState objectForKey: @"ForceField"]  
			forKey: @"ForceField"];
		[encoder encodeObject: [currentState objectForKey: @"NonbondedCalculationMethod"]  
			forKey: @"NonbondedCalculationMethod"];
		[encoder encodeObject: [currentState objectForKey: @"InactiveTerms"]  
			forKey: @"InactiveTerms"];
		[encoder encodeObject: [currentState objectForKey: @"CustomTerms"]  
			forKey: @"CustomTerms"];
		[encoder encodeObject: encodedState forKey: @"TermPotentials"];
		[encoder encodeBool: kineticUpdate forKey: @"KineticUpdate"];
		[encoder encodeObject: scheduler forKey: @"Scheduler"];
	}
	else
		[NSException raise: NSInvalidArgumentException
			format: @"%@ class does not support non keyed coding", [self class]];

}

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

Accessors 

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

- (NSDictionary*) termPotentials
{
	NSMutableDictionary* dict, *potentials;
	NSEnumerator* potentialEnum;	
	id pot;
	double value;

	dict = [NSMutableDictionary dictionaryWithCapacity: 1];
	potentials = [currentState valueForKey: @"TermPotentials"];
	potentialEnum = [potentials keyEnumerator];	
	while(pot = [potentialEnum nextObject])
	{
		value = *(double*)[[potentials valueForKey: pot] pointerValue];
		[dict setObject: [NSNumber numberWithDouble: value] forKey: pot];
	}

	return dict;
}

- (NSDictionary*) allEnergies
{
	id dict;
	
	dict = [self termPotentials];
	[dict setObject: [NSNumber numberWithDouble: total] 
					forKey: @"TotalEnergy"];
	[dict setObject: [NSNumber numberWithDouble: kinetic]
				forKey: @"KineticEnergy"];
	[dict setObject: [NSNumber numberWithDouble: potential]
				forKey: @"PotentialEnergy"];
	[dict setObject: [NSNumber numberWithDouble: temperature]
				forKey: @"Temperature"];
	[dict setObject: [NSNumber numberWithDouble: time]
				forKey: @"Time"];

	return dict;
}

- (double) totalEnergy
{
	return total;
}

- (double) potentialEnergy
{
	return potential;
}

- (double) kineticEnergy
{
	return kinetic;
}

- (double) temperature
{
	return temperature;
}

- (double) time
{
	return time;
}

- (int) degreesOfFreedom
{
	return DOF;
}

- (BOOL) isKineticUpdate
{
	return kineticUpdate;
}

- (void) setKineticUpdate: (BOOL) value
{
	kineticUpdate = value;
}

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

Environment Setters

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

- (void) setTimeStep: (int) tstep
{
	timestep = tstep;	
}

- (void) setUpdateInterval: (int) updateInt
{
	update_interval = updateInt;
}

/** Dependant Setters **/

- (void) set_DOF: (NSValue*) aValue
{
	DOF =  3*no_of_atoms - 3;	

	if(DOF <= 0)
	{
		NSLog(@"There are 0 degrees of freedom!\n \
		This is most likely because you have only 1 atom in the simulation.\n \
		Since translational degrees of freedom have been removed the atom is essentially frozen.\n \
		There are no forces acting on the atom and it has no momentum therefore a simulation is not possible.\n"); 
		exit(2);
	}

	[self updateDependantsOfKey: @"_DOF"];
}

- (void) set_KineticToTemperature: (NSValue*) aValue
{
	ke_2_temp =  (2*KB_1)/(DOF);
}

@end


static inline double CalcKineticEnergy(double** coords, double** velocity, int no_of_particles)
{
	register int i, j;
	double * vhold;
	double en, enhold;
	
	en = 0;

	for(i=0; i<no_of_particles; i++)
	{	
		vhold = velocity[i];
		for(enhold = 0,j=0; j< 3; j++)
			enhold += *(vhold + j)* *(vhold + j);

		en += enhold*coords[i][3];
	}
	en = en*0.5;

	return en;
}


