package yi.stat.em;

import java.util.HashMap;
import yi.stat.StatUtil;
import java.util.Iterator;

/**
 * Normal GaussianMixtureModel assumes that each sample has the same weight of
 * 1.  GMMOnWeightedData allows having a weight with each sample.
 */

public class ComponentialGMM
//    extends GaussianMixtureModel
    extends AlignedGaussianMixtureModel
    implements MixtureComponential {

  /**
   * Remember the weight of each data sample
   */
  HashMap m_dataWeight = new HashMap();

  /**
   * Initialize the Gaussian Mixture Model with number of components specified.
   * At this time, the mixture model is not ready for training yet since no
   * data is given.
   * @param m
   */
  public ComponentialGMM(int m) {
    m_numComponents = m;
  }

  /**
   * Add a new data sample (with weight)
   * @param data the data
   * @param weight the weight of the sample data
   */
  public void add(double[] data, double weight) {
    m_dataWeight.put(data, new Double(weight));
  }

  /**
   * Learn the model.  At this point, assume all the training data have been
   * added.  So first use all the data samples to set up the variables.  Then,
   * get the initial Gaussian components.
   */
  public void learnTheModel() {
    // get all the data samples.
    m_data = yi.util.data.Conversion.to2DDoubleArray(m_dataWeight.keySet());

//    // set up the minimums.
//    m_dataMinimums = yi.stat.StatUtil.minimumOnCol(m_data);
//
//    // set up the maximums.
//    m_dataMaximums = yi.stat.StatUtil.maximumOnCol(m_data);

    // initialize the components if it is not done yet
    if (m_initComponents == null || m_initComponents.size() == 0) {
      initializeComponents();
    }

    // Trigger learning as before
    super.learnTheModel();
  }

  /**
   * Calculate the mean of the weighted samples
   * @return
   * @throws ModelSingularException
   */
  public double[] getMeans() throws ModelSingularException {
    if (m_data == null || m_data.length == 0 || m_data[0] == null) {
      throw new ModelSingularException("There is no sample in this model");
    }
    int n = m_data.length;
    int dim = m_data[0].length;
    double[] mean = new double[dim];
    double weight_sum = 0;
    for (int i = 0; i < n; i++) { // for each data sample
      // get the weight of that data sample
      double weight = ( (Double) m_dataWeight.get(m_data[i])).doubleValue();
      // accumulate
      for (int j = 0; j < dim; j++) {
        mean[j] += m_data[i][j] * weight;
      }
      weight_sum += weight;
    }
    for (int i = 0; i < dim; i++)
      mean[i] /= weight_sum;
    return mean;
  }

  /**
   * Return the weight/probability sum of all samples
   * @return
   */
  public double getProbabilitySum() {
    double sum = 0;
    for (Iterator i = m_dataWeight.values().iterator(); i.hasNext(); ) {
      sum += ((Double)i.next()).doubleValue();
    }
    return sum;
  }

  /**
   * Initialize the components.  A simple method is used here: pick up m sample
       * points randomly as the means of the components; use 1/10 of global variance
   * on each dimension to construct the covariance matrix for each component.
   */
  protected void initializeComponents() {

    // n samples.
    int n = m_data.length;

    // it doesn't work if the number of components is greater than the number of samples
    if (m_numComponents > n) {
      System.err.println("The number of data samples, " + n +
                         ", is less than the specified number of components, " +
                         m_numComponents + ".  The component number is reduced");
      m_numComponents = n;
    }

    // randomly pick up m samples as the means.
    int[] random_index = StatUtil.randomize(n);

    // get the global variance on each dimension
    double[] variances = StatUtil.varianceOnCol(m_data);
    int dim = variances.length;
    double[][] covariance_matrix = new double[dim][dim];
    for (int i = 0; i < dim; i++)
      covariance_matrix[i][i] = variances[i] / 10;

    for (int i = 0; i < m_numComponents; i++) {
      GaussianComponent g = (GaussianComponent) createEMComponent();
      g.setMeans(m_data[random_index[i]]);
      g.setCovariance(covariance_matrix);
      m_initComponents.put(new Integer(i), g);
      m_weights.put(new Integer(i), new Double(1.0 / m_numComponents));
    }
  }

  /**
   * This method overrides that in MixtureModel because the weight of the sample
   * data is not always 1 as in MixtureModel
   * @param components
   * @param model_probs
   * @param data
   * @return
   */
  protected HashMap getDensitiesUnderComponents(HashMap components,
                                         HashMap model_probs, double[] data) {


    // get the probabilities under different components assuming the sum of the
    // data sample is 1
    HashMap probs = super.getDensitiesUnderComponents(components, model_probs, data);

    // adding weight

    // if the data is not present in the data->weight map, then assuming the
    // weight is 1
    Double weight = (Double)m_dataWeight.get(data);
    if (weight == null) return probs;

    double w = weight.doubleValue();

    for (Iterator iter = probs.keySet().iterator(); iter.hasNext();) {
      Object key = iter.next();
      double prob = ( (Double) probs.get(key)).doubleValue();
      probs.put(key, new Double(prob * w));
    }
    return probs;
  }

  /**
   * Calculate the likelihood function of the model parameters.  The weights
   * of different data sample must be considered:
   * p(x_1^w_1, x_2^w_2, ... | parameter) => log(.) = sum{p(x_i|parameter)*w_i}
   * @param components
   * @param weights
   * @param data
   * @return
   */
  protected double getParameterLogLikelihood(ComponentHashMap components, HashMap weights, double[][] data) throws ModelSingularException{
    double like = 0;
    for (int i = 0; i < data.length; i++) {
      double[] xk = data[i];
      Double weight = (Double)m_dataWeight.get(xk);
      if (weight == null) throw new ModelSingularException("Is this data from training samples?");
      double w = weight.doubleValue();
      double d = getDensity(components, weights, xk);
      if (d <1e-300) d=1e-300;
      like += Math.log(d) * w;
    }
    return like;
  }

  /**
   * Recalculate the weights of each component and update m_weights
   * @param new_components
   */
  protected void recalculateWeights(ComponentHashMap components) {
    double total_weight = 0;


    for (Iterator iter = components.keySet().iterator(); iter.hasNext(); ) {
      Object key = iter.next();
      MixtureComponential component = (MixtureComponential) components.get(key);
      double prob = component.getProbabilitySum();
      m_weights.put(key, new Double(prob));
      total_weight += prob;
    }

    // normalize
    for (Iterator iter = components.keySet().iterator(); iter.hasNext(); ) {
      Object key = iter.next();
      double weight = ((Double)m_weights.get(key)).doubleValue();
      m_weights.put(key, new Double(weight/total_weight));
    }
  }
}