package yi.stat.em;

import yi.stat.kmeans.*;
import yi.stat.hist.*;

/**
 * Create GaussianMixtureModel
 */

public class GaussianMixtureModelFactory {
  public GaussianMixtureModelFactory() {
  }

  /**
   * Create a Gaussian mixture model using EM algorithm.  EM is initialized by
   * several steps of K-Means.  The EM algorithm will be restarted several
   * times and the best result is kept.
   *
   * <p>
   * Since EM is initialized by K-Means, there are some parameters for K-Means
   * algorithm:
   *   <code>scale</code>,
   *   <code>kmeans_restart_times</code>,
   *   and <code>kmeans_steps</code>.
   * </p>
   *
   * <p>
   * since k-means tends to generate empty clusters, to have <code>m</code>
   * valid initial clusters for EM, it is better to use a "k" bigger than the
   * desired "m".  The strategy here is to increase k by "bigger_k_ratio" if the
   * cluster number generated by k-means is smaller than m.  If the cluster
   * number from k-means initial is bigger than m, then only the m biggest
   * cluster will be kept for EM initialization.
   * </p>
   * @param m
   * @param data
   * @param scale since K-Means is sensitive to the scale on different
   *              dimension, scaling on data on different dimension can be
   *              important
   * @param restart_times The EM algorithm will be restarted several times and
   *              the best result is kept.
   * @param gmm_break_threshold the break trehsold for EM algorithm.  See
   *              {@link #BREAK_THRESHOLD}
   * @param kmeans_restart_times k-means restarts "kmeans_restart_times" to
   *              find the best one
   * @param kmeans_steps run k-means for "kmeans_steps" iterations to get the
   *              initialized clusters for EM
   * @param bigger_k_ratio since k-means tends to generate empty clusters, to
   *              have <code>m</code> valid initial clusters for EM, it is
   *              better to use a "k" bigger than the desired "m".  The strategy
   *              here is to increase k by "bigger_k_ratio" if the cluster
   *              number generated by k-means is smaller than m.
   * @param fixer to fix singular covariance matrix
   * @return
   * @see yi.stat.kmeans.KMeans
   */
  public GaussianMixtureModel createGaussianMixtureModel(
      int m, double[][] data, double[] scale,
      int restart_times, double gmm_break_threshold,
      int kmeans_restart_times, int kmeans_steps, double bigger_k_ratio,
      SingularGaussianModelFixable fixer) {

    /** @todo k-means or partition? */
    // k-means
    int[] labels;

    /**
     * If the data is not enough
     */
    if (m>=data.length) {
      m = data.length;
      bigger_k_ratio = 1;
      System.err.println("data is not enough: set bigger_k_ratio = 1");
    }

    int m_copy = m;

    double best_score = 0;
    GaussianMixtureModel best_gmm = null;
    double[] scores = new double[restart_times];
    for (int times = 0; times < restart_times; times++) {
      System.out.println("GMM construction # " + (times + 1));
      System.out.println("data dim = " + data[0].length);

      // set initial k bigger than m
      int k = m_copy;

      while (true) {
        KMeans k_means = new KMeans();

        // initial k-means with bigger k
        labels = k_means.clustering(k, data[0].length, data, scale,
                                    kmeans_restart_times, kmeans_steps, true);

        int ncluster = k_means.getClusterNum();

        // if no required to have bigger k
        if (bigger_k_ratio <= 1.0) {
          m = ncluster;
          break;
        }

        if (ncluster < m) { // if the cluster number generated is less than the
          // desired "m", increase k and re-cluster
          System.out.println("k-means cluster num = " + ncluster +
                             ".  will increase");
//          k = (int) Math.ceil(k * (double)m / (double)ncluster * bigger_k_ratio); // if the generated cluster is less than m, increase k
          k = (int) Math.ceil((double)k * (double)m / (double)ncluster); // if the generated cluster is less than m, increase k
        }
        else if (ncluster > m) { // if the generated clusters are more than
          System.out.println("k-means cluster num = " + ncluster +
                             ".  will use only the biggest ones");
          // desired, use only the biggest ones.
          double[][] means = k_means.means;
          double[][] picked_means = new double[m][];

          // pick up the biggest clusters
          SimpleHistogram hist = new SimpleHistogram(0, ncluster - 1);
          hist.putToBin(labels);
          int[] counts = hist.getHistogram();
          for (int i = 0; i < m; i++) {
            // find the highest now
            int highest = counts[0];
            int highest_index = 0;
            for (int j = 1; j < ncluster; j++) {
              if (counts[j] > highest) {
                highest = counts[j];
                highest_index = j;
              }
            }
            picked_means[i] = means[highest_index];
            // so that it wouldn't be picked again
            counts[highest_index] = 0;
          }

          // do k-means again using the picked means
          KMeans k_means2 = new KMeans();
          // scale the means (since k-means didn't scale means
          labels = k_means2.clustering(data, scale, m, picked_means, 1, false);
          break;
        }
        else {
          break;
        }
      }
      System.out.println("component num = " + m);
      GaussianMixtureModel gmm = createGaussianMixtureModel(m, data, labels);
      gmm.BREAK_THRESHOLD = gmm_break_threshold;
      gmm.setSigularFixer(fixer);
      try {
        double score = gmm.getParameterLogLikelihood(data);
        scores[times] = score;
        if (best_gmm == null || score > best_score) {
          best_score = score;
          best_gmm = gmm;
        }
      } catch (ModelSingularException e) {
        e.printStackTrace();
      }
    }


    yi.util.io.OutputResult.print(System.out, scores);

    return best_gmm;
  }

  /**
   * Create a Gaussian mixture model which uses a splitting and vanishing
   * algorithm during training to learn the number of mixture components.
   * @param data
   * @param mcvr
   * @param mcsr
   * @return
   */
  public GaussianMixtureModel createSplittingAndVanishingGaussianMixtureModel(
      int m, double[][] data, double[] scale,
      int restart_times, double gmm_break_threshold,
      int kmeans_restart_times, int kmeans_steps, double bigger_k_ratio,
      SingularGaussianModelFixable fixer,
      double mcvr, double mcsr) {
    GaussianMixtureModel gmm = createGaussianMixtureModel(m, data, scale,
        restart_times, gmm_break_threshold,
        kmeans_restart_times, kmeans_steps, bigger_k_ratio, fixer);
    gmm.m_initComponents = gmm.m_rstComponents;
    gmm.SPLITTING_AND_VANISHING = true;
    gmm.MCSR = mcsr;
    gmm.MCVR = mcvr;
    gmm.learnTheModel();
    return gmm;
  }


  GaussianMixtureModel createGaussianMixtureModel(int m, double[][] data, int[] init_labels) {
    return new GaussianMixtureModel(m, data, init_labels);
  }

}