package yi.stat.em;

import java.util.HashMap;
import java.util.*;
import yi.util.io.*;
import yi.iu.processor.color.*;
import java.text.*;

/**
 * Mixture of Gaussian Mixtures
 */

public class GMMMixtureModel extends ControlledIEGMM{

  /**
   * Number of components in each GMM (Gaussian mixture).  Ideally, this
   * variable should be declared as non-static one.  However, it will be used
   * in the method of createEMComponent() which is a virtual method and will be
   * called from the initial method of the parent classs.  At that time, the
   * instance variables in this class are not initialized yet.  So here declare
   * it as static (class) variable, which is certainly an ugly solution.
   */
  public int m_numComponentsInGMM = 8;

  public double m_breakThresholdInGMM = 1e-7;

  public GMMMixtureModel() {
  }

  public GMMMixtureModel(HashMap feature_map, HashMap cluster_map, double outer_break_t, int numComponentsInGMM, double breakThresholdInGMM){
    super(feature_map, cluster_map);

    BREAK_THRESHOLD = outer_break_t;

    // when call the constructor of its parents, mixture components are created
    // at that time, "m_numComponentsInGMM" and "m_breakThresholdInGMM" are not
    // initialized at all, so here they are
    m_numComponentsInGMM = numComponentsInGMM;
    m_breakThresholdInGMM = breakThresholdInGMM;
    for (Iterator i = m_initComponents.keySet().iterator(); i.hasNext(); ) {
      ComponentialGMM gmm = (ComponentialGMM)m_initComponents.get(i.next());
      gmm.m_numComponents = m_numComponentsInGMM;
      gmm.BREAK_THRESHOLD = m_breakThresholdInGMM;
    }
  }


  /**
   * Use Gaussian Mixture Models as components.  Since using the mixture model
   * as a component.  It will take a long time if each component learn very
   * precisely.  Remember the learning processes will be performed for each
   * component in each iteration.  So the break_threshold should be relaxed to
   * allow a early stop.
   * @return
   */
  MixtureComponential createEMComponent() {
    ComponentialGMM c =new ComponentialGMM(m_numComponentsInGMM);
    c.BREAK_THRESHOLD = m_breakThresholdInGMM;
//    c.SPLITTING_AND_VANISHING = true;
//    c.MCVR = 3;
//    c.MCSR = 1.5;

    return c;
  }

  /**
   * When training by iterations, a new set of components are created in each
   * iteration.  This is for create it
   * @param last_components
   * @return
   */
  protected ComponentHashMap createNewComponents(ComponentHashMap last_components) {
//    ComponentHashMap new_gmms = new ComponentHashMap();
//    // for mixture of mixture, it should not learn the gmm components from
//    // the beginning, but we can just go from the last step learning
//    for (Iterator c = last_components.keySet().iterator(); c.hasNext(); ) {
//      Object comp_key = c.next();
//      ComponentialGMM gmm = (ComponentialGMM)last_components.get(comp_key);
//      ComponentialGMM new_gmm = new ComponentialGMM(gmm.getRstNumComponents());
//      new_gmm.BREAK_THRESHOLD = gmm.BREAK_THRESHOLD;
//      new_gmm.m_initComponents = (ComponentHashMap)gmm.m_rstComponents.clone();
//      new_gmm.m_rstComponents = (ComponentHashMap)gmm.m_rstComponents.clone();
//      new_gmm.m_weights = (HashMap)gmm.m_weights.clone();
//      new_gmms.put(comp_key, new_gmm);
//    }
//    return new_gmms;

    // clear the data saved in each component, then it is ready for the
    // next iteration
    for (Iterator c = last_components.keySet().iterator(); c.hasNext(); ) {
      Object comp_key = c.next();
      ComponentialGMM gmm = (ComponentialGMM)last_components.get(comp_key);
      gmm.m_data = null;
      gmm.m_dataWeight.clear();
      gmm.m_initComponents.clear();
      gmm.m_initComponents = gmm.m_rstComponents;
    }
    return last_components;
  }

  /**
   * The iteration step is made of E-step and M-step
   */
  protected void iterationSteps() {
    try{

      // iteration statistical info
      LinkedList likelihoods = new LinkedList();

      learnComponentModels(m_initComponents);
      ComponentHashMap last_components = (ComponentHashMap) m_initComponents.
          clone();
      if (this instanceof GMMMixtureModel) {
       int a=0;
      }
      double last_likelihood = getParameterLogLikelihood(last_components,
          m_weights, m_data);

      double best_likelihood = last_likelihood;
      ComponentHashMap best_components = last_components;
      if (this instanceof GMMMixtureModel) {
        System.out.println("iter 0, likelihood = " + last_likelihood);
        for (Iterator o = m_initComponents.keySet().iterator(); o.hasNext(); ) {
          Object obj = o.next();
          MixtureModel mm = (MixtureModel)last_components.get(obj);
          int[][] gmmmm_gc_color_means = new int[MAXIMUM_GMMMIXTURE_ITER][];
          gmmmm_gc_color_means[0] = mm.getRstMeansInColor("hsi");
          m_gmmmm_gc_color_means.put(obj, gmmmm_gc_color_means);
          double[][] gmmmm_gc_weights = new double[MAXIMUM_GMMMIXTURE_ITER][];
          gmmmm_gc_weights[0] = mm.getRstComponentWeightsInSortedArray();
          m_gmmmm_gc_weights.put(obj, gmmmm_gc_weights);
        }
      }
      else {
        likelihoods.addLast(new Double(last_likelihood));
      }
      double like = 0;

      // remember mean history
      addToMeanHist(mean_hist, getComponentMeans(m_initComponents));
      // see the histogram of each component (Gaussian)
      if (DRAW_COMPONET) {
//      drawComponentHistogram(m_initComponents, 0);
      }

      /** @todo timing */
      // timing
      long time_e = 0, time_m = 0, time_i = 0, time_l = 0;

      // init a empty component set
      ComponentHashMap new_components = createNewComponents(last_components);
      int iteration_step = 1;
      do {

        long time1 = System.currentTimeMillis();

        // E-step
        for (int i = 0; i < m_data.length; i++) {
          double[] xk = m_data[i];
          // p(wi | xk)
          HashMap probs = getDensitiesUnderComponents(last_components,
              m_weights, xk);

          Iterator prob_iter = probs.keySet().iterator();
          while (prob_iter.hasNext()) {
            Object key = prob_iter.next();
            double prob = ( (Double) probs.get(key)).doubleValue(); // p(wi|xk)
            if (Double.isNaN(prob)) {
              System.err.println("data xk got NaN prob for component \"" + key +
                                 "\"");
              getDensitiesUnderComponents(last_components,
                                          m_weights, xk);
            }
            new_components.add(key, xk, prob);
          }
        }

        long time2 = System.currentTimeMillis();

        // M-step
        // recaculate the obj_probs (weights for each component)
        recalculateWeights(new_components);

        long time3 = System.currentTimeMillis();

        // check the integrity of the new models
        learnComponentModels(new_components);

        // remember the mean
        addToMeanHist(mean_hist, getComponentMeans(new_components));

        long time4 = System.currentTimeMillis();

        // check the likelihood of the models
        // break if it it the time
        like = getParameterLogLikelihood(new_components, m_weights, m_data);

        System.err.println("iter = " + iteration_step + ", likelihood = " + like);
//        for (Iterator j = new_components.keySet().iterator(); j.hasNext();) {
//          Object obj = j.next();
//          System.out.println("component = " + obj);
//          System.out.println((MixtureComponential)new_components.get(obj));
//        }

        if (like < -1e10 || Double.isNaN( (like))) {
          System.err.println("likelihood is undesirable: " + like);
//          {
//            double like2 = 0;
//            for (int i = 0; i < m_data.length; i++) {
//              double[] xk = m_data[i];
//              double d = getDensity(new_components, m_weights, xk);
//              OutputResult.print(System.err, xk);
//              System.err.println("density = " + d);
//              like2 += Math.log(d);
//              System.err.println("like = " + like2);
//              if (like2 < -1e6 || Double.isNaN( (like2))) {
//                getDensity(new_components, m_weights, xk);
//              }
//            }
//          }
//          break;
        }

        if (this instanceof GMMMixtureModel) {

          if (like > best_likelihood) {
            best_likelihood = like;
            best_components = new_components;
          }

          getParameterLogLikelihood(new_components, m_weights, m_data);
          System.out.println("outer iter = " + iteration_step + ", likelihood = " +
                             like);
          for (Iterator o = m_initComponents.keySet().iterator(); o.hasNext(); ) {
            Object obj = o.next();
            MixtureModel mm = (MixtureModel)new_components.get(obj);
            int[][] gmmmm_gc_color_means = (int[][]) m_gmmmm_gc_color_means.get(
                obj);
            if (iteration_step >= gmmmm_gc_color_means.length) break;
            gmmmm_gc_color_means[iteration_step] = mm.getRstMeansInColor("hsi");
            double[][] gmmmm_gc_weights = (double[][]) m_gmmmm_gc_weights.get(
                obj);
            gmmmm_gc_weights[iteration_step] =
                mm.getRstComponentWeightsInSortedArray();
          }
          if (iteration_step >= MAXIMUM_GMMMIXTURE_ITER && like >= best_likelihood ||
              iteration_step >= MAXIMUM_GMMMIXTURE_ITER * 2) {
//            // USE THE ONE WITH BEST LIKELIHOOD
//            new_components = best_components;
//            System.out.println("use the one with best likelihood of " + best_likelihood);

//            for (Iterator o = m_initComponents.keySet().iterator(); o.hasNext(); ) {
//              Object obj = o.next();
//              int[][] gmmmm_gc_color_means = (int[][]) m_gmmmm_gc_color_means.
//                  get(obj);
//              double[][] gmmmm_gc_weights = (double[][]) m_gmmmm_gc_weights.get(
//                  obj);
//              OutputResult.print(obj.toString() + ".gc.color.means",
//                                 gmmmm_gc_color_means);
//              int[] color_means_demo = ColorUtil.createSpectrum(
//                  gmmmm_gc_color_means,
//                  gmmmm_gc_color_means[0].length, 50, 50, 0xffffffff);
//              OutputResult.outputGIFImage(obj.toString() + ".gc.color.means.gif",
//                                          color_means_demo,
//                                          50 * gmmmm_gc_color_means[0].length,
//                                          50 * gmmmm_gc_color_means.length);
//              OutputResult.print(obj.toString() + ".gc.weights",
//                                 gmmmm_gc_weights);
//            }
            System.out.println("GMM mixture iteration maximum set to " +
                               MAXIMUM_GMMMIXTURE_ITER + ", so jump out");
            break;
          }
        }
        else {
          likelihoods.addLast(new Double(like));
        }

        if (Math.abs(like - last_likelihood) <
            BREAK_THRESHOLD * Math.abs(last_likelihood)) {

          for (Iterator o = m_initComponents.keySet().iterator(); o.hasNext(); ) {
            Object obj = o.next();
            int[][] gmmmm_gc_color_means = (int[][]) m_gmmmm_gc_color_means.
                get(obj);
            int[] gmmmm_gc_color_mean = gmmmm_gc_color_means[iteration_step];
            int[] color_means_demo = ColorUtil.createSpectrum(
                gmmmm_gc_color_mean, 50, 50);
            OutputResult.outputGIFImage(obj.toString() + ".gc.color.means.gif",
                                        color_means_demo,
                                        50 * gmmmm_gc_color_mean.length,
                                        50);
          }

          if (SPLITTING_AND_VANISHING) {
            Random r = new Random();
            boolean changed = false;
            // use a splitting and vanishing algorithm during training to learn
            // the number of mixture components.  Two thresholds are defined,
            // a mixture-coefficient vanishing ratio (MCVR), and a
            // mixture-coefficient splitting ration (MCSR).  Under a M-component
            // mixture, with component probabilities p_m, if p_m < 1/(M*MCVR),
            // then the m-th component will vanish.   If p_m>MCSR/M, that
            // component will split.

            int current_comp_num = new_components.size();
            // first, vanishing
            // second, splitting
            for (Iterator c=new_components.keySet().iterator(); c.hasNext();) {
              Object comp_key = c.next();
              // get the weight of the component
              double weight = ((Double)m_weights.get(comp_key)).doubleValue();
              if (weight < 1 / (current_comp_num * MCVR)) {
                System.out.println("vanish component " + comp_key + " , weight = " + weight);
                c.remove();
                m_weights.remove(comp_key);
                changed = true;
                continue;
              }
              if (weight > MCSR / current_comp_num) {
                int new_comp_index = getNumComponents() + (m_splitCount++);
                MixtureComponential comp = (MixtureComponential)new_components.get(comp_key);
                MixtureComponential new_comp = createEMComponent();
                if (new_comp instanceof GaussianComponent) {
                  GaussianComponent gc = (GaussianComponent)new_comp;
                  double[] new_means = (double[])comp.getMeans();
                  double[][] cov = (double[][])((GaussianComponent)comp).getCovariance().clone();
                  for (int d=0; d<new_means.length; d++) {
                    new_means[d] += r.nextGaussian() * Math.sqrt(cov[d][d]);
                  }
                  gc.setMeans(new_means);
                  gc.setCovariance(cov);
                  // add the new component
                  System.out.println("split component " + comp_key + " to " + new_comp_index + ", weight = " + weight/2);
                  new_components.put(new Integer(new_comp_index), new_comp);
                  m_weights.put(new Integer(new_comp_index), new Double(weight/2));
                  m_weights.put(comp_key, new Double(weight/2));
                  break;
                } else {
                  System.err.println("splitting and vanishing algorithm only works for Gaussian components");
                  break;
                }
              }
            }
            if (!changed) break;
          } else {
              break;
          }
        }
        last_likelihood = like;

        long time5 = System.currentTimeMillis();

        iteration_step++;
        if (iteration_step > MAX_ITERATION_STEP) {
          System.err.println("iteration steps greater than MAX_ITERATION_STEP " +
                             MAX_ITERATION_STEP);
          break;
        }
        last_components = new_components;
        new_components = createNewComponents(last_components);

        if (SHOW_PROGRESS) {
          System.out.println(like);
        }

        time_e += time2 - time1;
        time_m += time3 - time2;
        time_i += time4 - time3;
        time_l += time5 - time4;

      }
      while (true);
//      System.out.println("likelihood for the parameters = " + like);
//
//      System.out.println("timing e=:" + (time_e / (double) iteration_step) +
//                         ", m=" + (time_m / (double) iteration_step) +
//                         ", checkModelsIntegrity = " +
//                         (time_i / (double) iteration_step) +
//                         ", getParameterLikelihood = " +
//                         (time_l / (double) iteration_step));
//
//      // show the mean hist
//      showMeanHist(mean_hist, "mean_hist.ppm");

      // keep the stable means
      m_rstComponents = new_components;
      for (Iterator c = m_rstComponents.keySet().iterator(); c.hasNext(); ) {
        Object comp_key = c.next();
        ComponentialGMM gmm = (ComponentialGMM)last_components.get(comp_key);
        System.out.println("component number of component " + comp_key + "\t" + gmm.getRstNumComponents());
      }


//      if (SHOW_PROGRESS)
//        System.out.println();
      double[] likelihood_iters = yi.util.data.Conversion.toDoubleArray(
          likelihoods);
      OutputResult.print(System.out, likelihood_iters);
    }
    catch (ModelSingularException e) {
      e.printStackTrace();
    }
  }

  /**
   * Check the integrity of each components.  Fix if possible
   * @param new_components
   * @param old_components
   */
  protected void learnComponentModels(HashMap components) {
    int i_index = 0;
    int index_to_monitor = -1;
    for (Iterator i = components.keySet().iterator(); i.hasNext(); i_index++) {
      Object key = i.next();
      if (i_index == index_to_monitor) {
        i_index = index_to_monitor;
      }

      MixtureComponential component = (MixtureComponential)components.get(key);
      if (component instanceof ComponentialGMM) {
        System.out.println("learn Component model on " + key);
      }
      try {
//        int pre_num_comp = 0;
//        if (component instanceof ComponentialGMM) {
//          pre_num_comp = ((ComponentialGMM)component).getRstNumComponents();
//        }
        component.learnTheModel();
        if (component instanceof ComponentialGMM) {
          int post_num_comp = ((ComponentialGMM)component).getRstNumComponents();
          System.out.println("component " + key + ", index = " + i_index + " has " + post_num_comp + "components");
        }
      } catch (ModelSingularException e) {
        // e.printStackTrace();
        double weight = ((Double)m_weights.get(key)).doubleValue();
        double prob_sum = component.getProbabilitySum();
        System.err.print("\"" + key + "\"" + " singular ");
        final NumberFormat nf = NumberFormat.getInstance();
        nf.setMaximumFractionDigits(2);
        System.err.println("(weight = " + nf.format(weight * 100) + "%, prob_sum = " + nf.format(prob_sum) + ")" + ", removed");
        i.remove();
        continue;
      }
    }
  }

}