package yi.iu.processor.clc;

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

import yi.iu.processor.color.*;
import yi.iu.processor.line.*;
import yi.iu.processor.ort.*;
import yi.util.config.*;
import yi.util.io.*;
import yi.stat.em.SimplifiedEMCumulator;

/**
 * Color Consistent Line Cluster.  Given an image, first find the dominant color,
 * find the color pair (in terms of dominant colors) of each line, and return
 * the lines and their color pair.  To tolerance the errors, it is possible to
 * find multiple color pairs for a line.  As the result, the returned lines may
 * have duplicates, and the number of the lines is possibly more than that of
 * detected.
 */
public class CLC_Color extends CLC_Clustering
{

  /**
   * When consider the color of one side of a line, the pixels of that color
   * must be greater than this threshold to be qualified.
   */
  public double LINE_COLOR_QUALIFY_THRESHOLD = 1.0/3;
  /**
   * When consider the color of one side of a line, the consideration window
   * is shifted sevearl time to tolerant the noises.  This value is how many
   * steps the windows shift.
   */
  public int LINE_COLOR_WIN_SHIFT = 1;
  /**
   * When consider the color of one side of a line, the consideration window
   * width is not 1-pixel tolerant the noises.  This value is the window width
   */
  public int  LINE_COLOR_WIN_WIDTH = 3;

  public final static int POST_RECLASSIFY_NONE = 0;
  public final static int POST_RECLASSIFY_GAUSSIAN = 2;
  public final static int POST_RECLASSIFY_EM = 3;
  public int POST_RECLASSIFY_METHOD = POST_RECLASSIFY_EM;
  public double MIN_COLOR_VARIANCE = 5 * 5;



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

  public final static int COLOR_SPACE_CLUSTERING = 1;
  public final static int CP_SPACE_CLUSTERING = 2;
  public int COLOR_PAIR_CLUSTERING_METHOD = COLOR_SPACE_CLUSTERING;


  public boolean DEMO_COLOR_CLC = true;
  public boolean DEMO_COLOR_CLC_REGION = false;

  public String[] getParameterNames() {
    return super.getParameterNames(new String[] {
      "DEMO_COLOR_CLC",
      "DEMO_COLOR_CLC_REGION",
      "LINE_COLOR_WIN_SHIFT",
      "LINE_COLOR_WIN_WIDTH",
      "LINE_COLOR_QUALIFY_THRESHOLD",
      "COLOR_PAIR_CLUSTERING_METHOD",
      "POST_RECLASSIFY_METHOD"
    });
  }

  public LineClusterSets getClusters(LineClusterSet lines, int[] rgb, int width, int height) {

    if (COLOR_PAIR_CLUSTERING_METHOD == CP_SPACE_CLUSTERING) {
      return getClusters2(lines,rgb, width, height);
    }


    // new line cluster sets
    LineClusterSets line_clusters = new LineClusterSets(width, height);

    // get the color pairs for the lines
    LineColorPairs lcp = new LineColorPairs(rgb, width, height);

    int it_count = 0;

    for (Iterator i = lines.iterator(); i.hasNext(); it_count++) {
      ORTLine l = (ORTLine)i.next();

      Collection color_pairs = lcp.getLineColorPairs(l, LINE_COLOR_QUALIFY_THRESHOLD, LINE_COLOR_WIN_SHIFT, LINE_COLOR_WIN_WIDTH);
      Iterator j = color_pairs.iterator();
      while (j.hasNext()) {
        ColorPair cp = (ColorPair)j.next();
        line_clusters.add(cp, l);
      }
    }

    // post re-assign
    if (POST_RECLASSIFY_METHOD == POST_RECLASSIFY_GAUSSIAN) {
      postReClassify(lines, line_clusters, rgb, width, height);
    } else if (POST_RECLASSIFY_METHOD == POST_RECLASSIFY_EM) {
      postReClassify2(lines, line_clusters, rgb, width, height);
    }

    // filter the result
    line_clusters.filter(
      CLCConfig.getMinQualifiedLineLen(),
      CLCConfig.getMinQualifiedCLCLineLenSum(width, height));

    // inherit the properties from the parent (clustered from)
    line_clusters.inheritProperties(lines);

    return line_clusters;
  }


  void postReClassify(LineClusterSet lines, LineClusterSets line_clusters, int[] rgb, int width, int height) {

    int clusterNumber = line_clusters.size();

    SimplifiedEMCumulator[] em_models = new SimplifiedEMCumulator[clusterNumber];
    double MAX_COLOR_VARIANCE = 7.5*7.5; //about 20/2.63
    for (int i=0; i<clusterNumber; i++) em_models[i] = new SimplifiedEMCumulator(MIN_COLOR_VARIANCE, MAX_COLOR_VARIANCE);

    // for each cluster, construct its data model
    int it_count = 0;
    for (Iterator i=line_clusters.keySet().iterator(); i.hasNext(); it_count++) {
      ColorPair cp =(ColorPair)i.next();
      LineClusterSet line_cluster = (LineClusterSet)line_clusters.get(cp);
//      System.out.println("cluster#" + it_count);
//      System.out.println(cp);

      int j_count = 0;
      for (Iterator j=line_cluster.iterator(); j.hasNext(); j_count++) {
        ORTLine line = (ORTLine)j.next();
        // for mean colors on the line and shifted lines
        int[][] colors = line.rgbMeanColors(1, rgb, width, height);
        // find which pair is the best match to the color pair of this cluster
        ColorPair best_match = new ColorPair();
        double d = cp.bestMatch(colors, best_match);
//        System.out.println(" " + j_count + ": " + d + " " + best_match);
        int[] best_x = ColorUtil.RGB2IntArray(best_match.x);
        int[] best_y = ColorUtil.RGB2IntArray(best_match.y);
        double[] data = new double[] {
            best_x[0], best_x[1], best_x[2],
            best_y[0], best_y[1], best_y[2]
        };
        em_models[it_count].add(data, 1);
      }
//      System.out.println();
    }

//    it_count = 0;
//    for (Iterator i=line_clusters.keySet().iterator(); i.hasNext(); it_count++) {
//      ColorPair cp =(ColorPair)i.next();
//
//      System.out.println("model #" + it_count);
//      try {
//        System.out.println(cp);
//        OutputResult.print(System.out, em_models[it_count].getMeans());
//        em_models[it_count].printSigmas(2.63);
//      }catch (Exception e) {
//
//      }
//    }

    for (Iterator i= lines.iterator(); i.hasNext();) {
      ORTLine line = (ORTLine)i.next();
      int[][] colors = line.rgbMeanColors(1, rgb, width, height);
      int j=0;
      for (Iterator it = line_clusters.keySet().iterator(); it.hasNext(); j++) {
        ColorPair cp =(ColorPair)it.next();
        LineClusterSet line_cluster = (LineClusterSet)line_clusters.get(cp);
        ColorPair best_match = new ColorPair();
        if (ColorPair.belongsTo(colors, em_models[j], best_match, 2.63)) {
          if (line_cluster != null && !line_cluster.contains(line)) {
//            if (j==70) {
//              System.out.println("post merge line " + line + " to line cluster " + j);
//              try {
//                double[] means = em_models[j].getMeans();
//                System.out.println("([" + means[0] + ", " + means[1] + ", " +
//                                   means[2] + "), (" +
//                                   means[3] + ", " + means[4] + ", " + means[5] +
//                                   ")]");
//              } catch (Exception e) {
//
//              }
//              System.out.println(best_match);
//            }
          }
          line_clusters.add(cp, line);
        }
      }
    }
  }


  void postReClassify2(LineClusterSet lines, LineClusterSets line_clusters, int[] rgb, int width, int height) {

    int clusterNumber = line_clusters.size();

    SimplifiedEMCumulator[] em_models = new SimplifiedEMCumulator[clusterNumber];
    MIN_COLOR_VARIANCE = 2*2;
    double MAX_COLOR_VARIANCE = -1; //about 20/2.63
    for (int i=0; i<clusterNumber; i++) {
      em_models[i] = new SimplifiedEMCumulator(MIN_COLOR_VARIANCE,
                                               MAX_COLOR_VARIANCE);
    }

    // for each cluster, construct its data model
    int it_count = 0;
    LineClusterSet[] old_clusters = new LineClusterSet[clusterNumber];
    LineClusterSet[] new_clusters = new LineClusterSet[clusterNumber];
    for (Iterator i=line_clusters.keySet().iterator(); i.hasNext(); it_count++) {
      ColorPair cp =(ColorPair)i.next();
      LineClusterSet line_cluster = (LineClusterSet)line_clusters.get(cp);
      old_clusters[it_count] = line_cluster;
      new_clusters[it_count] = line_cluster;
//      System.out.println("cluster#" + it_count);
//      System.out.println(cp);

      int j_count = 0;
      for (Iterator j=line_cluster.iterator(); j.hasNext(); j_count++) {
        ORTLine line = (ORTLine)j.next();
        // for mean colors on the line and shifted lines
        int[][] colors = line.rgbMeanColors(1, rgb, width, height);
        // find which pair is the best match to the color pair of this cluster
        ColorPair best_match = new ColorPair();
        double d = cp.bestMatch(colors, best_match);
//        System.out.println(" " + j_count + ": " + d + " " + best_match);
        int[] best_x = ColorUtil.RGB2IntArray(best_match.x);
        int[] best_y = ColorUtil.RGB2IntArray(best_match.y);
        double[] data = new double[] {
            best_x[0], best_x[1], best_x[2],
            best_y[0], best_y[1], best_y[2]
        };
        em_models[it_count].add(data, 1);
      }
//      System.out.println();
    }

    HashMap models = new HashMap();
    HashMap model_probs = new HashMap();
    for (int i=0; i<clusterNumber; i++) {
      models.put(Integer.toString(i), em_models[i]);
      double prob = (double)em_models[i].getCount()/(double)lines.size();
      model_probs.put(Integer.toString(i), new Double(prob));
//      try {
//        em_models[i].printSigmas();
//      } catch (Exception e) {
//        e.printStackTrace();
//      }
//      System.out.println(prob);
    }

    it_count = 0;
    for (Iterator i= lines.iterator(); i.hasNext();it_count++) {
      ORTLine line = (ORTLine)i.next();
      int[][] colors = line.rgbMeanColors(1, rgb, width, height);
      HashSet belongs = ColorPair.belongsTo(colors, models, model_probs, 0.4);
//      System.out.print("line #" + it_count + ":" + belongs.size() + "  :");
      for (Iterator it = belongs.iterator(); it.hasNext();) {
        int index = Integer.parseInt((String)it.next());
        new_clusters[index].add(line);
//        System.out.print(" " + index);
      }
//      System.out.println();
    }

    it_count = 0;
    for (Iterator i=line_clusters.keySet().iterator(); i.hasNext(); it_count++) {
      ColorPair cp = (ColorPair) i.next();
      line_clusters.put(cp, new_clusters[it_count]);
    }

//    for (int i=0; i<clusterNumber; i++) {
//      System.out.print("cluster #" + i + ": " + new_clusters[i].size());
//      double prob = (double) em_models[i].getCount() / (double) lines.size();
//      System.out.println("  prob = " + prob);
//    }
  }

  public LineClusterSets getClusters2(LineClusterSet lines, int[] rgb, int width, int height) {
    CLC_KMeansColorPair k_means = new CLC_KMeansColorPair();
    LineClusterSets line_clusters = k_means.getClusters(lines, rgb, width, height);

    // filter the result
    line_clusters.filter(
      CLCConfig.getMinQualifiedLineLen(),
      CLCConfig.getMinQualifiedCLCLineLenSum(width, height));

    // inherit the properties from the parent (clustered from)
    line_clusters.inheritProperties(lines);
    return line_clusters;
  }


  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);
    long startTime = System.currentTimeMillis();

    //find all the lines in an image
    yi.iu.processor.ort.DPU_Fex dp = new DPU_Fex();
    Vector[] segs= dp.getLinesAndArcs(pixels, d.width, d.height);
    Vector lines= segs[0];
    LineClusterSet lcs = new LineClusterSet(lines, d.width, d.height);

    // cluster
    CLC_Color clc_color = new CLC_Color();
    LineClusterSets lcss = clc_color.getReusableClusteres(lcs, pixels, d.width, d.height);

    long endTime = System.currentTimeMillis();
    System.out.println("running time: " + (endTime-startTime));

    // illustration each line cluster
    {
      Iterator i= lcss.keySet().iterator();
      int label = 1;
      while (i.hasNext()) {
        ColorPair cp = (ColorPair)i.next();
        LineClusterSet line_cluster = (LineClusterSet)lcss.get(cp);
        // draw
        int[] lineimg= yi.util.ImageUtil.makeBackground(pixels);
        line_cluster.draw(lineimg, d.width, d.height, 0xffffffff);
        // color spectrum
        int[] colors = new int[] {cp.x, cp.y};
        int color_w = 40, color_h = 50;
        int[] spectrum = yi.iu.processor.color.ColorUtil.createSpectrum(colors, color_w, color_h);
        for (int y=0; y<color_h; y++) {
          for (int x=0;x<color_w*colors.length; x++) {
            lineimg[y*d.width+x] = spectrum[y*color_w*colors.length+x];
          }
        }
        // make a simple color index
//        OutputResult.outputColorImage(img_name + ".cclc." + cp_index, lineimg, d.width, d.height, true);
        OutputResult.outputColorImage(img_name + ".cclc" +
                                      "." + clc_color.LINE_COLOR_WIN_SHIFT +
                                      "." + clc_color.LINE_COLOR_WIN_WIDTH,
                                      lineimg, d.width, d.height, true);

//        // output the region
//        lineimg= yi.util.ImageUtil.makeBackground(pixels);
//        Region r = line_cluster.getRegion();
//        Region.drawRegions(lineimg, d, r, colors[0]);
//        OutputResult.outputColorImage(img_name + ".clc_region", lineimg, d.width, d.height, true);
        label ++;
      }
    }


    // demo the line clusters
    if (clc_color.DEMO_COLOR_CLC) {
      int[] lineimg = lcss.demo(d.width, d.height);
      OutputResult.outputGIFImage(img_name + ".color_clcs.gif", lineimg, d.width, d.height);
    }

//    // demo the line cluster regions
//    {
//      Iterator i= lcss.values().iterator();
//      int[] lineimg= new int[pixels.length];
//      int label = 1;
//      while (i.hasNext()) {
//        LineClusterSet line_cluster = (LineClusterSet)i.next();
//        // draw region
//        Region r = line_cluster.getRegion();
//        Region.drawRegions(lineimg, d, r, label++);
//      }
//      lineimg = yi.util.ImageUtil.makeNewImageFromLabel(lineimg);
//      OutputResult.outputColorImage(img_name + ".clc_regions.", lineimg, d.width, d.height, true);
//    }

    // demo the line cluster regions (only the biggest)
    if (clc_color.DEMO_COLOR_CLC_REGION) {
      int[] lineimg = lcss.demoRegions(d.width, d.height);
      OutputResult.outputGIFImage(img_name + ".clc_regions.gif", lineimg, d.width, d.height);
    }

    System.exit(0);
  }
}