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

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

#include <iomanip>

#include <set>

#include "boundary.hpp"
#include "ellipsoid.hpp"
#include "linkedcell.hpp"
#include "tesselation.hpp"
#include "vector.hpp"
#include "verbose.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.;
  double Matrix[NDIM*NDIM];
  double InverseLength[3];
  double psi,theta,phi; // euler angles in ZX'Z'' convention

  //cout << Verbose(3) << "Begin of SquaredDistanceToEllipsoid" << endl;

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

  // 1. translate coordinate system so that ellipsoid center is in origin
  helper.CopyVector(&x);
  helper.SubtractVector(&EllipsoidCenter);
  RefPoint.CopyVector(&helper);
  //cout << Verbose(4) << "Translated given point is at " << RefPoint << "." << endl;

  // 2. transform coordinate system by inverse of rotation matrix and of diagonal matrix
  psi = EllipsoidAngle[0];
  theta = EllipsoidAngle[1];
  phi = EllipsoidAngle[2];
  Matrix[0] = cos(psi)*cos(phi) - sin(psi)*cos(theta)*sin(phi);
  Matrix[1] = -cos(psi)*sin(phi) - sin(psi)*cos(theta)*cos(phi);
  Matrix[2] = sin(psi)*sin(theta);
  Matrix[3] = sin(psi)*cos(phi) + cos(psi)*cos(theta)*sin(phi);
  Matrix[4] = cos(psi)*cos(theta)*cos(phi) - sin(psi)*sin(phi);
  Matrix[5] = -cos(psi)*sin(theta);
  Matrix[6] = sin(theta)*sin(phi);
  Matrix[7] = sin(theta)*cos(phi);
  Matrix[8] = cos(theta);
  helper.MatrixMultiplication(Matrix);
  helper.Scale(InverseLength);
  //cout << Verbose(4) << "Transformed RefPoint is at " << helper << "." << endl;

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

  // 4. transform back the constructed intersection point
  psi = -EllipsoidAngle[0];
  theta = -EllipsoidAngle[1];
  phi = -EllipsoidAngle[2];
  helper.Scale(EllipsoidLength);
  Matrix[0] = cos(psi)*cos(phi) - sin(psi)*cos(theta)*sin(phi);
  Matrix[1] = -cos(psi)*sin(phi) - sin(psi)*cos(theta)*cos(phi);
  Matrix[2] = sin(psi)*sin(theta);
  Matrix[3] = sin(psi)*cos(phi) + cos(psi)*cos(theta)*sin(phi);
  Matrix[4] = cos(psi)*cos(theta)*cos(phi) - sin(psi)*sin(phi);
  Matrix[5] = -cos(psi)*sin(theta);
  Matrix[6] = sin(theta)*sin(phi);
  Matrix[7] = sin(theta)*cos(phi);
  Matrix[8] = cos(theta);
  helper.MatrixMultiplication(Matrix);
  //cout << Verbose(4) << "Intersection is at " << helper << "." << endl;

  // 5. determine distance between backtransformed point and x
  distance = RefPoint.DistanceSquared(&helper);
  //cout << Verbose(4) << "Squared distance between intersection and RefPoint is " << distance << "." << endl;

  return distance;
  //cout << Verbose(3) << "End of SquaredDistanceToEllipsoid" << endl;
};

/** 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.x[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;
    }
  }

  //cout << "Current summed distance is " << SumDistance << "." << endl;
  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(ofstream *out, Vector *set, int N, Vector *EllipsoidCenter, double *EllipsoidLength, double *EllipsoidAngle)
{
  int status = GSL_SUCCESS;
  *out << Verbose(2) << "Begin of FitPointSetToEllipsoid " << endl;
  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->x[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->x[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);
        }
        *out << setprecision(3) << Verbose(4) << "Converged fit at: " << *EllipsoidCenter << ", lengths " << EllipsoidLength[0] << ", " << EllipsoidLength[1] << ", " << EllipsoidLength[2] << ", angles " << EllipsoidAngle[0] << ", " << EllipsoidAngle[1] << ", " << EllipsoidAngle[2] << " with summed distance " << s->fval << "." << endl;
      }

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

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

  } else {
    *out << Verbose(3) << "Not enough points provided for fit to ellipsoid." << endl;
    return false;
  }
  *out << Verbose(2) << "End of FitPointSetToEllipsoid" << endl;
  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(ofstream *out, class Tesselation *T, class LinkedCell *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;
  LinkedNodes *List = NULL;
  *out << Verbose(2) << "Begin of PickRandomPointSet" << endl;

  // allocate array
  if (x == NULL) {
    x = new Vector[PointsToPick];
  } else {
    *out << "WARNING: Given pointer to vector array seems already allocated." << endl;
  }

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

    *out << Verbose(2) << "LC Intervals:";
    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;
      *out << " [" << Nlower[i] << "," << Nupper[i] << "] ";
    }
    *out << endl;

    // 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]++) {
          List = LC->GetCurrentCell();
          PointsLeft += List->size();
        }
    *out << Verbose(2) << "There are " << PointsLeft << " atoms in this neighbourhood." << endl;
    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 = (rand() % PointsLeft);
      current = PickedAtomNrs.find(index);  // not present?
      if (current == PickedAtomNrs.end()) {
        //*out << Verbose(2) << "Picking atom nr. " << index << "." << endl;
        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]++) {
          List = LC->GetCurrentCell();
//          *out << Verbose(2) << "Current cell is " << LC->n[0] << ", " << LC->n[1] << ", " << LC->n[2] << " with No. " << LC->index << " containing " << List->size() << " points." << endl;
          if (List != NULL) {
//            if (List->begin() != List->end())
//              *out << Verbose(2) << "Going through candidates ... " << endl;
//            else
//              *out << Verbose(2) << "Cell is empty ... " << endl;
            for (LinkedNodes::iterator Runner = List->begin(); Runner != List->end(); Runner++) {
              if ((current != PickedAtomNrs.end()) && (*current == index)) {
                Candidate = (*Runner);
                *out << Verbose(2) << "Current picked node is " << **Runner << " with index " << index << "." << endl;
                x[PointsPicked++].CopyVector(Candidate->node);    // we have one more atom picked
                current++;    // next pre-picked atom
              }
              index++;  // next atom nr.
            }
//          } else {
//            *out << Verbose(2) << "List for this index not allocated!" << endl;
          }
        }
    *out << Verbose(2) << "The following points were picked: " << endl;
    for (size_t i=0;i<PointsPicked;i++)
      *out << Verbose(2) << x[i] << endl;
    if (PointsPicked == PointsToPick)  // break out of loop if we have all
      break;
  } while(1);

  *out << Verbose(2) << "End of PickRandomPointSet" << endl;
};

/** 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(ofstream *out, 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;
  *out << Verbose(2) << "Begin of PickRandomPointSet" << endl;

  // allocate array
  if (x == NULL) {
    x = new Vector[PointsToPick];
  } else {
    *out << "WARNING: Given pointer to vector array seems already allocated." << endl;
  }

  if (List != NULL)
    for (PointMap::iterator Runner = List->begin(); Runner != List->end(); Runner++) {
      threshold = 1. - (double)(PointsToPick - PointsPicked)/(double)PointsLeft;
      value = (double)rand()/(double)RAND_MAX;
      //*out << Verbose(3) << "Current node is " << *Runner->second->node << " with " << value << " ... " << threshold << ": ";
      if (value > threshold) {
        x[PointsPicked].CopyVector(Runner->second->node->node);
        PointsPicked++;
        //*out << "IN." << endl;
      } else {
        //*out << "OUT." << endl;
      }
      PointsLeft--;
    }
  *out << Verbose(2) << "The following points were picked: " << endl;
  for (size_t i=0;i<PointsPicked;i++)
    *out << Verbose(3) << x[i] << endl;

  *out << Verbose(2) << "End of PickRandomPointSet" << endl;
};

/** 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(ofstream *out, class Tesselation *T, class LinkedCell *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;
  *out << Verbose(0) << "Begin of FindDistributionOfEllipsoids" << endl;

  // 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.AddVector(Runner->second->node->node);
  Center.Scale(1./T->PointsOnBoundaryCount);
  *out << Verbose(1) << "Center is at " << Center << "." << endl;

  // 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--) {
    *out << Verbose(1) << "Determining data set " << number << " ... " << endl;
    // pick the point set
    x = NULL;
    //PickRandomPointSet(out, T, LCList, x, N);
    PickRandomNeighbouredPointSet(out, 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;
    }
    //*out << Verbose(2) << "MinDistance " << MinDistance << ", MaxDistance " << MaxDistance << "." << endl;
    EllipsoidCenter.CopyVector(&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(out, x, N, &EllipsoidCenter, &EllipsoidLength[0], &EllipsoidAngle[0])) {
      *out << Verbose(1) << "Picking succeeded!" << endl;
      // output obtained parameter set
      output << number << "\t";
      for (int i=0;i<3;i++)
        output << setprecision(9) << EllipsoidCenter.x[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
      *out << Verbose(1) << "Picking failed!" << endl;
      number++;
    }
    delete[](x);  // free allocated memory for point set
  }
  // close output and finish
  output.close();

  *out << Verbose(0) << "End of FindDistributionOfEllipsoids" << endl;
};