//Source file: D:\\Projects\\src\\yi\\stat\\em\\ControlledIEWGMM.java

package yi.stat.em;

import java.util.*;
import yi.util.data.*;
import yi.ds.hash.*;


/**
 * Not every feature vector has probabilities of 1.  If in the image containing
 * that feature vector, there are n objects, but m features, each feature vector
 * will only get weight of n/m.  This makes sure that the features from one image,
 * maybe there are several but all from one object, dominant the distribution of
 * the values of the object.
 */
public class ControlledIEWGMM extends ControlledIEGMM
{
  /**
   * remember the weight of each feature
   */
  protected DoubleScoreHashMap m_fvWeight = new DoubleScoreHashMap();


  public ControlledIEWGMM(HashMap feature_map, HashMap cluster_map)
  {
    // save all the feature vectors (from all the images)
    LinkedList all_data = new LinkedList();

    int object_num = 0;

    // for each image
    for (Iterator i= feature_map.keySet().iterator(); i.hasNext(); ) {
      String img = (String)i.next();
      Collection feature_list = (Collection)feature_map.get(img);
      all_data.addAll(feature_list);
      Collection cluster_list = (Collection)cluster_map.get(img);

      // weight for each feature vector
      double weight = 1.0 / (double)feature_list.size();
      object_num += cluster_list.size();
      // add every feature vector to every cluster
      // for each cluster (object)
      for (Iterator c = cluster_list.iterator(); c.hasNext(); ) {
        String cluster_name = (String)c.next();
        if (!m_initComponents.containsKey(cluster_name)) {
          GaussianComponent g = (GaussianComponent)createEMComponent();
          m_initComponents.put(cluster_name, g);
        }
        GaussianComponent cluster = (GaussianComponent)m_initComponents.get(cluster_name);
        for (Iterator f = feature_list.iterator(); f.hasNext(); ) {
          double[] data = (double[]) f.next();
          /** @todo the init probability ? */
          cluster.add(data, weight);
        }
      }
    }

    m_data = Conversion.to2DDoubleArray(all_data);
    m_numComponents = m_initComponents.size();

    // calculate the weight of each cluster
    Iterator iter = m_initComponents.keySet().iterator();
    double psum = 0;
    while (iter.hasNext()) {
      Object key = iter.next();
      GaussianComponent cumulator = (GaussianComponent) m_initComponents.get(key);
      double prob = cumulator.getProbabilitySum() / object_num;
      psum += prob;
      m_weights.put(key, new Double(prob));
    }

    m_clusterMap = cluster_map;

    for (Iterator i=feature_map.keySet().iterator(); i.hasNext(); ) {
      String img_name = (String)i.next();
      Collection feature_vect_list = (Collection)feature_map.get(img_name);
      Collection cluster_list = (Collection)cluster_map.get(img_name);
      double weight = cluster_list.size() / (double)feature_vect_list.size();
      for (Iterator j=feature_vect_list.iterator(); j.hasNext();) {
        Object fv = j.next();
        m_fv2imgMap.put(fv, img_name);
        m_fvWeight.putScore(fv, weight);
      }
    }

  }

  /**
   * Return the probabilities of the data under the different models represented
   * by the GMM.  This is the E-Step in EM algorithm.
   * The candidate clusters only are those objects described in the image.
   * Also adjust the probabilities by the weight of each feature vector
   * @param cumulators
   * @param model_probs
   * @param data
   * @return
   */
  protected HashMap getDensitiesUnderComponents(HashMap cumulators,
                                         HashMap model_probs, double[] data) {
    Collection cand_objs = (Collection)m_clusterMap.get(m_fv2imgMap.get(data));
    if (cand_objs == null) { // because this is a testing data. its objects is un-known
      return super.getDensitiesUnderComponents(cumulators, model_probs, data);
    }

    Iterator iter = cand_objs.iterator();
    HashMap probs = new HashMap();
    double prob_sum = 0;
    int component_num = 0;
    while (iter.hasNext()) {
      Object key = iter.next();
      if (!cumulators.containsKey(key)) {// the object may be removed due to matrix sigular
//        System.out.println(key.toString() + " has removed");
        continue;
      }
      GaussianComponent cumulator = (GaussianComponent) cumulators.get(key);
      try {
        double prob = cumulator.getDensity(data);
        double model_prob = ( (Double) model_probs.get(key)).doubleValue();
        prob *= model_prob;
        if (prob > 0) {
          probs.put(key, new Double(prob));
          prob_sum += prob;
          component_num ++;
        }
      }
      catch (GaussianModelSingularException e) {
        e.printStackTrace();
      }
    }

    double weight = 1;
    if (m_fvWeight.containsKey(data)) {
      weight = m_fvWeight.getScore(data);
    }

    // normalize
    iter = probs.keySet().iterator();
    while (iter.hasNext()) {
      Object key = iter.next();
      if (component_num == 0) probs.put(key, new Double(0));
      double prob = ( (Double) probs.get(key)).doubleValue();
      probs.put(key, new Double(prob / prob_sum * weight));
    }
    return probs;
  }

  /**
   * Calculate the log likelihood function of the model parameters
   * @param accumulators
   * @param weights
   * @param data
   * @return
   */
  protected double getParameterLogLikelihood(ComponentHashMap cumulators, HashMap weights, double[][] data) {
    Collection cand_objs = (Collection)m_clusterMap.get(m_fv2imgMap.get(data[0]));
    double weight = m_fvWeight.getScore(data[0]);

    double ws = 0;
    for (Iterator i=cumulators.keySet().iterator(); i.hasNext();) {
      Object key = i.next();
      double model_prob = ( (Double) weights.get(key)).doubleValue();
      ws += model_prob;
    }
    System.out.println("weight of sum =" + ws);

    double like = 0;
    for (int i = 0; i < data.length; i++) {
      double[] xk = data[i];

      Iterator iter = cand_objs.iterator();
      double prob_sum = 0;
      double component_weight_sum = 0;
      int component_num = 0;
      while (iter.hasNext()) {
        Object key = iter.next();
        if (!cumulators.containsKey(key)) {// the object may be removed due to matrix sigular
//        System.out.println(key.toString() + " has removed");
          continue;
        }
        GaussianComponent cumulator = (GaussianComponent) cumulators.get(key);
        try {
          double prob = cumulator.getDensity(xk);
          double model_prob = ( (Double) weights.get(key)).doubleValue();
          prob_sum += prob * model_prob;
          component_weight_sum += model_prob;
          component_num ++;
        }
        catch (GaussianModelSingularException e) {
          e.printStackTrace();
        }
      }

      if (component_num == 0) return 0;

      double d = prob_sum / component_weight_sum;
      if (d <1e-300) d=1e-300;
      like += Math.log(d)* weight;
    }
    return like;
  }

}
