package yi.iu.ft;

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.ui.ImageFrame;

/**
 * Class for generating region-based color features. <p>
 *
 * ------------------------------------------------------------------- <p>
 *
 * General options when processing an image from the command-line: <p>
 *
 * -i imagename <br>
 * Name of the image file to be processed. (required) <p>
 *
 * -o outputname <br>
 * Name of the feature output file.  You can specify it or use the default<p>
 *
 * -d <br>
 * Display the running result in a demo frame. <p>
 *
 * -w flags<br>
 * Specify whether the intermediate results will be written.  Eg. flag = 1101
 * means the 1st, 3rd, 4th intermediate results will be written.  In this
 * application, only 3 intermediate results are available: <br>
 * <ul>
 * <li>0: labels after k-means</li>
 * <li>1: labels after connected component</li>
 * <li>2: image with pixel using the mean color of the region</li>
 * </ul>
 *
 * ------------------------------------------------------------------- <p>
 *
 * Title:        Object Image Generation
 * Description:  This project provides the required utilities to construct Object Images.  An Object Image is an image containing one or several predefined object sketch.  This kind of image is suitable to be used as experiment image when testing a learning algorithm.
 * Copyright:    Copyright (c) 2002
 * Company:      University of Washington
 * @author Yi Li
 * @version 1.0
 */

public class FTExRegionColor extends FTExRegion{
    public static int K = 16;

    /**
     * initialize region-based color feature of an image file
     */
    public FTExRegionColor(String fileName) throws IOException{
        super(fileName);
    }

    /**
     * initialize region-based color feature of image.  Image ID must be
     * provided
     */
    public FTExRegionColor(String name, Image im) {
        imageID = name;
        d = new Dimension();
        pixels = ImageReader.getPixels(im, d);
    }

    /**
     * Return the color region features of a given image.  Each color feature
     * is in the formant of (R, G, B).   R, G, and B are integers of [0,255]
     *
     * [EFFECTS]
     *   the labels are sequentialized, and not necessarily be 0-15
     * The procedure:
     *   1. K-Means cluster in HSI color space
     *   2. find connected component (region)
     *   3. re-calculate the mean color for each region
     *   4. return image feature regions
     *
     * @param im image data
     * @param intermediateRsts if it is not null, it wil be used to return the
     *                         the intermediate results
     *        [0]: labels after K-means cluster in CY-HSI color space
     *        [1]: labels after connected component
     *        [2]: new image with pixels taking the mean color of their region
     *
     * @return an array of IFRegion
     */
    public static IFRegion[] extract(Image im, Object[] intermediateRsts) {
        // read image data
        Dimension d = new Dimension();
        int[] p = ImageReader.getPixels(im, d);

        // K-means cluster in CY-HSI color space
        // return labels from 0
        KMCluster c =  CYColorKMCluster.getCluster(K,
                            KMCluster.LINEAR, p);
        int[] labels = c.getLabels();

        // make the labels sequential
        int newn = Label.sequential(labels, null, false);
        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);

        // find connect component to make the regions
        ConnectComponent cc = new ConnectComponent();
        // find connected components (these are also sequentially labeled)
        // 0 should not be treated as background, but a region
        // return label from 0
        int[] compLabels = cc.compactLabeling(labels, d, false);
        int max = cc.getMaxLabel();
        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 = getMeanColors(p, 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
        return IFRegion.getIFRegions(compLabels, max, d, ifs);
    }

    /**
     * Calculate mean colors for regions
     * return rgb colors
     *
     * @param rgb
     * @param labels
     * @param maxLabel
     * @return
     */
    public static int[][] getMeanColors(int[] rgb, int[] labels, int maxLabel)
    {
        int[][] means = new int[maxLabel+1][3];

        // calculate the sum and count for every region
        long[][] sum  = new long[maxLabel+1][3];
        int[] count   = new int[maxLabel+1];
        for (int i=0; i<rgb.length; i++) {
            sum[labels[i]][0] += (rgb[i]>>16) & 0xff;
            sum[labels[i]][1] += (rgb[i]>> 8) & 0xff;
            sum[labels[i]][2] += rgb[i] & 0xff;
            count[labels[i]] ++;
        }

        for (int i=0; i<=maxLabel; i++) {
            means[i][0] = (int)(sum[i][0]/count[i]);
            means[i][1] = (int)(sum[i][1]/count[i]);
            means[i][2] = (int)(sum[i][2]/count[i]);
        }

        sum = null;
        count = null;

        return means;
    }

    public static void main(String[] args) {
      try {
        // get image name
        String image_name = ReadArgu.getOption('i', args);
        if (image_name.length() == 0) {
            System.err.println("image_name must be supplied with flag i");
            System.exit(-1);
        }

        // read image
        Image im = ImageReader.readImageFile(image_name);
        Dimension d =  new Dimension();
        int[] pixels = ImageReader.getPixels(im, d);
        Object[] intermediateRsts = new Object[3];
        IFRegion[] regions = FTExRegionColor.extract(im, intermediateRsts);

        // output the description of the region features
        String output_name = ReadArgu.getOption('o', args);
        if (output_name.length() > 0) {
          if (output_name.compareTo("1") == 0) {
            output_name = image_name + ".FTExRegionColor.xml";
            System.out.println("output file: " + output_name);
          }
          try {
            Writer out = new BufferedWriter(new FileWriter(output_name));
            IFRegion.outputIFRegions(out, image_name, regions);
            out.close();
          } catch (IOException e) {
            e.printStackTrace();
            System.exit(-1);
          }
        }

        // write the intermediate results
        String interRstFlag = ReadArgu.getOption('w', args);
        if (interRstFlag.length() > 0) {
          // if want to output the label image after k-means
          if (intermediateRsts.length > 0 &&
            ReadArgu.getCombinedFlag(interRstFlag, 0)) {
              int[] label = (int[])intermediateRsts[0];
              int[] rst= yi.util.ImageUtil.makeNewImageFromLabel(label);
              OutputResult.outputGIFImage(
                image_name + ".FTExRegionColor.0.gif",
                rst, d.width, d.height
              );
          }
          // if want to output the label image after connect component
          if (intermediateRsts.length > 1 &&
            ReadArgu.getCombinedFlag(interRstFlag, 1)) {
              int[] label = (int[])intermediateRsts[1];
              int[] rst= yi.util.ImageUtil.makeNewImageFromLabel(label);
              OutputResult.outputGIFImage(
                image_name + ".FTExRegionColor.1.gif",
                rst, d.width, d.height
              );
          }
          // if want to output the image with pixels representing the mean color
          // of the region
          if (intermediateRsts.length > 2 &&
            ReadArgu.getCombinedFlag(interRstFlag, 2)) {
              int[] rst = (int[])intermediateRsts[2];
              OutputResult.outputGIFImage(
                image_name + ".FTExRegionColor.2.gif",
                rst, d.width, d.height
              );
          }
        }

        // draw the regions
        if (ReadArgu.getOption('d', args).length() > 0) {
          IFRegion.markIFRegions(pixels, d, regions);
  //        DisplayFrame f = new DisplayFrame(pixels, d.width, d.height);
          ImageFrame f = new ImageFrame(pixels, d.width, d.height);
          f.pack();
          f.setVisible(true);
        }
        else {
          System.exit(0);
        }
      }catch (IOException e) {
        e.printStackTrace();
      }
    }
}