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

/**
	SCAAS implements boundary conditions for water in a
	spherical box.
	/note It could be adapted to other solvents with just
	the radial constraint.
**/

#define RADIAL_FORCE_CONSTANT 20*BOND_FACTOR
#define GAMMA 0.0085
#define FULL_GAMMA 0.061
#define POLARISATION_FORCE_CONSTANT 6.9*BOND_FACTOR

@implementation SCAAS

/**
This returns the COM of the molecule with respect to the centre
of the solvent sphere
**/

- (void) _putCOMOfMolecule: (int*) molecule in: (Vector3D*) centre
{
	int i, j, k;
	int atom;
	double** coordinates;

	coordinates = coordinatesMatrix->matrix;

	//find centre of mass
	
	for(i=0; i<3; i++)
		centre->vector[i] = 0;

	for(j=0; j<atoms_per_molecule; j++)
	{
		atom = molecule[j];
		for(k=0; k<3; k++)
			centre->vector[k] += coordinates[atom][k]*coordinates[atom][3];
	}

	for(i=0; i<3; i++)
		centre->vector[i] = centre->vector[i]/solvent_mass;

	for(i=0; i<3; i++)
		centre->vector[i] -= cavityCentre->vector[i];	

	Ad3DVectorLength(centre);
}

//\note this will only work for 3point water
//we need to know the specific geometry to find the dipole
//of a molecule
//if we have a rigid molecule however this is not needed as
//the dipole vector only rotates

- (double) _calculatePolarisationAngleOf: (int*)molecule 
		withCenterOfMass: (Vector3D*) CoM 
		dipoleVector: (Vector3D*) dipole;
{
	int j;
	double angle, charge;	
	double** coordinates;

	coordinates = coordinatesMatrix->matrix;

	//assumes atoms are 0 H H 
	//N.B. we should check this ...
	//charge is half the dipole charge magnitude (which is the oxygen charge)

	charge = coordinates[molecule[0]][5];

	for(j=0; j<3; j++)
	{
		dipole->vector[j] = coordinates[molecule[1]][j] + 
					coordinates[molecule[2]][j] - 
					2*coordinates[molecule[0]][j];
		dipole->vector[j] *= charge;
	} 
	
	angle = Ad3DDotProduct(dipole, CoM);
	Ad3DVectorLength(dipole);
	angle = acos(angle/(dipole->length*CoM->length));

	return angle;	
}

- (void) _initRandomForceGenerator
{
	twister = gsl_rng_alloc(gsl_rng_mt19937);

	/**\note Im not sure if there should be a three in this...
	If the center of mass has three degrees of freedom (which it should)*/

	/**\note The SCAAS paper is slightly ambiguous as to the form of the 
	random force. The paper uses y' for the adjusted friction constant
	(GAMMA here for the radial force) and yº for the real friction constant. 
	It then uses both A'y and Aºy for the random force. It seems logical to
	assume that A'y is calculated used y' however the paper says <(A'y)²> should
	be set using eq.6 which is the formula for Aºy. We are going to assume that
	this means use eq.6 substituting y' for yº **/
	
	variance = 2*KB*3*GAMMA/solvent_mass;	
}

- (void) _cleanUpSystem
{
	[memoryManager freeIntMatrix: solventIndexMatrix];
	free(radial_distance);
	gsl_rng_free(twister);
}

//\note Will there be a solute system and solvent system? How do we decide which is
//which. Should a solvent "know" what solute it contains etc. etc. etc.
//maybe the solvent should keep references to the system inserted into it ...
//also have to be aware there may be no solute

-(void) _initialisationForSystem
{
	int i, j;
	AdMatrix* soluteCoordinates;
	
	//retrieve and calculate the neccessary data for performing SCAAS

	solventSystem = [system valueForKey:@"solventSystem"]; 
	soluteSystem = [system valueForKey:@"soluteSystem"];

	coordinatesMatrix = [[solventSystem valueForKey: @"coordinates"] pointerValue];
	velocitiesMatrix = [[solventSystem valueForKey: @"velocities"] pointerValue];
	accelerationsMatrix = [[solventSystem valueForKey: @"accelerations"] pointerValue];
	cavityCentre = [[solventSystem dataSource] cavityCentre];

	//the number of solvent molecules that had to be ommitted due to the solute

	occlusion_factor = [[solventSystem valueForKeyPath:@"dataSource.numberOccludedMolecules"] intValue];
	atoms_per_molecule = [[solventSystem valueForKeyPath: @"dataSource.atomsPerMolecule"] intValue];
	no_solvent_atoms = coordinatesMatrix->no_rows;
	
	//the formula for calculating the equilibrium radial distance for each surface
	//atom can be written as ((alpha*j + beta)*solvent_mass)^1/3
	//\note incorporate the solvent mass here

	alpha = 3/(M_PI*4*solvent_density);
	beta = alpha*(occlusion_factor - 1);

	no_solvent_molecules = no_solvent_atoms/atoms_per_molecule;

	//create a matrix of rows = no solvent molecules and columns = atoms per molecule 

	solventIndexMatrix = [memoryManager allocateIntMatrixWithRows: no_solvent_molecules 
				withColumns: atoms_per_molecule];

	for(i=0; i < no_solvent_molecules; i++)
		for(j=0; j< atoms_per_molecule; j++)
			solventIndexMatrix->matrix[i][j] = atoms_per_molecule*i + j;

	//create an array to hold the radial distance of each molecule
	//and the positon of its centre of mass.
	//we only create the radial array here since we decide which atoms are
	//in the surface region based on this. We will need to create a
	//polarization array for the surface molecules each time we update them.
  	//(Since we dont know how many molecules are in the surface region and
	//we dont want to keep calculating them all every step)
	
	radial_distance = (Vector3D*)malloc(no_solvent_molecules*sizeof(Vector3D)); 

	//Find the mass of one molecule
	
	for(solvent_mass =0, i=0; i< atoms_per_molecule; i++)
		solvent_mass += coordinatesMatrix->matrix[i][3];
		
	[self _initRandomForceGenerator];
	
	//find the charge of the solute

	soluteCoordinates = [[soluteSystem valueForKey: @"coordinates"] pointerValue];
	
	solute_charge = 0.0;
	for(i=0; i<soluteCoordinates->no_rows; i++)
		solute_charge += soluteCoordinates->matrix[i][5];

	GSPrintf(stderr, @"The solute charge is %f\n", solute_charge);

	if(fabs(solute_charge) > 0.001)	
	{
		isChargedSolute = true;
		GSPrintf(stderr, @"Charged Solute\n");
	}
}

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

Object Creation

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

- (void) _initialiseDependants
{
	KBT = targetTemperature*KB;
}

- (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)
		{
			surface_region = 1.5;
			targetTemperature = 300;
			
			//\note these two variables should be retrieved from the solvent system
			//this goes for below aswell

			inner_sphere = 20 - 1.5;
			solvent_density = 1000*DENSITY_FACTOR;
		}
		else
		{
			surface_region = [[environment valueForKey: @"BoundarySize"] 
						doubleValue];

			inner_sphere = [[environment valueForKey: @"SolvationSphereRadius"] 
						doubleValue];	
			inner_sphere -= surface_region;

			solvent_density = [[environment valueForKey: @"SolventDensity"] 
						doubleValue];
			solvent_density *= DENSITY_FACTOR;

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

		}
		[self _initialiseDependants];
		comparison_pt = AdIndexSorter;
	}
	
	return self;	
}

- (void) dealloc
{
	[memoryManager freeIntMatrix: solventIndexMatrix];
	free(radial_distance);
	gsl_rng_free(twister);
	[super dealloc];
}

- (void) setSystem: (id) object
{
	if(system != nil)
		[self _cleanUpSystem];

	system = object;
	[self _initialisationForSystem];
}

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

SCAAS methods

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

- (void) _setupSCAAS
{
	int i, j;
	int molecule;
	
	//calculate the radial distance of each molecule	

	for(i=0; i<no_solvent_molecules; i++)
		[self _putCOMOfMolecule: solventIndexMatrix->matrix[i] in: &radial_distance[i]];

	//calculate sigma for the random force

	sigma = sqrt(variance*targetTemperature);

	//how many molecules are inside the surface region?

	for(no_surface_molecules = 0, i=0; i < no_solvent_molecules; i++)
		if(radial_distance[i].length > inner_sphere)
			no_surface_molecules++;

	inside_count = no_solvent_molecules - no_surface_molecules;

	//copy the indices of the surface molecules to an array and sort them based
	//on their radial distance
	
	radial_sorter = (Sort*)malloc(no_surface_molecules*sizeof(Sort)); 

	for(i=0, j=0; i<no_solvent_molecules; i++)
		if(radial_distance[i].length > inner_sphere)
		{
			radial_sorter[j].index = i;
			radial_sorter[j].property = radial_distance[i].length;
			j++;
		}

	qsort(radial_sorter, no_surface_molecules, sizeof(Sort), comparison_pt); 
	
	//create an array to hold the polarization angles and dipole vectors

	polarisation_angles = (double*)malloc(no_surface_molecules*sizeof(double));
	dipoles = (Vector3D*)malloc(no_surface_molecules*sizeof(Vector3D));

	//use radial sorter to find surface molecule indexes
	//this way  polarisation_angles[i] will refer to the molecule radial_sorter[i].index

	for(i=0; i<no_surface_molecules;i++)
	{
		molecule = radial_sorter[i].index;
		polarisation_angles[i] = [self _calculatePolarisationAngleOf: solventIndexMatrix->matrix[molecule]
						withCenterOfMass: &radial_distance[molecule]
						dipoleVector: &dipoles[i]];
	}

	//sort the array

	polarisation_sorter = (Sort*)malloc(no_surface_molecules*sizeof(Sort)); 

	for(i=0; i<no_surface_molecules; i++)
	{
		polarisation_sorter[i].index = i;
		polarisation_sorter[i].property = polarisation_angles[i];
	}

	qsort(polarisation_sorter, no_surface_molecules, sizeof(Sort), comparison_pt);
}	

- (void) _applyRadialConstraint
{
	int i, j, k, molecule_count;
	int molecule_no, atom_no;
	double** velocities, **accelerations, **coordinates;
	double total_force, fraction, constraintDistance, ran_force[3]; 
	Vector3D unit_vector; 		//for holding the unit_vector in the direction of the force
	Vector3D velocity;

	//dereference the matrix pointers

	velocities = velocitiesMatrix->matrix;
	accelerations = accelerationsMatrix->matrix;
	coordinates = coordinatesMatrix->matrix;

	//calculate the radial constraint force on each atom of each molecule

	for(i = 0, molecule_count=inside_count; i < no_surface_molecules; i++, molecule_count++)
	{
		/**
		for each molecule we need to calculate the force on each
		of its constituent atoms
		Each of these forces is in the same direction as the total
		force.
		First calculate the total force on the centre of mass of the
		molecule and then assign it to the consituent atoms
		The total force acts along the position vector of the centre of mass
		in the opposite direction
		**/

		//There are no_surface_molecules in radial_sorter. Well go through them from start to finish

		molecule_no = radial_sorter[i].index;
		AdGet3DUnitVector(&radial_distance[molecule_no], &unit_vector);

		constraintDistance =  cbrt((alpha*molecule_count + beta)*solvent_mass);

		//multiply by -1 to change direction

		total_force = -1*RADIAL_FORCE_CONSTANT*(radial_sorter[i].property - constraintDistance);

		//generate a value for the random force
		
		for(j=0; j<3; j++)
			ran_force[j] = solvent_mass*gsl_ran_gaussian(twister, sigma);
		
		for(j = 0; j < atoms_per_molecule; j++)
		{
			//which atom are we at 

			atom_no = solventIndexMatrix->matrix[molecule_no][j];

			//unit vector should be the unit vector in the direction of the centre of
			//mass - so we multiply the total force by -1 above.
			
			for(k=0; k<3; k++)
				accelerations[atom_no][k] += total_force*unit_vector.vector[k];

			//apply the random force
				
			for(k=0; k<3; k++)
				accelerations[atom_no][k] += ran_force[k];
		}

		//now apply the frictional force against the direction of motion of the centre of mass	
		//find the velocity of the centre of mass of this molecule

		for(j=0; j<3; j++)
			velocity.vector[j] = 0;

		for(k=0; k<3; k++)
		{
			for(j=0; j<atoms_per_molecule; j++)
			{
				velocity.vector[k] += 
					velocities[solventIndexMatrix->matrix[molecule_no][j]][k]*
					coordinates[solventIndexMatrix->matrix[molecule_no][j]][3];
			}
			velocity.vector[k] = velocity.vector[k]/solvent_mass;
		}

		Ad3DVectorLength(&velocity);	
		AdGet3DUnitVector(&velocity, &unit_vector);

		total_force = GAMMA*solvent_mass*velocity.length;

		for(j = 0; j < atoms_per_molecule; j++)
		{
			atom_no = solventIndexMatrix->matrix[molecule_no][j];
			for(k=0; k<3; k++)
				accelerations[atom_no][k] -= total_force*unit_vector.vector[k];
		}
	}
	
}

-(void) _applyPolarisationConstraint
{
	int i, j, k;
	int molecule, index, atom;
	double constraint_angle, cos_constraint, cos_actual, actual_angle;
	double langevin;
	double force_mag, factor, A, B;
	double** velocities, **accelerations, **coordinates;
	Vector3D Force;
	Vector3D *Dipole, *Center;

	velocities = velocitiesMatrix->matrix;
	accelerations = accelerationsMatrix->matrix;
	coordinates = coordinatesMatrix->matrix;
	
	for(i=0; i<no_surface_molecules; i++)
	{
		//the index into the polarisation_angle array and the dipole array
		
		index = polarisation_sorter[i].index;

		//the index into the molecule matrix
		
		molecule = radial_sorter[index].index;

		Dipole = &dipoles[index];
		Center = &radial_distance[molecule];

		cos_constraint = 1 + (1 - 2*(i + 1))/(double)no_surface_molecules;
		constraint_angle = acos(cos_constraint);
		
		//if the soluted is charged we need to shift the constraint angles
	
		if(isChargedSolute)
		{
			langevin = Dipole->length*solute_charge;
			langevin /= (Center->length*Center->length)*(1 + Center->length)*KBT;
			langevin = 1/tanh(langevin) - 1/langevin;
			constraint_angle = constraint_angle - 3*langevin*sin(constraint_angle)/2;
		}

		actual_angle = polarisation_angles[index];
		cos_actual = cos(actual_angle);
		force_mag = -1*POLARISATION_FORCE_CONSTANT*(actual_angle - constraint_angle);
		
		factor = -1/(sqrt(1 - cos_actual*cos_actual))*1/(Dipole->length*Center->length);

		for(j=0; j<atoms_per_molecule; j++)
		{
			atom = solventIndexMatrix->matrix[molecule][j];
			
			A = coordinates[atom][3]/solvent_mass - 
				coordinates[atom][5]*cos_actual*Center->length/Dipole->length;
			B = coordinates[atom][5] - 
				coordinates[atom][3]*cos_actual*Dipole->length/(Center->length * solvent_mass);

			for(k=0; k<3; k++)
			{
				Force.vector[k] = force_mag*factor*(A*Dipole->vector[k] + B*Center->vector[k]);
				accelerations[atom][k] += Force.vector[k];
			}
		}
	}
}


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

Public Methods

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

//For the moment we will do the polarization and radial
//constraints seperatly thought later we can probably combine
//the frictional and random parts of both in one loop

- (void) applyBoundaryConditions
{
	NSDebugLLog(@"SCAAS", @"Setting Up SCAAS");
	[self _setupSCAAS];
	NSDebugLLog(@"SCAAS", @"Applying radial constraint");
	[self _applyRadialConstraint];
	NSDebugLLog(@"SCAAS", @"Applying polarisation constraint");
	[self _applyPolarisationConstraint];
	free(radial_sorter);
	free(polarisation_angles);
	free(polarisation_sorter);
	free(dipoles);
}



@end
