// ivector/agglomerative-clustering.cc

// Copyright  2017-2018  Matthew Maciejewski
//                 2018  David Snyder
//                 2019  Dogan Can

// See ../../COPYING for clarification regarding multiple authors
//
// 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
//
// THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED
// WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE,
// MERCHANTABLITY OR NON-INFRINGEMENT.
// See the Apache 2 License for the specific language governing permissions and
// limitations under the License.

#include <algorithm>
#include "ivector/agglomerative-clustering.h"

namespace kaldi {

void AgglomerativeClusterer::Cluster() {
  if (num_points_ > first_pass_max_points_)
    ClusterTwoPass();
  else
    ClusterSinglePass();
}

void AgglomerativeClusterer::ClusterSinglePass() {
  InitializeClusters(0, num_points_);
  ComputeClusters(min_clusters_);
  AssignClusters();
}

void AgglomerativeClusterer::ClusterTwoPass() {
  // This is the first pass loop. We divide the input into equal size subsets
  // making sure each subset has at most first_pass_max_points_ points. Then, we
  // cluster the points in each subset separately until a stopping criterion is
  // reached. We set the minimum number of clusters to 10 * min_clusters_ for
  // each subset to avoid early merging of most clusters that would otherwise be
  // kept separate in single pass clustering.
  BaseFloat num_points = static_cast<BaseFloat>(num_points_);
  int32 num_subsets = ceil(num_points / first_pass_max_points_);
  int32 subset_size = ceil(num_points / num_subsets);
  for (int32 n = 0; n < num_points_; n += subset_size) {
    InitializeClusters(n, std::min(n + subset_size, num_points_));
    ComputeClusters(min_clusters_ * 10);
    AddClustersToSecondPass();
  }

  // We swap the contents of the first and second pass data structures so that
  // we can use the same method to do second pass clustering.
  clusters_map_.swap(second_pass_clusters_map_);
  active_clusters_.swap(second_pass_active_clusters_);
  cluster_cost_map_.swap(second_pass_cluster_cost_map_);
  queue_.swap(second_pass_queue_);
  count_ = second_pass_count_;

  // This is the second pass. It moves through the queue merging clusters
  // determined in the first pass until a stopping criterion is reached.
  ComputeClusters(min_clusters_);

  AssignClusters();
}

uint32 AgglomerativeClusterer::EncodePair(int32 i, int32 j) {
  if (i < j)
    return (static_cast<uint32>(i) << 16) + static_cast<uint32>(j);
  else
    return (static_cast<uint32>(j) << 16) + static_cast<uint32>(i);
}

std::pair<int32, int32> AgglomerativeClusterer::DecodePair(uint32 key) {
  return std::make_pair(static_cast<int32>(key >> 16),
                        static_cast<int32>(key & 0x0000FFFFu));
}

void AgglomerativeClusterer::InitializeClusters(int32 first, int32 last) {
  KALDI_ASSERT(last > first);
  clusters_map_.clear();
  active_clusters_.clear();
  cluster_cost_map_.clear();
  queue_ = QueueType();  // priority_queue does not have a clear method

  for (int32 i = first; i < last; i++) {
    // create an initial cluster of size 1 for each point
    std::vector<int32> ids;
    ids.push_back(i);
    AhcCluster *c = new AhcCluster(i + 1, -1, -1, ids);
    clusters_map_[i + 1] = c;
    active_clusters_.insert(i + 1);

    // propagate the queue with all pairs from the cost matrix
    for (int32 j = i + 1; j < last; j++) {
      BaseFloat cost = costs_(i, j);
      uint32 key = EncodePair(i + 1, j + 1);
      cluster_cost_map_[key] = cost;
      if (cost <= threshold_)
        queue_.push(std::make_pair(cost, key));
    }
  }
}

void AgglomerativeClusterer::ComputeClusters(int32 min_clusters) {
  while (active_clusters_.size() > min_clusters && !queue_.empty()) {
    std::pair<BaseFloat, uint32> pr = queue_.top();
    int32 i, j;
    std::tie(i, j) = DecodePair(pr.second);
    queue_.pop();
    // check to make sure clusters have not already been merged
    if ((active_clusters_.find(i) != active_clusters_.end()) &&
        (active_clusters_.find(j) != active_clusters_.end())) {
      if (clusters_map_[i]->size + clusters_map_[j]->size <= max_cluster_size_)
        MergeClusters(i, j);
    }
  }
}

void AgglomerativeClusterer::MergeClusters(int32 i, int32 j) {
  AhcCluster *clust1 = clusters_map_[i];
  AhcCluster *clust2 = clusters_map_[j];
  // For memory efficiency, the first cluster is updated to contain the new
  // merged cluster information, and the second cluster is later deleted.
  clust1->id = ++count_;
  clust1->parent1 = i;
  clust1->parent2 = j;
  clust1->size += clust2->size;
  clust1->utt_ids.insert(clust1->utt_ids.end(), clust2->utt_ids.begin(),
                         clust2->utt_ids.end());
  // Remove the merged clusters from the list of active clusters.
  active_clusters_.erase(i);
  active_clusters_.erase(j);
  // Update the queue with all the new costs involving the new cluster
  std::set<int32>::iterator it;
  for (it = active_clusters_.begin(); it != active_clusters_.end(); ++it) {
    // The new cost is the sum of the costs of the new cluster's parents
    BaseFloat new_cost = cluster_cost_map_[EncodePair(*it, i)] +
                         cluster_cost_map_[EncodePair(*it, j)];
    uint32 new_key = EncodePair(*it, count_);
    cluster_cost_map_[new_key] = new_cost;
    BaseFloat norm = clust1->size * (clusters_map_[*it])->size;
    if (new_cost / norm <= threshold_)
      queue_.push(std::make_pair(new_cost / norm, new_key));
  }
  active_clusters_.insert(count_);
  clusters_map_[count_] = clust1;
  delete clust2;
}

void AgglomerativeClusterer::AddClustersToSecondPass() {
  // This method collects the results of first pass clustering for one subset,
  // i.e. adds the set of active clusters to the set of second pass active
  // clusters and computes the costs for the newly formed cluster pairs.
  std::set<int32>::iterator it1, it2;
  int32 count = second_pass_count_;
  for (it1 = active_clusters_.begin(); it1 != active_clusters_.end(); ++it1) {
    AhcCluster *clust1 = clusters_map_[*it1];
    second_pass_clusters_map_[++count] = clust1;

    // Compute new cluster pair costs
    for (it2 = second_pass_active_clusters_.begin();
         it2 != second_pass_active_clusters_.end(); ++it2) {
      AhcCluster *clust2 = second_pass_clusters_map_[*it2];
      uint32 new_key = EncodePair(count, *it2);

      BaseFloat new_cost = 0.0;
      std::vector<int32>::iterator utt_it1, utt_it2;
      for (utt_it1 = clust1->utt_ids.begin();
           utt_it1 != clust1->utt_ids.end(); ++utt_it1) {
         for (utt_it2 = clust2->utt_ids.begin();
              utt_it2 != clust2->utt_ids.end(); ++utt_it2) {
           new_cost += costs_(*utt_it1, *utt_it2);
         }
      }

      second_pass_cluster_cost_map_[new_key] = new_cost;
      BaseFloat norm = clust1->size * clust2->size;
      if (new_cost / norm <= threshold_)
        second_pass_queue_.push(std::make_pair(new_cost / norm, new_key));
    }

    // Copy cluster pair costs that were already computed in the first pass
    int32 count2 = second_pass_count_;
    for (it2 = active_clusters_.begin(); it2 != it1; ++it2) {
      uint32 key = EncodePair(*it1, *it2);
      BaseFloat cost = cluster_cost_map_[key];
      BaseFloat norm = clust1->size * (clusters_map_[*it2])->size;
      uint32 new_key = EncodePair(count, ++count2);
      second_pass_cluster_cost_map_[new_key] = cost;
      if (cost / norm <= threshold_)
        second_pass_queue_.push(std::make_pair(cost / norm, new_key));
    }
  }
  // We update second_pass_count_ and second_pass_active_clusters_ here since
  // above loop assumes they do not change while the loop is running.
  while (second_pass_count_ < count)
    second_pass_active_clusters_.insert(++second_pass_count_);
}

void AgglomerativeClusterer::AssignClusters() {
  assignments_->resize(num_points_);
  int32 label_id = 0;
  std::set<int32>::iterator it;
  // Iterate through the clusters and assign all utterances within the cluster
  // an ID label unique to the cluster. This is the final output and frees up
  // the cluster memory accordingly.
  for (it = active_clusters_.begin(); it != active_clusters_.end(); ++it) {
    ++label_id;
    AhcCluster *cluster = clusters_map_[*it];
    std::vector<int32>::iterator utt_it;
    for (utt_it = cluster->utt_ids.begin();
         utt_it != cluster->utt_ids.end(); ++utt_it)
      (*assignments_)[*utt_it] = label_id;
    delete cluster;
  }
}

void AgglomerativeCluster(
    const Matrix<BaseFloat> &costs,
    BaseFloat threshold,
    int32 min_clusters,
    int32 first_pass_max_points,
    BaseFloat max_cluster_fraction,
    std::vector<int32> *assignments_out) {
  KALDI_ASSERT(min_clusters >= 0);
  KALDI_ASSERT(max_cluster_fraction >= 1.0 / min_clusters);
  AgglomerativeClusterer ac(costs, threshold, min_clusters,
                            first_pass_max_points, max_cluster_fraction,
                            assignments_out);
  ac.Cluster();
}

}  // end namespace kaldi.