package yi.iu.processor.color;

import java.awt.*;
import java.util.*;

import yi.util.config.*;
import yi.util.io.*;
import yi.util.io.input.*;
import yi.util.io.output.*;
import yi.util.log.*;
import yi.iu.imgdb.learn.FeatureRepresentable;
import yi.iu.processor.color.reg.CIELabLabeledRgnFts;
import yi.iu.imgdb.learn.RepresentableFeatureGenerator;

/**
 * <p>
 * Class DPU_DominantColor finds the dominant colors of an image:
 * </p>
 *
 * <pre>
 * 1. RGB->CIEL*a*b*
 * 2. calculate the number of pixels in any colors
 *    (using hash-table to save memory, not good when lots of colors in the image)
 * 3. find the dominant colors
 * </pre>
 *
 * <p>
 * Options:
 * </p>
 *
 * <pre>
 * 1. Use class QuanLab to get the index of a CIELab color; QuanLab use 3 bin
 *    size arguments, m_lShrinkFactor, m_aShrinkFactor, and m_bShrinkFactor
 * 2. TCD, MAX_CLUSTER_NUM, MAX_CLUSTERED_PXL (see the comments)
 * </pre>
 *
 * <P=p>
 * Example:
 * </p>
 *
 * <code><pre>
 *   DPU_DominantColor dp = new DPU_DominantColor(); <br/>
 *   DS_DominantColor rst = dp.getReusableDominantColors(pixels, width, height); <br/>
 *   int[] map = rst.getColorMap(); // pixel->dominant color map <br/>
 *   Collection colors = rst.getDominantColors(); // dominant colors <br/>
 *   int   cNum = colors.size(); // number of dominant colors <br/>
 * </pre></code>
 *
 * @author Yi Li
 * @version 1.0
 */
public class DPU_DominantColor implements Configurable, RstReusable, RepresentableFeatureGenerator{

  /**
   * Minmum distance between two color clusters.  If the distance between two
   * clusters is less than this number, they will be merged.
   */
  public double MIN_CLUSTER_DIST = 9;

  /**
   * Maximum number of color clusters.
   */
  public int    MAX_CLUSTER_NUM = 20;

  /**
   * It is currently not used.
   * Maximum number of clustered pixels.  If the number of clustered pixels
   * are higher than this number, the clustering algorithm will stop.  This
   * value is represented as a percentage.
   */
//  public double MAX_CLUSTERED_PXL = 0.9;

  /**
   * It is currently not used.  This can be used to stop the clustering
   * algorithm.  It is the minimum size of the color bin qualified to consider.
   * If no color bin can satify this rule, then stop.
   */
//  public double MIN_CLR_BIN_SIZE = 0.18;

  /**
   * At the final step, every pixel will be assigned to a closest dominant color.
   * This is only true when their distance is not greater than this maximum
   * assignable distance.
   */
  public double MAX_ASSGN_DIST = 100;

  /**
   * A CIEL*a*b* color will be first indexed and put into a hashtable (as
   * histogram).  This is the length of a bin for component L
   */
  public int    L_SHRINK_FACTOR = 3;

  /**
   * A CIEL*a*b* color will be first indexed and put into a hashtable (as
   * histogram).  This is the length of a bin for component a*
   */
  public int    A_SHRINK_FACTOR = 3;

  /**
   * A CIEL*a*b* color will be first indexed and put into a hashtable (as
   * histogram).  This is the length of a bin for component b*
   */
  public int    B_SHRINK_FACTOR = 3;

  /**
   * Cluster merge distance will be increased during the clustering, but can't
   * be too big.  If higher than MIN_CLUSTER_DIST * MAX_EXTENDED_DIST, stop.
   */
  public double MAX_EXTENDED_DIST = 1.5;

  /**
   * The color is considered according to how many pixels have that color.  If
   * the number is lower than MAX_PXL_NUM_WORTH_CONSIDER, stop.
   */
  public int    MAX_PXL_NUM_WORTH_CONSIDER = 10;

  public int    LOG_LEVEL = 1;

  public int    REUSE  = DIU.INPUT_NONE;
  public int    OUTPUT = DOU.OUTPUT_NONE;

  int COLOR_HASHTABLE_INIT_SIZE = 100;

  public String[] getParameterNames() {
    return new String[] {
      "MIN_CLUSTER_DIST",
      "MAX_CLUSTER_NUM",
      "MAX_CLUSTERED_PXL",
      "MIN_CLR_BIN_SIZE",
      "MAX_ASSGN_DIST",
      "MAX_EXTENDED_DIST",
      "MAX_PXL_NUM_WORTH_CONSIDER",
      "L_SHRINK_FACTOR",
      "A_SHRINK_FACTOR",
      "B_SHRINK_FACTOR",
      "LOG_LEVEL",
      "REUSE",
      "OUTPUT"
    };
  }

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

  public DPU_DominantColor(double mincd, int maxcn, double maxad,
    int lsf, int asf, int bsf) {
    this();
    MIN_CLUSTER_DIST =mincd;
    MAX_CLUSTER_NUM  =maxcn;
    MAX_ASSGN_DIST   =maxad;
    L_SHRINK_FACTOR  = lsf;
    A_SHRINK_FACTOR  = asf;
    B_SHRINK_FACTOR  = bsf;
  }


  /**
   * The wrapper function of getDominantColors plus usable functionility
   * @param source Source image stored in integer array.  (0xffrrggbb format).
   * @param width image width
   * @param height image height
   * @return an instance of class DS_DominantColor.  (see the comments in class
   *         DS_DominantColor for details.)
   * @see #getReusableDominantColors
   */
  public DS_DominantColor getReusableDominantColors(int source[], int width, int height) {
      DS_DominantColor rst = (DS_DominantColor)reuse(REUSE);

      if (rst == null) {
        rst = getDominantColors(source, width, height);
      }

      output(OUTPUT, rst);

      return rst;
  }

  /**
   * The main process of DominantColor.  Given an image in integer (0xffrrggbb)
   * array, it return the dominant colors (in CIELab format).
   *
   * @param source Source image stored in integer array.  (0xffrrggbb format).
   * @param width image width
   * @param height image height
   * @return an instance of class DS_DominantColor.  (see the comments in class
   *         DS_DominantColor for details.)
   */
  public DS_DominantColor getDominantColors(int source[], int width, int height)
  {
      int i, j;

      // get the CIE L*a*b* color space
      LogHandler.log(this, LOG_LEVEL, 3, "Get the dominant colors");

      int[] lab= CIELab.rgb2lab(source);  // adjusted labs

      // clustering
      int[][] cluster= new int[MAX_CLUSTER_NUM][4];   //store the info of cluster
                                          //[0]-size [1]-sum of L [2]-sum of a
                                          //[3]-sum of b
      int cNum = clustering(lab, cluster);

      //calculate the mean color of each cluster
      double[][] cluster_color= new double[cNum][3];
      for (i=0;i<cNum;i++) {
          cluster_color[i][0]= cluster[i][1]/(double)cluster[i][0];
          cluster_color[i][1]= cluster[i][2]/(double)cluster[i][0];
          cluster_color[i][2]= cluster[i][3]/(double)cluster[i][0];
      }

      //assign each pixel to one of the dominant color clusters
      int[] label = new int[source.length];
      for (i=0;i<source.length;i++) {
          //find the distance of that color to all the existing clusters
          double close_d=MAX_ASSGN_DIST;    //the smallest distance
          int close_i=-1;     //the cluster with smallest distance
          for (j=0;j<cNum;j++) {
              double dist=CIELab.edistance(cluster[j][1]/(double)cluster[j][0],
                  cluster[j][2]/(double)cluster[j][0],
                  cluster[j][3]/(double)cluster[j][0],
                  CIELab.getL(lab[i]),
                  CIELab.getA(lab[i]),
                  CIELab.getB(lab[i])
                  );
              if (dist<=close_d) {
                  close_d=dist; close_i=j;
              }
          }
          //if no cluster close to this color, new a cluster
          if (close_i<0) {
              label[i]=cNum;
          }
          else {
              label[i]=close_i;
          }
      }

      lab=null;cluster=null;System.gc();
      DS_DominantColor rst = new DS_DominantColor(cluster_color, label, width, height);

      return rst;
  }

  /**
   * Cluster image pixels.  Image pixels are given in Lab format.  The algorithm
   * first finds the most import color (the color with maximum number of pixels)
   * unclustered, then tries to classify it to the existing clusters if they are
   * close enough.  If there is no such a close cluster, start a new cluster.
   * If the cluster number is greater than the predefine MAX_CLUSTER_NUM, loose
   * the minmum distance (among clusters) threshold, MIN_CLUSTER_DIST.  If there
   * is no import color (no one with large enough amount of pixels) or the
   * distance is relaxed too much, stops.
   *
   * @param lab     the Lab values of image pixels.
   * @param cluster return the resultant clusters.  This passed in parameter
   *                should be inited as int[MAX_CLUSTER_NUM][4] before calling
   *                the function.  The real number of clusters is returned as
   *                the return value of the function.  This array is used to
   *                return the details of each color cluster:
   *                [0]-size [1]-sum of L [2]-sum of a [3]-sum of b
   * @return        the real number of color clusters after clustering
   */
  int clustering(int[] lab, int[][] cluster) {
      // First, index the color values
      // and make a histogram with respect to those indexes
      // in a bin, a integer array used to store:
      // 0: the total number of color index falling into this bin
      // 1: the summary of value of L component of those lab colors
      // 2: the summary of value of a component of those lab colors
      // 3: the summary of value of b component of those lab colors
      // Because the histogram is sparse, it will be stored in a hashtable
      // instead of in a vector-like data structure
      LogHandler.log(this, LOG_LEVEL, 3, "Quantilization");

      LabIndexMaker lab_index =
        new LabIndexMaker(L_SHRINK_FACTOR, A_SHRINK_FACTOR, B_SHRINK_FACTOR);
      int indexSize= lab_index.size();

      // the hashtable used to store the color histogram
      Hashtable hist= new Hashtable(COLOR_HASHTABLE_INIT_SIZE);
      int i, index;
      for (i=0;i<lab.length;i++) {
          index= lab_index.getIndex(lab[i]);
          if (index>=indexSize) {
              System.err.println("  index:"+index+" >="+indexSize);
              continue;
          }
          Object key= new Integer(index);
          if (hist.containsKey(key)) {
              int[] info= (int[])hist.get(key);
              info[0]++;
              info[1]+=CIELab.getL(lab[i]);
              info[2]+=CIELab.getA(lab[i]);
              info[3]+=CIELab.getB(lab[i]);
          }
          else {
              int[] info={1,
                CIELab.getL(lab[i]),CIELab.getA(lab[i]),CIELab.getB(lab[i])};
              hist.put(key,info);
          }
      }

      LogHandler.log(this, LOG_LEVEL, 4,
        "non-empty=" + hist.size() + " %=" + hist.size()/(double)indexSize);

      //clustering
      LogHandler.log(this, LOG_LEVEL, 3, "Clustering");
      int cNum=0;                         //number of cluster
      int consideredNum=0;                //number of considered pixel

      int j=0;
      int loopC=0;    // total loop number of the clustering
      double Td= MIN_CLUSTER_DIST;
              // minmum distance between two cluster to not be merged
      while (true){
          //System.err.println("Considered .. "+consideredNum);

          // find the color with the most pixels
          int max=0;
          int[] max_info=null;
          for (Enumeration e = hist.elements() ; e.hasMoreElements() ;) {
              int[] info=(int[])e.nextElement();
              if (info[0]>max) {
                  max=info[0];max_info=info;
              }
          }
          if (max ==0) break;
          // if the max size of the colors is small enough, break from clustering
          if (max < MAX_PXL_NUM_WORTH_CONSIDER) {
            LogHandler.log(this, LOG_LEVEL, 4, "breaking on max = " + max);
            break;
          }

          // find the smallest distance of that color to all the existing clusters
          double close_d = Td;    //the smallest distance
          int    close_i=-1;      //the cluster with smallest distance
          for (j=0;j<cNum;j++) {
              double dist=CIELab.edistance(
                  cluster[j][1]/(double)cluster[j][0],
                  cluster[j][2]/(double)cluster[j][0],
                  cluster[j][3]/(double)cluster[j][0],
                    max_info[1]/(double)max_info[0],
                    max_info[2]/(double)max_info[0],
                    max_info[3]/(double)max_info[0]
                  );
              if (dist < close_d) {
                  close_d = dist; close_i = j;
              }
          }

          if (cNum < MAX_CLUSTER_NUM) {
            //if no cluster close to this color, start a new cluster
            if (close_i<0) {
                cluster[cNum]=(int[])(max_info.clone());
                cNum++;
                max_info[0]=0; //mark the color to considered
            }
            //otherwise combine with the closest cluster
            else {
                cluster[close_i][0]+=max_info[0];
                cluster[close_i][1]+=max_info[1];
                cluster[close_i][2]+=max_info[2];
                cluster[close_i][3]+=max_info[3];
                max_info[0]=0; //mark the color to considered
            }
            consideredNum+= max;
            loopC++;
          }

          // if the number of found clusters is greater than the predefined
          // value, MAX_CLUSTER_NUM, then try to merge the close clusters and
          // loose the merge threshold if necessary
          if (cNum>=MAX_CLUSTER_NUM) {
              cNum= merge(cluster, cNum, Td);
              // if after merge, number of clusters became smaller, continue;
              if (cNum < MAX_CLUSTER_NUM) break;
              // loose the merge threshold, and try again
              Td ++;
              if (Td > MIN_CLUSTER_DIST * MAX_EXTENDED_DIST) {
                LogHandler.log(this, LOG_LEVEL, 4, "breaking on Td = " + Td);
                break;
              }
              cNum= merge(cluster, cNum, Td);
              // after merge, the number of cluster may remain the same.
              // However, in the next step of clustering, after a new color
              // is considered and merged into a existing cluster, the
              // distances may change and new merge may happen
              LogHandler.log(this, LOG_LEVEL, 4,
                "considered pixels = "+consideredNum/(double)lab.length);
              LogHandler.log(this, LOG_LEVEL, 4,
                "use new Threshold for distance = "+Td);
          }
      }
      LogHandler.log(this, LOG_LEVEL, 3, "Done");
      LogHandler.log(this, LOG_LEVEL, 3, "total clusters = "+cNum);
      LogHandler.log(this, LOG_LEVEL, 3,
        "considered pixels = " + consideredNum/(double)lab.length);
      LogHandler.log(this, LOG_LEVEL, 3, "loop counter = "+loopC);

      //combine the close clusters
      cNum= merge(cluster, cNum, Td);
      LogHandler.log(this, LOG_LEVEL, 3, "after post-combine, cluster number is "+cNum);

      return cNum;
  }

  /**
   * Re-calculate the distance between each pair of clusters, merge them if the
   * distance is less than the threshold, Td.  This function may be called when
   * Td is gotten relaxed (bigger).  In addition, although the dist. between
   * color c1 and c2 satisfies the threshold (less than Td), after c3 is
   * classified as the same cluster of c1, the dist. between the new cluster and
   * c2 may not.  (the mean value of the new cluster changed)
   *
   * @param cluster the clusters. [MAX_CLUSTER_NUM][4]. [0]-size(number of
   *                pixels); [1]-sum of L; [2]-sum of a; [3]-sum of b.
   * @param cNum    the real number of clusters (the array cluster may have empty
   *                space at the end.
   * @param Td      merge threshold (Td stands for "threshold of distance")
   * @return        the number of clusters after merge
   */
  int merge(int[][] cluster, int cNum, double Td)
  {
      int i,j;
      int old_cNum;
      do {
        old_cNum = cNum;
        for (i=0;i<cNum-1;i++) {
            for (j=i+1;j<cNum;j++) {
                double dist=CIELab.edistance(
                    cluster[j][1]/(double)cluster[j][0],
                    cluster[j][2]/(double)cluster[j][0],
                    cluster[j][3]/(double)cluster[j][0],
                    cluster[i][1]/(double)cluster[i][0],
                    cluster[i][2]/(double)cluster[i][0],
                    cluster[i][3]/(double)cluster[i][0]
                    );
                if (dist<MIN_CLUSTER_DIST) {
                    //combine
                    cluster[i][0]+=cluster[j][0];
                    cluster[i][1]+=cluster[j][1];
                    cluster[i][2]+=cluster[j][2];
                    cluster[i][3]+=cluster[j][3];
                    //shift the clusters to eliminate the combined
                    for (int k=j;k<cNum-1;k++)
                        cluster[k]=cluster[k+1];
                    cNum--;
                }
            }
        }
      } while (cNum != old_cNum);
      return cNum;
  }

  public void setReuse(int reuse) {
    REUSE = reuse;
  }

  public Object reuse(int reuse_level) {
    if (reuse_level != DIU.INPUT_NONE) {
      try {
        DS_DominantColor rst = (DS_DominantColor)DIU.in(
            reuse_level,
            DOU.shortName(this.getClass().getName()),
            "DS_DominantColor",
            Class.forName("yi.iu.processor.color.DS_DominantColor"));
        return rst;
      }
      catch (Exception e) {
        LogHandler.log(this, LOG_LEVEL, 1, "Can't Reuse DS_DominantColor:\n" + e);
      }
    }
    return null;
  }

  public void output(int output_level, DS_DominantColor ds){
    if (output_level != DOU.OUTPUT_NONE) {
      try {
        DOU.out(
            output_level,
            DOU.shortName(this.getClass().getName()),
            DOU.shortName(ds.getClass().getName()),
            ds);
      } catch (DataOutputException e) {
        LogHandler.log(this, LOG_LEVEL, 1, e.toString());
      }
    }
  }

  /**
   * <p>
   * Class LabIndexMaker is to index a color instance in CIE L*a*b* space.  It will
   * return an integer as the index.
   * </p>
   *
   * <p>
   * A lab value is in the format of 0xllaaabbb, where
   * </p>
   *
   * <pre>
   * l:[0..100], a:[0..1000], b:[0..400]
   * </pre>
   *
   * <p>
   * An index will be very big if we leave lab as it is.  So first the values on
   * L, a, and b will be shrunken by a shrink factor, say 3; then the biggest
   * index will be only 1/9 of the one before
   * </p>
   *
   * <p>
   * Example: Suppose the shrink factors are 3 for every component of a Lab
   * color.  The ranges of Lab will be [0,33], [0,333], [0,133].  A color
   * (50, 200, 300) will be shrinked to (16, 66, 100), and the index
   * is 16*333*133+66*133+100
   * </p>
   */
  class LabIndexMaker {

      int m_lShrinkFactor = 3;
      int m_aShrinkFactor = 3;
      int m_bShrinkFactor = 3;

      int m_sizeL =  100/m_lShrinkFactor+1;
      int m_sizeA = 1000/m_aShrinkFactor+1;
      int m_sizeB =  400/m_bShrinkFactor+1;

      int m_sizeAB = m_sizeA * m_sizeB;

      public LabIndexMaker(int lShrinkFactor, int aShrinkFactor, int bShrinkFactor) {
        m_lShrinkFactor = lShrinkFactor;
        m_aShrinkFactor = aShrinkFactor;
        m_bShrinkFactor = bShrinkFactor;
      }

      /**
       * Given a lab color in format of 0xllaabb, return its index.
       * @param lab a lab color.
       * @return the index.
       */
      int getIndex(int lab) {
          int l=(lab>>24)&0xff;
          int a=(lab>>12)&0xfff;
          int b=(lab)&0xfff;
          l/=m_lShrinkFactor;a/=m_aShrinkFactor;b/=m_bShrinkFactor;

          return (l* m_sizeAB + a * m_sizeB + b);
      }

      /**
       * @return the shrinked size of lab space
       */
      int size() {
          return m_sizeL * m_sizeA * m_sizeB;
      }
  }

  public int[] getDemo(int[] pixels, int width, int height) {
    DS_DominantColor rst = getReusableDominantColors(pixels, width, height);

    int[] demo= new int[pixels.length];
    int[] map = rst.getColorMap();
    Collection colors = rst.getDominantColors();
    int   cNum = colors.size();

    // change to RGB colors
    int[] rgb_colors = new int[colors.size()];
    {
      Iterator i = colors.iterator();
      int j = 0;
      while (i.hasNext()) {
        rgb_colors[j++] = CIELab.lab2rgb((CIELab)i.next());
      }
    }

    // demo image
    for (int i=0;i<demo.length;i++) {
      if (map[i] >= cNum) demo[i]=0xffff0000;
      else demo[i] = rgb_colors[map[i]];
    }

//    // demo the color regions one by one
//    for (int i=0; i<cNum; i++) {
//      int[] region_img = new int[pixels.length];
//      int color = rgb_colors[i];
//      int bk_color = ColorUtil.inverseColor(color);
//      for (int j=0; j<region_img.length; j++) {
//        if (map[j] == i) region_img[j] = color;
//        else region_img[j] = bk_color;
//      }
//      String img_name = Configuration.getStringParameter("SOURCE_NAME", "unknown");
//      OutputResult.outputGIFImage(img_name + ".color." + i + ".gif", region_img, width, height);
//    }

    return demo;
  }

  public FeatureRepresentable generateFeatureRepresentable(String feature_type, int[] pixels, int width, int height) {
    DS_DominantColor ds = getDominantColors(pixels, width, height);
    int[] labels = ds.m_colorMap;

    ArrayList labs_list = ds.m_dominantColors;
    double[][] labs = new double[labs_list.size()][3];
    int i=0;
    for (Iterator it=labs_list.iterator(); it.hasNext(); i++) {
      CIELab lab = (CIELab)it.next();
      labs[i][0] = lab.L;
      labs[i][1] = lab.a;
      labs[i][2] = lab.b;
    }

    CIELabLabeledRgnFts rst = new CIELabLabeledRgnFts(labs, labels, width, height, width, height, false);
    return rst;
  }

  public static void main(String[] args) {
    // init configuration
    Configuration.init(args);

    // get the image name
    String img_name = ReadArgu.getOption('i', args);
    if (img_name == null) {
      System.err.println("Usage: -i image_name");
      System.exit(1);
    }

    // read in image
    Image image = null;
    Dimension d = new Dimension();
    int[] pixels = null;
    try {
      image = yi.util.io.ImageReader.readImageFile(img_name);
      pixels = yi.util.io.ImageReader.getPixels(image, d);
    }
    catch (java.io.IOException e) {
      System.err.println("Can't read image " + img_name);
      e.printStackTrace();
      System.exit(2);
    }


    // process
    Configuration.setParameter("SOURCE_NAME", img_name);
    DPU_DominantColor dp = new DPU_DominantColor();
    int[] rst_img = dp.getDemo(pixels, d.width, d.height);

    // output result
    OutputResult.outputGIFImage(img_name + "." + DOU.shortName(dp.getClass().getName()) + ".gif", rst_img, d.width, d.height);

    System.exit(0);
  }
}