// Copyright (c) 2013-2020, SIB - Swiss Institute of Bioinformatics and
//                          Biozentrum - University of Basel
// 
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
// 
//   http://www.apache.org/licenses/LICENSE-2.0
// 
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.


#include <promod3/core/cluster.hh>
#include <promod3/core/message.hh>

#include <queue>

namespace {

struct Cluster {

  Cluster() { }

  void Merge(const Cluster& other){
    members.insert(members.end(), other.members.begin(), other.members.end());
  }

  std::vector<uint> members;
};

// gives avg distance between two clusters
Real ClusterDistanceAVG(const Cluster& one, 
                        const Cluster& two, 
                        const ost::TriMatrix<Real>& dist_matrix) {
  typedef std::vector<uint>::const_iterator MyIter;
  Real sum  = 0.0;
  for (MyIter i = one.members.begin(); i != one.members.end(); ++i) {
    for (MyIter j = two.members.begin(); j != two.members.end(); ++j) {
      sum += dist_matrix.Get(*i, *j);
    }
  }
  return sum / (one.members.size() * two.members.size());
}

// gives max distance between two clusters
Real ClusterDistanceMAX(const Cluster& one, 
                        const Cluster& two, 
                        const ost::TriMatrix<Real>& dist_matrix) {
  typedef std::vector<uint>::const_iterator MyIter;
  Real max_dist = -std::numeric_limits<Real>::max();
  Real current_dist;
  for (MyIter i = one.members.begin(); i != one.members.end(); ++i) {
    for (MyIter j = two.members.begin(); j != two.members.end(); ++j) {
      current_dist = dist_matrix.Get(*i, *j);
      if (current_dist > max_dist) max_dist = current_dist;
    }
  }
  return max_dist;
}

// gives min distance between two clusters
Real ClusterDistanceMIN(const Cluster& one, 
                        const Cluster& two, 
                        const ost::TriMatrix<Real>& dist_matrix) {
  typedef std::vector<uint>::const_iterator MyIter;
  Real min_dist = std::numeric_limits<Real>::max();
  Real current_dist;
  for (MyIter i = one.members.begin(); i != one.members.end(); ++i) {
    for (MyIter j = two.members.begin(); j != two.members.end(); ++j) {
      current_dist = dist_matrix.Get(*i, *j);
      if (current_dist < min_dist) min_dist = current_dist;
    }
  }
  return min_dist;
}

// returns whether actual val is smaller than extreme val
bool NegativeCompare(Real actual_val, Real extreme_val) {
  return actual_val < extreme_val;
}

// returns whether actual val is larger than extreme val
bool PositiveCompare(Real actual_val, Real extreme_val) {
  return actual_val > extreme_val;
}

struct QueueEntry{

  QueueEntry() { }
  QueueEntry(int a, int b, Real d): first(a), second(b), dist(d) { }

  int first;
  int second;
  Real dist;
};

// Similar compare functions as above but specifically for the priority queue
// The operators are the other way around than above! this is because of the 
// sorting behaviour of the queue
bool NegativeCompareQueueEntries(const QueueEntry& a, const QueueEntry& b){
  return a.dist > b.dist; 
}

bool PositiveCompareQueueEntries(const QueueEntry& a, const QueueEntry& b){
  return a.dist < b.dist; 
}

} // anon ns


namespace promod3 { namespace core {

void DoClustering(const ost::TriMatrix<Real>& dist_matrix, 
                  DistanceMetric metric,
                  Real thresh, char direction, 
                  std::vector<std::vector<uint> >& output){

  output.clear();

  uint num = dist_matrix.GetSize();

  // do the stupid cases
  if(num == 0){
    return;
  }

  if(num == 1){
    std::vector<uint> super_big_cluster(1,0);
    output.push_back(super_big_cluster);
  }

  //let's get a function pointer to compare according to direction
  typedef bool (*cmp_f_ptr)(Real actual_val, Real other_val);
  cmp_f_ptr cmp_ptr;

  if (direction == '+') {
    cmp_ptr = &PositiveCompare;
  } else if (direction == '-') {
    cmp_ptr = &NegativeCompare;
  } else {
    throw promod3::Error("Direction must be '+' or '-'");
  }

  //let's get a function ptr to the right metric function
  typedef Real (*metric_f_ptr)(const Cluster& one, 
                               const Cluster& two, 
                               const ost::TriMatrix<Real>& dist_matrix);
  metric_f_ptr metric_ptr;

  switch (metric) {
    case AVG_DIST_METRIC: {
      metric_ptr = &ClusterDistanceAVG;
      break;
    }
    case MAX_DIST_METRIC: {
      metric_ptr = &ClusterDistanceMAX;
      break;
    }
    case MIN_DIST_METRIC: {
      metric_ptr = &ClusterDistanceMIN;
      break;
    }
    default: {
      metric_ptr = &ClusterDistanceAVG;
    }
  }

  // in agglomerative clustering we alway merge two clusters, we can therefore
  // represent this process using a binary tree, that has 2*num-1 nodes.
  // This is the upper limit for us because of the thresh value.
  // so we might stop merging before building up the whole tree
  // The idea is to build up sub trees from the bottom up and only
  // keep their roots (clusters that can be merged)
  std::vector<Cluster*> clusters(2*num - 1, NULL);
  for (uint i = 0; i < num; ++i) {
    clusters[i] = new Cluster;
    clusters[i]->members.push_back(i);
  }

  // generate a priority queue containing the distances between clusters
  typedef bool (*queue_cmp_f_ptr)(const QueueEntry&, const QueueEntry&);
  queue_cmp_f_ptr queue_cmp_ptr;
  if(direction == '+'){
    queue_cmp_ptr = &PositiveCompareQueueEntries;
  }
  else{
    queue_cmp_ptr = &NegativeCompareQueueEntries;
  } 
  std::priority_queue<QueueEntry,std::vector<QueueEntry>,queue_cmp_f_ptr> 
  distance_queue(queue_cmp_ptr);

  for(uint i = 0; i < num; ++i){
    for(uint j = i + 1; j < num; ++j){
      distance_queue.push(QueueEntry(i, j, dist_matrix.Get(i, j)));
    }
  }

  uint merge_iteration = 0;

  // there is a maximum of num - 1 merge events
  // after that, everything would be in the same cluster
  for( ; merge_iteration < num - 1; ++merge_iteration) {

    // search in priority queue until we get an element containing two
    // root clusters. Note, that we also find distance entries associated to 
    // clusters that have already been merged (and therefore deleted). 
    // We just ignore them and continue the search
    // Important: All clusters are merged into one cluster after num-1 merge steps. 
    // We're therefore guaranteed to find two clusters to merge at this point. 
    while((clusters[distance_queue.top().first] == NULL || 
           clusters[distance_queue.top().second] == NULL)){
      distance_queue.pop();
    }

    int cluster_a = distance_queue.top().first;
    int cluster_b = distance_queue.top().second;
    Real actual_distance = distance_queue.top().dist;
    distance_queue.pop();

    if(!(*cmp_ptr)(actual_distance, thresh)){
      // don't merge clusters anymore, they're all above/below thresh
      break;
    }

    // move cluster a to new location and merge cluster b into it
    int new_cluster_idx = num + merge_iteration;
    clusters[new_cluster_idx] = new Cluster(*clusters[cluster_a]);
    clusters[new_cluster_idx]->Merge(*clusters[cluster_b]);

    // delete the old clusters (we only keep the roots!)
    delete clusters[cluster_a];
    delete clusters[cluster_b];
    clusters[cluster_a] = NULL;
    clusters[cluster_b] = NULL;

    // calculate the distance from new cluster to all other root clusters
    for(int i = 0; i < new_cluster_idx; ++i){
      if(clusters[i] != NULL){
        Real distance = (*metric_ptr)(*clusters[i], *clusters[new_cluster_idx], 
                                      dist_matrix);
        distance_queue.push(QueueEntry(i, new_cluster_idx, distance));
      }
    }
  }

  //fill the output and clean up
  for (uint i = 0; i < num + merge_iteration; ++i) {
    if(clusters[i] != NULL){
      output.push_back(clusters[i]->members);
      delete clusters[i];
    }
  }
}

}} // ns
