package yi.iu.imgdb.learn;

import yi.util.config.*;
import yi.iu.processor.color.reg.*;
import yi.util.io.input.*;
import yi.iu.imgdb.desc.*;
import java.util.*;
import yi.util.io.*;
import yi.ds.*;
import yi.ds.fold.*;
import yi.ai.classifiers.*;
import org.jmat.gui.*;
import org.jmat.data.*;
import java.io.*;
import yi.iu.imgdb.learn.kmeans.*;
import yi.util.io.output.*;
import javax.swing.*;
import java.awt.event.*;
import java.awt.Color;
import java.awt.Dimension;
import yi.iu.processor.color.*;
import yi.ds.hash.*;
import yi.stat.em.*;

/**
 * Test EM like algorithm on image database
 */

public class ImgDBLearnApp {
  HashMap m_imgObjects;
  HashMap m_imgFeatures;

  public ImgDBLearnApp(HashMap imgObjects, HashMap imgFeatures) {
    m_imgObjects = imgObjects;
    m_imgFeatures = imgFeatures;

  }

  public ImgDBLearnApp(ImgDescManager img_manager, int scale, String data_processor_name) {

    m_imgObjects = new HashMap();
    m_imgFeatures = new HashMap();

    RstReusable data_processor = null;
    try {
      data_processor = (RstReusable) Class.forName(data_processor_name).
          newInstance();
    } catch (Exception e) {
      e.printStackTrace();
      return;
    }

    Set images = img_manager.getFiles();

    Iterator iter = images.iterator();
    while (iter.hasNext()) {
      String img_name = (String) iter.next();
//      System.out.print("reading : " + img_name);

      // get the object descriptions
      Collection descriptions = img_manager.getDesc(img_name);
      m_imgObjects.put(img_name, descriptions);
//      System.out.print(", descriptions #= " + descriptions.size());

      // get the image features
      String scaled_img_name = img_name;
      if (scale != 0) scaled_img_name = img_name + yi.iu.processor.resample.Resampler.RESIZED_FILE_SUFFIX + scale;
      Configuration.setParameter("SOURCE_NAME", scaled_img_name);
      Featurelizable data_set = (Featurelizable)data_processor.reuse(DIU.INPUT_XML);
      double[][] fvs = data_set.getFeatureVectors(0);
      Collection ft_vts = Arrays.asList(fvs);
      m_imgFeatures.put(img_name, ft_vts);
//      System.out.println(", feature vectors #= " + ft_vts.size());
    }
    removeObjWithoutEnoughInst(Math.min(images.size(), 30));

  }


  /**
   * Remove objects without enough instance.  Those objects (keywords) are not
   * good for learning, so just remove them
   */
  private void removeObjWithoutEnoughInst(int min_count) {
    System.out.println("remove objects without enough instances: threshold = "  + min_count);
    // count the appearence of objects
    CounterHashMap counter = new CounterHashMap();
    for (Iterator i=m_imgObjects.keySet().iterator(); i.hasNext(); ) {
      Collection objs = (Collection)m_imgObjects.get(i.next());
      counter.increaseCount(objs);
    }

    // find the objects to remove
    LinkedList obj_to_remove = new LinkedList();
    for (Iterator i=counter.keySet().iterator(); i.hasNext(); ) {
      String key = (String)i.next();
      int count = counter.getCount(key);
      if (count < min_count) {
        System.out.println(key + " #" + count + ", removed");
        obj_to_remove.add(key);
      }
      else {
        System.out.println(key + " #" + count + ", keep");
      }
    }

    // remove the objects from the descriptions
    for (Iterator i=m_imgObjects.keySet().iterator(); i.hasNext(); ) {
      Collection objs = (Collection)m_imgObjects.get(i.next());
      objs.removeAll(obj_to_remove);
    }

    List images = new ArrayList(m_imgObjects.keySet());
    Collections.sort(images);
    for (Iterator i=images.iterator(); i.hasNext(); ) {
      String key = (String)i.next();
      Collection objs = (Collection)m_imgObjects.get(key);
      System.out.print(key + ":: ");
      yi.util.io.OutputResult.printStringCollection(System.out, objs, ',');
      System.out.println();
    }
  }


//  /**
//   * Add training cases using all the images read.
//   * @return
//   */
//  protected ImgDBLearner addTrainingCases(String learner_name) {
//
//    ImgDBLearner learner = null;
//    try {
//      learner = (ImgDBLearner) Class.forName(learner_name).
//          newInstance();
//    } catch (Exception e) {
//      e.printStackTrace();
//      return null;
//    }
//
//    Set images = m_imgObjects.keySet();
//
//    Iterator iter = images.iterator();
//    while (iter.hasNext()) {
//      String img_name = (String)iter.next();
//      Collection descriptions = (Collection)m_imgObjects.get(img_name);
//      Collection ft_vts = (Collection)m_imgFeatures.get(img_name);
//      learner.addImageInst(img_name, descriptions, ft_vts);
//    }
//
//    return learner;
//
//  }
//
//  public HashMap getModel(String learner_name) {
//    ImgDBLearner learner =  addTrainingCases(learner_name);
//    return learner.getModels();
//  }

  /**
   * Use the given images.  Do training and testing on folds
   * @param learner_name
   * @param classifier_name
   * @param fold
   * @param models_from_folds to return the models from each fold.  Give null
   *        if the caller not interested in getting it.
   * @return the classification result.  For each object, the corresponding
   *         image scores are returned. eg. "tree"->{1.jpg->0.9; 2.jpg->0.5; ...}
   */
  public HashMapHashMap foldedLearn(String learner_name, String classifier_name, int fold, List models_from_folds) {

    // one classifier for all the folds
    ImgDBClassifier classifier = null;
    try {
      classifier = (ImgDBClassifier) Class.forName(
          classifier_name).newInstance();
    } catch (Exception e) {
      e.printStackTrace();
      return null;
    }


    // store the classification result
    HashMapHashMap classified = new HashMapHashMap();

    // folding...
    Set images = m_imgObjects.keySet();
    ArrayList folds = (new DataFolder()).nFold(images, fold);

    // for each fold
    for (int i=0; i<fold; i++) {
      // get the training set and testing set
      TwoLinkedListData two_list = (TwoLinkedListData)folds.get(i);
      Collection training = two_list.first;
      Collection testing = two_list.second;

      // new an instance of learner
      ObjImgModel learner = null;
      try {
        learner = (ObjImgModel) Class.forName(learner_name).
            newInstance();
      }
      catch (Exception e) {
        e.printStackTrace();
        return null;
      }

      // construct the learner
      for (Iterator iter = training.iterator(); iter.hasNext(); ) {
        String img_name = (String)iter.next();
        learner.addImageInst(img_name,
                             (Collection)m_imgObjects.get(img_name),
                             (Collection)m_imgFeatures.get(img_name));
      }
      // get the learning resultant models
      learner.learnTheModel();

      System.out.println("\n--------------------------------------");
      System.out.println(learner);

      // save the model from each fold for later.
      if (models_from_folds != null) {
        models_from_folds.add(learner);
      }

      // testing
      classifier.setModels(learner);
      for (Iterator iter = testing.iterator(); iter.hasNext(); ) {
        String img_name = (String)iter.next();
        DoubleScoreHashMap classify_rst = classifier.classify((Collection)m_imgFeatures.get(img_name));
        classified.put(img_name, classify_rst);
      }
    }
    return classified;
  }

  /**
   * Given the models from the folders.  Return a hashmap with object names as
   * key, and a list of means (double array) from different folds as value.  eg.
   * "tree"->{{0,0,0},{1,0,1},...}
   * @param models_from_folds
   * @return
   */
  public DoubleArrayListMap getModelMeans(List models_from_folds) {
    DoubleArrayListMap means_from_folds = new DoubleArrayListMap();
    // loop for each fold
    for (Iterator i= models_from_folds.iterator(); i.hasNext(); ) {
      // get the models (a mixture model) of current fold
      HashMap models = (HashMap)i.next();
      // loop for each object
      for (Iterator iter = models.keySet().iterator(); iter.hasNext(); ) {
        Object obj_name = iter.next();
        // find the model of that object
        GaussianComponent accumulator = (GaussianComponent) models.get(obj_name);
        try {
          double[] mean = accumulator.getMeans();
          means_from_folds.addToList(obj_name, mean);
        }
        catch (GaussianModelSingularException e) {
          e.printStackTrace();
        }
      }
    }
    return means_from_folds;
  }

//  /**
//   * Given the classified map, return the learning result analysis.  An entry
//   * of the map is like:
//   * 1.jpg -> {"tree"->0.9; "sky"->0.2}
//   * @param classified
//   * @return
//   */
//  public HashMap analysisClassifiedMap(HashMap classified) {
//    // the result
//    HashMap analysis = new HashMap();
//
//    // analysis each image
//    List images = new ArrayList(classified.keySet());
//    Collections.sort(images);
//    for (Iterator i= images.iterator(); i.hasNext(); ) {
//      Object key = i.next();
//      Collection objs = (Collection)m_imgObjects.get(key);
//      System.out.println(key);
//      System.out.println("objects:");
//      yi.util.io.OutputResult.printStringCollection(System.out, objs, ',');
//      DoubleScoreHashMap predicted = (DoubleScoreHashMap)classified.get(key);
//      System.out.println("classified: " + predicted);
//
//      Collection found = predicted.keySet();
//      ClassificationPerformance performance = ResultAnalysis.ResultAnalysis(objs, found);
//      analysis.put(key, performance);
//      System.out.println(performance);
//    }
//
//    // print a desc format output
//    System.out.println();
//    System.out.println();
//    System.out.println("desc format results:");
//    System.out.println();
//    for (Iterator i= images.iterator(); i.hasNext(); ) {
//      Object key = i.next();
//      Collection objs = (Collection)m_imgObjects.get(key);
//      DoubleScoreHashMap predicted = (DoubleScoreHashMap)classified.get(key);
//      System.out.println(key + ":: " + predicted.objsToString());
//    }
//
//    return analysis;
//  }
//

  /**
   * Given the object->scores map, return the analysis.  An entry
   * of the map is like:
   * "tree" -> {1.jpg->0.9; 2.jpg->0.5; ...}
   * @param classified
   * @return
   */
  public HashMap analysisObjScoreMap(HashMap obj_scores) {
    // the result
    HashMap rocs = new HashMap();

    // analysis each object
    List objects = new ArrayList(obj_scores.keySet());
    Collections.sort(objects);
    for (Iterator i= objects.iterator(); i.hasNext(); ) {
      Object the_obj = i.next();
      System.out.println("----------------------");
      System.out.println(the_obj.toString() + ":");
      DoubleScoreHashMap one_obj_score = (DoubleScoreHashMap)obj_scores.get(the_obj);
      List the_images = new ArrayList(one_obj_score.keySet());
      Collections.sort(the_images, one_obj_score);
      Collections.reverse(the_images); // sorted from highest score to lowest
      int[] right_pos = new int[the_images.size()+1], mis_neg = new int[the_images.size()+1]; // right positive, mis-classified neg
      right_pos[0] = 0; mis_neg[0] = 0;
      int pos = 0, neg = 0; // total number of pos and neg
      int index =0;
      for (Iterator j= the_images.iterator(); j.hasNext(); index++) {
        Object the_img = j.next();
        double score = one_obj_score.getScore(the_img);
        System.out.println(the_img.toString() + "= " + score);
        Collection objs = (Collection)m_imgObjects.get(the_img);
        if (objs.contains(the_obj)) {
          pos++;
          right_pos[index+1] = right_pos[index] + 1;
          mis_neg[index+1] = mis_neg[index];
        } else {
          neg++;
          mis_neg[index+1] = mis_neg[index] + 1;
          right_pos[index+1] = right_pos[index];
        }
      }
      double[][] roc = new double[right_pos.length][2]; // receiver operating curve
      for (int j=0; j<right_pos.length; j++) {
        roc[j][0] = right_pos[j] / (double)pos;
        roc[j][1] = mis_neg[j] / (double)neg;
      }
      rocs.put(the_obj, roc);
    }

    return rocs;
  }

  public static void drawROCs(HashMap rocs) {
    // the draw panel
    Plot2DPanel pp = new Plot2DPanel();

    // analysis each object
    List objects = new ArrayList(rocs.keySet());
    Collections.sort(objects);
    for (Iterator i= objects.iterator(); i.hasNext(); ) {
      Object obj = i.next();
      double[][] roc = (double[][])rocs.get(obj);
      Matrix m = new Matrix(roc);
      m.toPlot2DPanel(pp, obj.toString(), "LINE");
    }
    new ExistablePlotViewer("ROC", pp);

  }



  /**
   * The passed in map has image name as key and the value is another map
   * which has object as key and score as value. For example, one entry can be
   * 1.jpg -> {"tree"->0.9; "sky"->0.2}
   * classifiedMap2ObjScoreMap will change the map to have objects as key,
   * image->score maps as values. Like,
   * "tree" -> {1.jpg->0.9; 2.jpg->0.5; ...}
   * @param classified
   * @return
   */
  public HashMap classifiedMap2ObjScoreMap(HashMap classified) {
    List images = new ArrayList(classified.keySet());
    Collections.sort(images);
    HashMap obj_scores = new HashMap();
    for (Iterator i= images.iterator(); i.hasNext(); ) {
      Object the_img = i.next();
      DoubleScoreHashMap predicted = (DoubleScoreHashMap)classified.get(the_img);
      for (Iterator j= predicted.keySet().iterator(); j.hasNext();) {
        Object the_obj = j.next();
        double score = predicted.getScore(the_obj);

        // put the score in the corresponding hashset
        DoubleScoreHashMap one_obj_score = null;
        if (obj_scores.containsKey(the_obj)) {
          one_obj_score = (DoubleScoreHashMap)obj_scores.get(the_obj);
        } else {
          one_obj_score = new DoubleScoreHashMap();
          obj_scores.put(the_obj, one_obj_score);
        }
        one_obj_score.putScore(the_img, score);
      }
    }
    return obj_scores;
  }

  public void printObjScoreMap(HashMap obj_scores) {
    // now, print the object candidates according to the score
    List objects = new ArrayList(obj_scores.keySet());
    Collections.sort(objects);
    for (Iterator i= objects.iterator(); i.hasNext(); ) {
      Object the_obj = i.next();
      System.out.println("----------------------");
      System.out.println(the_obj.toString() + ":");
      DoubleScoreHashMap one_obj_score = (DoubleScoreHashMap)obj_scores.get(the_obj);
      List the_images = new ArrayList(one_obj_score.keySet());
      Collections.sort(the_images, one_obj_score);
      for (Iterator j= the_images.iterator(); j.hasNext();) {
        Object the_img = j.next();
        double score = one_obj_score.getScore(the_img);
        System.out.println(the_img.toString() + "= " + score);
      }
    }
  }

  public static void outputObjScoreMapHtml(HashMap obj_scores, String prefix) {
    // now, print the object candidates according to the score
    List objects = new ArrayList(obj_scores.keySet());
    Collections.sort(objects);
    for (Iterator i= objects.iterator(); i.hasNext(); ) {
      Object the_obj = i.next();
      try {
        PrintStream out = new PrintStream(
            new BufferedOutputStream(new FileOutputStream(prefix + the_obj + ".html")));
        // header
        out.println("<html><body>");
        out.println("<h3>Image sorted by their score on " + the_obj + "<h3><p>");

        DoubleScoreHashMap one_obj_score = (DoubleScoreHashMap)obj_scores.get(the_obj);
        List the_images = new ArrayList(one_obj_score.keySet());
        Collections.sort(the_images, one_obj_score);
        for (Iterator j= the_images.iterator(); j.hasNext();) {
          Object the_img = j.next();
          double score = one_obj_score.getScore(the_img);
          out.println(
              "<img src=\"" +
              the_img.toString() + "\"> <br> score= " + score + "<br>");
        }
        out.println("</body></html>");
        out.close();
      }
      catch (java.io.IOException e) {
        e.printStackTrace();
      }

    }
  }

  public final static String USAGE = "Usage: -i img_name [-s img_scale] -p data_processor_name -l learner_name -c classifier_name [-f fold_count] [-r reuse_file_name]";

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

    // get the image name
    String imgmanager_name = ReadArgu.getOption('i', args);
    if (imgmanager_name == null) {
      System.err.println(USAGE);
      System.exit(1);
    }

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

    // get the data processor name
    String dp_name = ReadArgu.getOption('p', args, "KMeansHSISeg");
    if (dp_name.compareToIgnoreCase("ColorRegoin") == 0) dp_name = "yi.iu.processor.color.reg.DPU_ColorRegion";
    if (dp_name.compareToIgnoreCase("KMeansHSISeg") == 0) dp_name = "yi.iu.processor.color.reg.merge.KMeansHSISeg";

    // get the learner name
    String learner_name = ReadArgu.getOption('l', args, "EMWithKMeanInits");
    if (learner_name.compareToIgnoreCase("KMeans") == 0) learner_name = "yi.iu.imgdb.learn.kmeans.KMeansLearner";
    if (learner_name.compareToIgnoreCase("EM") == 0) learner_name = "yi.iu.imgdb.learn.em.EMLearner";
    if (learner_name.compareToIgnoreCase("EMWithKMeanInits") == 0) learner_name = "yi.iu.imgdb.learn.em.EMLearnerWithKMeanInits";

    // get the classifier name
    String classifier_name = ReadArgu.getOption('c', args, "IndepEventSet");
    if (classifier_name.compareToIgnoreCase("Closest") == 0) classifier_name = "yi.iu.imgdb.learn.kmeans.ClosestKMeansClassifier";
    if (classifier_name.compareToIgnoreCase("Probability") == 0) classifier_name = "yi.iu.imgdb.learn.em.ProbabilityClassifier";
    if (classifier_name.compareToIgnoreCase("IndepEventSet") == 0) classifier_name = "yi.iu.imgdb.learn.em.IndepEventSetClassifier";

    // get the fold count
    int fold = 5;
    String fold_str = ReadArgu.getOption('f', args);
    if (fold_str != null) {
      fold = Integer.parseInt(fold_str);
    }

    // the file name for output/reuse analysis result
    String reuse_file = ReadArgu.getOption('r', args);

    HashMap rocs = null;
    // reuse the results
    if (reuse_file != null) {
      try {
        rocs = (HashMap)yi.util.io.input.DIU_Ser.input(reuse_file);
      } catch (IOException e) { // if can't find the one to reuse, it is ok
      } catch (Exception e) {
        e.printStackTrace();
      }
    }

    if (rocs == null) {

      // new image manager
      ImgDescManager img_manager = new ImgDescManager();
      img_manager.readFromFile(imgmanager_name);

      ImgDBLearnApp app = new ImgDBLearnApp(
          img_manager,
          scale,
          dp_name
          );

      // fold learn
      // To save the model means for each fold.  The key is an object, and the
      // value is a LinkedList having means from different fold-run
      LinkedList models = new LinkedList();
      HashMap classified = app.foldedLearn(learner_name, classifier_name, fold, models);

      // get the models from each fold
      DoubleArrayListMap means_from_fold = app.getModelMeans(models);
      Dimension d = new Dimension();
      int[] means_img = means_from_fold.drawDataAsRGB(d, 30, 30);
      OutputResult.outputGIFImage("means.gif", means_img, d.width, d.height);

      //analysis = app.analysisClassifiedMap(classified);
      HashMap obj_scores = app.classifiedMap2ObjScoreMap(classified);
      rocs = app.analysisObjScoreMap(obj_scores);
      ImgDBLearnApp.outputObjScoreMapHtml(obj_scores, "EM_learn_");


      if (reuse_file != null) {
        try {
          yi.util.io.output.DOU_Ser.output(reuse_file, rocs);
        } catch (IOException e) {
          e.printStackTrace();
        }
      }
    }

    ImgDBLearnApp.drawROCs(rocs);



//    Collection performances = analysis.values();
//    Matrix matrix = ClassificationPerformance.toHistogramMatrix(performances);
//    matrix.toPlot3DFrame("performance", "precision vs recall", Plot3DPanel.BAR);

  }

}