package yi.iu.imgdb.learn;

import yi.util.io.*;
import java.util.*;
import yi.util.data.*;
import yi.util.config.*;
import yi.iu.imgdb.desc.*;
import java.lang.reflect.*;
import yi.ds.hash.*;
import yi.ds.sort.*;
import yi.util.io.output.*;
import yi.util.io.input.*;
import yi.ds.*;

/**
 * Utilities to deal with simplicity image set (59,895 images)
 */

public class SimImageSetUtil {
  /**
   * Convert image list to image descriptions
   */
  public static class List2Desc {
    public static void main(String[] args) {
      try {
        String[] files = ReadArgu.readStrings("sim.cd.files.list");
        HashMap desc_map = new HashMap();
        int category_count = 0;
        for (int i = 0; i < files.length; i++) {
          String path = yi.util.io.File.getPath(files[i]);
          String file = yi.util.io.File.getShortName(files[i]);
          int cat = 0;
          if (desc_map.containsKey(path)) {
            cat = ((Integer)desc_map.get(path)).intValue();
          } else {
            cat = category_count++;
            desc_map.put(path, new Integer(cat));
          }
          System.out.println(files[i]+ ":: " + cat);
        }
      }catch (Exception e) {
        e.printStackTrace();
      }
    }
  }


  /**
   * Test on the simplicity images
   */
  public static class Train {
    public final static String USAGE = "Usage: -c category [-a model_adapter_name]";


    public static void main(String[] args) {

      try {
        // init configuration
        Configuration.init(args);

        String category = ReadArgu.getOption('c', args);
        if (category == null) {
          System.err.println(USAGE);
          System.exit(1);
        }

        //
        // create the image data center on fly
        //
        // select images for training (x6, x18, ... %12=6 reserved for testing)
        // pick up all the positive and 1000 negative images for this category
        // read the description file
        int neg_img_num = 1000;
        String desc_file = "sim.desc";
        String[] desc_lines = ReadArgu.readStrings(desc_file);
        LinkedList training_images = new LinkedList();
        LinkedList neg_images = new LinkedList();
        for (int i = 0; i < desc_lines.length; i++) {
          int sep = desc_lines[i].indexOf(":: ");
          String file_name = desc_lines[i].substring(0, sep);
          // find the index of this file (last two digitals).  always ends with ".jpg"
          String last_2dig_str = file_name.substring(file_name.length() - 6,
              file_name.length() - 4);
          int last_2dig = Integer.parseInt(last_2dig_str);
          if (last_2dig % 12 == 6) {
//            System.out.println("skip " + file_name);
            continue;
          }
          String cat = desc_lines[i].substring(sep + 3);
          if (cat.compareTo(category) == 0) {
            training_images.add(file_name);
          }
          else {
            neg_images.add(file_name);
          }
        }
        System.out.println("positive image num = " + training_images.size());
        System.out.println("negative image num = " + neg_images.size());
        System.out.println("pick up " + neg_img_num +
                           " negative image for training");
        Random r = new Random();
        for (int i = 0; i < neg_img_num; i++) {
          int next = r.nextInt(neg_images.size());
          String file = (String) neg_images.remove(next);
          training_images.add(file);
          System.out.println(file);
        }

        LinkedList obj_list = new LinkedList();
        obj_list.add(category);

        String rep_path = Configuration.getParameter("REPRESENTATION_PATH", "");
        String rep_str = "cp;tp;pc";
        RepresentationDescriptor[] represent_desc = RepresentationDescriptor.
            parse(rep_str);
        ImgDataCenter.VERBOSE = false;
        ImgDataCenter idc = new ImgDataCenter(rep_path, training_images,
                                                 obj_list, desc_file,
                                                 represent_desc);

        String learn_path = Configuration.getParameter("LEARNING_PATH", "");

        //
        // create the model adapter to train
        //
        // arguments to initialize the StatModelAdapter
        Class[] argumentTypes = {
            ImgDataCenter.class,
            LinkedList.class,
            LinkedList.class,
            LinkedList.class
        };
        String model_adapter_name = ReadArgu.getOption('a', args , "yi.iu.imgdb.learn.GMMNNAdapter");
        // the StatModelAdapter class
        Class model_adapter_class = Class.forName(model_adapter_name);
        // the constructor
        Constructor model_adapter_constructor =
                 model_adapter_class.getConstructor( argumentTypes );

          // create the model adapter
          Object[] arguments = {
              idc,
              training_images,
              null,
              null
          };


          StatModelAdapter model_adapter =
              (StatModelAdapter)model_adapter_constructor.newInstance(arguments);

          // Save the model
          String model_output_name = learn_path + StringUtil.getFileSuffixNameRemoved(desc_file) + "." +
              rep_str + "." +
              StringUtil.getClassShortName(model_adapter_name) + "." +
              category + "." + "mdl";
          model_adapter.writeToFile(model_output_name);
          System.out.println("the trained model is saved to " + model_output_name);

      } catch (Exception e) {
        e.printStackTrace();
      }
    }
  }


  /**
   * Make image data center for testing on the simplicity images
   */
  public static class MakeTestImageDataCenter {

    public static void main(String[] args) {

      try {
        // init configuration
        Configuration.init(args);

        //
        // create the image data center on fly
        //
        // select images for testing (x6, x18, ... %12=6)
        String desc_file = "sim.desc";
        String[] desc_lines = ReadArgu.readStrings(desc_file);
        LinkedList training_images = new LinkedList();
        for (int i = 0; i < desc_lines.length; i++) {
          int sep = desc_lines[i].indexOf(":: ");
          String file_name = desc_lines[i].substring(0, sep);
          // find the index of this file (last two digitals).  always ends with ".jpg"
          String last_2dig_str = file_name.substring(file_name.length() - 6,
              file_name.length() - 4);
          int last_2dig = Integer.parseInt(last_2dig_str);
          if (last_2dig % 12 != 6) {
            continue;
          }
          String cat = desc_lines[i].substring(sep + 3);
          training_images.add(file_name);
        }
        System.out.println("testing image num = " + training_images.size());

        String rep_path = Configuration.getParameter("REPRESENTATION_PATH", "");
        String rep_str = "cp;tp;pc";
        RepresentationDescriptor[] represent_desc = RepresentationDescriptor.
            parse(rep_str);
        ImgDataCenter idc = new ImgDataCenter(rep_path, training_images,
                                                 null, desc_file,
                                                 represent_desc);

        String idc_name = rep_path + StringUtil.getFileSuffixNameRemoved(desc_file) + "." +
            rep_str + "." + "idc";
        idc.writeToFile(idc_name);

        System.out.println("the testing image data center is saved to " + idc_name);

      } catch (Exception e) {
        e.printStackTrace();
      }
    }
  }


  /**
   * Experiments on the simplicity images
   */
  public static class Test {
    public final static String USAGE = "Usage: -c category [-a model_adapter_name]";


    public static void main(String[] args) {

      try {
        // init configuration
        Configuration.init(args);

        String category = ReadArgu.getOption('c', args);
        if (category == null) {
          System.err.println(USAGE);
          System.exit(1);
        }

//        //
//        // create the image data center on fly
//        //
//        // select images for testing (x6, x18, ... %12=6)
        String desc_file = "sim.desc";
//        String[] desc_lines = ReadArgu.readStrings(desc_file);
//        LinkedList training_images = new LinkedList();
//        for (int i = 0; i < desc_lines.length; i++) {
//          int sep = desc_lines[i].indexOf(":: ");
//          String file_name = desc_lines[i].substring(0, sep);
//          // find the index of this file (last two digitals).  always ends with ".jpg"
//          String last_2dig_str = file_name.substring(file_name.length() - 6,
//              file_name.length() - 4);
//          int last_2dig = Integer.parseInt(last_2dig_str);
//          if (last_2dig % 12 != 6) {
//            continue;
//          }
//          String cat = desc_lines[i].substring(sep + 3);
//          training_images.add(file_name);
//        }
//        System.out.println("testing image num = " + training_images.size());

        String rep_path = Configuration.getParameter("REPRESENTATION_PATH", "");
        String rep_str = "cp;tp;pc";
        RepresentationDescriptor[] represent_desc = RepresentationDescriptor.
            parse(rep_str);
//        ImgDataCenter idc = new ImgDataCenter(rep_path, training_images,
//                                                 null, desc_file,
//                                                 represent_desc);

        String idc_name = rep_path +
            StringUtil.getFileSuffixNameRemoved(desc_file) + "." +
            rep_str + "." + "idc";
        ImgDataCenter idc = ImgDataCenter.createFromFile(idc_name);
        LinkedList training_images = idc.getImageList();
        System.out.println("testing image num = " + training_images.size());

        String learn_path = Configuration.getParameter("LEARNING_PATH", "");
        String model_adapter_name = ReadArgu.getOption('a', args , "yi.iu.imgdb.learn.GMMNNAdapter");

        String model_name = learn_path + StringUtil.getFileSuffixNameRemoved(desc_file) + "." +
            rep_str + "." +
            StringUtil.getClassShortName(model_adapter_name) + "." +
            category + "." + "mdl";
        StatModelAdapter model_adapter = StatModelAdapter.createFromFile(
            model_name);

        // Test
        model_adapter.m_imgDataCenter = idc;
        StatTestRst rst = model_adapter.test(training_images);
        String rst_output_name = learn_path + StringUtil.getFileSuffixNameRemoved(desc_file) + "." +
            rep_str + "." +
            StringUtil.getClassShortName(model_adapter_name) + "." +
            category + "." + "rst";
        rst.writeToFile(rst_output_name);
        System.out.println("the test result is saved to " + rst_output_name);

      } catch (Exception e) {
        e.printStackTrace();
      }
    }
  }

  public static HashMap getImgObjScoreMap(StatTestRst rst, ImgDescManager idm) {

    // analysis each object
    List objects = new ArrayList(rst.keySet());
    Collections.sort(objects);
    int num_obj = objects.size();

    // images
    List images = new ArrayList(((HashMap)rst.get(objects.get(0))).keySet());
    Collections.sort(images);

    HashMap img_obj_scores = new HashMap(images.size());

    int i_index = 0;
    for (Iterator i = images.iterator(); i.hasNext(); i_index++) {
      String img = (String) i.next();
      System.out.println(img);
      double[] scores = new double[num_obj];
      for (Iterator o = objects.iterator(); o.hasNext(); ) {
        String obj = (String) o.next();
        DoubleScoreHashMap one_obj_score = new DoubleScoreHashMap(
            (HashMap) rst.get(obj));
        int label_ind = Integer.parseInt(obj);
        scores[label_ind] = one_obj_score.getScore(img);
      }
      int[] sorted_index = SortUtil.sort(scores, false);
      img_obj_scores.put(img, new ScoreWithSortedIndex(scores, sorted_index));
    }
    return img_obj_scores;
  }

  /**
   * generate a hashmap containinig the image->object scores
   * for each entry, the key is the image name, the value is a ScoreWithSortedIndex
   * object.  ScoreWithSortedIndex.score[obj_index] is the score for the object
   * and ScoreWithSortedIndex.sortedIndex is the sorted index
   */

  public static class GenerateImgObjScoreMap {
    public final static String USAGE = "Usage: -d img_desc.desc -r rst_file";

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

        // get the datacenter
        String data_center_name = ReadArgu.getOption('d', args);
        if (data_center_name == null) {
          System.err.println(USAGE);
          System.exit(1);
        }

        // init the datacenter
        String rep_path = Configuration.getParameter("REPRESENTATION_PATH", "");
        ImgDescManager idm = new ImgDescManager();
        idm.readFromFile(data_center_name);

       String rst_files = ReadArgu.getOption('r', args);
       if (rst_files == null) {
         System.out.println(USAGE);
         System.exit(1);
       }
       String learn_path = Configuration.getParameter("LEARNING_PATH", "");
       StatTestRst rst = StatTestRst.createFromFile(learn_path + rst_files);

       HashMap map = getImgObjScoreMap(rst, idm);
       String output_name = learn_path + StringUtil.getFileSuffixNameRemoved(rst_files) + ".obj2score.rst";
       DOU_Ser.output(output_name, map);
       System.out.println("obj2score map saved to " + output_name);
     }
     catch (Exception e) {
       e.printStackTrace();
     }
   }
  }

  public static class PrintImgObjScoreMap {
    public final static String USAGE = "Usage: -r obj2score.rst";

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

       String rst_files = ReadArgu.getOption('r', args);
       if (rst_files == null) {
         System.out.println(USAGE);
         System.exit(1);
       }
       String learn_path = Configuration.getParameter("LEARNING_PATH", "");
       HashMap rst = (HashMap)DIU_Ser.input(learn_path + rst_files);

       ArrayList images = new ArrayList(rst.keySet());
       Collections.sort(images);

       for (Iterator i=images.iterator(); i.hasNext(); ) {
         String img = (String)i.next();
         ScoreWithSortedIndex map = (ScoreWithSortedIndex)rst.get(img);
         double[] scores = map.scores;
         int[] sorted_index = map.sortedIndex;
         System.out.print(img);
         for (int j=0; j<5; j++) {
           System.out.print("\t" + sorted_index[j]);
           System.out.print("\t" + scores[sorted_index[j]]);
         }
         System.out.println();
       }

     }
     catch (Exception e) {
       e.printStackTrace();
     }
   }
  }

  public static HashMap getCategoryDescMap(String cat_desc_name) throws java.io.IOException{
    String[] cat_desc_list = yi.util.io.ReadArgu.readStrings(cat_desc_name);
    HashMap cat_desc_map = new HashMap();
    for (int i=0; i<cat_desc_list.length; i++) {
      int sep_pos = cat_desc_list[i].indexOf("::");
      String cat = cat_desc_list[i].substring(0, sep_pos);
      while (cat.endsWith(" ")) cat = cat.substring(0, cat.length() - 1);
      String desc = cat_desc_list[i].substring(sep_pos + 2);
      while (desc.startsWith(" ")) desc = desc.substring(1);
      cat_desc_map.put(cat, desc);
    }
    return cat_desc_map;
  }

  public static class PrintImgCategorizationInHtml {
    public final static String USAGE = "Usage: -r obj2score.rst -d img.desc -c category.desc";

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

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

        // init the datacenter
        ImgDescManager idm = new ImgDescManager();
        idm.readFromFile(img_desc_name);

        // get the image description
        String cat_desc_name = ReadArgu.getOption('c', args);
        if (cat_desc_name == null) {
          System.err.println(USAGE);
          System.exit(1);
        }
        HashMap cat_desc_map = getCategoryDescMap(cat_desc_name);

        // init the datacenter
        ImgDescManager cat_idm = new ImgDescManager();
        cat_idm.readFromFile(cat_desc_name);

       String rst_files = ReadArgu.getOption('r', args);
       if (rst_files == null) {
         System.out.println(USAGE);
         System.exit(1);
       }
       String learn_path = Configuration.getParameter("LEARNING_PATH", "");
       HashMap rst = (HashMap)DIU_Ser.input(learn_path + rst_files);

       ArrayList images = new ArrayList(rst.keySet());
       Collections.sort(images);
       String img_path =  Configuration.getParameter("IMAGE_PATH", "");

       for (Iterator i=images.iterator(); i.hasNext(); ) {
         String img = (String)i.next();
         ScoreWithSortedIndex map = (ScoreWithSortedIndex)rst.get(img);
         double[] scores = map.scores;
         int[] sorted_index = map.sortedIndex;
         String img_name = img_path + img;
         System.out.println("<img src=\"" + img_name + "\">" + "<br>");
         LinkedList labels = new LinkedList(idm.getDesc(img));
         String label = (String)labels.getFirst();
         int cat = Integer.parseInt(label);
         String cat_desc = (String)cat_desc_map.get(Integer.toString(cat));
         System.out.println("cat = " + cat + " (" + cat_desc + ")" + "<br>");
         boolean found = false;
         for (int j=0; j<5; j++) {
           int j_cat = sorted_index[j];
           if (j_cat == cat) {
             System.out.println("<font color=red>");
             found = true;
           }
           System.out.println("top " + (j + 1) + " = " +
                              sorted_index[j] +
                              " (" +
                              cat_desc_map.get(Integer.toString(sorted_index[j])) +
                              "), p = " +
                              scores[sorted_index[j]] + "<br>");
           if (j_cat == cat) {
             System.out.println("</font>");
           }
         }
         if (!found) {
           for (int j=5; j<sorted_index.length; j++) {
             int j_cat = sorted_index[j];
             if (j_cat != cat)
               continue;
             System.out.println("<font color=red>");
             System.out.println("top " + (j + 1) + " = " +
                                sorted_index[j] +
                                " (" +
                                cat_desc_map.get(Integer.toString(sorted_index[
                 j])) +
                                "), p = " +
                                scores[sorted_index[j]] + "<br>");
             System.out.println("</font>");
             break;
           }
         }
         System.out.println("<br>");
       }

     }
     catch (Exception e) {
       e.printStackTrace();
     }
   }
  }

  /**
   * For print out the category description (for example, only print those
   * %10 == 0
   */
  public static class PrintCategoryDescription {
    public final static String USAGE = "Usage: -c category.desc";

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

        // get the Category description
        String cat_desc_name = ReadArgu.getOption('c', args);
        if (cat_desc_name == null) {
          System.err.println(USAGE);
          System.exit(1);
        }
        HashMap cat_desc_map = getCategoryDescMap(cat_desc_name);

        for (int i=0; i<599; i+=10) {
          String cat = Integer.toString(i);
          Object desc = cat_desc_map.get(cat);
          System.out.println(cat + "\t" + desc);
        }

     }
     catch (Exception e) {
       e.printStackTrace();
     }
   }
  }


  /**
   * Jie1 Cheng2
   * @param n
   * @return
   */
  public static int jieCheng(int n) {
    int product = 1;
    for (int i = 1; i <= n; i++) {
      product *= i;
    }
    return product;
  }

  public static HashMap getP_j_k(CounterHashMap cat_stat_map, int n, int k) {
    HashMap p_j_k_map = new HashMap(cat_stat_map.size());

    for (Iterator l=cat_stat_map.keySet().iterator(); l.hasNext(); ) {
      String label = (String)l.next();
      int m = cat_stat_map.getCount(label); // appear m times in all the categories
      double[] p_j_k = new double[k+1];
      double p = m/(double)n;
      for (int j=0; j<=k; j++) {
        double sum = 0;
        for (int i=j; i<=k; i++) {
          sum+= jieCheng(k)/(double)jieCheng(i)/(double)jieCheng(k-i)*Math.pow(p,i)*Math.pow(1-p, k-i);
        }
        p_j_k[j] = sum;
      }
      p_j_k_map.put(label, p_j_k);
    }
    return p_j_k_map;
  }

  /**
   * Calculate p(j,k): a description words appears at least j times in k
   * categories randomly.
   *
   */
  public static class CalculateP_j_k {
    public final static String USAGE = "Usage: -c category.desc";

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

        // get the Category description
        String cat_desc_name = ReadArgu.getOption('c', args);
        if (cat_desc_name == null) {
          System.err.println(USAGE);
          System.exit(1);
        }
        ImgDescManager idm = new ImgDescManager();
        idm.readFromFile(cat_desc_name);
        CounterHashMap cat_stat_map = idm.getLabelStat();

        int k=5;
        int n=599;
        HashMap p_j_k_map = getP_j_k(cat_stat_map, n, k);
     }
     catch (Exception e) {
       e.printStackTrace();
     }
   }
  }

  public static class Annotation {
    public final static String USAGE = "Usage: -r obj2score.rst -c category.desc -d img.desc";

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

       String rst_files = ReadArgu.getOption('r', args);
       if (rst_files == null) {
         System.out.println(USAGE);
         System.exit(1);
       }
       String learn_path = Configuration.getParameter("LEARNING_PATH", "");
       HashMap rst = (HashMap)DIU_Ser.input(learn_path + rst_files);

       // get the image description
       String img_desc_name = ReadArgu.getOption('d', args);
       if (img_desc_name == null) {
         System.err.println(USAGE);
         System.exit(1);
       }
       // init the datacenter
       ImgDescManager idm = new ImgDescManager();
       idm.readFromFile(img_desc_name);

       // get the Category description
       String cat_desc_name = ReadArgu.getOption('c', args);
       if (cat_desc_name == null) {
         System.err.println(USAGE);
         System.exit(1);
       }
       ImgDescManager cat_desc = new ImgDescManager();
       cat_desc.readFromFile(cat_desc_name);
       CounterHashMap cat_stat_map = cat_desc.getLabelStat();

       HashMap cat_desc_map = getCategoryDescMap(cat_desc_name);

       int k=5;
       int n=599;
       HashMap p_j_k_map = getP_j_k(cat_stat_map, n, k);

       ArrayList images = new ArrayList(rst.keySet());
       Collections.sort(images);
       String img_path = Configuration.getParameter("IMAGE_PATH", "");

       double threshold = 0.0824; // those less than this threshold are considered as label

//       LinkedList ps = new LinkedList();

       for (Iterator i=images.iterator(); i.hasNext(); ) {
         String img = (String)i.next();
         ScoreWithSortedIndex map = (ScoreWithSortedIndex)rst.get(img);
         double[] scores = map.scores;
         int[] sorted_index = map.sortedIndex;
         CounterHashMap label_counter = new CounterHashMap();
         for (int j=0; j<k; j++) {
           int predicted_cat = sorted_index[j];
           HashSet labels = cat_desc.getDesc(Integer.toString(predicted_cat));
           label_counter.increaseCount(labels);
         }
         DoubleScoreHashMap selected_labels = new DoubleScoreHashMap();
         for (Iterator l=label_counter.keySet().iterator(); l.hasNext(); ) {
           String label = (String)l.next();
           int count = label_counter.getCount(label); // appear times
           double p = ((double[])p_j_k_map.get(label))[count];
//           ps.add(new Double(p));
           if (p < threshold) {
             // System.out.println("\t" + label + "\t" + count + "\t" + p);
             selected_labels.putScore(label, p);
           }
         }
         Object cat = idm.getDesc(img).toArray()[0];
         HashSet manual_labels = cat_desc.getDesc((String)cat);
         int predicted = 0;
         for (Iterator j=manual_labels.iterator(); j.hasNext(); ) {
           if (selected_labels.containsKey(j.next())) predicted++;
         }
         int un_predicted = manual_labels.size() - predicted;
         if (un_predicted > 1) continue;

         // sort the selected labels
         ArrayList sorted_labels = selected_labels.getObjectsByOrder();
         System.out.println("<img src=\"" + img_path + img + "\">" + "<br>");
         System.out.println("copy " + img_path + img + " D:/research/mypapers/phd/fig/sim.ann." + yi.util.io.File.getShortName(img) + "<br>");
         System.out.print("predicted: ");
         int j_index = 0;
         for (Iterator j=sorted_labels.iterator(); j.hasNext(); j_index++) {
           Object label = j.next();
           // System.out.println("\t" + label + "\t" + selected_labels.getScore(label));
           if (j_index != 0) System.out.print(", ");
           if (!manual_labels.contains(label)) System.out.print("<font color=red>");
           System.out.print(label);
           if (!manual_labels.contains(label)) System.out.print("</font>");
         }
         System.out.println("<br>");
         System.out.println("manual" + "[" + cat + "]: "+ cat_desc_map.get(cat) + "<br>");
         for (Iterator j=manual_labels.iterator(); j.hasNext(); ) {
           Object man = j.next();
           if (!selected_labels.containsKey(man))
             System.out.println("<font color=blue>" + man + "</font>" + "<br>");
         }
       }
//       double[] ps_array = yi.util.data.Conversion.toDoubleArray(ps);
//       Arrays.sort(ps_array);
//       System.out.println("total length = " + ps_array.length);
//       System.out.println("median = " + ps_array[ps_array.length/2]);
       // total length = 69264
       // median = 0.08235085231951043
     }
     catch (Exception e) {
       e.printStackTrace();
     }
   }
  }

}

/**
 * Good for save the object scores for an image.  The sorted index are cached
 * to save running time
 */
class ScoreWithSortedIndex implements java.io.Serializable{
  double[] scores;
  int[] sortedIndex;
  public ScoreWithSortedIndex(double[] scores, int[] sorted_index) {
    this.scores = scores;
    this.sortedIndex = sorted_index;
  }
}