/*
 * Project: MoleCuilder
 * Description: creates and alters molecular systems
 * Copyright (C)  2010-2012 University of Bonn. All rights reserved.
 * Please see the LICENSE file or "Copyright notice" in builder.cpp for details.
 */

/*
 * ellipsoid.cpp
 *
 *  Created on: Jan 20, 2009
 *      Author: heber
 */

// include config.h
#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

#include "CodePatterns/MemDebug.hpp"

#include <gsl/gsl_multimin.h>
#include <gsl/gsl_vector.h>

#include <iomanip>

#include <set>

#include "CodePatterns/Log.hpp"
#include "ellipsoid.hpp"
#include "LinearAlgebra/Vector.hpp"
#include "LinearAlgebra/RealSpaceMatrix.hpp"
#include "LinkedCell/linkedcell.hpp"
#include "Tesselation/BoundaryPointSet.hpp"
#include "Tesselation/boundary.hpp"
#include "Tesselation/tesselation.hpp"

#include "RandomNumbers/RandomNumberGeneratorFactory.hpp"
#include "RandomNumbers/RandomNumberGenerator.hpp"

/** Determines squared distance for a given point \a x to surface of ellipsoid.
 * \param x given point
 * \param EllipsoidCenter center of ellipsoid
 * \param EllipsoidLength[3] three lengths of half axis of ellipsoid
 * \param EllipsoidAngle[3] three rotation angles of ellipsoid
 * \return squared distance from point to surface
 */
double SquaredDistanceToEllipsoid(Vector &x, Vector &EllipsoidCenter, double *EllipsoidLength, double *EllipsoidAngle)
{
  Vector helper, RefPoint;
  double distance = -1.;
  RealSpaceMatrix Matrix;
  double InverseLength[3];
  double psi,theta,phi; // euler angles in ZX'Z'' convention

  //LOG(3, "Begin of SquaredDistanceToEllipsoid");

  for(int i=0;i<3;i++)
    InverseLength[i] = 1./EllipsoidLength[i];

  // 1. translate coordinate system so that ellipsoid center is in origin
  RefPoint = helper = x - EllipsoidCenter;
  //LOG(4, "Translated given point is at " << RefPoint << ".");

  // 2. transform coordinate system by inverse of rotation matrix and of diagonal matrix
  psi = EllipsoidAngle[0];
  theta = EllipsoidAngle[1];
  phi = EllipsoidAngle[2];
  Matrix.set(0,0, cos(psi)*cos(phi) - sin(psi)*cos(theta)*sin(phi));
  Matrix.set(1,0, -cos(psi)*sin(phi) - sin(psi)*cos(theta)*cos(phi));
  Matrix.set(2,0, sin(psi)*sin(theta));
  Matrix.set(0,1, sin(psi)*cos(phi) + cos(psi)*cos(theta)*sin(phi));
  Matrix.set(1,1, cos(psi)*cos(theta)*cos(phi) - sin(psi)*sin(phi));
  Matrix.set(2,1, -cos(psi)*sin(theta));
  Matrix.set(0,2, sin(theta)*sin(phi));
  Matrix.set(1,2, sin(theta)*cos(phi));
  Matrix.set(2,2, cos(theta));
  helper *= Matrix;
  helper.ScaleAll(InverseLength);
  //LOG(4, "Transformed RefPoint is at " << helper << ".");

  // 3. construct intersection point with unit sphere and ray between origin and x
  helper.Normalize(); // is simply normalizes vector in distance direction
  //LOG(4, "Transformed intersection is at " << helper << ".");

  // 4. transform back the constructed intersection point
  psi = -EllipsoidAngle[0];
  theta = -EllipsoidAngle[1];
  phi = -EllipsoidAngle[2];
  helper.ScaleAll(EllipsoidLength);
  Matrix.set(0,0, cos(psi)*cos(phi) - sin(psi)*cos(theta)*sin(phi));
  Matrix.set(1,0, -cos(psi)*sin(phi) - sin(psi)*cos(theta)*cos(phi));
  Matrix.set(2,0, sin(psi)*sin(theta));
  Matrix.set(0,1, sin(psi)*cos(phi) + cos(psi)*cos(theta)*sin(phi));
  Matrix.set(1,1, cos(psi)*cos(theta)*cos(phi) - sin(psi)*sin(phi));
  Matrix.set(2,1, -cos(psi)*sin(theta));
  Matrix.set(0,2, sin(theta)*sin(phi));
  Matrix.set(1,2, sin(theta)*cos(phi));
  Matrix.set(2,2, cos(theta));
  helper *= Matrix;
  //LOG(4, "Intersection is at " << helper << ".");

  // 5. determine distance between backtransformed point and x
  distance = RefPoint.DistanceSquared(helper);
  //LOG(4, "Squared distance between intersection and RefPoint is " << distance << ".");

  return distance;
  //LOG(3, "End of SquaredDistanceToEllipsoid");
};

/** structure for ellipsoid minimisation containing points to fit to.
 */
struct EllipsoidMinimisation {
  int N;      //!< dimension of vector set
  Vector *x;  //!< array of vectors
};

/** Sum of squared distance to ellipsoid to be minimised.
 * \param *x parameters for the ellipsoid
 * \param *params EllipsoidMinimisation with set of data points to minimise distance to and dimension
 * \return sum of squared distance, \sa SquaredDistanceToEllipsoid()
 */
double SumSquaredDistance (const gsl_vector * x, void * params)
{
  Vector *set= ((struct EllipsoidMinimisation *)params)->x;
  int N = ((struct EllipsoidMinimisation *)params)->N;
  double SumDistance = 0.;
  double distance;
  Vector Center;
  double EllipsoidLength[3], EllipsoidAngle[3];

  // put parameters into suitable ellipsoid form
  for (int i=0;i<3;i++) {
    Center[i] = gsl_vector_get(x, i+0);
    EllipsoidLength[i] = gsl_vector_get(x, i+3);
    EllipsoidAngle[i] = gsl_vector_get(x, i+6);
  }

  // go through all points and sum distance
  for (int i=0;i<N;i++) {
    distance = SquaredDistanceToEllipsoid(set[i], Center, EllipsoidLength, EllipsoidAngle);
    if (!isnan(distance)) {
      SumDistance += distance;
    } else {
      SumDistance = GSL_NAN;
      break;
    }
  }

  //LOG(0, "Current summed distance is " << SumDistance << ".");
  return SumDistance;
};

/** Finds best fitting ellipsoid parameter set in Least square sense for a given point set.
 * \param *out output stream for debugging
 * \param *set given point set
 * \param N number of points in set
 * \param EllipsoidParamter[3] three parameters in ellipsoid equation
 * \return true - fit successful, false - fit impossible
 */
bool FitPointSetToEllipsoid(Vector *set, int N, Vector *EllipsoidCenter, double *EllipsoidLength, double *EllipsoidAngle)
{
  int status = GSL_SUCCESS;
  LOG(2, "Begin of FitPointSetToEllipsoid ");
  if (N >= 3) { // check that enough points are given (9 d.o.f.)
    struct EllipsoidMinimisation par;
    const gsl_multimin_fminimizer_type *T = gsl_multimin_fminimizer_nmsimplex;
    gsl_multimin_fminimizer *s = NULL;
    gsl_vector *ss, *x;
    gsl_multimin_function minex_func;

    size_t iter = 0;
    double size;

    /* Starting point */
    x = gsl_vector_alloc (9);
    for (int i=0;i<3;i++) {
      gsl_vector_set (x, i+0, EllipsoidCenter->at(i));
      gsl_vector_set (x, i+3, EllipsoidLength[i]);
      gsl_vector_set (x, i+6, EllipsoidAngle[i]);
    }
    par.x = set;
    par.N = N;

    /* Set initial step sizes */
    ss = gsl_vector_alloc (9);
    for (int i=0;i<3;i++) {
      gsl_vector_set (ss, i+0, 0.1);
      gsl_vector_set (ss, i+3, 1.0);
      gsl_vector_set (ss, i+6, M_PI/20.);
    }

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

    s = gsl_multimin_fminimizer_alloc (T, 9);
    gsl_multimin_fminimizer_set (s, &minex_func, 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) {
        for (int i=0;i<3;i++) {
          EllipsoidCenter->at(i) = gsl_vector_get (s->x,i+0);
          EllipsoidLength[i] = gsl_vector_get (s->x, i+3);
          EllipsoidAngle[i] = gsl_vector_get (s->x, i+6);
        }
        LOG(4, setprecision(3) << "Converged fit at: " << *EllipsoidCenter << ", lengths " << EllipsoidLength[0] << ", " << EllipsoidLength[1] << ", " << EllipsoidLength[2] << ", angles " << EllipsoidAngle[0] << ", " << EllipsoidAngle[1] << ", " << EllipsoidAngle[2] << " with summed distance " << s->fval << ".");
      }

    } while (status == GSL_CONTINUE && iter < 1000);

    gsl_vector_free(x);
    gsl_vector_free(ss);
    gsl_multimin_fminimizer_free (s);

  } else {
    LOG(3, "Not enough points provided for fit to ellipsoid.");
    return false;
  }
  LOG(2, "End of FitPointSetToEllipsoid");
  if (status == GSL_SUCCESS)
    return true;
  else
    return false;
};

/** Picks a number of random points from a LC neighbourhood as a fitting set.
 * \param *out output stream for debugging
 * \param *T Tesselation containing boundary points
 * \param *LC linked cell list of all atoms
 * \param *&x random point set on return (not allocated!)
 * \param PointsToPick number of points in set to pick
 */
void PickRandomNeighbouredPointSet(class Tesselation *T, class LinkedCell_deprecated *LC, Vector *&x, size_t PointsToPick)
{
  size_t PointsLeft = 0;
  size_t PointsPicked = 0;
  int Nlower[NDIM], Nupper[NDIM];
  set<int> PickedAtomNrs;   // ordered list of picked atoms
  set<int>::iterator current;
  int index;
  TesselPoint *Candidate = NULL;
  LOG(2, "Begin of PickRandomPointSet");

  // allocate array
  if (x == NULL) {
    x = new Vector[PointsToPick];
  } else {
    ELOG(2, "Given pointer to vector array seems already allocated.");
  }

  RandomNumberGenerator &random = RandomNumberGeneratorFactory::getInstance().makeRandomNumberGenerator("mt19937", "uniform_int");
  // check that random number generator's bounds are ok
  ASSERT(random.min() == 0,
      "PickRandomNeighbouredPointSet: Chosen RandomNumberGenerator's min "
      +toString(random.min())+" is not 0!");
  ASSERT(random.max() >= LC->N[0],
      "PickRandomNeighbouredPointSet: Chosen RandomNumberGenerator's max "
      +toString(random.max())+" is too small"+toString(LC->N[0])
      +" for axis 0!");
  ASSERT(random.max() >= LC->N[1],
      "PickRandomNeighbouredPointSet: Chosen RandomNumberGenerator's max "
      +toString(random.max())+" is too small"+toString(LC->N[1])
      +" for axis 1!");
  ASSERT(random.max() >= LC->N[2],
      "PickRandomNeighbouredPointSet: Chosen RandomNumberGenerator's max "
      +toString(random.max())+" is too small"+toString(LC->N[2])
      +" for axis 2!");

  do {
    for(int i=0;i<NDIM;i++) // pick three random indices
      LC->n[i] = ((int)random() % LC->N[i]);
    LOG(2, "INFO: Center cell is " << LC->n[0] << ", " << LC->n[1] << ", " << LC->n[2] << ".");
    // get random cell
    const TesselPointSTLList *List = LC->GetCurrentCell();
    if (List == NULL) {  // set index to it
      continue;
    }
    LOG(2, "INFO: Cell index is No. " << LC->index << ".");

    if (DoLog(2)) {
      std::stringstream output;
      output << "LC Intervals:";
      for (int i=0;i<NDIM;i++)
        output << " [" << Nlower[i] << "," << Nupper[i] << "] ";
      LOG(2, output.str());
    }

    for (int i=0;i<NDIM;i++) {
      Nlower[i] = ((LC->n[i]-1) >= 0) ? LC->n[i]-1 : 0;
      Nupper[i] = ((LC->n[i]+1) < LC->N[i]) ? LC->n[i]+1 : LC->N[i]-1;
    }

    // count whether there are sufficient atoms in this cell+neighbors
    PointsLeft=0;
    for (LC->n[0] = Nlower[0]; LC->n[0] <= Nupper[0]; LC->n[0]++)
      for (LC->n[1] = Nlower[1]; LC->n[1] <= Nupper[1]; LC->n[1]++)
        for (LC->n[2] = Nlower[2]; LC->n[2] <= Nupper[2]; LC->n[2]++) {
          const TesselPointSTLList *List = LC->GetCurrentCell();
          PointsLeft += List->size();
        }
    LOG(2, "There are " << PointsLeft << " atoms in this neighbourhood.");
    if (PointsLeft < PointsToPick) {  // ensure that we can pick enough points in its neighbourhood at all.
      continue;
    }

    // pre-pick a fixed number of atoms
    PickedAtomNrs.clear();
    do {
      index = (((int)random()) % PointsLeft);
      current = PickedAtomNrs.find(index);  // not present?
      if (current == PickedAtomNrs.end()) {
        //LOG(2, "Picking atom Nr. " << index << ".");
        PickedAtomNrs.insert(index);
      }
    } while (PickedAtomNrs.size() < PointsToPick);

    index = 0; // now go through all and pick those whose from PickedAtomsNr
    PointsPicked=0;
    current = PickedAtomNrs.begin();
    for (LC->n[0] = Nlower[0]; LC->n[0] <= Nupper[0]; LC->n[0]++)
      for (LC->n[1] = Nlower[1]; LC->n[1] <= Nupper[1]; LC->n[1]++)
        for (LC->n[2] = Nlower[2]; LC->n[2] <= Nupper[2]; LC->n[2]++) {
          const TesselPointSTLList *List = LC->GetCurrentCell();
//          LOG(2, "Current cell is " << LC->n[0] << ", " << LC->n[1] << ", " << LC->n[2] << " with No. " << LC->index << " containing " << List->size() << " points.");
          if (List != NULL) {
//            if (List->begin() != List->end())
//              LOG(2, "Going through candidates ... ");
//            else
//              LOG(2, "Cell is empty ... ");
            for (TesselPointSTLList::const_iterator Runner = List->begin(); Runner != List->end(); Runner++) {
              if ((current != PickedAtomNrs.end()) && (*current == index)) {
                Candidate = (*Runner);
                LOG(2, "Current picked node is " << (*Runner)->getName() << " with index " << index << ".");
                x[PointsPicked++] = Candidate->getPosition();    // we have one more atom picked
                current++;    // next pre-picked atom
              }
              index++;  // next atom Nr.
            }
//          } else {
//            LOG(2, "List for this index not allocated!");
          }
        }
    LOG(2, "The following points were picked: ");
    for (size_t i=0;i<PointsPicked;i++)
      LOG(2, x[i]);
    if (PointsPicked == PointsToPick)  // break out of loop if we have all
      break;
  } while(1);

  LOG(2, "End of PickRandomPointSet");
};

/** Picks a number of random points from a set of boundary points as a fitting set.
 * \param *out output stream for debugging
 * \param *T Tesselation containing boundary points
 * \param *&x random point set on return (not allocated!)
 * \param PointsToPick number of points in set to pick
 */
void PickRandomPointSet(class Tesselation *T, Vector *&x, size_t PointsToPick)
{
  size_t PointsLeft = (size_t) T->PointsOnBoundaryCount;
  size_t PointsPicked = 0;
  double value, threshold;
  PointMap *List = &T->PointsOnBoundary;
  LOG(2, "Begin of PickRandomPointSet");

  // allocate array
  if (x == NULL) {
    x = new Vector[PointsToPick];
  } else {
    ELOG(2, "Given pointer to vector array seems already allocated.");
  }

  RandomNumberGenerator &random = RandomNumberGeneratorFactory::getInstance().makeRandomNumberGenerator("mt19937", "uniform_int");
  const double rng_min = random.min();
  const double rng_max = random.max();
  if (List != NULL)
    for (PointMap::iterator Runner = List->begin(); Runner != List->end(); Runner++) {
      threshold = 1. - (double)(PointsToPick - PointsPicked)/(double)PointsLeft;
      value = (double)random()/(double)(rng_max-rng_min);
      if (value > threshold) {
        x[PointsPicked] = (Runner->second->node->getPosition());
        PointsPicked++;
        //LOG(3, "Current node is " << *Runner->second->node << " with " << value << " ... " << threshold << ": IN.");
      } else {
        //LOG(3, "Current node is " << *Runner->second->node << " with " << value << " ... " << threshold << ": OUT.");
      }
      PointsLeft--;
    }
  LOG(2, "The following points were picked: ");
  for (size_t i=0;i<PointsPicked;i++)
    LOG(3, x[i]);

  LOG(2, "End of PickRandomPointSet");
};

/** Finds best fitting ellipsoid parameter set in least square sense for a given point set.
 * \param *out output stream for debugging
 * \param *T Tesselation containing boundary points
 * \param *LCList linked cell list of all atoms
 * \param N number of unique points in ellipsoid fit, must be greater equal 6
 * \param number of fits (i.e. parameter sets in output file)
 * \param *filename name for output file
 */
void FindDistributionOfEllipsoids(class Tesselation *T, class LinkedCell_deprecated *LCList, int N, int number, const char *filename)
{
  ofstream output;
  Vector *x = NULL;
  Vector Center;
  Vector EllipsoidCenter;
  double EllipsoidLength[3];
  double EllipsoidAngle[3];
  double distance, MaxDistance, MinDistance;
  LOG(0, "Begin of FindDistributionOfEllipsoids");

  // construct center of gravity of boundary point set for initial ellipsoid center
  Center.Zero();
  for (PointMap::iterator Runner = T->PointsOnBoundary.begin(); Runner != T->PointsOnBoundary.end(); Runner++)
    Center += (Runner->second->node->getPosition());
  Center.Scale(1./T->PointsOnBoundaryCount);
  LOG(1, "Center is at " << Center << ".");

  // Output header
  output.open(filename, ios::trunc);
  output << "# Nr.\tCenterX\tCenterY\tCenterZ\ta\tb\tc\tpsi\ttheta\tphi" << endl;

  // loop over desired number of parameter sets
  for (;number >0;number--) {
    LOG(1, "Determining data set " << number << " ... ");
    // pick the point set
    x = NULL;
    //PickRandomPointSet(T, LCList, x, N);
    PickRandomNeighbouredPointSet(T, LCList, x, N);

    // calculate some sensible starting values for parameter fit
    MaxDistance = 0.;
    MinDistance = x[0].ScalarProduct(x[0]);
    for (int i=0;i<N;i++) {
      distance = x[i].ScalarProduct(x[i]);
      if (distance > MaxDistance)
        MaxDistance = distance;
      if (distance < MinDistance)
        MinDistance = distance;
    }
    //LOG(2, "MinDistance " << MinDistance << ", MaxDistance " << MaxDistance << ".");
    EllipsoidCenter = Center;  // use Center of Gravity as initial center of ellipsoid
    for (int i=0;i<3;i++)
      EllipsoidAngle[i] = 0.;
    EllipsoidLength[0] = sqrt(MaxDistance);
    EllipsoidLength[1] = sqrt((MaxDistance+MinDistance)/2.);
    EllipsoidLength[2] = sqrt(MinDistance);

    // fit the parameters
    if (FitPointSetToEllipsoid(x, N, &EllipsoidCenter, &EllipsoidLength[0], &EllipsoidAngle[0])) {
      LOG(1, "Picking succeeded!");
      // output obtained parameter set
      output << number << "\t";
      for (int i=0;i<3;i++)
        output << setprecision(9) << EllipsoidCenter[i] << "\t";
      for (int i=0;i<3;i++)
        output << setprecision(9) << EllipsoidLength[i] << "\t";
      for (int i=0;i<3;i++)
        output << setprecision(9) << EllipsoidAngle[i] << "\t";
      output << endl;
    } else { // increase N to pick one more
      LOG(1, "Picking failed!");
      number++;
    }
    delete[](x);  // free allocated memory for point set
  }
  // close output and finish
  output.close();

  LOG(0, "End of FindDistributionOfEllipsoids");
};