package yi.iu.processor.color.reg;

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

import yi.iu.ft.*;
import yi.util.config.*;
import yi.util.io.input.*;
import yi.util.io.output.*;
import yi.util.log.*;
import java.awt.*;
import java.awt.image.*;
import java.io.*;
import java.util.*;
import yi.iu.processor.cluster.*;
import yi.iu.processor.morph.ConnectComponent;
import yi.util.io.ImageReader;
import yi.util.io.ReadArgu;
import yi.util.io.OutputResult;
import yi.util.label.Label;
import yi.iu.region.Region;
import yi.iu.processor.morph.*;
import yi.iu.ui.SegFrame;
import yi.iu.processor.filter.GaussianSmooth;
import yi.iu.processor.color.*;
import yi.iu.processor.color.quan.SmithColorQuantization;
import yi.iu.processor.resample.*;

public class DPU_ColorRegion implements Configurable, RstReusable{

  public int    K = 16;
  /**
   * Gussian smooth theta.  theta = GAUSSIAN_THETA * Math.sqrt(width*height).
   * Eg. width = 600, height = 400, theta = around 5
   */
  public double GAUSSIAN_THETA = 0.01;
  /**
   * filter threshold.  detail comparison see cv010
   */
  public double FILTER_THRESHOLD = 0.2;
  public int    REUSE  = DIU.INPUT_NONE;
  public int    OUTPUT = DOU.OUTPUT_XML;
  public int    REUSE_KMCLUSTER = DIU.INPUT_NONE;
  public int    OUTPUT_KMCLUSTER = DOU.OUTPUT_NONE;
  public boolean SKIP_EXISTING = false;

  /**
   * If a region is smaller than x*y*MINIMUM_REGION_SIZE, it will be discarded.
   */
  public double MINIMUM_REGION_SIZE = 0.025;

  /**
   * The regions are first filtered using closing and this is the size of the
   * disk used.
   */
  public int    CLOSING_DISK_SIZE = 5;

  public int RESAMPLER = Resampler.NEAREST_RESAMPLER;

  public boolean DEMO_FILTERING = false;

  public boolean DEMO_ELLIPSE_REGIONS = false;
  public boolean DEMO_COLOR_REGIONS = true;

  public String[] getParameterNames() {
    return new String[] {
      "REUSE",
      "OUTPUT",
      "REUSE_KMCLUSTER",
      "OUTPUT_KMCLUSTER",
      "SKIP_EXISTING",
      "K",
      "MINIMUM_REGION_SIZE",
      "CLOSING_DISK_SIZE",
      "FILTER_THRESHOLD",
      "RESAMPLER",
      "DEMO_FILTERING",
      "DEMO_COLOR_REGIONS",
      "DEMO_ELLIPSE_REGIONS"
    };
  }

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

  /**
   * Background pixles (those pixels can't be assigned to any region) are ignored
   * @param pixels
   * @param width
   * @param height
   * @return
   */
  public DS_RgbColorRegions getColorRegions(int[] pixels, int width, int height, double scale) {
    DS_RgbColorRegions rst = (DS_RgbColorRegions)reuse(REUSE);
    if (rst != null) {
      output(OUTPUT, rst);
      return rst;
    }

    Dimension d = new Dimension(width, height);

    // K-means cluster in CY-HSI color space
    // return labels from 0
    KMCluster c = (KMCluster)reuseKMCluster(REUSE_KMCLUSTER);
    if (c == null) c = CYColorKMCluster.getCluster(K, KMCluster.LINEAR, pixels);
    outputKMCluster(OUTPUT_KMCLUSTER, c);

//    for (int i=0;i<pixels.length;i++) pixels[i] = pixels[i]&0xff;
//    int[][] info= new int[1][];
//    info[0] = pixels;
//    KMCluster c = new KMCluster(K, KMCluster.LINEAR, info);

    int[] labels = c.getLabels();

//    DPU_DominantColor dc = new DPU_GongProminentColor();
//    DS_DominantColor rst = dc.getDominantColors(pixels, width, height);
//    int[] labels = rst.getColorMap();

//    SmithColorQuantization quan = new SmithColorQuantization();
//    int[] labels = quan.quantize(pixels, width, height, null);
//    int[] counts = new int[SmithColorQuantization.MAX_LABEL];
//    // filter small regions.  small regions are labeled -1.
//    for (int i=0;i<labels.length;i++) counts[labels[i]]++;
//    for (int i=0;i<counts.length;i++) {
//      if (counts[i]<MINIMUM_REGION_SIZE*width*height) counts[i]=-1;
//      else counts[i]=i;
//    }
//    for (int i=0;i<labels.length;i++) labels[i] = counts[labels[i]];


    // make the labels sequential
    int newn = Label.sequential(labels, null, false);

//    int[][] cluster_means = getMeanColors(pixels, labels, newn);
//    int[] cluster_img = yi.util.ImageUtil.makeNewImageFromLabel(labels);
//    OutputResult.outputColorImage("t_cluster.ppm",cluster_img,width, height);
//    int[] cluster_means_colors = new int[cluster_means.length];
//    for (int i=0;i<cluster_means.length;i++) {
//      cluster_means_colors[i] = 0xff000000 | (cluster_means[i][0]<<16) |
//      (cluster_means[i][1]<<8) |cluster_means[i][2];
//    }
//    for (int i=0;i<cluster_img.length;i++) cluster_img[i]=cluster_means_colors[labels[i]];
//    OutputResult.outputColorImage("t_cluster_img.ppm",cluster_img,width, height);


//    if (intermediateRsts != null && intermediateRsts.length > 0)
//      intermediateRsts[0] = labels;

    // test code to display the k-mean clustering result
//    int[] segimg = yi.util.ImageUtil.makeNewImageFromLabel_back(labels);
//    SegFrame sf = new SegFrame(segimg, d.width, d.height, labels);
//    sf.setVisible(true);

//    int target_label = 2;
//    for (int i=0; i<labels.length; i++) {
//      labels[i] = (labels[i]==target_label)?1:0;
//    }
//    Closing close = new Closing(Morph.DISK, 10);
//    labels = close.getProcessed(labels, d);
    labels = filter(labels, width, height); // label 0 is background, and the labels are sequentialized already

    int[] compLabels = labels;
    // demo the result
    if (DEMO_COLOR_REGIONS) {
      int[] segimg = yi.util.ImageUtil.makeNewImageFromLabel(compLabels);
      String img_name = Configuration.getStringParameter("SOURCE_NAME", "unknown");
      OutputResult.outputGIFImage(img_name + "_color_regions.gif", segimg, d.width, d.height);
    }

      int max = yi.stat.StatUtil.maximum(compLabels);
//    if (intermediateRsts != null && intermediateRsts.length > 1)
//      intermediateRsts[1] = compLabels;

    // test code to display the connected component labeling result
//        int[] segimg2 = yi.util.ImageUtil.makeNewImageFromLabel_back(compLabels);
//        SegFrame sf2 = new SegFrame(segimg2, d.width, d.height, compLabels);
//        sf2.setVisible(true);

    // recalculate the mean color for every region
    int[][] means = DS_RgbColorRegions.getMeanColors(pixels, compLabels, max);

    // construct Color Image Features (IFColor)
//    ImageFeature[] ifs = new ImageFeature[means.length];   // use 1 integer RGB
//    for (int i=0; i<means.length; i++) {
//      ifs[i] = new IFColor(means[i][0], means[i][1], means[i][2]);
//    }

//    int[] dspimg = new int[compLabels.length];
//    for (int i=0; i<compLabels.length; i++) {
//        dspimg[i] = ((IFColor)ifs[compLabels[i]]).getRGB();
//    }
//    if (intermediateRsts != null && intermediateRsts.length > 2)
//      intermediateRsts[2] = dspimg;

//        // test to display the regions with mean color
//        SegFrame f = new SegFrame("Color Regions", dspimg, d.width, d.height,
//            compLabels);
//        f.setVisible(true);

    // construct region features and return
//    IFRegion[] regions = IFRegion.getIFRegions(compLabels, max, d, ifs);
    Region[] regions = Region.getRegions(compLabels, max, d);
    Vector color_regions = new Vector();
    // !! the background is useless, so start with regions with label = 1
    for (int i=1; i<= max; i++) {
      color_regions.add(new RgbColorRegion(means[i][0], means[i][1], means[i][2], regions[i]));
    }
    rst = new DS_RgbColorRegions(width, height, color_regions, compLabels);
    rst.setScale(scale);

    output(OUTPUT, rst);
    return rst;
  }

  public Object reuse(int reuse_level) {
    if (reuse_level != DIU.INPUT_NONE) {
      try {
        DS_RgbColorRegions rst = (DS_RgbColorRegions)DIU.in(
            reuse_level,
            DOU.shortName(this.getClass().getName()),
            "DS_RgbColorRegions",
            Class.forName("yi.iu.processor.color.reg.DS_RgbColorRegions"));
        return rst;
      }
      catch (Exception e) {
        System.err.println("Can't Reuse DS_RgbColorRegions");
        e.printStackTrace();
      }
    }
    return null;
  }

  public void output(int output_level, DS_RgbColorRegions ds) {
    try {
      if (output_level != DOU.OUTPUT_NONE) {
        DOU.out(
            output_level,
            DOU.shortName(this.getClass().getName()),
            DOU.shortName(ds.getClass().getName()),
            ds);
      }
    } catch (DataOutputException e) {
      e.printStackTrace();
    }
  }

  final static int BACKGROUND = 0;
  /**
   * Filter regions (including morph filter and size filter).   Suppose the pass
   * in parameter labels don't include background label and resultant background
   * will be returned as label 0.
   * @param labels
   * @param width
   * @param height
   * @return
   */
  int[] filter(int[] labels, int width, int height) {
    int pnum = labels.length;

    // find the maximum label
    int maxlabels = 0;
    for (int i=0; i<pnum; i++) {
      if (labels[i] > maxlabels) maxlabels = labels[i];
    }

    // init the resultant array
    int[] filtered = new int[pnum];
    for (int i=0; i<pnum; i++) filtered[i] = BACKGROUND;

    // deal with each label seperatly
    int start_label = BACKGROUND;
    for (int i=0; i<=maxlabels; i++) {
      int[] binary = new int[pnum];
      for (int j=0; j<pnum; j++) binary[j] = (labels[j]==i)?1:BACKGROUND;
      if (DEMO_FILTERING) {
        int[] binary_img = yi.util.ImageUtil.makeNewImageFromBinary(binary);
        String img_name = Configuration.getStringParameter("SOURCE_NAME", "unknown");
        yi.util.io.OutputResult.outputGIFImage(img_name + "." + i + ".bin.gif", binary_img, width, height);
      }
      binary = filterBinary(binary, width, height);
      if (DEMO_FILTERING) {
        int[] filtered_img = yi.util.ImageUtil.makeNewImageFromLabel(binary);
        String img_name = Configuration.getStringParameter("SOURCE_NAME", "unknown");
        yi.util.io.OutputResult.outputGIFImage(img_name + "." + i + ".filtered.gif", filtered_img, width, height);
      }

      int max_label = 0;
      for (int j=0; j<pnum; j++) {
        if (binary[j]!=0 && labels[j]==i) filtered[j] = start_label + binary[j];
        if (binary[j]>max_label) max_label = binary[j];
      }
//      System.out.println("label="+i+",max_label="+max_label+",start_label="+start_label);
      start_label += max_label;
//      segimg = yi.util.ImageUtil.makeNewImageFromLabel_back(filtered);
//      OutputResult.outputColorImage("t"+i+"_2.ppm", segimg, width, height);
    }

    return filtered;
  }

  int gau_count=0;

  /**
   * Do morph processing on a binary region image.  Try to do closing ...
   * In "binary" parameter, 0 is background, 1 is the region.  return 0 for
   * background
   * @param binary
   * @param width
   * @param height
   * @return
   */
  int[] filterBinary(int[] binary, int width, int height) {
    int pnum = binary.length;
//    Closing close = new Closing(Morph.DISK, CLOSING_DISK_SIZE);
//    binary = close.getProcessed(binary, new Dimension(width, height));
    GaussianSmooth gau = new GaussianSmooth(GAUSSIAN_THETA*Math.sqrt(width*height));
    for (int i=0;i<binary.length;i++) if (binary[i]!=0) binary[i]=255;
    binary = gau.getProcessed(binary, new Dimension(width,height));
//    int[] gau_img = yi.util.ImageUtil.makeNewImageFromGray(binary);
//    String count_str = "00"+gau_count;
//    count_str = count_str.substring(count_str.length()-2);
//    OutputResult.outputGrayImage("gau_theta_5_"+count_str+".pgm",binary, width, height);
//    gau_count++;
    int binary_threshold = (int)(255 * FILTER_THRESHOLD);
    for (int i=0;i<binary.length;i++) {
      binary[i] = (binary[i]>binary_threshold)?1:0;
    }


    // connected component
    ConnectComponent cc = new ConnectComponent();
    // find connected components (these are also sequentially labeled)
    // 0 should be treated as background
    int[] cc_labels = cc.compactLabeling(binary, new Dimension(width, height), true);

    // count component size
    int max_cc_label = cc.getMaxLabel();
    int[] sizes = new int[max_cc_label + 1];
    for (int i=0; i<pnum; i++) sizes[cc_labels[i]]++;

    // filter small component
    int size_threshold = (int)(MINIMUM_REGION_SIZE * width * height);
    for (int i=1; i<=max_cc_label; i++) {
      if (sizes[i] < size_threshold) sizes[i] = BACKGROUND;
      else sizes[i] = i;
    }
    sizes[BACKGROUND]=0;  // background remain the same
    for (int i=0; i<pnum; i++) cc_labels[i] = sizes[cc_labels[i]];
    Label.sequential(cc_labels,null,true);

    return cc_labels;
  }

  public int[] getDemo(int[] pixels, int width, int height)
  {
    DS_RgbColorRegions regions = getColorRegions(pixels, width, height, 1.0);
    int[] labels = regions.getLabels();
//    int[] segimg = yi.util.ImageUtil.makeNewImageFromLabel_back(labels);
//    SegFrame sf = new SegFrame(segimg, width, height, labels);
//    sf.setVisible(true);
    int[] rst= yi.util.ImageUtil.makeNewImageFromLabel(labels);

    return rst;
  }

  public void outputKMCluster(int output_level, KMCluster c) {
    if (output_level != DOU.OUTPUT_NONE) {
      try{
        DOU.out(
            output_level,
            DOU.shortName(this.getClass().getName()),
            "KMCluster",
            c);
      } catch (DataOutputException e) {
        e.printStackTrace();
      }
    }
  }

  public KMCluster reuseKMCluster(int reuse_level) {
    if (reuse_level != DIU.INPUT_NONE) {
      try {
        KMCluster rst = (KMCluster)DIU.in(
            reuse_level,
            DOU.shortName(this.getClass().getName()),
            "KMCluster",
            Class.forName("yi.iu.processor.cluster.KMCluster"));
        return rst;
      }
      catch (Exception e) {
        System.err.println("Can't Reuse KMCluster");
        e.printStackTrace();
      }
    }
    return null;
  }

  /**
   * Output the regions (a region per image)
   */
  public static class DemoRegions {
    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_ColorRegion dp = new DPU_ColorRegion();
      DS_RgbColorRegions ds = dp.getColorRegions(pixels, d.width, d.height, 1.0);
      ds.demoRegions(img_name, pixels, 200, 0xff000000, 0xffffffff);

      System.exit(0);
    }
  }


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

    // 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(img_name);
    Dimension d = new Dimension();
    int[] pixels = yi.util.io.ImageReader.getPixels(image, d);

    // rescale the image if necessary
    if (scale != 1) {
      DPU_ColorRegion dummy = new DPU_ColorRegion();
      ResampledDescriptor resampled = Resampler.scaleDown(pixels, d,
          dummy.RESAMPLER, scale);
      double the_scale = 1.0;
      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);
      yi.util.ImageUtil.setPixels(pixels);
      yi.util.ImageUtil.setDimension(d);

      DPU_ColorRegion dp = new DPU_ColorRegion();
      long startTime = System.currentTimeMillis();
      DS_RgbColorRegions ds = dp.getColorRegions(pixels, d.width, d.height, 1.0);
      long endTime = System.currentTimeMillis();
      System.out.println("running time: " + (endTime-startTime));

      if (dp.DEMO_ELLIPSE_REGIONS) {
        int[] rst_img = ds.getDemo(pixels, d.width, d.height);
        OutputResult.outputGIFImage(img_name + "." + DOU.shortName(ds.getClass().getName()) + ".gif", rst_img, d.width, d.height);
      }

//      int[] rst_img = dp.getDemo(pixels, d.width, d.height);
//      OutputResult.outputColorImage(img_name + "." + DOU.shortName(dp.getClass().getName()) + "_quan.ppm", rst_img, d.width, d.height);

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

    System.exit(0);
  }
}