//
// postprocess.cpp
//     The actual TriplClust post-processing.
//
// Author:  Christoph Dalitz
// Date:    2022-11-07
// License: see ../LICENSE
//

#include "postprocess.h"
#include "orthogonallsq.h"
#include "kdtree/kdtree.hpp"
#include "util.h"

#include <iostream>
#include <fstream>
#include <string>
#include <limits>
#include <math.h>       /* sqrt */




void writeClusterID(const std::vector<size_t> & correct, std::vector< std::vector<size_t> > & assignedIds);
void writeIDsInCloud(PointCloud &cloud, const std::vector< std::vector<size_t> > & assignedIds, size_t* next_label, size_t old_cluster_id);
int correctOverSegmentation(PointCloud &cloud, std::vector< std::vector<size_t> > & assignedIds);
int expandStart(PointCloud &cloud, std::vector< std::vector<size_t> > & assignedIds);


/** @file */


/**
 * @brief Processes entire point cloud and reassigns cluster id labels, if appropriate.
 *
 * @param cloud Pointcloud generated by triplclust
 * @param min_depth Minimum number of points making a branch
 * @param verbose Verbosity level
 * @return RC Number of split up clusters
 */
int process_pointcloud(PointCloud &cloud, int min_depth/*=25*/, int verbose/*=0*/)
{
  int nchanged = 0;

  // collect all cluster labels in cloud
  std::set<size_t> ids;
  Point p;
  //for(Point p : cloud) {
  for (std::vector<Point>::iterator iterator = cloud.begin(); iterator != cloud.end(); ++iterator) {
    p = *iterator;
    p.cluster_ids = iterator->cluster_ids;
    // concat cluster id set to ids
    ids.insert(p.cluster_ids.begin(), p.cluster_ids.end());
  }

  size_t last_label_begin = *ids.rbegin();
  size_t next_label = *ids.rbegin()+1;

  PointCloud temp;
  std::vector<size_t> temp_ids;

  // for each cluster label, call process_cluster
  size_t id;
  for (std::set<long unsigned int>::iterator iterator = ids.begin(); iterator != ids.end(); ++iterator) {
    id = *iterator;
    temp_ids.clear();
    for (unsigned int index = 0; index < cloud.size(); ++index) {
      // check if id is in Point in cloud
      if(cloud[index].cluster_ids.find(id) != cloud[index].cluster_ids.end()) {
        temp_ids.push_back(index);
      }
    }
    // call process_cluster with the cloud subset with cluster=id
    nchanged += process_cluster(cloud, temp_ids, id, &next_label, min_depth, verbose);
  }

  nchanged = (next_label-1)-last_label_begin;

  return nchanged;

}


// processes cluster with given id in point cloud
// and reassign cluster id labels, if appropriate
//   cloud:      *full* point cloud with label cluster_id
//   indexes:    the index of the point cloud elements in the cluster
//   cluster_id: cluster id that is to be processed
//   next_label: next free label for new ids
//   min_depth:  minimum number of points making a branch
//   verbose:    verbosity level
//   RC:         number of split up clusters (1 or 0)
int process_cluster(PointCloud &cloud, std::vector<size_t> indexes, size_t cluster_id, size_t* next_label, size_t min_depth, int verbose)
{
  int nchanged = 0;
  bool compound = false; // is trajectory made of many trajectories

  Graph graph(cloud, indexes);   // Constructor constructs MST

  std::vector< std::vector<size_t> > segments;  // will contain all segments in a trajectory

  // remove edges at branches with depth min_depth
  if(graph.removeBranches(min_depth)) {   // is compound trajectory
    compound = true;
    // remove edges considerably longer than neighborhood edges
    graph.removeLongEdges();
  }

  // if (verbose > 1): visualize MST with removed edges in different color
  if(verbose>1) {
    graph.generateGnuplot(true, "debug-cluster-" + std::to_string(cluster_id));
  }

  if(compound) {
    segments = graph.assignNewCurveIds();

    // correct oversegmentation with extrapolation by orthogonal_lsq
    correctOverSegmentation(cloud, segments);

    // extrapolate curves to start (points belong to more than one curve?)
    expandStart(cloud, segments);

    // adjust cluster id labels
    writeIDsInCloud(cloud, segments, next_label, cluster_id);

  }
  return nchanged;
}




/**
 * @brief Corrects Oversegmentation after removing edges from the MST.
 *
 * For every Segment:
 *
 * Compute PCA for last 4th of the segment. Then, for every follower segment
 * (meaning the last index (+tolerance) of the leader is greater than the first index of the follower),
 * calculate the orthogonal distance from the first point to the previously calculated PCA.
 *
 * The leader segment is assigned the index of the follower with the minimum distance.
 * If there are no followers the segment will not change.
 *
 * Note: The first index is the earliest in time.
 *
 * Note: If a segment only has followers with a distance greater than max_distance, the followers will be ignored.
 *
 * @param cloud Pointcloud from triplclust.
 * @return assignedIds: Contains the segments and will be modified.
 * @return int Number of Segments that were removed in this step.
 */
int correctOverSegmentation(PointCloud &cloud, std::vector< std::vector<size_t> > & assignedIds) {
  if(assignedIds.size() <= 1) {
    return 0;
  }
  int nchanged = assignedIds.size();
  int size_of_cluster=0;
  PointCloud pc;
  Point a;
  Point b;
  unsigned int max_cluster_id = assignedIds.size()-1;
  size_t tolerance = 10; //10
  double min_lsq_distance=std::numeric_limits<float>::max();
  double lsq_distance=0;
  unsigned int min_lsq_distance_cluster_id=max_cluster_id+1;
  std::vector<size_t> result;
  double max_distance = 10;


  for(unsigned int cluster_id = 0; cluster_id <= max_cluster_id; cluster_id++) {
    pc.clear();
    min_lsq_distance = std::numeric_limits<float>::max();
    lsq_distance=0;
    min_lsq_distance_cluster_id=max_cluster_id+1;

    size_of_cluster = assignedIds[cluster_id].size();

    // get the last 4th of the points in the cluster and push onto pc
    for(int i = (int)(3*size_of_cluster/4); i<size_of_cluster; i++) {
      pc.push_back(cloud[assignedIds[cluster_id][i]]);
    }
    // calc PCA
    orthogonal_lsq(pc, &a, &b);

    for(int i = (int)(3*size_of_cluster/4); i<size_of_cluster; i++) {
      max_distance += orthogonal_lsq_distance(&a, &b, &pc.front());
    }

    max_distance = 20*(max_distance / ((size_of_cluster-1)-(int)(3*size_of_cluster/4)));


    // get all possible follower segments and get min
    for(unsigned int i = 0; i <= max_cluster_id; i++) {
      if(i != cluster_id) {
        if((assignedIds[cluster_id].back() + tolerance) > (assignedIds[i].front())) {
          lsq_distance = orthogonal_lsq_distance(&a,&b,&cloud[assignedIds[i].front()]);
          if((lsq_distance < max_distance) && (lsq_distance < min_lsq_distance)) {
            min_lsq_distance = lsq_distance;
            min_lsq_distance_cluster_id = i;
          }
        }
      }
    }
    if(min_lsq_distance_cluster_id != (max_cluster_id+1)) nchanged++;
    result.push_back(min_lsq_distance_cluster_id);
  }

  writeClusterID(result, assignedIds);

  if((assignedIds.size() == 0) || ((assignedIds.size()-1) > max_cluster_id)) {
    std::cout<<"Something has gone wrong while trying to correct oversegmentation"<<std::endl;
    exit(1);
  }

  return nchanged-assignedIds.size();
}


/**
 * @brief Extends other segments into the first segment, if appropriate.
 *
 * For every segment that isn't the first segment (Cluster 0):
 *
 * * Calculate first Component of PCA G for first fifth of segment. The position vector is the middle/average.
 * * Calculate the orthogonal distances of the PCA to the points. Calculate mean and
 * * standard deviation of the distances.
 * *
 * *
 * * For every Point in Cluster 0 that precedes the follower segment:
 * * * Calculate orthogonal distance to G
 * * * if distance in [mean-2*standardDeviation, mean+2*standardDeviation]
 * * * * assign Point to Cluster in the follower segment
 *
 * note: if a point in cluster 0 is assigned to another cluster, the point will be in both cluster 0 and in the new cluster!
 *
 * @param cloud Pointcloud from triplclust.
 * @return assignedIds: Contains the segments and will be modified.
 * @return int Returns Number of individual points from Cluster 0 that were assigned to another segment.
 */
int expandStart(PointCloud &cloud, std::vector< std::vector<size_t> > & assignedIds) {

  if(assignedIds.size() <= 1) {
    return 0;
  }

  int nchanged = 0;
  int size_of_cluster=0;
  PointCloud pc;
  Point a;
  Point a_middle;
  Point b;
  Point b0;
  std::vector<double> distances;
  double mean = 0.0;
  double standardDeviation = 0.0;
  double upperBound = 0.0; // mean + 2*sd
  double lowerBound = 0.0; // mean - 2*sd
  double current_distance = 0.0;

  // For every cluster that isnt the first cluster:
  for(unsigned int cluster_id = 1; cluster_id <= assignedIds.size()-1; cluster_id++) {
    if(assignedIds[cluster_id].size() == 0) break;
    pc.clear();
    size_of_cluster = assignedIds[cluster_id].size();
    // Add every Point in first 5th of cluster to pc.
    for(int i = 0; i<(int)(size_of_cluster/5); i++) {
      pc.push_back(cloud[assignedIds[cluster_id][i]]);
    }
    // calc PCA
    orthogonal_lsq(pc, &a, &b);

    Point p;
    for (std::vector<Point>::iterator iterator = pc.begin(); iterator != pc.end(); ++iterator) {
      p = *iterator;
      distances.push_back(orthogonal_lsq_distance(&a,&b,&p));
    }

    // avoid division by zero
    if(distances.size() <= 1) {
      std::cout<<"Empty distances\n";
      break;
    }

    double d;
    for (std::vector<double>::iterator iterator = distances.begin(); iterator != distances.end(); ++iterator) {
      d = *iterator;
      mean += d;
    }

    mean /= distances.size();

    for (std::vector<double>::iterator iterator = distances.begin(); iterator != distances.end(); ++iterator) {
      d = *iterator;
      standardDeviation += ((d - mean)*(d - mean));
    }

    standardDeviation *= 1/((double)distances.size()-1);
    standardDeviation = sqrt(standardDeviation);

    upperBound = mean+(2*standardDeviation);
    lowerBound = mean-(2*standardDeviation);



    for(unsigned int i = 0; i<assignedIds[0].size(); i++) {
      if(assignedIds[0][i] > assignedIds[cluster_id][0]) {
        current_distance = orthogonal_lsq_distance(&a, &b, &cloud[assignedIds[0][i]]);
        if((current_distance <= upperBound)    &&    (current_distance >= lowerBound) ) {
            assignedIds[cluster_id].push_back(assignedIds[0][i]);

            //!!!!!! uncomment following lines to erase Indexes from Cluster 0 when they are assigned to another Cluster.

            //assignedIds[0].erase(assignedIds[0].begin() + i);
            //i--;
            //size_of_cluster0--;
            nchanged++;
        }
      }
    }
  }
  return nchanged;
}


/**
 * @brief Adds oversegmented Clusters together and modifies assignedIds appropriately. Every index in correct corresponds to an original cluster and the content assigns this original cluster to another one (or none).
 *
 * Helper function of correctOversegmentation(). Checks every element in param correct. If it is not
 * assignedIds.size() then correctOversegmenation() has deemed it a part of another segment.
 * The corresponding Vector (by index) in assignedIds is written to the index specified in correct.
 * Finally, empty vectors in assignedIds are removed.
 * Also keeps track of Cluster 0 and ensures that Cluster 0 stays in front of assignedIds.
 *
 * @param correct:
 * @returns assignedIds
 */
void writeClusterID(const std::vector<size_t> & correct, std::vector< std::vector<size_t> > & assignedIds) {

  if(correct.size()>assignedIds.size()) {
    // dont correct anything
    std::cout<<"writeClusterID. Invalid input, Oversegmentation will not be corrected.\n";
    return;
  }

  unsigned int zero_location=0;

  for(unsigned int i = 0; i<correct.size(); i++) {
    if(correct[i] < assignedIds.size()) {
      if((i == 0) || (i == zero_location)) zero_location = correct[i];
      if(assignedIds[correct[i]].front() > assignedIds[i].front()) {
        assignedIds[correct[i]].insert( assignedIds[correct[i]].end(), assignedIds[i].begin(), assignedIds[i].end() );
      } else {
        assignedIds[correct[i]].insert( assignedIds[correct[i]].begin(), assignedIds[i].begin(), assignedIds[i].end() );
      }
      assignedIds[i].clear();
    }
  }

  if(zero_location < assignedIds.size()) {
    std::vector< std::vector<size_t> >::iterator it = assignedIds.begin() + zero_location;
    std::rotate(assignedIds.begin(), it, it + 1);
  } else {
    std::cout<<"writeClusterID: zero_location impossible.\n";
  }

  // remove all items that are empty vectors = all vectors that were appended to another vector due to being oversegmented
  std::vector< std::vector<size_t> > help = assignedIds;
  assignedIds.clear();
  for (size_t i=0; i<help.size(); i++) {
    if (!help[i].empty()) assignedIds.push_back(help[i]);
  }

}

/**
 * @brief Assigns the final Cluster IDs to the pointcloud. Increments next_label as needed.
 * Assigns cluster 0 the old cluster ID.
 * @param cloud Output
 * @param assignedIds Vector in which every element represents one new cluster identified by the indexes from the original pointcloud.
 * @param next_label Next Label to be assigned to a new cluster.
 */
void writeIDsInCloud(PointCloud &cloud, const std::vector< std::vector<size_t> > & assignedIds, size_t* next_label, size_t old_cluster_id) {
  // skip cluster 0, so that cluster 0 keeps old cluster id
  for(size_t cluster_id = 1; cluster_id<assignedIds.size(); cluster_id++) {
    for(unsigned int i = 0; i < assignedIds[cluster_id].size(); i++) {
      cloud[assignedIds[cluster_id][i]].cluster_ids.erase(old_cluster_id);
      cloud[assignedIds[cluster_id][i]].cluster_ids.insert(*next_label);
    }
    (*next_label)++;
  }
}
