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

package yi.stat.em;

import java.util.*;
import yi.util.io.*;
import java.io.*;


/**
 * Not only the initial step is controlled, but also the Expection step is
 * controlled: each feature vector will only be classified to the known object
 * cluster of that image.
 */
public class ControlledIEGMM extends ControlledInitGMM
{
  /**
   * Map from feature vector to image name.  This information should be saved
   * since one featuer vector will be only assigned to the clusters
   * corresponding to the objects contained in the corresponding image
   */
  protected HashMap m_fv2imgMap = new HashMap();

  /**
   * Map from image name to a collection of its objects
   */
  protected HashMap m_clusterMap;

  protected ControlledIEGMM() {

  }

  /**
   * Besides others, remember the image->objects map, and feature->object map.
   * So that when assigning a featuer vector, it will be only assigned to the
   * clusters corresponding to the objects contained in the corresponding image
   * @param feature_map
   * @param cluster_map
   */
  public ControlledIEGMM(HashMap feature_map, HashMap cluster_map)
  {
    super(feature_map, cluster_map);
    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);
      for (Iterator j=feature_vect_list.iterator(); j.hasNext();) {
        Object fv = j.next();
        m_fv2imgMap.put(fv, img_name);
      }
    }
  }

  /**
   * 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 are only those objects described in the image.
   * @param cumulators
   * @param model_probs
   * @param data
   * @return
   */
  protected HashMap getDensitiesUnderComponents(HashMap components,
                                         HashMap model_probs, double[] data) {
    // find the objects described in the corresponding image
    Collection cand_objs = (Collection)m_clusterMap.get(m_fv2imgMap.get(data));

    // for a testing data, its objects is un-known.  Just use the usual method
    if (cand_objs == null) {
      return super.getDensitiesUnderComponents(components, model_probs, data);
    }

    HashMap probs = new HashMap();
    double prob_sum = 0;
    int component_num = 0;

    // caculate p(data|component)
    for (Iterator iter = cand_objs.iterator(); iter.hasNext();) {
      Object key = iter.next();
      MixtureComponential component = (MixtureComponential) components.get(key);
      if (!components.containsKey(key)) {// the object may be removed due to matrix sigular
        continue;
      }
      try {
        double prob = component.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 (ModelSingularException e) {
        e.printStackTrace();
      }
    }

    // normalize
    for (Iterator iter = probs.keySet().iterator(); iter.hasNext();) {
      Object key = iter.next();
      if (component_num == 0) { // all the corresponding components are removed?
        probs.put(key, new Double(0));
      }
      double prob = ( (Double) probs.get(key)).doubleValue();
      probs.put(key, new Double(prob / prob_sum));
    }
    return probs;
  }

  /**
   * Calculate the log likelihood function of the model parameters
   * @param accumulators
   * @param weights
   * @param data
   * @return
   */
  protected double getParameterLogLikelihood(ComponentHashMap components, HashMap weights, double[][] data) {

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

      Collection cand_objs = (Collection)m_clusterMap.get(m_fv2imgMap.get(xk));
//      OutputResult.print(new PrintWriter(System.out),cand_objs);
//      System.out.println("#obj=" + cand_objs.size());
      double prob_sum = 0;
      double component_weight_sum = 0;
      int component_num = 0;
      for (Iterator iter = cand_objs.iterator(); iter.hasNext();) {
        Object key = iter.next();
        if (!components.containsKey(key)) {// the object may be removed due to matrix sigular
          continue;
        }
        MixtureComponential component = (MixtureComponential) components.get(key);
        try {
          double prob = component.getDensity(xk);
          double model_prob = ( (Double) weights.get(key)).doubleValue();
          prob_sum += prob * model_prob;
          component_weight_sum += model_prob;
          component_num ++;
        }
        catch (ModelSingularException 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);
    }
    return like;
  }


}
