package yi.iu.imgdb.learn.kmeans;

import java.util.*;
import yi.ds.*;
import yi.math.*;
import yi.iu.imgdb.learn.em.*;
import yi.iu.imgdb.learn.*;

/**
 * KMeansLearner implements the object feature learning idea suggested by Jeff
 * Bilms.  The detailed idea is documented in
 * Projects\learning2\doc\learning.doc
 *
 * The algorithm contains two steps: boot step and iteration step.
 *
 * See main() for a simple example.
 */

public class KMeansModel implements ObjImgModel{

  /**
   * Store the init accumulators for the objects.  The key is the name of the
   * object and the value is the corresponding accumulator (DoubleArrayCumulator)
   */
  AccumulatorHashMap m_initAccumulators = new AccumulatorHashMap();

  /**
   * Store the final accumulators for the objects after iteration step.  The key
   * is the name of the object and the value is the corresponding accumulator
   * (DoubleArrayCumulator)
   */
  AccumulatorHashMap m_rstAccumulators = new AccumulatorHashMap();

  /**
   * Store the mapping from image to its object list.  The key is the name of
   * the image and the value is a String collection
   */
  HashMap m_imgObjList = new HashMap();

  /**
   * Store the mapping from image to its feature vector list.  The key is the
   * name of the image and the value is a double[] collection
   */
  HashMap m_imgFvList = new HashMap();

  public double BREAK_THRESHOLD = 1e-4;

  public final static int BREAK_ON_MEAN_DIFF = 0;
  public final static int BREAK_ON_STEP_COUNT = 1;
  public int BREAK_COND = BREAK_ON_MEAN_DIFF;
  /**
   * If break condition (BREAK_COND) is set to BREAK_ON_STEP_COUNT, this number
   * decides it breaks after how many steps.  See iterationSteps() for more
   * information
   */
  public int break_on_step_count = 5;

  public boolean SHOW_PROGRESS = true;


  public KMeansModel() {
  }

  /**
   * Add a new image instance in the init step.
   * @param img_name the image name
   * @param obj_names the names of the ojects contained in the image
   * @param ft_vect the feature vectors of the image.  The type is a collection
   * of double[]
   */
  public void addImageInst(String img_name, Collection obj_names, Collection ft_vects) {
    Iterator f_iter = ft_vects.iterator();
    while (f_iter.hasNext()) {
      double[] ft_vect = (double[])f_iter.next();
      Iterator o_iter = obj_names.iterator();
      while (o_iter.hasNext()) {
        m_initAccumulators.accumulate((String)o_iter.next(), ft_vect);
      }
    }
    // save the information
    m_imgObjList.put(img_name, obj_names);
    m_imgFvList.put(img_name, ft_vects);
  }

  public HashMap getMeans() {
    if (m_rstAccumulators == null) {
      learnTheModel();
    }
    return m_rstAccumulators.getMeans();
  }

  /**
   * This method triggers the iteration steps.  It returns the resultant model
   * for each object calculated in the iteration step.
   * @return The key is the name of the object and the value is the
   * corresponding model
   */
  public void learnTheModel() {
    // go iteration steps
    iterationSteps();

  }

  /**
   * The iteration steps described in Projects\learning2\doc\learning.doc
   */
  void iterationSteps() {

    HashMap last_means = m_initAccumulators.getMeans();

    // init a empty accumulator set
    AccumulatorHashMap new_accumulators = new AccumulatorHashMap();
    int iteration_step = 0;
    do {
      // assign each feature vector to its cloest mean
      Iterator i_iter = m_imgFvList.keySet().iterator();
      while (i_iter.hasNext()) {
        String img_name = (String)i_iter.next();
        Collection ft_vects = (Collection)m_imgFvList.get(img_name);
        Iterator ft_iter = ft_vects.iterator();
        while (ft_iter.hasNext()) {
          double[] fv = (double[]) ft_iter.next();
          // find the cloest one
          String obj = getClosest(last_means, fv);
          new_accumulators.accumulate(obj, fv);
        }
      }

      // test the break condition
      if (BREAK_COND == BREAK_ON_MEAN_DIFF) {
        // if the means didn't change, then exit the iterations
        HashMap new_means = new_accumulators.getMeans();
        double change_sum = getSumOfDistance(last_means, new_means); // the sum of differences between last_means and new_means
        last_means = new_means;
//        System.out.println("step: " + iteration_step + ", changes = " +
//                           change_sum);
        if (change_sum <= BREAK_THRESHOLD * last_means.size())
          break;
      } else if (BREAK_COND == BREAK_ON_STEP_COUNT) {
//        System.out.println("step: " + iteration_step);
        if (iteration_step >= break_on_step_count) {
          break;
        }
      }
      if (SHOW_PROGRESS && (iteration_step%10 == 0)) System.out.print(".");
      iteration_step++;
    } while (true);

    if (SHOW_PROGRESS) System.out.println();

    // keep the stable means
    m_rstAccumulators = new_accumulators;
  }

  /**
   * Return the summary of the distance between the corresponding vectors in the
   * two HashMap.  If there is no corresponding vector for a key, then vector 0
   * is assumed.
   * @param vectors1
   * @param vectors2
   * @return
   */
  private double getSumOfDistance(HashMap vectors1, HashMap vectors2) {
    double diff_sum = 0; // the sum of differences between last_means and new_means

    // loop through the keys in vectors 1
    Iterator iter = vectors1.keySet().iterator();
    while (iter.hasNext()) {
      Object key = iter.next();
      double[] v1 = (double[])vectors1.get(key);
      if (!vectors2.containsKey(key)) {
        diff_sum += VectorUtil.lenth(v1);
      } else {
        double[] v2 = (double[]) vectors2.get(key);
        diff_sum += VectorUtil.EuclideanDistance(v1, v2);
      }
    }

    // loop through the keys in vectors 2
    iter = vectors2.keySet().iterator();
    while (iter.hasNext()) {
      Object key = iter.next();
      if (!vectors1.containsKey(key)) {
        double[] v2 = (double[])vectors2.get(key);
        diff_sum += VectorUtil.lenth(v2);
      }
    }
    return diff_sum;
  }

  /**
   * Return the name of closest mean.  Users is responsible for making sure that
   * the means is not empty
   * @param means
   * @param value
   * @return
   */
  private String getClosest(HashMap means, double[] value) {
    double closest_d = -1;
    String closest_obj = null;

    Iterator iter = means.keySet().iterator();
    while (iter.hasNext()) {
      String key = (String)iter.next();
      double[] mean = (double[])means.get(key);
      double dist = yi.math.VectorUtil.EuclideanDistance(mean, value);
      if (closest_obj == null || dist < closest_d) {
        closest_d = dist;
        closest_obj = key;
      }
    }
    return closest_obj;
  }

  public String toString() {
    StringBuffer str = new StringBuffer();
    if (m_rstAccumulators != null) {
      HashMap means = m_rstAccumulators.getMeans();
      Iterator iter = means.keySet().iterator();
      while (iter.hasNext()) {
        String key = (String) iter.next();
        double[] mean = (double[]) means.get(key);
        str.append("<<" + key + ">>\nmean: ");
        for (int j = 0; j < mean.length; j++) {
          str.append(mean[j]);
          str.append(" ");
        }
        str.append('\n');
      }
    }
    return str.toString();
  }

  public static class TestOnGaussianMixture {
    public static void main(String[] args) {
      KMeansModel learner = new KMeansModel();

      String str_sun = "Sun";
      String str_tree = "Tree";
      String str_water = "Water";

      String[] objects = new String[] {
          str_sun, str_tree, str_water};

      // means for each object
      double[][] means = {
          {180, 50, 50},  // sun
          {20, 150, 50},  // tree
          {20, 20, 200}   // water
      };

      // variances
      double[][] variances = {
          {20, 30, 30},   // sune
          {15, 50, 30},   // tree
          {15, 15, 50}    // water
      };

      // probability for each object
      double[] obj_probs = {
          0.5, 0.8, 0.6};

      EMTestDataGenerator generator = new EMTestDataGenerator(objects,
          obj_probs, means, variances, new Random(0));

      int img_num = 200;
      for (int i = 0; i < img_num; i++) {
        String img_name = "img" + i;
        EMTestData data = generator.generate();
        System.out.println(img_name + "\n" + data);
        learner.addImageInst(img_name, data.objects, data.features);
      }

      learner.learnTheModel();
      System.out.println(learner);
    }

  }

  public static void main(String[] args) {
    KMeansModel learner = new KMeansModel();
    String str_sun = "Sun";
    String str_water = "Water";
    String str_tree = "Tree";

    double[] v_sun = new double[]{255, 0, 0};
    double[] v_water = new double[]{0, 0, 255};
    double[] v_tree = new double[]{0, 255, 0};

    String[] objects = new String[]{str_sun, str_water, str_tree};
    double[][] vectors = new double[][] {v_sun, v_water, v_tree};

    Collection[] objs = new Collection[8];
    Collection[] vects = new Collection[8];
    for (int i=0; i<8; i++) {
      objs[i] = new HashSet();
      vects[i] = new LinkedList();

      for (int k=0; k<objects.length; k++) {
        if (((i>>k) & 1) != 0) {
          objs[i].add(objects[k]);
          vects[i].add(vectors[k]);
        }
      }
      // init steps
      learner.addImageInst("img" + i, objs[i], vects[i]);
    }

    learner.learnTheModel();
    System.out.println(learner);
  }

}