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 yi.stat.em.*;
import java.awt.Dimension;
import yi.ds.hash.*;
import yi.stat.em.ModelSingularException;


/**
 * This is the learning approach in which each object is modelled as a Gaussian
 * Mixture Model
 */

public class ImgDBLearnApp2 extends ImgDBLearnApp{

  public ImgDBLearnApp2(HashMap imgObjects, HashMap imgFeatures) {
    super(imgObjects, imgFeatures);

  }

  public ImgDBLearnApp2(ImgDescManager img_manager, int scale, String data_processor_name) {
    super(img_manager, scale,  data_processor_name);
  }


  /**
   * Use the given images.  Do training and testing on folds
   * @param learner_name
   * @param classifier_name
   * @param fold
   * @param verify_frac how many obj-images will be used to verify
   * @param models_from_folds to return the trained model.  Give null, if not
   *        interested in the results.  Entrys like:
   *        "tree"->{models_from_fold_1, 2, ...}
   * @return the classification result. For each object, the corresponding
   *         image scores are returned. eg. "tree"->{1.jpg->9.8; 2.jpg->7.5; ...}
   */
  public HashMapHashMap foldedLearn(String learner_name, String classifier_name, int fold,
                          double verify_frac, LinkedListHashMap models_from_folds) {

    HashMapHashMap obj_img_scores = new HashMapHashMap();

    // first, get all the objects need to be learned
    Set objects = new HashSet();
    for (Iterator i = m_imgObjects.keySet().iterator(); i.hasNext(); ) {
      objects.addAll((Collection)m_imgObjects.get(i.next()));
    }

    // debug only
    //objects.clear();
    //objects.add("cherry trees");
    //objects.add("building");
    //objects.add("flowers");


    // learn the objects one by one
    for (Iterator it = objects.iterator(); it.hasNext(); ) {
      String obj = (String)it.next();

      System.out.println();
      System.out.println("training on \"" + obj + "\"" + "\n");

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

      int total_false_pos = 0, total_false_neg = 0;
      int total_pos = 0, total_neg = 0;
      HashMap right_pos_score = new HashMap(),
          right_neg_score = new HashMap(),
          false_neg_score = new HashMap(),
          false_pos_score = new HashMap();

      // for each fold
      for (int i=0; i<fold; i++) {

        System.out.println("fold #" + i + "\n");

        // get the training set and testing set
        TwoLinkedListData two_list = (TwoLinkedListData) folds.get(i);
        Collection training = two_list.first;
        Collection testing = two_list.second;

        System.out.println("learning # = " + training.size() + ", " +
                           "testing # = " + testing.size() + "\n");

        // seperate positive training instances and negative training instance
        Vector negative = new Vector();
        Vector positive = new Vector();
        for (Iterator j = training.iterator(); j.hasNext(); ) {
          Object img = j.next();
          Collection objs = (Collection) m_imgObjects.get(img);
          if (objs.contains(obj)) {
            positive.add(img);
          }
          else {
            negative.add(img);
          }
        }

        System.out.println("positive learning # = " + positive.size() + ", " +
                           "negative learning # = " + negative.size() + "\n");

        // seperate part of the positive training instance for verification
        int verify_n = (int) (positive.size() * verify_frac);
        Vector positive_v = new Vector();
        // seperate the part of training and the part for verificatoin
        for (int j = 0; j < verify_n; j++) {
          Object v = positive.remove(0);
          positive_v.add(v);
        }

        System.out.println("Training set: " + positive.size());
        OutputResult.printSortedCollection(System.out, positive);
        System.out.println("Positive Verifying set: " + positive_v.size());
        OutputResult.printSortedCollection(System.out, positive_v);
        System.out.println("Negative Verifying set: " + negative.size());
        OutputResult.printSortedCollection(System.out, negative);

        // now, use those in positive ones for training
        Collection training_fts = new Vector();
        for (Iterator j = positive.iterator(); j.hasNext(); ) {
          Collection fts = (Collection) m_imgFeatures.get(j.next());
          training_fts.addAll(fts);
        }
        double[][] data = yi.util.data.Conversion.to2DDoubleArray(training_fts);

        System.out.println("total feature vectors #= " + data.length + "\n");

        // increase kernels one by one
        double last_score = -Double.MAX_VALUE;
        double last_threshold = 0;
        GaussianMixtureModel best_gmm = null;
        for (int m = 1; true; m++) {

          System.out.println("kernel # = " + m + "\n");

          // training
          GaussianMixtureModel gmm = new GaussianMixtureModel(m, data, null, null); // suppose the data is within [0,1)
          System.out.println(gmm);

          // verify on positive cases
          Collections.sort(positive_v);
          double total_score = 1;
          double[] pscores = new double[verify_n]; // store the scores of positive cases
          for (int j = 0; j < verify_n; j++) {
            String img = (String) positive_v.elementAt(j);
            double score = 0;
            try {
              score = getScore(gmm, img);
            } catch (ModelSingularException e) {
              e.printStackTrace();
            }
            score = Math.log(score);
            pscores[j] = score;
//            System.out.println("+" + img + " = " + score);
            total_score += score;
          }

          // verify on negative cases
          Collections.sort(negative);
          double[] nscores = new double[negative.size()]; // store the scores of negative cases
          for (int j = 0; j < negative.size(); j++) {
            String img = (String) negative.elementAt(j);
            double score = 0;
            try {
              score = getScore(gmm, img);
            } catch (ModelSingularException e) {
              e.printStackTrace();
            }

            score = Math.log(score);
            nscores[j] = score;
//            System.out.println("-" + img + " = " + score);
            total_score -= score;
          }

              /** @todo this commented out part is to learn the optimal threshold */
//          // analysize the scores
//          // first put all the candidate threshold together
//          Vector cand = new Vector();
//          for (int j=0; j<verify_n; j++) {
//            cand.add(new Double(pscores[j]));
//          }
//          for (int j=0; j<negative.size(); j++) {
//            cand.add(new Double(nscores[j]));
//          }
//          // sort the scores
//          Collections.sort(cand);
//          // find the optimal threshld
//          double last_t = ((Double)cand.elementAt(0)).doubleValue();
//          int best_j = 0;
//          for (best_j=0; best_j<cand.size(); best_j++) {
//            double t = ((Double)cand.elementAt(best_j)).doubleValue();
//            int false_pos = 0;
//            for (int k=0; k<verify_n; k++) {
//              if (pscores[k] < t) false_pos++;
//            }
//            int false_neg = 0;
//            for (int k=0; k<negative.size(); k++) {
//              if (nscores[k] >= t) false_neg++;
//            }
//            // at the begining, since the threshold is too small, so false_pos
//            // should be 0; while false_neg = #of false cases.
//            // so the search should stop when the false_pos ratio is greater
//            // than false_neg ratio
//            double false_pos_ratio = false_pos / (double)verify_n;
//            double false_neg_ratio = false_neg / (double)negative.size();
//            System.out.println("t = " + t + ": " + false_pos + "/" + verify_n + "," +
//                               false_neg + "/" + negative.size());
//            if (false_pos_ratio >= false_neg_ratio) break;
//          }
//          double best_threshold = (best_j == 0) ?
//              ((Double)cand.elementAt(0)).doubleValue() :
//              (((Double)cand.elementAt(best_j - 1)).doubleValue() +
//               ((Double)cand.elementAt(best_j)).doubleValue()) / 2;
//
//          // recalculate the false pos and false neg for the optimal threshold
//          int false_pos = 0;
//          for (int k=0; k<verify_n; k++) {
//            if (pscores[k] < best_threshold) false_pos++;
//          }
//          int false_neg = 0;
//          for (int k=0; k<negative.size(); k++) {
//            if (nscores[k] >= best_threshold) false_neg++;
//          }
//
//          System.out.println("best threshold = " + best_threshold + ": " + false_pos + "/" + verify_n + "," +
//                               false_neg + "/" + negative.size());

          System.out.println();

          System.out.println("total score = " + total_score);
          System.out.println();
          // compare the current score to the last one.  If not better, just quit
          if (total_score > last_score) {
            last_score = total_score;
            best_gmm = gmm;
                /** @todo this commented out part is to learn the optimal threshold */
//            last_threshold = best_threshold;
          }
          else {
            break;
          }
        }

        if (models_from_folds != null) {
          models_from_folds.addToList(obj, best_gmm);
        }

        for (Iterator j = testing.iterator(); j.hasNext(); ) {
          Object img = j.next();
          double score = 0;
          try {
            score = getScore(best_gmm, img);
          } catch (ModelSingularException e) {
            e.printStackTrace();
          }

          score = Math.log(score);
          obj_img_scores.put(obj, img, new Double(score));
        }
      }


        /** @todo this comment out part is for testing using threshold */
//        // seperate positive test instances and negative test instance
//        Vector negative_test = new Vector();
//        Vector positive_test = new Vector();
//        for (Iterator j = testing.iterator(); j.hasNext(); ) {
//          Object img = j.next();
//          Collection objs = (Collection) m_imgObjects.get(img);
//          if (objs.contains(obj)) {
//            positive_test.add(img);
//          }
//          else {
//            negative_test.add(img);
//          }
//        }
//
//        // test on positive cases
//        Collections.sort(positive_test);
//        int false_negative = 0;
//        double[] pscores = new double[positive_test.size()];  // store the scores of positive cases
//        for (int j=0; j<positive_test.size(); j++) {
//          String img = (String)positive_test.elementAt(j);
//          double score = getScore(best_gmm, img);
//          score = Math.log(score);
//          pscores[j] = score;
//          if (score >= last_threshold) {
//            right_pos_score.put(img, new Double(score));
//            System.out.print(" "); // correctly classified
//          } else {
//            false_neg_score.put(img, new Double(score));
//            System.out.print("*");
//            false_negative ++;
//          }
//          System.out.println("+" + img + " = " + score);
//        }
//
//        // verify on negative cases
//        Collections.sort(negative_test);
//        int false_positive = 0;
//        double[] nscores = new double[negative_test.size()];  // store the scores of negative cases
//        for (int j=0; j<negative_test.size(); j++) {
//          String img = (String)negative_test.elementAt(j);
//          double score = getScore(best_gmm, img);
//          score = Math.log(score);
//          nscores[j] = score;
//          if (score < last_threshold) {
//            right_neg_score.put(img, new Double(score));
//            System.out.print(" "); // correctly classified
//          } else {
//            false_pos_score.put(img, new Double(score));
//            System.out.print("*");
//            false_positive ++;
//          }
//          System.out.println("-" + img + " = " + score);
//        }
//
//        System.out.println("\n--------------------------------------");
//        System.out.println(best_gmm);
//        System.out.println("threshold = " + last_threshold);
//        System.out.println("Test result: False Negative = " +
//                           false_negative + "/" + positive_test.size() + ", "
//                           + "False Positive = " +
//                           false_positive + "/" + negative_test.size()
//                           );
//        System.out.println("--------------------------------------\n");
//
//        total_false_pos += false_positive;
//        total_false_neg += false_negative;
//        total_pos += positive_test.size();
//        total_neg += negative_test.size();
//      }
//      System.out.println("\n======================================");
//      System.out.println("Test result: False Negative = " +
//                         total_false_neg + "/" + total_pos + ", "
//                         + "False Positive = " +
//                         total_false_pos + "/" + total_neg
//                         );
//      System.out.println("======================================\n");
//      System.out.println("Right Positive:");
//      OutputResult.printSortedMap(System.out, right_pos_score);
//      System.out.println("False Negative:");
//      OutputResult.printSortedMap(System.out, false_neg_score);
//      System.out.println("Right Negative:");
//      OutputResult.printSortedMap(System.out, right_neg_score);
//      System.out.println("False Positive:");
//      OutputResult.printSortedMap(System.out, false_pos_score);
//      System.out.println("======================================\n");

    }
    return obj_img_scores;

  }

  private double getScore(GaussianMixtureModel gmm, Object img) throws ModelSingularException{
    Collection fts = (Collection)m_imgFeatures.get(img);
    double p_on = 1;
    for (Iterator i=fts.iterator(); i.hasNext();) {
      double[] f = (double[])i.next();
      double p = gmm.getDensity(f);
      p_on*= p;
    }
    return p_on;
  }

  public DoubleArrayListMap getModelMeans(List models_from_folds) {
    DoubleArrayListMap means_from_folds = new DoubleArrayListMap();
    // loop for each fold
    int i_index = 0;
    for (Iterator i= models_from_folds.iterator(); i.hasNext(); i_index++) {
      // construct a key
      String fold_str = "fold" + i_index;
      // get the models (a mixture model) of current fold
      GaussianMixtureModel models = (GaussianMixtureModel)i.next();
      HashMap means = models.getComponentMeans();

      // loop for each model
      ArrayList model_list = new ArrayList(means.keySet());
      Collections.sort(model_list);
      for (Iterator iter = model_list.iterator(); iter.hasNext(); ) {
        Object model_name = iter.next();
        double[] mean = (double[]) means.get(model_name);
        means_from_folds.addToList(fold_str, mean);
      }
    }
    return means_from_folds;
  }

  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, "GMMEM");
    if (learner_name.compareToIgnoreCase("GMMEM") == 0) learner_name = "yi.stat.em.GMMEM";

    // get the classifier name
    String classifier_name = ReadArgu.getOption('c', args, "Closest");
    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";

    // get the fold count
    int fold = 2;
    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 analysis = null;
    // reuse the results
    if (reuse_file != null) {
      try {
        analysis = (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 (analysis == null) {

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

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

      // fold learn
      LinkedListHashMap models_from_folds = new LinkedListHashMap();
      HashMapHashMap obj_img_scores = app.foldedLearn(learner_name,
          classifier_name, fold, 1.0/3, models_from_folds);

      // get the models from each fold
      ArrayList objs = new ArrayList(models_from_folds.keySet());
      Collections.sort(objs);
      for (Iterator i=objs.iterator(); i.hasNext(); ) {
        // the object
        Object obj = i.next();
        LinkedList models = (LinkedList)models_from_folds.get(obj);
        DoubleArrayListMap means_from_fold = app.getModelMeans(models);
        Dimension d = new Dimension();
        int[] means_img = means_from_fold.drawDataAsRGB(d, 30, 30);
        OutputResult.outputGIFImage(obj.toString() + "_means.gif", means_img, d.width, d.height);
      }
    }
  }

}
