package yi.stat.kmeans;

import java.util.*;

/**
 * A new implmentation of K-Means (replace yi.ai.cluster.KMCluster).  The last
 * version uses integer, while the current version will use double.  Other
 * differeneces include:
 *
 * <ul>
 *   <li>
 *   scaling on each dimension leads to different meanings of Eulidean distances
 *   of two multiple-dimensional data.  So scaling the data in each dimension to
 *   1 in the last implementation is not alwasys a good choice.  This new
 *   implementation allows the user to specify the "scales" on each dimension
 *   </li>
 *   <li>
 *   add a new restart method to restart K-Means several times to get the best
 *   clustering since the results depend on the initial condition a lot.
 *   </li>
 * </ul>
 */
public class KMeans
    implements java.io.Serializable {

  public boolean OUTPUT_RESTART = false;

  Random r = new Random();

  /**
   * number of cluster after clustering (may be different from the specified K)
   */
  public int k = 0;

  /**
   * classification result for each data point
   */
  public int[] label = null;

  /**
   * Resultant means
   */
  public double[][] means = null;

  /**
   * classification score: sum of squared Euclidean distances between data
   * points and their closest means
   */
  public double score = 0;

  public KMeans() {
  }

  /**
   * Initialize with specified random seed (useful in debugging)
   * @param r specified random seed (useful in debugging)
   */
  public KMeans(Random r) {
    this.r = r;
  }

  /**
   * Return k-means clustering result.  Every data is given a label
   * @return the data->cluster index map
   */
  public int[] getLabels() {
    return this.label;
  }

  /**
   * Return the resultant cluster number
   * @return the resultant cluster number
   */
  public int getClusterNum() {
    return k;
  }


  /**
   * Clustering data uses k-means with "restart_num" restart.
   *
   * Since Euclidean distance is used to calculate the distance between data
   * points, it is important that data values on different dimension are scaled
   * properly.  The data will be re-scaled as specified in "scales".  If "null"
   * is passed in, then no rescaling will be performed.
   * @param k as in "K"-means
   * @param d data dimension
   * @param data the data to be clustered.  First dimension is the different
   *             data points and the second dimension is different components
   *             of a data point (as for multiple-dimensional data)
   * @param scales The data will be re-scaled as specified in "scales".  If
   *               "null" is passed in, then no rescaling will be performed.
   * @param restart_num the K-Means algorithm runs "restart_num" times and the
   *                    the best result is kept
   * @param max_step the K-Means will stop after "max_step" iterations.  If a
   *                 non-positive value is passed in, the K-means will not stop
   *                 until converged.
   * @param no_empty_label remove empty labels in the results
   * @return the cluster label for each data.  <code>lable[x] = y </code> means
   *         the x-th data point is classified to the y-th cluster
   */
  public int[] clustering(int k, int d, double[][] data, double[] scales,
                          int restart_num, int max_step, boolean no_empty_label) {

    // rescale the data on each dimension
    double[][] data_copy = yi.util.data.CloneUtil.clone(data);
    if (scales != null) {
      for (int i = 0; i < data.length; i++) {
        for (int j = 0; j < d; j++) {
          data_copy[i][j] *= scales[j];
        }
      }
    }

    double[] maximums = yi.stat.StatUtil.maximumOnCol(data_copy);
    double[] minimums = yi.stat.StatUtil.minimumOnCol(data_copy);

    // run "restart_num" times
    double best_score = Double.MAX_VALUE;
    int[] best_labels = null;
    double[][] best_means = null;

    double[] scores = new double[restart_num];
    for (int times = 0; times < restart_num; times++) {

      // initialize the means
      double[][] means = initializeMeans(k, d, minimums, maximums);

      // iteration
      int[] labels = iteration(data_copy, k, means, max_step);
      double score = clusteringScore(data_copy, labels, means);

      scores[times] = score;

      if (score < best_score) {
        best_score = score;
        best_labels = labels;
        best_means = means;
      }
    }

    if (OUTPUT_RESTART) {
      System.out.print(k);
      for (int i=0; i<scores.length; i++) System.out.print(" " + scores[i]);
      System.out.println(" " + best_score);
    }

    // scale back the means
    if (scales != null) {
      for (int i = 0; i < best_means.length; i++) {
        for (int j = 0; j < d; j++) {
          best_means[i][j] /= scales[j];
        }
      }
    }

    // remember the results
    this.k = k;
    this.label = best_labels;
    this.means = best_means;
    this.score = best_score;

    if (no_empty_label) {
      k = eliminateEmptyCluster();
    }

    return this.label;
  }

  /**
   * K-Means clustering on the given data and given means
   * @param data the data to be clustered.  First dimension is the different
   *             data points and the second dimension is different components
   *             of a data point (as for multiple-dimensional data)
   * @param k as in "K"-means
   * @param means inital means
   * @param max_step the K-Means will stop after "max_step" iterations.  If a
   *                 non-positive value is passed in, the K-means will not stop
   *                 until converged.
   * @param no_empty_label remove empty labels in the results
   * @return the cluster label for each data.  <code>lable[x] = y </code> means
   *         the x-th data point is classified to the y-th cluster
   */
  public int[] clustering(double[][] data, int k, double[][] means, int max_step, boolean no_empty_label) {
    return clustering(data, null, k, means, max_step, no_empty_label);
  }

  /**
   * K-Means clustering on the given data and given means
   * @param data the data to be clustered.  First dimension is the different
   *             data points and the second dimension is different components
   *             of a data point (as for multiple-dimensional data)
   * @param scales
   * @param k as in "K"-means
   * @param means inital means
   * @param max_step the K-Means will stop after "max_step" iterations.  If a
   *                 non-positive value is passed in, the K-means will not stop
   *                 until converged.
   * @param no_empty_label remove empty labels in the results
   * @return the cluster label for each data.  <code>lable[x] = y </code> means
   *         the x-th data point is classified to the y-th cluster
   */
  public int[] clustering(double[][] data, double[] scales, int k, double[][] means, int max_step, boolean no_empty_label) {
    // rescale the data on each dimension
    double[][] data_copy = yi.util.data.CloneUtil.clone(data);
    double[][] means_copy = yi.util.data.CloneUtil.clone(means);
    if (scales != null) {
      for (int i = 0; i < data.length; i++) {
        for (int j = 0; j < data[i].length; j++) {
          data_copy[i][j] *= scales[j];
        }
      }
      for (int i = 0; i < means.length; i++) {
        for (int j = 0; j < means[i].length; j++) {
          means_copy[i][j] *= scales[j];
        }
      }
    }

    // iteration
    int[] labels = iteration(data_copy, k, means_copy, max_step);
    double score = clusteringScore(data_copy, labels, means_copy);

    // scale back the means
    if (scales != null) {
      for (int i = 0; i < means_copy.length; i++) {
        for (int j = 0; j < means_copy[i].length; j++) {
          means_copy[i][j] /= scales[j];
        }
      }
    }

    // remember the results
    this.k = k;
    this.label = labels;
    this.means = means_copy;
    this.score = score;

    if (no_empty_label) {
      k = eliminateEmptyCluster();
    }
    return this.label;
  }

  /**
   * Iteration step of K-Means
   * @param data the data to cluster
   * @param k as in K-Means
   * @param means initial means and will be changed to the resultant means
   * @param max_step the k-means will quit after "max_step" iterations.  If a
   *                 non-postive (&lt;=0) value is passed in, then k-means will
   *                 continue until converged.
   * @return the cluster label for each data.  <code>lable[x] = y </code> means
   *         the x-th data point is classified to the y-th cluster
   */
  protected int[] iteration(double[][] data, int k, double[][] means, int max_step) {

    /**
     * size of data
     */
    int size = data.length;

    /**
     * dimension of data vector
     */
    int d = data[0].length; //number of elements of the information vector (hsi, =3)

    /**
     * classification result
     */
    label = new int[data.length];

    int i, j, m, n, temp;
    int ic = 1;
    boolean changed = true;

    //used when re-calculate the means before the iteration step
    double[][] sum = new double[k][d];
    int[] num = new int[k];

    /**
     * step counter
     */
    int steps = 0;
    while (max_step <= 0 || steps < max_step) {
      steps++;
//      System.out.println("step #" + steps);

      changed = false;
      //find the minimum distance for each vector
      for (i = 0; i < size; i++) {
        double min_dist = yi.math.VectorUtil.EuclideanDistance(data[i], means[0]);
        int min_dist_index = 0; //remember the index of the mean with the minimum distance
        for (m = 1; m < k; m++) {
          double dist = yi.math.VectorUtil.EuclideanDistance(data[i], means[m]);
          if (min_dist > dist) {
            min_dist = dist;
            min_dist_index = m;
          }
        }
        // assign the data to the closest cluster
        if (min_dist_index != label[i]) {
          label[i] = min_dist_index;
          changed = true;
        }
      }
      if (!changed)
        break;

      //recalculate the means
      for (m = 0; m < k; m++)
        java.util.Arrays.fill(sum[m], 0);
      java.util.Arrays.fill(num, 0);
      for (i = 0; i < size; i++) {
        temp = label[i];
        for (n = 0; n < d; n++)
          sum[temp][n] += data[i][n];
        num[temp]++;
      }

      for (m = 0; m < k; m++) {
        if (num[m] == 0)
          continue;
        for (n = 0; n < d; n++) {
          means[m][n] = sum[m][n] / num[m];
        }
      }
    }

    return label;
  }

  /**
   * Remove empty label.  If there is no data classified to cluster
   * <code>c</code>, then that cluster will be removed.  The index of cluster
   * greater than <code>c</code> will be decreased by 1.
   * @param label the data cluster label to process
   * @param k the number of clusters
   * @return the new number of label
   */
  protected int eliminateEmptyCluster() {
    boolean[] empty = new boolean[k];
    java.util.Arrays.fill(empty, true);

    // find out which clusters are empty
    for (int i = 0; i < label.length; i++) {
      empty[label[i]] = false;
    }

    // set up the old-label -> new-label map
    int[] map = new int[k];
    int count = 0;
    for (int i = 0; i < empty.length; i++) {
      if (!empty[i]) {
        map[i] = count++;
      }
    }
    k = count;

    // relabel the result
    for (int i = 0; i < label.length; i++) {
      label[i] = map[label[i]];
    }

    // change the resultant means accordingly
    double[][] new_means = new double[k][];
    count = 0;
    for (int i = 0; i < empty.length; i++) {
      if (!empty[i]) {
        new_means[count++] = means[i];
      }
    }
    means = new_means;
    return count;
  }

  /**
   * Generate initial means randomly by unified distribution.  The distribution
   * of the data on each dimension is important for the random sampling.  This
   * method takes a simple approach which assuming the data values on each
   * dimension is unified distributed.  If "minimums" and/or "maximums" are
   * given, then they will be used, or "0"s and/or "1"s will be assumed.
   * @param k as in K-Means
   * @param d data dimension
   * @param minimums minimum value on each data dimension. If 0's, give null
   * @param maximums maximum value on each data dimension. If 1's, give null
   * @return initial means.  First dimension is the data points and the second
   *         dimension is the data dimensions.
   */
  protected double[][] initializeMeans(int k, int d, double[] minimums, double[] maximums) {
    double[][] means = new double[k][d];

    // default value of minimums is 0
    if (minimums == null) {
      minimums = new double[d];
      for (int i = 0; i < d; i++)
        minimums[i] = 0;
    }

    // default value of maximums is 1
    if (maximums == null) {
      maximums = new double[d];
      for (int i = 0; i < d; i++)
        maximums[i] = 1;
    }

    double[] range = new double[d];
    for (int i = 0; i < d; i++)
      range[i] = maximums[i] - minimums[i];

    for (int i = 0; i < k; i++) {
      for (int j = 0; j < d; j++) {
        means[i][j] = r.nextDouble() * range[j] + minimums[j];
      }
    }
    return means;
  }

  /**
   * Return the sum of the squared Euclidean distance between each data point
   * to its classified mean.
   * @param data the data to be clustered
   * @param labels the cluster label for each data
   * @param means the means of each cluster
   * @return the sum of the squared Euclidean distance between each data point
   *         to its classified mean.
   */
  protected double clusteringScore(double[][] data, int[] labels, double[][] means) {
    double dist_sum = 0;
    for (int i = 0; i < data.length; i++) {
      dist_sum +=
          yi.math.VectorUtil.squaredEuclideanDistance(data[i], means[label[i]]);
    }
    return dist_sum;

  }
}