/*   []----------------------------------------[]
      |            NewtonDynamics.cpp          |
     []----------------------------------------[]
      |                                        |
      |   AUTHOR:    MFSomers 2005.            |
      |   USE:       Implements a classical    |
      |              dynamics framework class  |
      |              to be used for different  |
      |              integrators of the Newton |
      |              equation of motion ...    |
      |                                        |
     []----------------------------------------[]
*/

// Copyright (C) 2005 M.F. Somers, Theoretical Chemistry Department, Leiden University
//
// This is free software; you can redistribute it and/or modify it under the terms of
// the GNU Lesser General Public License as published by the Free Software Foundation.
//
// http://www.gnu.org/licenses/lgpl.txt

#include "Global.h"
#include "Vector.h"
#include "Constants.h"
#include "ClassicalDynamics.h"

#include "NewtonDynamics.h"
#include "DistanceConstraint.h"

/* -------------------------------------------------- NewtonDynamics ----- */

int NewtonDynamics::RunTrajectory( double Start, double Step, int nSteps, ConstraintListPointer Constraints )
{
 DistanceConstraintPointer C;
 
 // first we setup the constraint tables indices to particles once at the start of a trajectory
 
 if( Constraints != NULL )
  for( int CIndex = 0; CIndex < Constraints -> size(); ++CIndex )
   if( ( C = (DistanceConstraintPointer)((*Constraints)[ CIndex ]) ) != NULL )
    if( ( C -> Type == AGeneralDistanceConstraint ) && ( C -> Distance > 0.0 ) )
     for( int n = 0; n < C -> GetParticleListSize(); ++n )   
      if( C -> ParticleIDs[ n ] ) 
       C -> ParticleIndices[ n ] = IsParticleIncluded( C -> ParticleIDs[ n ] );
  
 return( ClassicalDynamics::RunTrajectory( Start, Step, nSteps, Constraints ) );
}

/* ----------------------------------------------------------------------- */
/* This algorithm implements the Velocity Verlet scheme with distance      */
/* constraints being applied by the RATTLE algorithm. See the refs:        */
/*                                                                         */
/* H. C. Andersen, Journal of Computational Physics 52, 24-34 (1983)       */
/* http://www.CCL.net/cca/software/SOURCES/FORTRAN/allen-tildesley-book/f.09.shtml */

int NewtonDynamics::Integrator( double Step, double Start, ConstraintListPointer Constraints )
{
 ParticleListPointer CurrentState = TheParticleListList[ 1 ];
 ParticleListPointer NewState = TheParticleListList[ 0 ];
 int                 Done;
   
 if( CurrentState == NULL ) return ( 0 );
 if( NewState == NULL ) return( 0 ); 
 if( Step < CONST_EPSILON ) return( 0 ); 

 #pragma omp parallel
 {
  DistanceConstraintPointer C;
  ParticlePointer           TheParticle, TheSecondParticle;           // local private variables
  int                       n, TheParticleIndex, TheSecondParticleIndex, Size;
  Vector                    v, dv, Force, Momentum;
  double                    Mass;
 
  // First part of the velocity verlet scheme, updating positions a full time step but the momenta half a time step. 

  Size = CurrentState -> size();

  #pragma omp for schedule( static )
  for( n = 0; n < Size; ++n )
   if( ( ( TheParticle = (*CurrentState)[ n ] ) != NULL ) && ( ( TheSecondParticle = (*NewState)[ n ] ) != NULL ) )
   {
    Mass = TheParticle -> Mass;                       
 
    if( Mass != 0.0 )
    {
      Mass = Step / Mass;
    
     Force = TheParticle -> TheTotalForce();
     Momentum = TheParticle -> Momentum;
     
     // first update positions for the full time step      
     v = TheParticle -> Position;
    
     dv = Momentum * Mass;        // get first order change
     v += dv; 
    
     dv = 0.5 * Step * Force * Mass;    // and second order
     v += dv;
        
     TheSecondParticle -> Position = v;
    
     // now update velocities/momenta half a time step   
     v = Momentum;
   
     dv = 0.5 * Step * Force;
     v += dv;
    
     TheSecondParticle -> Momentum = v;

     TheSecondParticle -> TheInternalFlag = 1;          // also signal particle was moved
    }
   } 
  
  
  // Now start doing the constraints if they can be done, the indices have been setup already (RATTLE A)
  if( Constraints != NULL )
  {
   int Itteration = 0;
 
   Size = Constraints -> size();
  
   do
   {
    #pragma omp barrier

    Done = 1; 

    #pragma omp for schedule( static )
    for( int CIndex = 0; CIndex < Size; ++CIndex )
     if( ( C = (DistanceConstraintPointer)((*Constraints)[ CIndex ]) ) != NULL )
      if( ( C -> Type == AGeneralDistanceConstraint ) && ( C -> Distance > 0.0 ) && ( Itteration < C -> MaxItterations ) )
      {
       double DistanceSqr = ( C -> Distance ) * ( C -> Distance );
           
       for( int PIndex = n = 0; n < C -> NrOfParticlesIncluded() - 1; ++n )
       {
        while( ( TheParticleIndex = C -> ParticleIndices[ PIndex ] ) < 0 ) ++PIndex;          // get first and second particle 
        if( (*CurrentState)[ TheParticleIndex ]-> Mass == 0.0 ) continue;                     // and make sure they have a mass
       
        ++PIndex; while( ( TheSecondParticleIndex = C -> ParticleIndices[ PIndex ] ) < 0 ) ++PIndex;
        if( (*CurrentState)[ TheSecondParticleIndex ]-> Mass == 0.0 ) continue;
       
        TheParticle       = (*NewState)[ TheParticleIndex ];
        TheSecondParticle = (*NewState)[ TheSecondParticleIndex ];
        
        // if one of these particles has been moved because of constraints or whatever
        if( ( TheParticle -> TheInternalFlag > 0 ) || ( TheSecondParticle -> TheInternalFlag > 0 ) )
        {
 
         Vector PAB = TheParticle -> Position - TheSecondParticle -> Position;
  	 double AbsPABSqr = AbsSqr( PAB );
	 double DiffSqr = DistanceSqr - AbsPABSqr;

 	 if( fabs( DiffSqr ) > ( 2.0 * (C -> Tolerance) * DistanceSqr ) )
	 {
		
      	  Vector RAB = (*CurrentState)[ TheParticleIndex ] -> Position - (*CurrentState)[ TheSecondParticleIndex ] -> Position;
	  double RABDotPAB = Dot( RAB, PAB );
	 
	  double InvMassA = ( 1.0 / TheParticle -> Mass );
	  double InvMassB = ( 1.0 / TheSecondParticle -> Mass );
	 
	  double GAB = 0.5 * DiffSqr / ( ( InvMassA + InvMassB ) * RABDotPAB );
	 
 	  dv = GAB * RAB;
	 
	  TheParticle -> Position       += dv * InvMassA;
	  TheSecondParticle -> Position -= dv * InvMassB;

	  dv = dv / Step;
	 
	  TheParticle -> Momentum       += dv;
	  TheSecondParticle -> Momentum -= dv;
	 
	  // signal we need to continue	 
	  Done = 0;
	 
	  // and that these two particles have been moved so convergence checks are again performed
	  TheParticle -> TheInternalFlag = 2;
	  TheSecondParticle -> TheInternalFlag = 2;
	 } 
	
	
 	 TheParticle -> TheInternalFlag--;
 	 TheSecondParticle -> TheInternalFlag--;
        }
       
       }
    
      }
   
    Itteration++;    

    #pragma omp barrier
   } while( !Done );
  
  }

 
  // then update the momenta another half a time step according to the Velocity Verlet

  Size = NewState -> size();

  #pragma omp for schedule( static )
  for( n = 0; n < Size; ++n )
   if( ( TheSecondParticle = (*NewState)[ n ] ) != NULL )
   {
    v  = TheSecondParticle -> Momentum;
   
    dv = 0.5 * Step * TheSecondParticle -> TheTotalForce(); 
    v += dv;
    
    TheSecondParticle -> Momentum = v;

    TheSecondParticle -> TheInternalFlag = 1;        // also signal this particle was 'moved'
   }   
  
  
  // and finally we start doing the velocity constraints if they can be done (RATTLE B)
  if( Constraints != NULL )
  {
   int Itteration = 0;
   
   Size = Constraints -> size();
  
   do
   {
    #pragma omp barrier

    Done = 1; 

    #pragma omp for schedule( static )
    for( int CIndex = 0; CIndex < Size; ++CIndex )
     if( ( C = (DistanceConstraintPointer)((*Constraints)[ CIndex ]) ) != NULL )
      if( ( C -> Type == AGeneralDistanceConstraint ) && ( C -> Distance > 0.0 ) && ( Itteration < C -> MaxItterations ) )
      {
       double DistanceSqr = ( C -> Distance ) * ( C -> Distance );
         
       for( int PIndex = n = 0; n < C -> NrOfParticlesIncluded() - 1; ++n )
       {
        while( ( TheParticleIndex = C -> ParticleIndices[ PIndex ] ) < 0 ) ++PIndex;          // get first and second particle 
        if( (*CurrentState)[ TheParticleIndex ]-> Mass == 0.0 ) continue;                     // and make sure they have a mass
       
        ++PIndex; while( ( TheSecondParticleIndex = C -> ParticleIndices[ PIndex ] ) < 0 ) ++PIndex;
        if( (*CurrentState)[ TheSecondParticleIndex ]-> Mass == 0.0 ) continue;
       
        TheParticle       = (*NewState)[ TheParticleIndex ];
        TheSecondParticle = (*NewState)[ TheSecondParticleIndex ];
       
        // if one of these particles has been moved because of constraints or whatever
        if( ( TheParticle -> TheInternalFlag > 0 ) || ( TheSecondParticle -> TheInternalFlag > 0 ) )
        {
 	 double InvMassA = ( 1.0 / TheParticle -> Mass );
	 double InvMassB = ( 1.0 / TheSecondParticle -> Mass );
              
         Vector RAB = TheParticle -> Position - TheSecondParticle -> Position;
	 Vector VAB = InvMassA * TheParticle -> Momentum - InvMassB * TheSecondParticle -> Momentum;
	 double RABDotVAB = Dot( RAB, VAB );
	
	 double GAB = -1.0 * RABDotVAB / ( ( InvMassA + InvMassB ) * DistanceSqr );
	
	 if( fabs( GAB ) > C -> Tolerance )
	 {
		
     	  dv = GAB * RAB;
	 
	  TheParticle -> Momentum       += dv;
	  TheSecondParticle -> Momentum -= dv;
	 
	  // signal we need to continue	 
	  Done = 0;
	 
	  // and that these two particles have been moved so convergence checks are again performed
	  TheParticle -> TheInternalFlag = 2;
	  TheSecondParticle -> TheInternalFlag = 2;
	 }  
	
	
	 TheParticle -> TheInternalFlag--;
	 TheSecondParticle -> TheInternalFlag--;
        }
       
       }
    
      }
   
    Itteration++;    

    #pragma omp barrier
   } while( !Done );

  } 
  
 }  // end of parallel region
 
 return( 1 );
}

/* ----------------------------------------------------------------------- */

int NewtonDynamics::IntegratorConservationCheck( double Step )
{
 return( 1 );
}

/* ----------------------------------------------------------------------- */

int NewtonDynamics::TrajectoryContinuationCheck( double At, double Step, int nStep, int nSteps )
{
 return( 1 );
}

/* ----------------------------------------------------------------------- */