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.FeatureRepresentationList;
import yi.iu.processor.color.reg.*;
import yi.iu.processor.resample.*;
import yi.iu.imgdb.learn.FeatureRepresentable;
import yi.iu.imgdb.learn.RepresentableFeatureGenerator;

/**
 * Reference:
 * "Advancing Content-based Image Retrieval By Exploiting Image Color and Region Features"
 * Yihong Gong
 */
public class DPU_GongProminentColor extends DPU_DominantColor implements RepresentableFeatureGenerator{

  public int RESAMPLER = Resampler.NEAREST_RESAMPLER;
  public int    SEED_NUM = 20;
  /**
   * If found two color regions are similar to merge, then check whether their
   * radius with this threshold.  If one of the them has a greater radius (than
   * this threshold), then the merge will be dropped.
   */
  public double MAX_MERGABLE_RADIUS = 6;

  public boolean DEMO_STEP = false;

  public String[] getParameterNames() {
    return new String[] {
      "MIN_CLUSTER_DIST",
      "SEED_NUM",
      "MAX_MERGABLE_RADIUS",
      "L_SHRINK_FACTOR",
      "A_SHRINK_FACTOR",
      "B_SHRINK_FACTOR",
      "LOG_LEVEL",
      "REUSE",
      "OUTPUT"
    };
  }

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

  /**
   * The main process of DominantColor.  Given an image in integer (0xffrrggbb)
   * array, it return the dominant colors.
   *
   * @param source Source image stored in integer array.  (0xffrrggbb format).
   * @param d Dimension of the image.
   * @param label Output the image in terms of the label of dominant colors.
   *              The dominant colors are returned as an array.  The array
   *              index is used to represent to which dominant color a pixel
   *              is clustered.  If there are N dominant colors, then the
   *              valid values of this label array are 0,..,N.  If a pixel is
   *              assigned to no dominant color, a N is assigned to it.
   * @return An array storing the dominant colors.  Each dominant color is
   *         a 3-length double array having the values of L, a, b.
   */
  public DS_DominantColor getDominantColors(int source[], int width, int height)
  {
      DS_DominantColor rst = (DS_DominantColor)reuse(REUSE);
      if (rst != null) {
        output(OUTPUT, rst);
        return rst;
      }

      // 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

      // index the color values
      // and make a histogram with respect to those index
      // 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
      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();

      Hashtable hist= new Hashtable(COLOR_HASHTABLE_INIT_SIZE);
      for (int i=0;i<lab.length;i++) {
          int 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);

      // find seeds
      // 1. Project the input image into the HVC color space. Initialize the
      //    seed set to empty SEEDS = 0, and let k = 1.
      // 2. In the HVC color space, find a color Zk that has the largest pixel
      //    count, and that has not been marked as a seed. If SEEDS is empty,
      //    add Zk to SEEDS; otherwise, compare Zk with each Zj in SEEDS. If Zk
      //    satisfies
      //        min(E(Zk,Zj)) > Tsd
      //    add Zk to SEEDS. In the case that Zk is added to SEEDS, mark Zk as
      //    a seed in the color space, and increment k by 1.
      // 3. If k >Ts, or no colors with a non-zero pixel count are left in the
      //    color space, go to Color Clustering; otherwise go to Step 2.

      LogHandler.log(this, LOG_LEVEL, 3, "find seeds");
      Vector  seeds  = new Vector(SEED_NUM);
      HashSet marked = new HashSet();
      do {
        // find a color Zk that has the largest pixel count and that has not
        // been marked as a seed
        int max=0;
        int[] max_info=null;
        Integer max_key = null;
        for (Enumeration i = hist.keys(); i.hasMoreElements();) {
          Integer key  = (Integer)i.nextElement();
          int[]   info = (int[])hist.get(key);
          if (info[0]>max && !marked.contains(key)) {
            max=info[0];max_info=info;max_key=key;
          }
        }
        if (max ==0) break;

        // mark it anyway
        marked.add(max_key);

        // If Zk satisfies min(E(Zk,Zj)) > Tsd
        boolean satisfied = true;
        for (Iterator i=seeds.iterator(); i.hasNext();) {
          double[] seed = (double[])i.next();
          double dist=CIELab.distance(
              seed[0],
              seed[1],
              seed[2],
              max_info[1]/(double)max,
              max_info[2]/(double)max,
              max_info[3]/(double)max
          );
          if (dist <= MIN_CLUSTER_DIST) {
            satisfied = false;
            break;
          }
        }

        // add to seeds
        if (satisfied) {
          seeds.add(new double[] {
            max_info[1]/(double)max,
            max_info[2]/(double)max,
            max_info[3]/(double)max});
        }
      } while (seeds.size() < SEED_NUM);

      LogHandler.log(this, LOG_LEVEL, 4, "found seeds # : " + seeds.size());

      // demonstrate the original seeds
      if (DEMO_STEP) {
        String img_name = Configuration.getStringParameter("SOURCE_NAME", "undef");
        demo_seeds(seeds, img_name + ".seed0.gif");
      }

      // Color Clustering:
      // 1. For each seed Zj in SEEDS, form a empty cluster Sj with radius
      //    R(Sj) = 0.0.
      // 2. In the HVC color space, for each color Pr with a non-zero pixel
      //    count, find a seed Zx in SEEDS that satisfies
      //        E(Pr, Zx) = min{E(Pr, Zj)} 1<=j<=k
      //    and add Pr to cluster Sx.
      // 3. For each cluster Sx, update its seed Zx and calculate its radius
      //    R(Sx).  The updated Zx is equivalent to the center color (averaged
      //    color) of cluster Sx, while R(Sx) equals the averaged NBS distance
      //    of all the colors included in Sx to the cluster center Zx.
      // 4. For any cluster pairs Si and Sj with their center distance and
      //    radiuses satisfying:
      //        ENBS(Zi, Zj) <= Tsd
      //        R(Si) < Tcr; R(Sx) < Tcr
      //    merge Si and Sj to form a new cluster Si' with its center
      //    recalculated.  Decrement the cluster counter k by 1.
      // 5. If no merge occurs in Step 4, terminate the operation; otherwise go
      //    to Step 2.
      do{
        // step 2
        ArrayList[] clusters = new ArrayList[seeds.size()];
        for (int i=0; i<clusters.length; i++) clusters[i] = new ArrayList();
        for (Enumeration e = hist.elements(); e.hasMoreElements();) {
          int[] c = (int[])e.nextElement();
          double[] cc = new double[] {
                            c[1]/(double)c[0],
                            c[2]/(double)c[0],
                            c[3]/(double)c[0]};
          // find the closest seed
          double close_d = 0xffff;    //the smallest distance
          int    close_i  = -1;      //the cluster with smallest distance
          for (int j=0; j<seeds.size(); j++) {
            double[] seed = (double[])seeds.elementAt(j);
            double dist = CIELab.distance(
                            seed[0], seed[1], seed[2],
                            cc[0], cc[1], cc[2]);
            if (dist < close_d) {
                  close_d = dist; close_i = j;
            }
          }
          clusters[close_i].add(c);
        }

        // For each cluster Sx, update its seed Zx and calculate its radius R(Sx).
        double[][] average = new double[clusters.length][];
        double[]   radius  = new double[clusters.length];
        for (int i=0; i<clusters.length; i++) {
          average[i] = average(clusters[i]);
          radius[i] = radius(average[i], clusters[i]);
        }

        // merge
        // if a cluster j is merged to another one the average[j] will be set to null
        boolean didMerge = false;
        for (int i=0; i<clusters.length-1; i++) {
          if (average[i]==null) continue;
          for (int j=i+1; j<clusters.length; j++) {
            if (average[j]==null) continue;
            if (CIELab.distance(average[i][0], average[i][1], average[i][2],
                                average[j][0], average[j][1], average[j][2]) >
                MIN_CLUSTER_DIST) continue;
            if (radius[i] >= MAX_MERGABLE_RADIUS || radius[j] >= MAX_MERGABLE_RADIUS) {
              continue;
            }
            merge(clusters[i], clusters[j]);
            average[i] = average(clusters[i]);
            radius[i] = radius(average[i], clusters[i]);
            average[j] = null;
            didMerge = true;
          }
        }

        // construct new seeds
        seeds.clear();
        for (int i=0; i<average.length; i++) {
          if (average[i]!=null) seeds.add(average[i].clone());
        }

        if (!didMerge) break;
      } while (true);

      // demo the seeds after merging
      if (DEMO_STEP) {
        String img_name = Configuration.getStringParameter("SOURCE_NAME", "undef");
        demo_seeds(seeds, img_name + ".seed1.gif");
      }

      // assign the pixel to dominant colors
      int[] label = new int[source.length];
      for (int i=0;i<source.length;i++) {
          double close_d=0xffff;    //the smallest distance
          int close_i=-1;     //the cluster with smallest distance
          for (int j=0; j<seeds.size(); j++) {
            double[] seed = (double[])seeds.elementAt(j);
            double dist=CIELab.distance(seed[0], seed[1], seed[2],
                CIELab.getL(lab[i]),
                CIELab.getA(lab[i]),
                CIELab.getB(lab[i])
                );
            if (dist<=close_d) {
                close_d=dist; close_i=j;
            }
          }
          label[i]=close_i;
      }


      lab=null;hist=null;System.gc();
      double[][] dc = new double[seeds.size()][];
      for (int i= 0; i< seeds.size(); i++) dc[i] = (double[])seeds.elementAt(i);
      rst = new DS_DominantColor(dc, label, width, height);

      return rst;
  }

  /**
   * Calculate the average of a color cluster.  Each element in the arraylist
   * is a int[4], with size of pixel stored in [0], sum of l, a, b stored in
   * [1],[2],[3]
   * @param a
   * @return
   */
  private double[] average(ArrayList a) {
    double[] sum = new double[4];
    for (Iterator i = a.iterator(); i.hasNext();) {
      int[] data = (int[])i.next();
      sum[0] += data[0];
      sum[1] += data[1];
      sum[2] += data[2];
      sum[3] += data[3];
    }
    return new double[] {sum[1]/sum[0],sum[2]/sum[0],sum[3]/sum[0]};
  }

  private double radius(double[] average, ArrayList a) {
    double sum = 0;
    int count = 0;
    for (Iterator i = a.iterator(); i.hasNext();) {
      int[] data = (int[])i.next();
      double dist = CIELab.squaredDistance(average[0], average[1], average[2],
                        data[1]/(double)data[0], data[2]/(double)data[0], data[3]/(double)data[0]);
      sum += dist*data[0];
      count += data[0];
    }
    return Math.sqrt(sum/count);
  }

  private void merge(ArrayList a, ArrayList b) {
    a.addAll(b);
    b.clear();
  }

  /**
   * Demo the colors of the seeds
   * @param seeds: a vector of double[] corresponding to l, a, b, respectively
   */
  private void demo_seeds(Vector  seeds, String file_name) {
    // save the RGB colors of the seeds
    int[] colors = new int[seeds.size()];

    // convert each seed color from lab to rgb
    int i_index = 0;
    for (Iterator i=seeds.iterator(); i.hasNext(); i_index++){
      double[] info = (double[])i.next();
      double l = info[0];
      double a = info[1];
      double b = info[2];
      colors[i_index] = yi.iu.processor.color.CIELab.lab2rgb(l, a, b);
    }
    int spectrum_blk_w = 20;
    int spectrum_blk_h = 20;
    int[] spectrum = yi.iu.processor.color.ColorUtil.createSpectrum(colors, spectrum_blk_w, spectrum_blk_h);
    yi.util.io.OutputResult.outputGIFImage(file_name, spectrum, spectrum_blk_w*colors.length, spectrum_blk_h);
  }

  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_GongProminentColor dp = new DPU_GongProminentColor();
    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);
  }

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

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

  public CIELabLabeledRgnFts getCIELabLabeledRgnFts(int[] pixels, int width, int height, int old_w, int old_h) {
    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, old_w, old_h, false);
    return rst;
  }

  public CIELabLabeledRgnFts getReusableCIELabLabeledRgnFts(int[] pixels, int width, int height, int old_w, int old_h)
  {

     CIELabLabeledRgnFts rst = (CIELabLabeledRgnFts)reuseCIELabLabeledRgnFts(REUSE);

     if (rst == null) {
       rst = getCIELabLabeledRgnFts(pixels, width, height, old_w, old_h);
     }

     outputCIELabLabeledRgnFts(OUTPUT, rst);

     return rst;
  }

  /**
   * Generate the CIELabLabeledRgnFts (region features with labels)
   */
  public static class GenCIELabLabeledRgnFts {

    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);
     }

     String from_path = ReadArgu.getOption('f', args, "");
     String to_path = ReadArgu.getOption('t', args, "");


     // get the scale
     String scale_str = ReadArgu.getOption('s', args);
     int scale = 1;
     if (scale_str != null) {
       scale = Integer.parseInt(scale_str);
     }

     // read in image
     Image image = yi.util.io.ImageReader.JAIReadImage(from_path + img_name);
     Dimension d = new Dimension();
     int[] pixels = yi.util.io.ImageReader.getPixels(image, d);
     Dimension old_d = (Dimension)d.clone();

     // rescale the image if necessary
     double the_scale = 1.0;
     if (scale != 1) {
       DPU_GongProminentColor dummy = new DPU_GongProminentColor();
       ResampledDescriptor resampled = Resampler.scaleDown(pixels, d,
           dummy.RESAMPLER, scale);
       if (resampled != null) {
         pixels = resampled.m_resampled;
         d = resampled.m_newDimension;
         the_scale = resampled.m_scale;
         img_name = img_name + Resampler.RESIZED_FILE_SUFFIX + scale;
       }
     }

     try {
       // process
       Configuration.setParameter("SOURCE_NAME", img_name);
       Configuration.setParameter("DEST_PATH", to_path);
       yi.util.ImageUtil.setPixels(pixels);
       yi.util.ImageUtil.setDimension(d);

       long startTime = System.currentTimeMillis();
       DPU_GongProminentColor processor = new DPU_GongProminentColor();
       CIELabLabeledRgnFts ds = processor.getReusableCIELabLabeledRgnFts(
           pixels, d.width, d.height, old_d.width, old_d.height);
       long endTime = System.currentTimeMillis();
       System.out.println("running time: " + (endTime-startTime));

       int[] demo = ds.colorDemo(null);
       yi.util.io.File.makeDirectory(to_path, false);
       yi.util.io.OutputResult.outputGIFImage(to_path + img_name + ".dc.rgn.gif", demo, d.width,
                                              d.height, false);

     } catch (Exception e) {
       System.err.println("KMeansHSISeg Terminated.");
       e.printStackTrace();
     }

     System.exit(0);
    }

  }
  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;
  }

}