/*
 * molecule_geometry.cpp
 *
 *  Created on: Oct 5, 2009
 *      Author: heber
 */

#include "atom.hpp"
#include "bond.hpp"
#include "config.hpp"
#include "element.hpp"
#include "helpers.hpp"
#include "leastsquaremin.hpp"
#include "log.hpp"
#include "memoryallocator.hpp"
#include "molecule.hpp"

/************************************* Functions for class molecule *********************************/


/** Centers the molecule in the box whose lengths are defined by vector \a *BoxLengths.
 * \param *out output stream for debugging
 */
bool molecule::CenterInBox()
{
  bool status = true;
  const Vector *Center = DetermineCenterOfAll();
  double *M = ReturnFullMatrixforSymmetric(cell_size);
  double *Minv = InverseMatrix(M);

  // go through all atoms
  ActOnAllVectors( &Vector::SubtractVector, Center);
  ActOnAllVectors( &Vector::WrapPeriodically, (const double *)M, (const double *)Minv);

  Free(&M);
  Free(&Minv);
  delete(Center);
  return status;
};


/** Bounds the molecule in the box whose lengths are defined by vector \a *BoxLengths.
 * \param *out output stream for debugging
 */
bool molecule::BoundInBox()
{
  bool status = true;
  double *M = ReturnFullMatrixforSymmetric(cell_size);
  double *Minv = InverseMatrix(M);

  // go through all atoms
  ActOnAllVectors( &Vector::WrapPeriodically, (const double *)M, (const double *)Minv);

  Free(&M);
  Free(&Minv);
  return status;
};

/** Centers the edge of the atoms at (0,0,0).
 * \param *out output stream for debugging
 * \param *max coordinates of other edge, specifying box dimensions.
 */
void molecule::CenterEdge(Vector *max)
{
  Vector *min = new Vector;

//  Log() << Verbose(3) << "Begin of CenterEdge." << endl;
  atom *ptr = start->next;  // start at first in list
  if (ptr != end) {   //list not empty?
    for (int i=NDIM;i--;) {
      max->x[i] = ptr->x.x[i];
      min->x[i] = ptr->x.x[i];
    }
    while (ptr->next != end) {  // continue with second if present
      ptr = ptr->next;
      //ptr->Output(1,1,out);
      for (int i=NDIM;i--;) {
        max->x[i] = (max->x[i] < ptr->x.x[i]) ? ptr->x.x[i] : max->x[i];
        min->x[i] = (min->x[i] > ptr->x.x[i]) ? ptr->x.x[i] : min->x[i];
      }
    }
//    Log() << Verbose(4) << "Maximum is ";
//    max->Output(out);
//    Log() << Verbose(0) << ", Minimum is ";
//    min->Output(out);
//    Log() << Verbose(0) << endl;
    min->Scale(-1.);
    max->AddVector(min);
    Translate(min);
    Center.Zero();
  }
  delete(min);
//  Log() << Verbose(3) << "End of CenterEdge." << endl;
};

/** Centers the center of the atoms at (0,0,0).
 * \param *out output stream for debugging
 * \param *center return vector for translation vector
 */
void molecule::CenterOrigin()
{
  int Num = 0;
  atom *ptr = start;  // start at first in list

  Center.Zero();

  if (ptr->next != end) {   //list not empty?
    while (ptr->next != end) {  // continue with second if present
      ptr = ptr->next;
      Num++;
      Center.AddVector(&ptr->x);
    }
    Center.Scale(-1./Num); // divide through total number (and sign for direction)
    Translate(&Center);
    Center.Zero();
  }
};

/** Returns vector pointing to center of all atoms.
 * \return pointer to center of all vector
 */
Vector * molecule::DetermineCenterOfAll() const
{
  atom *ptr = start->next;  // start at first in list
  Vector *a = new Vector();
  Vector tmp;
  double Num = 0;

  a->Zero();

  if (ptr != end) {   //list not empty?
    while (ptr->next != end) {  // continue with second if present
      ptr = ptr->next;
      Num += 1.;
      tmp.CopyVector(&ptr->x);
      a->AddVector(&tmp);
    }
    a->Scale(1./Num); // divide through total mass (and sign for direction)
  }
  return a;
};

/** Returns vector pointing to center of gravity.
 * \param *out output stream for debugging
 * \return pointer to center of gravity vector
 */
Vector * molecule::DetermineCenterOfGravity()
{
  atom *ptr = start->next;  // start at first in list
  Vector *a = new Vector();
  Vector tmp;
  double Num = 0;

  a->Zero();

  if (ptr != end) {   //list not empty?
    while (ptr->next != end) {  // continue with second if present
      ptr = ptr->next;
      Num += ptr->type->mass;
      tmp.CopyVector(&ptr->x);
      tmp.Scale(ptr->type->mass);  // scale by mass
      a->AddVector(&tmp);
    }
    a->Scale(-1./Num); // divide through total mass (and sign for direction)
  }
//  Log() << Verbose(1) << "Resulting center of gravity: ";
//  a->Output(out);
//  Log() << Verbose(0) << endl;
  return a;
};

/** Centers the center of gravity of the atoms at (0,0,0).
 * \param *out output stream for debugging
 * \param *center return vector for translation vector
 */
void molecule::CenterPeriodic()
{
  DeterminePeriodicCenter(Center);
};


/** Centers the center of gravity of the atoms at (0,0,0).
 * \param *out output stream for debugging
 * \param *center return vector for translation vector
 */
void molecule::CenterAtVector(Vector *newcenter)
{
  Center.CopyVector(newcenter);
};


/** Scales all atoms by \a *factor.
 * \param *factor pointer to scaling factor
 */
void molecule::Scale(const double ** const factor)
{
  atom *ptr = start;

  while (ptr->next != end) {
    ptr = ptr->next;
    for (int j=0;j<MDSteps;j++)
      ptr->Trajectory.R.at(j).Scale(factor);
    ptr->x.Scale(factor);
  }
};

/** Translate all atoms by given vector.
 * \param trans[] translation vector.
 */
void molecule::Translate(const Vector *trans)
{
  atom *ptr = start;

  while (ptr->next != end) {
    ptr = ptr->next;
    for (int j=0;j<MDSteps;j++)
      ptr->Trajectory.R.at(j).Translate(trans);
    ptr->x.Translate(trans);
  }
};

/** Translate the molecule periodically in the box.
 * \param trans[] translation vector.
 * TODO treatment of trajetories missing
 */
void molecule::TranslatePeriodically(const Vector *trans)
{
  double *M = ReturnFullMatrixforSymmetric(cell_size);
  double *Minv = InverseMatrix(M);

  // go through all atoms
  ActOnAllVectors( &Vector::SubtractVector, trans);
  ActOnAllVectors( &Vector::WrapPeriodically, (const double *)M, (const double *)Minv);

  Free(&M);
  Free(&Minv);
};


/** Mirrors all atoms against a given plane.
 * \param n[] normal vector of mirror plane.
 */
void molecule::Mirror(const Vector *n)
{
  ActOnAllVectors( &Vector::Mirror, n );
};

/** Determines center of molecule (yet not considering atom masses).
 * \param center reference to return vector
 */
void molecule::DeterminePeriodicCenter(Vector &center)
{
  atom *Walker = start;
  double *matrix = ReturnFullMatrixforSymmetric(cell_size);
  double *inversematrix = InverseMatrix(cell_size);
  double tmp;
  bool flag;
  Vector Testvector, Translationvector;

  do {
    Center.Zero();
    flag = true;
    while (Walker->next != end) {
      Walker = Walker->next;
#ifdef ADDHYDROGEN
      if (Walker->type->Z != 1) {
#endif
        Testvector.CopyVector(&Walker->x);
        Testvector.MatrixMultiplication(inversematrix);
        Translationvector.Zero();
        for (BondList::const_iterator Runner = Walker->ListOfBonds.begin(); Runner != Walker->ListOfBonds.end(); (++Runner)) {
         if (Walker->nr < (*Runner)->GetOtherAtom(Walker)->nr) // otherwise we shift one to, the other fro and gain nothing
            for (int j=0;j<NDIM;j++) {
              tmp = Walker->x.x[j] - (*Runner)->GetOtherAtom(Walker)->x.x[j];
              if ((fabs(tmp)) > BondDistance) {
                flag = false;
                Log() << Verbose(0) << "Hit: atom " << Walker->Name << " in bond " << *(*Runner) << " has to be shifted due to " << tmp << "." << endl;
                if (tmp > 0)
                  Translationvector.x[j] -= 1.;
                else
                  Translationvector.x[j] += 1.;
              }
            }
        }
        Testvector.AddVector(&Translationvector);
        Testvector.MatrixMultiplication(matrix);
        Center.AddVector(&Testvector);
        Log() << Verbose(1) << "vector is: ";
        Testvector.Output();
        Log() << Verbose(0) << endl;
#ifdef ADDHYDROGEN
        // now also change all hydrogens
        for (BondList::const_iterator Runner = Walker->ListOfBonds.begin(); Runner != Walker->ListOfBonds.end(); (++Runner)) {
          if ((*Runner)->GetOtherAtom(Walker)->type->Z == 1) {
            Testvector.CopyVector(&(*Runner)->GetOtherAtom(Walker)->x);
            Testvector.MatrixMultiplication(inversematrix);
            Testvector.AddVector(&Translationvector);
            Testvector.MatrixMultiplication(matrix);
            Center.AddVector(&Testvector);
            Log() << Verbose(1) << "Hydrogen vector is: ";
            Testvector.Output();
            Log() << Verbose(0) << endl;
          }
        }
      }
#endif
    }
  } while (!flag);
  Free(&matrix);
  Free(&inversematrix);

  Center.Scale(1./(double)AtomCount);
};

/** Transforms/Rotates the given molecule into its principal axis system.
 * \param *out output stream for debugging
 * \param DoRotate whether to rotate (true) or only to determine the PAS.
 * TODO treatment of trajetories missing
 */
void molecule::PrincipalAxisSystem(bool DoRotate)
{
  atom *ptr = start;  // start at first in list
  double InertiaTensor[NDIM*NDIM];
  Vector *CenterOfGravity = DetermineCenterOfGravity();

  CenterPeriodic();

  // reset inertia tensor
  for(int i=0;i<NDIM*NDIM;i++)
    InertiaTensor[i] = 0.;

  // sum up inertia tensor
  while (ptr->next != end) {
    ptr = ptr->next;
    Vector x;
    x.CopyVector(&ptr->x);
    //x.SubtractVector(CenterOfGravity);
    InertiaTensor[0] += ptr->type->mass*(x.x[1]*x.x[1] + x.x[2]*x.x[2]);
    InertiaTensor[1] += ptr->type->mass*(-x.x[0]*x.x[1]);
    InertiaTensor[2] += ptr->type->mass*(-x.x[0]*x.x[2]);
    InertiaTensor[3] += ptr->type->mass*(-x.x[1]*x.x[0]);
    InertiaTensor[4] += ptr->type->mass*(x.x[0]*x.x[0] + x.x[2]*x.x[2]);
    InertiaTensor[5] += ptr->type->mass*(-x.x[1]*x.x[2]);
    InertiaTensor[6] += ptr->type->mass*(-x.x[2]*x.x[0]);
    InertiaTensor[7] += ptr->type->mass*(-x.x[2]*x.x[1]);
    InertiaTensor[8] += ptr->type->mass*(x.x[0]*x.x[0] + x.x[1]*x.x[1]);
  }
  // print InertiaTensor for debugging
  Log() << Verbose(0) << "The inertia tensor is:" << endl;
  for(int i=0;i<NDIM;i++) {
    for(int j=0;j<NDIM;j++)
      Log() << Verbose(0) << InertiaTensor[i*NDIM+j] << " ";
    Log() << Verbose(0) << endl;
  }
  Log() << Verbose(0) << endl;

  // diagonalize to determine principal axis system
  gsl_eigen_symmv_workspace *T = gsl_eigen_symmv_alloc(NDIM);
  gsl_matrix_view m = gsl_matrix_view_array(InertiaTensor, NDIM, NDIM);
  gsl_vector *eval = gsl_vector_alloc(NDIM);
  gsl_matrix *evec = gsl_matrix_alloc(NDIM, NDIM);
  gsl_eigen_symmv(&m.matrix, eval, evec, T);
  gsl_eigen_symmv_free(T);
  gsl_eigen_symmv_sort(eval, evec, GSL_EIGEN_SORT_ABS_DESC);

  for(int i=0;i<NDIM;i++) {
    Log() << Verbose(1) << "eigenvalue = " << gsl_vector_get(eval, i);
    Log() << Verbose(0) << ", eigenvector = (" << evec->data[i * evec->tda + 0] << "," << evec->data[i * evec->tda + 1] << "," << evec->data[i * evec->tda + 2] << ")" << endl;
  }

  // check whether we rotate or not
  if (DoRotate) {
    Log() << Verbose(1) << "Transforming molecule into PAS ... ";
    // the eigenvectors specify the transformation matrix
    ActOnAllVectors( &Vector::MatrixMultiplication, (const double *) evec->data );
    Log() << Verbose(0) << "done." << endl;

    // summing anew for debugging (resulting matrix has to be diagonal!)
    // reset inertia tensor
    for(int i=0;i<NDIM*NDIM;i++)
      InertiaTensor[i] = 0.;

    // sum up inertia tensor
    ptr = start;
    while (ptr->next != end) {
      ptr = ptr->next;
      Vector x;
      x.CopyVector(&ptr->x);
      //x.SubtractVector(CenterOfGravity);
      InertiaTensor[0] += ptr->type->mass*(x.x[1]*x.x[1] + x.x[2]*x.x[2]);
      InertiaTensor[1] += ptr->type->mass*(-x.x[0]*x.x[1]);
      InertiaTensor[2] += ptr->type->mass*(-x.x[0]*x.x[2]);
      InertiaTensor[3] += ptr->type->mass*(-x.x[1]*x.x[0]);
      InertiaTensor[4] += ptr->type->mass*(x.x[0]*x.x[0] + x.x[2]*x.x[2]);
      InertiaTensor[5] += ptr->type->mass*(-x.x[1]*x.x[2]);
      InertiaTensor[6] += ptr->type->mass*(-x.x[2]*x.x[0]);
      InertiaTensor[7] += ptr->type->mass*(-x.x[2]*x.x[1]);
      InertiaTensor[8] += ptr->type->mass*(x.x[0]*x.x[0] + x.x[1]*x.x[1]);
    }
    // print InertiaTensor for debugging
    Log() << Verbose(0) << "The inertia tensor is:" << endl;
    for(int i=0;i<NDIM;i++) {
      for(int j=0;j<NDIM;j++)
        Log() << Verbose(0) << InertiaTensor[i*NDIM+j] << " ";
      Log() << Verbose(0) << endl;
    }
    Log() << Verbose(0) << endl;
  }

  // free everything
  delete(CenterOfGravity);
  gsl_vector_free(eval);
  gsl_matrix_free(evec);
};


/** Align all atoms in such a manner that given vector \a *n is along z axis.
 * \param n[] alignment vector.
 */
void molecule::Align(Vector *n)
{
  atom *ptr = start;
  double alpha, tmp;
  Vector z_axis;
  z_axis.x[0] = 0.;
  z_axis.x[1] = 0.;
  z_axis.x[2] = 1.;

  // rotate on z-x plane
  Log() << Verbose(0) << "Begin of Aligning all atoms." << endl;
  alpha = atan(-n->x[0]/n->x[2]);
  Log() << Verbose(1) << "Z-X-angle: " << alpha << " ... ";
  while (ptr->next != end) {
    ptr = ptr->next;
    tmp = ptr->x.x[0];
    ptr->x.x[0] =  cos(alpha) * tmp + sin(alpha) * ptr->x.x[2];
    ptr->x.x[2] = -sin(alpha) * tmp + cos(alpha) * ptr->x.x[2];
    for (int j=0;j<MDSteps;j++) {
      tmp = ptr->Trajectory.R.at(j).x[0];
      ptr->Trajectory.R.at(j).x[0] =  cos(alpha) * tmp + sin(alpha) * ptr->Trajectory.R.at(j).x[2];
      ptr->Trajectory.R.at(j).x[2] = -sin(alpha) * tmp + cos(alpha) * ptr->Trajectory.R.at(j).x[2];
    }
  }
  // rotate n vector
  tmp = n->x[0];
  n->x[0] =  cos(alpha) * tmp +  sin(alpha) * n->x[2];
  n->x[2] = -sin(alpha) * tmp +  cos(alpha) * n->x[2];
  Log() << Verbose(1) << "alignment vector after first rotation: ";
  n->Output();
  Log() << Verbose(0) << endl;

  // rotate on z-y plane
  ptr = start;
  alpha = atan(-n->x[1]/n->x[2]);
  Log() << Verbose(1) << "Z-Y-angle: " << alpha << " ... ";
  while (ptr->next != end) {
    ptr = ptr->next;
    tmp = ptr->x.x[1];
    ptr->x.x[1] =  cos(alpha) * tmp + sin(alpha) * ptr->x.x[2];
    ptr->x.x[2] = -sin(alpha) * tmp + cos(alpha) * ptr->x.x[2];
    for (int j=0;j<MDSteps;j++) {
      tmp = ptr->Trajectory.R.at(j).x[1];
      ptr->Trajectory.R.at(j).x[1] =  cos(alpha) * tmp + sin(alpha) * ptr->Trajectory.R.at(j).x[2];
      ptr->Trajectory.R.at(j).x[2] = -sin(alpha) * tmp + cos(alpha) * ptr->Trajectory.R.at(j).x[2];
    }
  }
  // rotate n vector (for consistency check)
  tmp = n->x[1];
  n->x[1] =  cos(alpha) * tmp +  sin(alpha) * n->x[2];
  n->x[2] = -sin(alpha) * tmp +  cos(alpha) * n->x[2];

  Log() << Verbose(1) << "alignment vector after second rotation: ";
  n->Output();
  Log() << Verbose(1) << endl;
  Log() << Verbose(0) << "End of Aligning all atoms." << endl;
};


/** Calculates sum over least square distance to line hidden in \a *x.
 * \param *x offset and direction vector
 * \param *params pointer to lsq_params structure
 * \return \f$ sum_i^N | y_i - (a + t_i b)|^2\f$
 */
double LeastSquareDistance (const gsl_vector * x, void * params)
{
  double res = 0, t;
  Vector a,b,c,d;
  struct lsq_params *par = (struct lsq_params *)params;
  atom *ptr = par->mol->start;

  // initialize vectors
  a.x[0] = gsl_vector_get(x,0);
  a.x[1] = gsl_vector_get(x,1);
  a.x[2] = gsl_vector_get(x,2);
  b.x[0] = gsl_vector_get(x,3);
  b.x[1] = gsl_vector_get(x,4);
  b.x[2] = gsl_vector_get(x,5);
  // go through all atoms
  while (ptr != par->mol->end) {
    ptr = ptr->next;
    if (ptr->type == ((struct lsq_params *)params)->type) { // for specific type
      c.CopyVector(&ptr->x);  // copy vector to temporary one
      c.SubtractVector(&a);   // subtract offset vector
      t = c.ScalarProduct(&b);           // get direction parameter
      d.CopyVector(&b);       // and create vector
      d.Scale(&t);
      c.SubtractVector(&d);   // ... yielding distance vector
      res += d.ScalarProduct((const Vector *)&d);        // add squared distance
    }
  }
  return res;
};

/** By minimizing the least square distance gains alignment vector.
 * \bug this is not yet working properly it seems
 */
void molecule::GetAlignvector(struct lsq_params * par) const
{
    int np = 6;

   const gsl_multimin_fminimizer_type *T =
     gsl_multimin_fminimizer_nmsimplex;
   gsl_multimin_fminimizer *s = NULL;
   gsl_vector *ss;
   gsl_multimin_function minex_func;

   size_t iter = 0, i;
   int status;
   double size;

   /* Initial vertex size vector */
   ss = gsl_vector_alloc (np);

   /* Set all step sizes to 1 */
   gsl_vector_set_all (ss, 1.0);

   /* Starting point */
   par->x = gsl_vector_alloc (np);
   par->mol = this;

   gsl_vector_set (par->x, 0, 0.0);  // offset
   gsl_vector_set (par->x, 1, 0.0);
   gsl_vector_set (par->x, 2, 0.0);
   gsl_vector_set (par->x, 3, 0.0);  // direction
   gsl_vector_set (par->x, 4, 0.0);
   gsl_vector_set (par->x, 5, 1.0);

   /* Initialize method and iterate */
   minex_func.f = &LeastSquareDistance;
   minex_func.n = np;
   minex_func.params = (void *)par;

   s = gsl_multimin_fminimizer_alloc (T, np);
   gsl_multimin_fminimizer_set (s, &minex_func, par->x, ss);

   do
     {
       iter++;
       status = gsl_multimin_fminimizer_iterate(s);

       if (status)
         break;

       size = gsl_multimin_fminimizer_size (s);
       status = gsl_multimin_test_size (size, 1e-2);

       if (status == GSL_SUCCESS)
         {
           printf ("converged to minimum at\n");
         }

       printf ("%5d ", (int)iter);
       for (i = 0; i < (size_t)np; i++)
         {
           printf ("%10.3e ", gsl_vector_get (s->x, i));
         }
       printf ("f() = %7.3f size = %.3f\n", s->fval, size);
     }
   while (status == GSL_CONTINUE && iter < 100);

  for (i=0;i<(size_t)np;i++)
    gsl_vector_set(par->x, i, gsl_vector_get(s->x, i));
   //gsl_vector_free(par->x);
   gsl_vector_free(ss);
   gsl_multimin_fminimizer_free (s);
};