/*
  Copyright (C) 2010 The ESPResSo project
  Copyright (C) 2002,2003,2004,2005,2006,2007,2008,2009,2010 Max-Planck-Institute for Polymer Research, Theory Group, PO Box 3148, 55021 Mainz, Germany
  
  This file is part of ESPResSo.
  
  ESPResSo 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 3 of the License, or
  (at your option) any later version.
  
  ESPResSo 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 General Public License for more details.
  
  You should have received a copy of the GNU General Public License
  along with this program.  If not, see <http://www.gnu.org/licenses/>. 
*/
/** \file rotation.c  Molecular dynamics integrator for rotational motion.
 *
 *  A velocity Verlet <a HREF="http://ciks.cbt.nist.gov/~garbocz/dpd1/dpd.html">algorithm</a>
 *  using quaternions is implemented to tackle rotational motion. A random torque and a friction
 *  term are added to provide the constant NVT conditions. Due to this feature all particles are
 *  treated as 3D objects with 3 translational and 3 rotational degrees of freedom if ROTATION
 *  flag is set in \ref config.h "config.h".
*/

#include <mpi.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include "utils.h"
#include "integrate.h"
#include "interaction_data.h"
#include "particle_data.h"
#include "communication.h"
#include "grid.h"
#include "cells.h"
#include "verlet.h"
#include "rotation.h"
#include "ghosts.h"
#include "forces.h"
#include "p3m.h"
#include "thermostat.h"
#include "initialize.h"

/****************************************************
 *                     DEFINES
 ***************************************************/
/**************** local variables  *******************/

#ifdef ROTATION

/** moment of inertia. Currently we define the inertia tensor here to be constant.
    If it is not spherical the angular velocities have to be refined several times
    in the \ref convert_torques_propagate_omega. Also the kinetic energy in file
    \ref statistics.c is calculated assuming that I[0] =  I[1] =  I[2] = 1  */
static double I[3] = { 1, 1, 1};

/** \name Privat Functions */
/************************************************************/
/address@hidden/

/** define rotation matrix A for a given particle */
static void define_rotation_matrix(Particle *p, double A[9]);

/** define first and second time derivatives of a quaternion, as well
    as the angular acceleration. */
static void define_Qdd(Particle *p, double Qd[4], double Qdd[4], double S[3], double Wd[3]);

/address@hidden/

/** convert quaternions to the director */
void convert_quat_to_quatu(double quat[4], double quatu[3])
{
  /* director */
  quatu[0] = 2*(quat[1]*quat[3] + quat[0]*quat[2]);
  quatu[1] = 2*(quat[2]*quat[3] - quat[0]*quat[1]);
  quatu[2] =   (quat[0]*quat[0] - quat[1]*quat[1] -
		quat[2]*quat[2] + quat[3]*quat[3]); 
}

#ifdef DIPOLES

/** convert a dipole moment to quaternions and dipolar strength  */
int convert_dip_to_quat(double dip[3], double quat[4], double *dipm)
{
  double dm;
  // Calculate magnitude of dipole moment
  dm = sqrt(dip[0]*dip[0] + dip[1]*dip[1] + dip[2]*dip[2]);
  *dipm = dm;
  convert_quatu_to_quat(dip,quat);

  return 0;
}



/** convert quaternion director to the dipole moment */
void convert_quatu_to_dip(double quatu[3], double dipm, double dip[3])
{
  /* dipole moment */
  dip[0] = quatu[0]*dipm;
  dip[1] = quatu[1]*dipm;
  dip[2] = quatu[2]*dipm;
}

#endif
/** Convert director to quaternions */
int convert_quatu_to_quat(double d[3], double quat[4])
{
  double d_xy, dm;
  double theta2, phi2;
  double costhe2,sinthe2,cosphi2,sinphi2;


  // Calculate magnitude of the given vector
  dm = sqrt(d[0]*d[0] + d[1]*d[1] + d[2]*d[2]);

  // The vector needs to be != 0 to be converted into a quaternion
  if (dm < ROUND_ERROR_PREC) {
    return 1;
  }
  else {
    // Calculate angles 
    d_xy = sqrt(d[0]*d[0] + d[1]*d[1]);
    // If dipole points along z axis:
    if (d_xy == 0){
      // We need to distinguish between (0,0,d_z) and (0,0,d_z)
      if (d[2]>0)
       theta2 = 0;
      else
       theta2 = PI/2.;
      phi2 = 0;
    }
    else {
      // Here, we take care of all other directions
      //Here we suppose that theta2 = 0.5*theta and phi2 = 0.5*(phi - PI/2),
      //where theta and phi - angles are in spherical coordinates
      theta2 = 0.5*acos(d[2]/dm);
      if (d[1] < 0) phi2 = -0.5*acos(d[0]/d_xy) - PI*0.25;
      else phi2 = 0.5*acos(d[0]/d_xy) - PI*0.25;
    }

    // Calculate the quaternion from the angles; mjw below is slightly faster unless compiler does this
     costhe2 = cos(theta2);
     sinthe2 = sin(theta2);
     cosphi2 = cos(phi2);
     sinphi2 = sin(phi2);
     quat[0] =  costhe2 * cosphi2;
     quat[1] = -sinthe2 * cosphi2;
     quat[2] = -sinthe2 * sinphi2;
     quat[3] =  costhe2 * sinphi2;
/*   quat[0] =  cos(theta2) * cos(phi2);
     quat[1] = -sin(theta2) * cos(phi2);
     quat[2] = -sin(theta2) * sin(phi2);
     quat[3] =  cos(theta2) * sin(phi2); */
  }
  return 0;
}

/** Here we use quaternions to calculate the rotation matrix which
    will be used then to transform torques from the laboratory to
    the body-fixed frames */  
void define_rotation_matrix(Particle *p, double A[9])
{
  A[0 + 3*0] = p->r.quat[0]*p->r.quat[0] + p->r.quat[1]*p->r.quat[1] -
               p->r.quat[2]*p->r.quat[2] - p->r.quat[3]*p->r.quat[3] ;
  A[1 + 3*1] = p->r.quat[0]*p->r.quat[0] - p->r.quat[1]*p->r.quat[1] +
               p->r.quat[2]*p->r.quat[2] - p->r.quat[3]*p->r.quat[3] ;
  A[2 + 3*2] = p->r.quat[0]*p->r.quat[0] - p->r.quat[1]*p->r.quat[1] -
               p->r.quat[2]*p->r.quat[2] + p->r.quat[3]*p->r.quat[3] ;

  A[0 + 3*1] = 2*(p->r.quat[1]*p->r.quat[2] + p->r.quat[0]*p->r.quat[3]);
  A[0 + 3*2] = 2*(p->r.quat[1]*p->r.quat[3] - p->r.quat[0]*p->r.quat[2]);
  A[1 + 3*0] = 2*(p->r.quat[1]*p->r.quat[2] - p->r.quat[0]*p->r.quat[3]);

  A[1 + 3*2] = 2*(p->r.quat[2]*p->r.quat[3] + p->r.quat[0]*p->r.quat[1]);
  A[2 + 3*0] = 2*(p->r.quat[1]*p->r.quat[3] + p->r.quat[0]*p->r.quat[2]);
  A[2 + 3*1] = 2*(p->r.quat[2]*p->r.quat[3] - p->r.quat[0]*p->r.quat[1]);
}

/** calculate the second derivative of the quaternion of a given particle
    as well as Wd vector which is the angular acceleration of this particle */
void define_Qdd(Particle *p, double Qd[4], double Qdd[4], double S[3], double Wd[3])
{
  double S1;

  /* calculate the first derivative of the quaternion */
  Qd[0] = 0.5 * ( -p->r.quat[1] * p->m.omega[0] -
                   p->r.quat[2] * p->m.omega[1] -
		   p->r.quat[3] * p->m.omega[2] );

  Qd[1] = 0.5 * (  p->r.quat[0] * p->m.omega[0] -
                   p->r.quat[3] * p->m.omega[1] +
		   p->r.quat[2] * p->m.omega[2] );

  Qd[2] = 0.5 * (  p->r.quat[3] * p->m.omega[0] +
                   p->r.quat[0] * p->m.omega[1] -
		   p->r.quat[1] * p->m.omega[2] );

  Qd[3] = 0.5 * ( -p->r.quat[2] * p->m.omega[0] +
                   p->r.quat[1] * p->m.omega[1] +
		   p->r.quat[0] * p->m.omega[2] );

  /* calculate the second derivative of the quaternion */
  

// substituted by mjw
#ifdef ROTATIONAL_INERTIA
	I[0]=p->p.rinertia[0];
	I[1]=p->p.rinertia[1];
	I[2]=p->p.rinertia[2];
#endif
  Wd[0] =  (p->f.torque[0] + p->m.omega[1]*p->m.omega[2]*(I[1]-I[2]))/I[0];
  Wd[1] =  (p->f.torque[1] + p->m.omega[2]*p->m.omega[0]*(I[2]-I[0]))/I[1];
  Wd[2] =  (p->f.torque[2] + p->m.omega[0]*p->m.omega[1]*(I[0]-I[1]))/I[2];

/*  
  Wd[0] =  (p->f.torque[0] + p->m.omega[1]*p->m.omega[2]*(I[1]-I[2]))/I[0]/p->p.rinertia[0];
  Wd[1] =  (p->f.torque[1] + p->m.omega[2]*p->m.omega[0]*(I[2]-I[0]))/I[1]/p->p.rinertia[1];
  Wd[2] =  (p->f.torque[2] + p->m.omega[0]*p->m.omega[1]*(I[0]-I[1]))/I[2]/p->p.rinertia[2];
#else
  Wd[0] =  (p->f.torque[0] + p->m.omega[1]*p->m.omega[2]*(I[1]-I[2]))/I[0];
  Wd[1] =  (p->f.torque[1] + p->m.omega[2]*p->m.omega[0]*(I[2]-I[0]))/I[1];
  Wd[2] =  (p->f.torque[2] + p->m.omega[0]*p->m.omega[1]*(I[0]-I[1]))/I[2];
#endif */

  S1 = Qd[0]*Qd[0] + Qd[1]*Qd[1] + Qd[2]*Qd[2] + Qd[3]*Qd[3];

  Qdd[0] = 0.5 * ( -p->r.quat[1] * Wd[0] -
		    p->r.quat[2] * Wd[1] -
	  	    p->r.quat[3] * Wd[2] ) - p->r.quat[0] * S1;

  Qdd[1] = 0.5 * (  p->r.quat[0] * Wd[0] -
		    p->r.quat[3] * Wd[1] +
		    p->r.quat[2] * Wd[2] ) - p->r.quat[1] * S1;
  
  Qdd[2] = 0.5 * (  p->r.quat[3] * Wd[0] +
		    p->r.quat[0] * Wd[1] -
		    p->r.quat[1] * Wd[2] ) - p->r.quat[2] * S1;
  
  Qdd[3] = 0.5 * ( -p->r.quat[2] * Wd[0] +
                    p->r.quat[1] * Wd[1] +
		    p->r.quat[0] * Wd[2] ) - p->r.quat[3] * S1;

  S[0] = S1;
  S[1] = Qd[0]*Qdd[0]  + Qd[1]*Qdd[1]  + Qd[2]*Qdd[2]  + Qd[3]*Qdd[3];
  S[2] = Qdd[0]*Qdd[0] + Qdd[1]*Qdd[1] + Qdd[2]*Qdd[2] + Qdd[3]*Qdd[3];
}

/** propagate angular velocities and quaternions \todo implement for
       fixed_coord_flag */
void propagate_omega_quat()
{
  Particle *p;
  Cell *cell;
  int c,i,j, np;
  double lambda, dt2, dt4, dtdt, dtdt2;

  dt2 = time_step*0.5;
  dt4 = time_step*0.25;
  dtdt = time_step*time_step;
  dtdt2 = dtdt*0.5;

  INTEG_TRACE(fprintf(stderr,"%d: propagate_omega_quat:\n",this_node));

  for (c = 0; c < local_cells.n; c++) {
    cell = local_cells.cell[c];
    p  = cell->part;
    np = cell->n;
    for(i = 0; i < np; i++) {
	  double Qd[4], Qdd[4], S[3], Wd[3];
	  
	  #ifdef VIRTUAL_SITES
	   // Virtual sites are not propagated in the integration step
	   if (ifParticleIsVirtual(&p[i]))
	    continue;
	  #endif

	  define_Qdd(&p[i], Qd, Qdd, S, Wd);
	  
	  lambda = 1 - S[0]*dtdt2 - sqrt(1 - dtdt*(S[0] + time_step*(S[1] + dt4*(S[2]-S[0]*S[0]))));

	  for(j=0; j < 3; j++){
	    p[i].m.omega[j]+= dt2*Wd[j];
	  }
	  ONEPART_TRACE(if(p[i].p.identity==check_id) fprintf(stderr,"%d: OPT: PV_1 v_new = (%.3e,%.3e,%.3e)\n",this_node,p[i].m.v[0],p[i].m.v[1],p[i].m.v[2]));

	  p[i].r.quat[0]+= time_step*(Qd[0] + dt2*Qdd[0]) - lambda*p[i].r.quat[0];
	  p[i].r.quat[1]+= time_step*(Qd[1] + dt2*Qdd[1]) - lambda*p[i].r.quat[1];
	  p[i].r.quat[2]+= time_step*(Qd[2] + dt2*Qdd[2]) - lambda*p[i].r.quat[2];
	  p[i].r.quat[3]+= time_step*(Qd[3] + dt2*Qdd[3]) - lambda*p[i].r.quat[3];
	  // Update the director
	  convert_quat_to_quatu(p[i].r.quat, p[i].r.quatu);
#ifdef DIPOLES
	  // When dipoles are enabled, update dipole moment
	  convert_quatu_to_dip(p[i].r.quatu, p[i].p.dipm, p[i].r.dip);
#endif
	

      ONEPART_TRACE(if(p[i].p.identity==check_id) fprintf(stderr,"%d: OPT: PPOS p = (%.3f,%.3f,%.3f)\n",this_node,p[i].r.p[0],p[i].r.p[1],p[i].r.p[2]));
    }
  }
}

/** convert the torques to the body-fixed frames and propagate angular velocities */
void convert_torques_propagate_omega()
{
  Particle *p;
  Cell *cell;
  int c,i, np, times;
  double dt2, tx, ty, tz;

  dt2 = time_step*0.5;
  INTEG_TRACE(fprintf(stderr,"%d: convert_torques_propagate_omega:\n",this_node));
  for (c = 0; c < local_cells.n; c++) {
    cell = local_cells.cell[c];
    p  = cell->part;
    np = cell->n;
    for(i = 0; i < np; i++) {
      double A[9];
      define_rotation_matrix(&p[i], A);

      tx = A[0 + 3*0]*p[i].f.torque[0] + A[0 + 3*1]*p[i].f.torque[1] + A[0 + 3*2]*p[i].f.torque[2];
      ty = A[1 + 3*0]*p[i].f.torque[0] + A[1 + 3*1]*p[i].f.torque[1] + A[1 + 3*2]*p[i].f.torque[2];
      tz = A[2 + 3*0]*p[i].f.torque[0] + A[2 + 3*1]*p[i].f.torque[1] + A[2 + 3*2]*p[i].f.torque[2];

      
      if ( thermo_switch & THERMO_LANGEVIN ) {
#if defined (VIRTUAL_SITES) && defined(THERMOSTAT_IGNORE_NON_VIRTUAL)
       if (!ifParticleIsVirtual(&p[i]))
#endif
       {
	friction_thermo_langevin_rotation(&p[i]);

	p[i].f.torque[0]+= tx;
	p[i].f.torque[1]+= ty;
	p[i].f.torque[2]+= tz;
       }
      } else {
	p[i].f.torque[0] = tx;
	p[i].f.torque[1] = ty;
	p[i].f.torque[2] = tz;
      }
    
      ONEPART_TRACE(if(p[i].p.identity==check_id) fprintf(stderr,"%d: OPT: SCAL f = (%.3e,%.3e,%.3e) v_old = (%.3e,%.3e,%.3e)\n",this_node,p[i].f.f[0],p[i].f.f[1],p[i].f.f[2],p[i].m.v[0],p[i].m.v[1],p[i].m.v[2]));
    
/*
#ifdef ROTATIONAL_INERTIA
	  p[i].m.omega[0]+= dt2*p[i].f.torque[0]/p[i].p.rinertia[0]/I[0];
	  p[i].m.omega[1]+= dt2*p[i].f.torque[1]/p[i].p.rinertia[1]/I[1];
	  p[i].m.omega[2]+= dt2*p[i].f.torque[2]/p[i].p.rinertia[2]/I[2];
#else
	  p[i].m.omega[0]+= dt2*p[i].f.torque[0]/I[0];
	  p[i].m.omega[1]+= dt2*p[i].f.torque[1]/I[1];
	  p[i].m.omega[2]+= dt2*p[i].f.torque[2]/I[2];
#endif */
// substituted by mjw
#ifdef ROTATIONAL_INERTIA
	I[0]=p->p.rinertia[0];
	I[1]=p->p.rinertia[1];
	I[2]=p->p.rinertia[2];
#endif
	  p[i].m.omega[0]+= dt2*p[i].f.torque[0]/I[0];
	  p[i].m.omega[1]+= dt2*p[i].f.torque[1]/I[1];
	  p[i].m.omega[2]+= dt2*p[i].f.torque[2]/I[2];

	  /* if the tensor of inertia is isotropic, the following refinement is not needed.
	     Otherwise repeat this loop 2-3 times depending on the required accuracy */
	  for(times=0;times<=4;times++) { 
	    double Wd[3];

/*
#ifdef ROTATIONAL_INERTIA
	    Wd[0] = (p[i].m.omega[1]*p[i].m.omega[2]*(I[1]-I[2]))/I[0]/p[i].p.rinertia[0];
	    Wd[1] = (p[i].m.omega[2]*p[i].m.omega[0]*(I[2]-I[0]))/I[1]/p[i].p.rinertia[1]; 
	    Wd[2] = (p[i].m.omega[0]*p[i].m.omega[1]*(I[0]-I[1]))/I[2]/p[i].p.rinertia[2];
#else
	    Wd[0] = (p[i].m.omega[1]*p[i].m.omega[2]*(I[1]-I[2]))/I[0];
	    Wd[1] = (p[i].m.omega[2]*p[i].m.omega[0]*(I[2]-I[0]))/I[1]; 
	    Wd[2] = (p[i].m.omega[0]*p[i].m.omega[1]*(I[0]-I[1]))/I[2];
#endif */
// substituted by mjw
	    Wd[0] = (p[i].m.omega[1]*p[i].m.omega[2]*(I[1]-I[2]))/I[0];
	    Wd[1] = (p[i].m.omega[2]*p[i].m.omega[0]*(I[2]-I[0]))/I[1]; 
	    Wd[2] = (p[i].m.omega[0]*p[i].m.omega[1]*(I[0]-I[1]))/I[2];
 
	    p[i].m.omega[0]+= dt2*Wd[0];
	    p[i].m.omega[1]+= dt2*Wd[1];
	    p[i].m.omega[2]+= dt2*Wd[2];
	  }
	
      
      ONEPART_TRACE(if(p[i].p.identity==check_id) fprintf(stderr,"%d: OPT: PV_2 v_new = (%.3e,%.3e,%.3e)\n",this_node,p[i].m.v[0],p[i].m.v[1],p[i].m.v[2]));
      
    }
  }
}

/** convert the torques to the body-fixed frames before the integration loop */
void convert_initial_torques()
{
  Particle *p;
  Cell *cell;
  int c,i, np;
  double tx, ty, tz;

  INTEG_TRACE(fprintf(stderr,"%d: convert_initial_torques:\n",this_node));
  for (c = 0; c < local_cells.n; c++) {
    cell = local_cells.cell[c];
    p  = cell->part;
    np = cell->n;
    for(i = 0; i < np; i++) {
      double A[9];
      define_rotation_matrix(&p[i], A);

      tx = A[0 + 3*0]*p[i].f.torque[0] + A[0 + 3*1]*p[i].f.torque[1] + A[0 + 3*2]*p[i].f.torque[2];
      ty = A[1 + 3*0]*p[i].f.torque[0] + A[1 + 3*1]*p[i].f.torque[1] + A[1 + 3*2]*p[i].f.torque[2];
      tz = A[2 + 3*0]*p[i].f.torque[0] + A[2 + 3*1]*p[i].f.torque[1] + A[2 + 3*2]*p[i].f.torque[2];

      if ( thermo_switch & THERMO_LANGEVIN ) {
      
	friction_thermo_langevin_rotation(&p[i]);
	p[i].f.torque[0]+= tx;
	p[i].f.torque[1]+= ty;
	p[i].f.torque[2]+= tz;
      } else {
	p[i].f.torque[0] = tx;
	p[i].f.torque[1] = ty;
	p[i].f.torque[2] = tz;
      }

      ONEPART_TRACE(if(p[i].p.identity==check_id) fprintf(stderr,"%d: OPT: SCAL f = (%.3e,%.3e,%.3e) v_old = (%.3e,%.3e,%.3e)\n",this_node,p[i].f.f[0],p[i].f.f[1],p[i].f.f[2],p[i].m.v[0],p[i].m.v[1],p[i].m.v[2]));
    }
  }
}

/** convert torques from the body-fixed frames to space-fixed coordinates */
void convert_torques_body_to_space(Particle *p, double *torque)
{
  double A[9];
  define_rotation_matrix(p, A);

  torque[0] = A[0 + 3*0]*p->f.torque[0] + A[1 + 3*0]*p->f.torque[1] + A[2 + 3*0]*p->f.torque[2];
  torque[1] = A[0 + 3*1]*p->f.torque[0] + A[1 + 3*1]*p->f.torque[1] + A[2 + 3*1]*p->f.torque[2];
  torque[2] = A[0 + 3*2]*p->f.torque[0] + A[1 + 3*2]*p->f.torque[1] + A[2 + 3*2]*p->f.torque[2];

}

void convert_omega_body_to_space(Particle *p, double *omega)
{
  double A[9];
  define_rotation_matrix(p, A);

  omega[0] = A[0 + 3*0]*p->m.omega[0] + A[1 + 3*0]*p->m.omega[1] + A[2 + 3*0]*p->m.omega[2];
  omega[1] = A[0 + 3*1]*p->m.omega[0] + A[1 + 3*1]*p->m.omega[1] + A[2 + 3*1]*p->m.omega[2];
  omega[2] = A[0 + 3*2]*p->m.omega[0] + A[1 + 3*2]*p->m.omega[1] + A[2 + 3*2]*p->m.omega[2];

}


/** Multiply two quaternions */
void multiply_quaternions(double a[4], double b[4], double result[4])
{
 // Formula from http://www.j3d.org/matrix_faq/matrfaq_latest.html
 result[0] = a[0]*b[0] - a[1]*b[1] - a[2]*b[2] - a[3]*b[3];
 result[1] = a[0] * b[1] + a[1] * b[0] + a[2] * b[3] - a[3] * b[2];
 result[2] = a[0] * b[2] + a[2] * b[0] + a[3] * b[1] - a[1] * b[3]; 
 result[3] = a[0] * b[3] + a[3] * b[0] + a[1] * b[2] - a[2] * b[1];
}


#endif