// ivector/agglomerative-clustering.h
  
  // 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.
  
  #ifndef KALDI_IVECTOR_AGGLOMERATIVE_CLUSTERING_H_
  #define KALDI_IVECTOR_AGGLOMERATIVE_CLUSTERING_H_
  
  #include <vector>
  #include <queue>
  #include <set>
  #include <unordered_map>
  #include <functional>
  #include "base/kaldi-common.h"
  #include "matrix/matrix-lib.h"
  #include "util/stl-utils.h"
  
  namespace kaldi {
  
  /// AhcCluster is the cluster object for the agglomerative clustering. It
  /// contains three integer IDs: its own ID and the IDs of its "parents", i.e.
  /// the clusters that were merged to form it. It also contains the size (the
  /// number of points in the cluster) and a vector of the IDs of the utterances
  /// contained in the cluster.
  struct AhcCluster {
    int32 id,
      parent1,
      parent2,
      size;
    std::vector<int32> utt_ids;
    AhcCluster(int32 id, int32 p1, int32 p2, std::vector<int32> utts)
        : id(id), parent1(p1), parent2(p2), utt_ids(utts) {
      size = utts.size();
    }
  };
  
  /// The AgglomerativeClusterer class contains the necessary mechanisms for the
  /// actual clustering algorithm.
  class AgglomerativeClusterer {
   public:
    AgglomerativeClusterer(
        const Matrix<BaseFloat> &costs,
        BaseFloat threshold,
        int32 min_clusters,
        int32 first_pass_max_points,
        BaseFloat max_cluster_fraction,
        std::vector<int32> *assignments_out)
        : costs_(costs), threshold_(threshold), min_clusters_(min_clusters),
          first_pass_max_points_(first_pass_max_points),
          assignments_(assignments_out) {
      num_points_ = costs.NumRows();
  
      // The max_cluster_size_ is a hard limit on the number points in a cluster.
      // This is useful for handling degenerate cases where some outlier points
      // form their own clusters and force everything else to be clustered
      // together, e.g. when min-clusters is provided instead of a threshold.
      max_cluster_size_ = ceil(num_points_ * max_cluster_fraction);
  
      // The count_, which is used for identifying clusters, is initialized to
      // num_points_ because cluster IDs 1..num_points_ are reserved for input
      // points, which are the initial set of clusters.
      count_ = num_points_;
  
      // The second_pass_count_, which is used for identifying the initial set of
      // second pass clusters and initializing count_ before the second pass, is
      // initialized to 0 and incremented whenever a new cluster is added to the
      // initial set of second pass clusters.
      second_pass_count_ = 0;
    }
  
    // Clusters points. Chooses single pass or two pass algorithm.
    void Cluster();
  
    // Clusters points using single pass algorithm.
    void ClusterSinglePass();
  
    // Clusters points using two pass algorithm.
    void ClusterTwoPass();
  
   private:
    // Encodes cluster pair into a 32bit unsigned integer.
    uint32 EncodePair(int32 i, int32 j);
    // Decodes cluster pair from a 32bit unsigned integer.
    std::pair<int32, int32> DecodePair(uint32 key);
    // Initializes the clustering queue with singleton clusters
    void InitializeClusters(int32 first, int32 last);
    // Does hierarchical agglomerative clustering
    void ComputeClusters(int32 min_clusters);
    // Adds clusters created in first pass to second pass clusters
    void AddClustersToSecondPass();
    // Assigns points to clusters
    void AssignClusters();
    // Merges clusters with IDs i and j and updates cost map and queue
    void MergeClusters(int32 i, int32 j);
  
    const Matrix<BaseFloat> &costs_;  // cost matrix
    BaseFloat threshold_;  // stopping criterion threshold
    int32 min_clusters_;  // minimum number of clusters
    int32 first_pass_max_points_;  // maximum number of points in each subset
    std::vector<int32> *assignments_;  // assignments out
  
    int32 num_points_;  // total number of points to cluster
    int32 max_cluster_size_;  // maximum number of points in a cluster
    int32 count_;  // count of first pass clusters, used for identifying clusters
    int32 second_pass_count_;  // count of second pass clusters
  
    // Priority queue using greater (lowest costs are highest priority).
    // Elements contain pairs of cluster IDs and their cost.
    typedef std::pair<BaseFloat, uint32> QueueElement;
    typedef std::priority_queue<QueueElement, std::vector<QueueElement>,
      std::greater<QueueElement>  > QueueType;
    QueueType queue_, second_pass_queue_;
  
    // Map from cluster IDs to cost between them
    std::unordered_map<uint32, BaseFloat> cluster_cost_map_;
    // Map from cluster ID to cluster object address
    std::unordered_map<int32, AhcCluster*> clusters_map_;
    // Set of unmerged cluster IDs
    std::set<int32> active_clusters_;
  
    // Map from second pass cluster IDs to cost between them
    std::unordered_map<uint32, BaseFloat> second_pass_cluster_cost_map_;
    // Map from second pass cluster ID to cluster object address
    std::unordered_map<int32, AhcCluster*> second_pass_clusters_map_;
    // Set of unmerged second pass cluster IDs
    std::set<int32> second_pass_active_clusters_;
  };
  
  /** This is the function that is called to perform the agglomerative
   *  clustering. It takes the following arguments:
   *   - A matrix of all pairwise costs, with each row/column corresponding
   *      to an utterance ID, and the elements of the matrix containing the
   *      cost for pairing the utterances for its row and column
   *   - A threshold which is used as the stopping criterion for the clusters
   *   - A minimum number of clusters that will not be merged past
   *   - A maximum fraction of points that can be in a cluster
   *   - A vector which will be filled with integer IDs corresponding to each
   *      of the rows/columns of the score matrix.
   *
   *  The basic algorithm is as follows:
   *  \code
   *      while (num-clusters > min-clusters && smallest-merge-cost <= threshold)
   *          if (size-of-new-cluster <= max-cluster-size)
   *              merge the two clusters with lowest cost
   *  \endcode
   *
   *  The cost between two clusters is the average cost of all pairwise
   *  costs between points across the two clusters.
   *
   *  The algorithm takes advantage of the fact that the sum of the pairwise
   *  costs between the points of clusters I and J is equiavlent to the
   *  sum of the pairwise costs between cluster I and the parents of cluster
   *  J. In other words, the total cost between I and J is the sum of the
   *  costs between clusters I and M and clusters I and N, where
   *  cluster J was formed by merging clusters M and N.
   *
   *  If the number of points to cluster is larger than first-pass-max-points,
   *  then clustering is done in two passes. In the first pass, input points are
   *  divided into contiguous subsets of size at most first-pass-max-points and
   *  each subset is clustered separately. In the second pass, the first pass
   *  clusters are merged into the final set of clusters.
   *
   */
  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);
  
  }  // end namespace kaldi.
  
  #endif  // KALDI_IVECTOR_AGGLOMERATIVE_CLUSTERING_H_