/*
   Project: Adun

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

   Author: Michael Johnston

   Created: 2005-06-23 11:06:55 +0200 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 "LangevinSimulator.h"

inline static void generateLDForceMatrix(AdMatrix*, double, double, gsl_rng*, AdMatrix*);

@implementation LangevinSimulator

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

Object Creation and Maintainence

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

- (void) _createForceMatrices
{
	int numberOfAtoms, i;
	AdMatrix* matrix;
	id subsystem;
	NSEnumerator* subsystemEnum;

	[forceMatrices release];
	forceMatrices = [NSMutableDictionary new];

	subsystemEnum = [subsystems objectEnumerator];
	while(subsystem = [subsystemEnum nextObject])
	{
		NSDebugLLog(@"LangevinSimulator", @"Creating Force Matrix for %@\n",
			[subsystem systemName]);
		
		numberOfAtoms = [subsystem numberOfAtoms];
		matrix = (AdMatrix*)[memoryManager allocateMatrixWithRows: numberOfAtoms 
						withColumns: 3];
		[forceMatrices setObject: [NSValue valueWithPointer: matrix]
			forKey: [subsystem systemName]];
	}
}

- (void) handleChangeInSystemStatus: (NSNotification*) aNotification
{
	AdMatrix* forceMatrix;
	NSEnumerator* matrixEnum;
	NSDictionary* infoDict;
	NSString* previousStatus, *currentStatus;
	id value;

	//We are interested in systems who change from "active" or to "active"
	
	infoDict = [aNotification userInfo];
	previousStatus = [infoDict objectForKey:@"previousStatus"];
	currentStatus = [infoDict objectForKey:@"currentStatus"];
	if(!([previousStatus isEqual: @"Active"] || 
		[currentStatus isEqual: @"Active"]))
	{
		return;
	}
	
	//Reacquire the current active subsystems and update the
	//forceMatrices dictionary
	
	[subsystems release];
	subsystems = [system systemsOfType: @"Standard"
			withStatus: @"Active"];
	[subsystems retain];		

	//free the previous force matrices (if any)

	matrixEnum = [forceMatrices keyEnumerator];
	while(value = [matrixEnum nextObject])
	{
		forceMatrix = [[forceMatrices valueForKey: value]
					pointerValue];
		if(forceMatrix != NULL)
			[memoryManager freeMatrix: forceMatrix];
	}
	[forceMatrices release];
 
	//create new ones
	[self _createForceMatrices];

}

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

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

- (id) initWithEnvironment: (id) object observe: (BOOL) value
{
	if(self = [super initWithEnvironment: object observe: value])
	{
		memoryManager = [AdMemoryManager appMemoryManager];

		if(environment != nil)
			seed = [[environment valueForKey: @"Seed"] intValue];	
		else
			seed = 238919;		
		
		//FIXME: this should be read from the environment
	
		gamma = 0.05;
	
		twister = gsl_rng_alloc(gsl_rng_mt19937);
		gsl_rng_set(twister, seed);
		variance = 2*KB*target_temperature;
	}	

	return self;
}

- (void) dealloc
{
	NSEnumerator* matrixEnum;
	AdMatrix* forceMatrix;
	id value;

	matrixEnum = [forceMatrices keyEnumerator];
	while(value = [matrixEnum nextObject])
	{
		forceMatrix = [[forceMatrices valueForKey: value]
					pointerValue];
		[memoryManager freeMatrix: forceMatrix];
	}
	[forceMatrices release];

	gsl_rng_free(twister);
	[super dealloc];
}

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

Public Methods

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

- (void) production
{
	register int j, k;
	double holder;
	id subsystem;
	AdMatrix* accelerations, *velocities, *coordinates, *matrix;
	double **coords, **accel, **vel;
	NSEnumerator* subsystemEnum;

	pool = [[NSAutoreleasePool alloc] init];

	GSPrintf(stderr, @"Initial energy calculation\n");
	//Zero the accelerations
	
	subsystemEnum = [subsystems objectEnumerator];
	while(subsystem = [subsystemEnum nextObject])
		[subsystem zeroAccelerations];

	[forceField calculateAccelerations];
	
	subsystemEnum = [subsystems objectEnumerator];
	while(subsystem = [subsystemEnum nextObject])
		[subsystem zeroAccelerations];

	[system update];

	//generate initial langevin random force matrix
	
	subsystemEnum = [subsystems objectEnumerator];
	while(subsystem = [subsystemEnum nextObject])
	{
		coordinates = [[subsystem coordinates] pointerValue];
		matrix = [[forceMatrices objectForKey: [subsystem systemName]] pointerValue];
		generateLDForceMatrix(matrix, variance, gamma, twister, coordinates);
	}		

	GSPrintf(stderr, @"Begining Simulation\n");
	for(currentStep=0; currentStep< no_of_steps; currentStep++)
	{
		//update the protein velocity a half a step 
		
		subsystemEnum = [subsystems objectEnumerator];
		while(subsystem = [subsystemEnum nextObject])
		{
			coordinates = [[subsystem coordinates] pointerValue];
			coords = coordinates->matrix;
			accelerations = [[subsystem accelerations] pointerValue];
			accel = accelerations->matrix;
			velocities = [[subsystem velocities] pointerValue];
			vel = velocities->matrix;
			matrix = [[forceMatrices objectForKey: [subsystem systemName]] pointerValue];

			for(j=0; j<velocities->no_rows; j++)
				for(k=0; k<3; k++)
					vel[j][k] += timefac*(accel[j][k] - 
							gamma*vel[j][k] + 
							matrix->matrix[j][k]);
			
			//update the position vector
			for(j=0; j<velocities->no_rows; j++)
				for(k=0; k<3; k++)
					coords[j][k] += vel[j][k]*time_step;
		}
		//calculate the  accelerations

		[forceField calculateAccelerations];
		
		subsystemEnum = [subsystems objectEnumerator];
		while(subsystem = [subsystemEnum nextObject])
		{	
			coordinates = [[subsystem coordinates] pointerValue];
			coords = coordinates->matrix;
			accelerations = [[subsystem accelerations] pointerValue];
			accel = accelerations->matrix;
			velocities = [[subsystem velocities] pointerValue];
			vel = velocities->matrix;
			matrix = [[forceMatrices objectForKey: [subsystem systemName]] pointerValue];

			//generate a new langevin random force matrix
			
			generateLDForceMatrix(matrix, variance, gamma, twister, coordinates);
		
			//update the velocity by another half a step
		
			for(j=0; j<velocities->no_rows; j++)
			{
				holder = 2/(2 + gamma*time_step);
				for(k=0; k<3; k++)
				{
					vel[j][k] += timefac*(accel[j][k] + matrix->matrix[j][k]);
					vel[j][k] = holder*vel[j][k];
				}
			}
		}
		
		[system frameUpdate];
		[timer increment];
		
		if(endSimulation)
			break;
	}	
	
	[pool release];
}

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

Accessors

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

- (void) setGamma: (double) value
{
	gamma = value;
}

- (double) gamma
{
	return gamma;
}

- (void) setSystem: (id) object 
{
	NSEnumerator* matrixEnum;
	AdMatrix* forceMatrix;
	id value;
	
	[super setSystem: object];

	//free the previous force matrices (if any)

	matrixEnum = [forceMatrices keyEnumerator];
	while(value = [matrixEnum nextObject])
	{
		forceMatrix = [[forceMatrices valueForKey: value]
					pointerValue];
		if(forceMatrix != NULL)
			[memoryManager freeMatrix: forceMatrix];
	}
	[forceMatrices release];
 
	//create new ones
	[self _createForceMatrices];
}

- (void) setTargetTemperature: (double) aValue
{
	[super setTargetTemperature: aValue];
	variance = 2*KB*target_temperature;
}

@end

inline static void generateLDForceMatrix(AdMatrix* matrix, 
	double variance, 
	double gamma, 
	gsl_rng* twister, 
	AdMatrix* coordinates)
{
	//declarations
	
	register int i, j;
	double sigma, holder;
	
	for (i=0; i < coordinates->no_rows; i++)
	{
		holder = variance*gamma*coordinates->matrix[i][4];
		sigma = sqrt(holder);

		for(j=0; j < 3; j++)
			matrix->matrix[i][j] = gsl_ran_gaussian(twister, sigma);
	}
}

