package yi.iu.processor.color.reg.merge;

import yi.iu.img.*;
import java.awt.*;
import java.util.*;
import yi.stat.*;
import yi.ds.sum.*;
import yi.ds.*;
import yi.stat.hist.*;
import yi.iu.processor.morph.*;
import yi.util.config.*;

/**
 * ClosestCCMerge merges the connected components.  A connected component (CC)
 * could be merged to a neighbor CC if they are close enough
 */

public class ClosestCCMerge implements Configurable{

  public boolean DEMO_MERGE = false;
  public boolean DEMO_STEPS = false;
  public boolean OUTPUT_REGIONS = false;

  private int m_steps = 0;

  public double MERGE_THRESHOLD = 9;

  /**
   * In merging, the mean of the distance between pixels in two regions may be
   * biggger than that within those regions.  For example, in region 1, the
   * average distance is 4, in region 2, the average distance is 5.  On the
   * boundary the average distance between pixels in region 1 and those in
   * region 2, say, is 6.  The average distance within the regions before
   * merge is:
   *
   * d_1 = (4 * n_1 + 5 * n_2) / (n_1 + n_2)
   *
   * n_1 is the connection in region 1, n_2 is the connection in region 2
   *
   * The average distance on the boundary is
   *
   * d_2 = 6
   *
   * The change ratio is d_2 / d_1.  If the number is too big, the merge is
   * invalid.
   */
  public double MEAN_DIST_CHANGE_RATIO = 1.25;

  public String[] getParameterNames()
  {
   return new String[] {
       "DEMO_MERGE",
       "DEMO_STEPS",
       "OUTPUT_REGIONS",
       "MERGE_THRESHOLD",
       "MEAN_DIST_CHANGE_RATIO",
   };
  }

  public ClosestCCMerge()
  {
   Configuration.configParameters(this);
  }

  /**
   * Return the distances of the neighboring CC.  For example, the neighboring
   * CC of a CC and the distance between them.
   * @param cc_labels
   * @param diffs
   * @param d
   * @return the array is indexed by the label of cc.  Each element of the array
   *         is a HashMap in which the key is the lable of the neighoring cc and
   *         value is the number of neighboring pixels and the total distance.
   *         For example,
   *         rst[3] = {4=>{4, 20.5}, 6=>{2, 15.3}} means cc3 is neighboring cc4
   *         and they have 4 neighboring pixels and the total distance is 20.5;
   *         cc3 is neighboring cc6 and they have 2 neighboring pixels and the
   *         total distance is 15.3
   */
  public IndexedDoubleSummationHashMap[] getNeighborCCDist(int[] cc_labels, DifferentialImage diffs, Dimension d) {
    // the number of labels
    int label_num = yi.stat.StatUtil.maximum(cc_labels) + 1;

    // the results
    IndexedDoubleSummationHashMap[] results = new IndexedDoubleSummationHashMap[label_num];
    for (int i=0; i<results.length; i++) results[i] = new IndexedDoubleSummationHashMap();

    // sum
    for (int y=0, ptr=0; y<d.height; y++) {
      for (int x=0; x<d.width; x++, ptr++) {
        // horizontal
        if (x < d.width - 1) {
          double dist = diffs.getHorizontalDistance(x, y);
          results[cc_labels[ptr]].putToSum(cc_labels[ptr+1], dist);
          results[cc_labels[ptr+1]].putToSum(cc_labels[ptr], dist);
        }
        // vertical
        if (y < d.height - 1) {
          double dist = diffs.getVerticalDistance(x, y);
          results[cc_labels[ptr]].putToSum(cc_labels[ptr+d.width], dist);
          results[cc_labels[ptr+d.width]].putToSum(cc_labels[ptr], dist);
        }
      }
    }
    return results;
  }




  /**
   * Pick up which region this region should be merged to.  First the size of
   * this region must be larger than
   * @param cc_sizes the sizes of all the regions
   * @param dists stores the distances from this region to others
   * @param i
   * @param dist_threshold
   * @return the lable of region this region should be merged to.  Return the
   *         label of itself if no other region can satisfy the requirement
   */
  private int pickupLabelToMerge(int[] cc_sizes, IndexedDoubleSummationHashMap dists,
                                 int i, double dist_threshold) {
    int to = i;
    double min_dist = dist_threshold;

    for (Iterator it = dists.keySet().iterator(); it.hasNext(); ) {

      // the label of the region merged to
      Integer key = (Integer) it.next();
      int merge_to = key.intValue();

      // if the label is the same
      if (merge_to == i)
        continue;

      // only merge to a bigger region
      if (cc_sizes[merge_to] < cc_sizes[i])
        continue;

      // when their sizes are same, low labeled region will be merged
      if (cc_sizes[merge_to] == cc_sizes[i] && i > merge_to)
        continue;


      // get the distance
      DoubleSummationStruct dist_sum = (DoubleSummationStruct) dists.get(key);

      // if the connecting pixels are too small compared to its area
//      if (dist_sum.counter < Math.sqrt(cc_sizes[i]) * 0.5) continue;

      double dist = dist_sum.getAverage();
      // the distance must be less than the dist_threshold
      if (dist < min_dist) {
        min_dist = dist;
        to = merge_to;
      }
    }

    if (i != to) {

//      System.out.println("picked: " + i + "->" + to + ", d=" + min_dist +
//                         ", size=" + cc_sizes[i] + "," + cc_sizes[to]);
    }

    return to;

  }

  public void merge(int[] labels, DifferentialImage diffs, Dimension d, double dist_threshold, int size_threshold) {
    for (int i=1; true; i++) {
      System.out.println("merge #" + i);
      boolean merged = mergeOnce(labels, diffs, d, dist_threshold, size_threshold);
      if (!merged) break;
      if (DEMO_MERGE) {
        int[] labels_demo = yi.util.ImageUtil.makeNewImageFromLabel(labels, false);
        yi.util.io.OutputResult.outputGIFImage("merge.ccc." + i + ".s.gif",
                                               labels_demo, d.width,
                                               d.height, false);
      }
    }
  }
/*
  public boolean mergeOnce(int[] labels, DifferentialImage diffs, Dimension d, double dist_threshold, int size_threshold) {

    ConnectComponent cc = new ConnectComponent();
    int[] cc_labels = cc.compactLabeling(labels, d, false);

    IndexedDoubleSummationHashMap[] cc_dists = getNeighborCCDist(cc_labels, diffs, d);

    // the number of labels
    int label_num = cc_dists.length;

    // the size of each cc
    SimpleHistogram hist = new SimpleHistogram(cc_labels);
    int[] cc_sizes = hist.getHistogram();

    // merge
    boolean will_merge = false;
    MergeMap merge_map = new MergeMap();
    for (int i=0; i<label_num; i++) {

      // the size must less than the size_threshold
      if (cc_sizes[i] >= size_threshold) continue;

      IndexedDoubleSummationHashMap dists = cc_dists[i];

      int to = pickupLabelToMerge(cc_sizes, dists, i, dist_threshold);
      if (to != i) {
        merge_map.put(i, to);
        will_merge = true;
      }
    }

    if (will_merge) {

      // the cc_labels -> labels map
      int[] cc_label_map = new int[label_num];
      for (int i = 0; i < cc_labels.length; i++)
        cc_label_map[cc_labels[i]] = labels[i];

        // merge result
      for (int i = 0; i < cc_labels.length; i++) {
        labels[i] = cc_label_map[merge_map.get(cc_labels[i])];
      }
    }

    return will_merge;
  }
      */

  public boolean mergeOnce(int[] labels, DifferentialImage diffs, Dimension d, double dist_threshold, int size_threshold) {

    ConnectComponent cc = new ConnectComponent();
    int[] cc_labels = cc.compactLabeling(labels, d, false);

    IndexedDoubleSummationHashMap[] cc_dists = getNeighborCCDist(cc_labels, diffs, d);

    // the number of labels
    int label_num = cc_dists.length;

    // the size of each cc
    SimpleHistogram hist = new SimpleHistogram(cc_labels);
    int[] cc_sizes = hist.getHistogram();

    // merge
    boolean will_merge = false;
    MergeMap merge_map = new MergeMap();
    for (int i=0; i<label_num; i++) {

      // the size must less than the size_threshold
      if (cc_sizes[i] >= size_threshold) continue;

      IndexedDoubleSummationHashMap dists = cc_dists[i];

      int to = pickupLabelToMerge(cc_sizes, dists, i, dist_threshold);
      if (to != i) {
        merge_map.put(i, to);
        will_merge = true;
      }
    }

    if (will_merge) {

      // the cc_labels -> labels map
      int[] cc_label_map = new int[label_num];
      for (int i = 0; i < cc_labels.length; i++)
        cc_label_map[cc_labels[i]] = labels[i];

        // merge result
      for (int i = 0; i < cc_labels.length; i++) {
        labels[i] = cc_label_map[merge_map.get(cc_labels[i])];
      }
    }

    return will_merge;
  }

  public boolean oneStepMerge(int[] pixels, int[] labels, DifferentialImage diffs, Dimension d, int small_region_threshold) {

    int width = d.width, height = d.height;

    // find the connected components
    ConnectComponent cc = new ConnectComponent();
    int[] cc_labels = cc.compactLabeling(labels, d, false);

    if (OUTPUT_REGIONS) {
      yi.util.io.OutputResult.outputImageRegions("rgn." + m_steps, pixels, cc_labels,
                                                 width,
                                                 height);
    }

    // mask the big regions
    NeighboringCCInfo[] ncc_infos = NeighboringCCInfo.getNeighboringCCInfos(
        cc_labels, diffs, d, small_region_threshold, -1);

    double[][] lab = new double[cc_labels.length][3];
    yi.iu.processor.color.CIELab.rgb2lab(pixels, lab);
    double[][] lab_mean = yi.util.label.Label.getMeans(lab, cc_labels);

    ClosestNeighborWithCloseMeanPicker crp = new ClosestNeighborWithCloseMeanPicker();
    Dimension closest = crp.findClosestRelation(ncc_infos, lab_mean, MERGE_THRESHOLD);

    // closest.height is the label to be merged, closest.width is the label
    // merged to
    if (closest == null) {
      return false;
    }

    int closest_from = closest.width;
    int closest_to = closest.height;

    if (DEMO_MERGE) {

      int[] two_labels = generateTwoLabels(cc_labels, closest_from,
                                           closest_to);
      int[] image = yi.util.ImageUtil.getPixels();
      int[] two_labels_demo = yi.util.ImageUtil.makeNewColorImageFromLabel(two_labels, image, true, "lab");
      yi.util.io.OutputResult.outputGIFImage("merge", two_labels_demo,
                                             d.width,
                                             d.height, true);
      int[] marked = yi.util.ImageUtil.markRegionBoundary(image, two_labels, d.width, d.height, 0xffff0000, true);
      yi.util.io.OutputResult.outputColorImage("merge_b", marked,
                                             d.width,
                                             d.height, true);

      NeighboringCCInfo[] ncc = NeighboringCCInfo.getNeighboringCCInfos(
          two_labels, diffs, d, small_region_threshold, -1);
      for (int k=0; k<ncc.length; k++) {
        for (Iterator l=ncc[k].keySet().iterator(); l.hasNext(); ) {
          Object key = l.next();
          DoubleSummationStruct dist_sum = (DoubleSummationStruct) ncc[k].get(
              key);
          System.out.println(Integer.toString(k) + "," + key + "=" + dist_sum.total + "/" + dist_sum.counter + "~" + dist_sum.getAverage());
        }
      }
//      getDistBetweenTwoRegions(two_labels, pixels, d.width, d.height);

    }

    //mergeLabels(labels, d, closest.height, closest.width, cc_labels, mask);
    mergeLabels(labels, cc_labels, closest_from, closest_to);

    if (DEMO_MERGE) {
      int[] two_labels = generateTwoLabels(cc_labels, closest_from,
                                           closest_to);
      int[] image = yi.util.ImageUtil.getPixels();
      int[] two_labels_demo = yi.util.ImageUtil.makeNewColorImageFromLabel(two_labels, image, true, "lab");
      yi.util.io.OutputResult.outputGIFImage("merge", two_labels_demo,
                                             d.width,
                                             d.height, true);
    }
    return true;

  }

//  public double getDistBetweenTwoRegions(int[] label, int[] image, int width, int height) {
//    double[][] lab = new double[label.length][3];
//    yi.iu.processor.color.CIELab.rgb2lab(image, lab);
//    int num_label = yi.stat.StatUtil.maximum(label) + 1;
//
//    yi.stat.hist.SimpleHistogram hist = new yi.stat.hist.SimpleHistogram(label);
//    System.out.println(hist);
//
//    double sum=0;
//    int count=0;
//    for (int y=0; y<height; y++) {
//      for (int x=0; x<width; x++) {
//        int p=y*width+x;
//        if (x<width-1 && ((label[p]==1&&label[p+1]==2) || (label[p]==2&&label[p+1]==1))) {
//          sum += yi.math.VectorUtil.EuclideanDistance(lab[p], lab[p+1]);
//          count++;
//        }
//        if (y < height - 1 && ((label[p] == 1 && label[p + width] == 2) ||
//            (label[p] == 2 && label[p + width] == 1))) {
//          sum += yi.math.VectorUtil.EuclideanDistance(lab[p], lab[p + width]);
//          count++;
//        }
//      }
//    }
//    System.out.println("sum=" + sum + ", count=" + count + ", avg=" + (sum/count));
//    return sum/count;
//  }


  public boolean efficientOneStepMerge(int[] pixels, int[] labels, DifferentialImage diffs, Dimension d, int small_region_threshold) {
    int width = d.width, height = d.height;

    // find the connected components
    ConnectComponent cc = new ConnectComponent();
    int[] cc_labels = cc.compactLabeling(labels, d, false);

    if (OUTPUT_REGIONS) {
      yi.util.io.OutputResult.outputImageRegions("rgn."+m_steps, pixels, cc_labels,
                                                 width,
                                                 height);
    }

    // mask the big regions

    NeighboringCCInfo[] ncc_infos = NeighboringCCInfo.getNeighboringCCInfos(
        cc_labels, diffs, d, small_region_threshold, -1);

    EfficientClosestRelationPicker crp = new EfficientClosestRelationPicker();
    LinkedList merges = crp.findClosestRelation(ncc_infos, MERGE_THRESHOLD);

    // closest.height is the label to be merged, closest.width is the label
    // merged to
    if (merges == null || merges.size() == 0) {
      return false;
    }

    //mergeLabels(labels, d, closest.height, closest.width, cc_labels, mask);
    Merge merge = new Merge(ncc_infos.length);
    for (Iterator i=merges.iterator(); i.hasNext(); ) {
      Dimension closest = (Dimension)i.next();
      int closest_from = closest.width;
      int closest_to = closest.height;
      merge.merge(closest_from, closest_to);
    }
    int[] merge_map = merge.result();
    mergeLabels(labels, cc_labels, merge_map);

    return true;

  }

  class MergeMap extends HashMap {
    void put(int from, int to) {
      put(new Integer(from), new Integer(to));
    }
    int get(int from) {
      Integer key = new Integer(from);
      if (!containsKey(key)) return from;
      else return ((Integer)get(key)).intValue();
    }
  }

  /**
   * Pick up the closest ones by their distances between boundary pixels
   */
  public class ClosestNeighborPicker {
    public Dimension findClosestRelation(NeighboringCCInfo[] ncc_infos, double threshold) {
      int size = ncc_infos.length;
      double IMPOSSIBLE_BIG_DIST = threshold;
      double highest_score = IMPOSSIBLE_BIG_DIST;
      int highest_score_from = 0, highest_score_to = 0;

      int cclabel_num = ncc_infos.length;
      for (int i = 0; i < cclabel_num; i++) {

        NeighboringCCInfo dists = ncc_infos[i];

        for (Iterator it = dists.keySet().iterator(); it.hasNext(); ) {

          // the label of the region merged to
          Integer key = (Integer) it.next();
          int merge_to = key.intValue();
          if (i==merge_to) continue;
          if (dists.size > ncc_infos[merge_to].size)
            continue;

          DoubleSummationStruct dist_sum = (DoubleSummationStruct) dists.get(
              key);

          double score = dist_sum.getAverage();

          DoubleSummationStruct from_dist_sum = (DoubleSummationStruct)dists.get(new Integer(i));
          DoubleSummationStruct to_dist_sum = (DoubleSummationStruct)ncc_infos[merge_to].get(new Integer(merge_to));
          double average_dist = (from_dist_sum.total + to_dist_sum.total)/
              (from_dist_sum.counter + to_dist_sum.counter);

          if (score < threshold && score < average_dist * MEAN_DIST_CHANGE_RATIO) {
//          if (score < highest_score) {
            highest_score = score;
            highest_score_from = i;
            highest_score_to = merge_to;
          }
        }
      }

      if (highest_score >= threshold)
        return null;

      System.out.println("closest: " + highest_score_from + "," +
                         highest_score_to +
                         " = " + highest_score);
      return new Dimension(highest_score_from, highest_score_to);
    }

  }

  /**
   * Pick up the closest ones by their distances between boundary pixels.  Those
   * regions must have a close mean
   */
  public class ClosestNeighborWithCloseMeanPicker {
    public Dimension findClosestRelation(NeighboringCCInfo[] ncc_infos, double[][] means, double t_mean) {
      int size = ncc_infos.length;
      double IMPOSSIBLE_BIG_DIST = Double.MAX_VALUE;
      double highest_score = IMPOSSIBLE_BIG_DIST;
      int highest_score_from = 0, highest_score_to = 0;

      int cclabel_num = ncc_infos.length;
      for (int i = 0; i < cclabel_num; i++) {

        NeighboringCCInfo dists = ncc_infos[i];

        for (Iterator it = dists.keySet().iterator(); it.hasNext(); ) {

          // the label of the region merged to
          Integer key = (Integer) it.next();
          int merge_to = key.intValue();
          if (i==merge_to) continue;
          if (dists.size > ncc_infos[merge_to].size)
            continue;

          DoubleSummationStruct dist_sum = (DoubleSummationStruct) dists.get(
              key);

          double dist_bound = dist_sum.getAverage();
          DoubleSummationStruct from_dist_sum = (DoubleSummationStruct)dists.get(new Integer(i));
          DoubleSummationStruct to_dist_sum = (DoubleSummationStruct)ncc_infos[merge_to].get(new Integer(merge_to));
          double average_dist = (from_dist_sum.total + to_dist_sum.total)/
              (from_dist_sum.counter + to_dist_sum.counter);

          double mean_dist = yi.math.VectorUtil.EuclideanDistance(means[i], means[key.intValue()]);


          if (mean_dist < t_mean && dist_bound < average_dist * MEAN_DIST_CHANGE_RATIO) {
//          if (score < highest_score) {
            highest_score = dist_bound;
            highest_score_from = i;
            highest_score_to = merge_to;
          }
        }
      }

      if (highest_score >= Double.MAX_VALUE)
        return null;

      System.out.println("closest: " + highest_score_from + "," +
                         highest_score_to +
                         " = " + highest_score);
      return new Dimension(highest_score_from, highest_score_to);
    }

  }

  public class EfficientClosestRelationPicker {
    /**
     * Return a list of Dimension (class Dimension is used to represent the merges)
     * @param ncc_infos
     * @param threshold
     * @return
     */
    public LinkedList findClosestRelation(NeighboringCCInfo[] ncc_infos, double threshold) {
      int size = ncc_infos.length;

      LinkedList merges = new LinkedList();

      int cclabel_num = ncc_infos.length;
      for (int i = 0; i < cclabel_num; i++) {

        NeighboringCCInfo dists = ncc_infos[i];

        for (Iterator it = dists.keySet().iterator(); it.hasNext(); ) {

          // the label of the region merged to
          Integer key = (Integer) it.next();
          int merge_to = key.intValue();
          if (i==merge_to) continue;
          if (dists.size > ncc_infos[merge_to].size)
            continue;

          DoubleSummationStruct dist_sum = (DoubleSummationStruct) dists.get(
              key);

          double score = dist_sum.getAverage();

          DoubleSummationStruct from_dist_sum = (DoubleSummationStruct)dists.get(new Integer(i));
          DoubleSummationStruct to_dist_sum = (DoubleSummationStruct)ncc_infos[merge_to].get(new Integer(merge_to));
          double average_dist = (from_dist_sum.total + to_dist_sum.total)/
              (from_dist_sum.counter + to_dist_sum.counter);

          if (score < threshold && score < average_dist * 2) {
            merges.add(new Dimension(i, merge_to));
            System.out.println("merge: " + i + "," +
                               merge_to +
                               " = " + score + " / " + average_dist);
          }
        }
      }

      return merges;
    }

  }


  private int[] generateTwoLabels(int[] labels, int l1, int l2) {
    int[] demo = new int[labels.length];
    for (int i=0; i<demo.length; i++) {
      if (labels[i] == l1) demo[i] = 1;
      else if (labels[i] == l2) demo[i] = 2;
      else demo[i] = 0;
    }
    return demo;
  }
  public void mergeLabels(int[] labels, int[] cc_labels,
                          int from, int to) {
    // cc label -> label map
    int cc_labels_num = yi.stat.StatUtil.maximum(cc_labels) + 1;
    int[] cc_labels_map = new int[cc_labels_num];
    for (int i=0; i<labels.length; i++) cc_labels_map[cc_labels[i]] = labels[i];

    int label_to = cc_labels_map[to];

    for (int i=0; i<labels.length; i++) {
      if (cc_labels[i] == from) {
        labels[i] = label_to;
        cc_labels[i] = to;
      }
    }
  }

  public void mergeLabels(int[] labels, int[] cc_labels,
                          int[] merge_maps) {
    // cc label -> label map
    int cc_labels_num = yi.stat.StatUtil.maximum(cc_labels) + 1;
    int[] cc_labels_map = new int[cc_labels_num];
    for (int i=0; i<labels.length; i++) cc_labels_map[cc_labels[i]] = labels[i];

    for (int i=0; i<labels.length; i++) {
      labels[i] = cc_labels_map[merge_maps[cc_labels[i]]];
    }
  }

  public int[] seg(int[] pixels, int[] labels, Dimension d, int small_region_threshold) {

    int max = yi.stat.StatUtil.maximum(labels);
//    DifferentialImage diffs = new HSIDifferentialImage(pixels, d);
    DifferentialImage diffs = new CIELabDifferentialImage(pixels, d);

    while (true) {
      m_steps++;
      System.out.println("step# = " + m_steps);
      if (DEMO_STEPS) {
//        int[] rst = yi.util.ImageUtil.makeNewImageFromLabel(labels, max, false);
        int[] image = yi.util.ImageUtil.getPixels();
        int[] rst = yi.util.ImageUtil.makeNewColorImageFromLabel(labels, image, true, "lab");
        yi.util.io.OutputResult.outputGIFImage("merge", rst, d.width,
                                               d.height, true);
      }

      if (!oneStepMerge(pixels, labels, diffs, d, small_region_threshold))
        break;
//      if (!efficientOneStepMerge(pixels, labels, diffs, d, small_region_threshold))
//        break;
    }
    if (DEMO_STEPS) {
//      int[] rst = yi.util.ImageUtil.makeNewImageFromLabel(labels, max, false);
      int[] image = yi.util.ImageUtil.getPixels();
      int[] rst = yi.util.ImageUtil.makeNewColorImageFromLabel(labels, image, true, "lab");
      yi.util.io.OutputResult.outputGIFImage("merge", rst, d.width,
                                             d.height, true);
    }

    return labels;

  }

}