package yi.iu.imgdb.learn;

import java.util.*;
import yi.util.io.OutputResult;
import yi.stat.em.*;
import yi.stat.nn.*;
import yi.util.config.Configuration;
import yi.util.io.input.*;
import java.io.*;

/**
 * Train on multiple objects the same time.  In GMMNNAdapter, we learn the
 * object models one by one.  The another strategy is to learn the GMM from the
 * features of all the object, then construct NN with multiple output.
 */

public class GMMNNAllObjAdapter extends GMMNNAdapter {

  public final static String ALL_OBJ_NAME = "All";

  public GMMNNAllObjAdapter(ImgDataCenter data_center, LinkedList img_list,
                      LinkedList ft_list, LinkedList obj_list) {
    super(data_center, img_list, ft_list, obj_list);
  }

  /**
   * Learn the GMM-NN models
   * @param data_center
   * @param img_list
   * @param ft_list
   * @param obj_list
   */
  public void learn(ImgDataCenter data_center, LinkedList img_list,
                      LinkedList ft_list, LinkedList obj_list) {

    // remove "bg" since it is not useful here
    for (Iterator o = m_objectList.iterator(); o.hasNext(); ) {
      String object = (String) o.next();
      if (object.compareTo("bg") == 0) {
        o.remove();
      }
    }
    System.out.println("objects to learn: ");
    OutputResult.print(System.out, m_objectList);

    // learn the objects all together
    System.out.print("learning on ");
    for (Iterator o = m_objectList.iterator(); o.hasNext(); ) {
      String object = (String) o.next();
      System.out.print(object + " ");
    }
    System.out.println();

    // for each feature, learn a gmm
    for (Iterator f = m_featureList.iterator(); f.hasNext(); ) {
      String feature = (String) f.next();
      try {
        learnGaussianMixtureModelOnFeature(ALL_OBJ_NAME, feature, img_list, M);
      }
      catch (ModelSingularException e) {
        e.printStackTrace();
      }

    }

    constructAbstractRegionModel(ALL_OBJ_NAME, img_list);
    double[][] scores_concat = getFeatureVectors(img_list, ALL_OBJ_NAME, null);

    if (OUTPUT_NN_DATA) {
      System.out.println("score on pos writen to " + ALL_OBJ_NAME + ".nn.dat");
      yi.util.io.OutputResult.print(ALL_OBJ_NAME.toString() + ".nn.dat",
                                    scores_concat,
                                    scores_concat[0].length,
                                    scores_concat.length);
    }

    // for experiemnt
    LinkedList test_img = StatModelFoldExperiment.testing;
    double[][] test_scores = null;
    if (test_img != null) {
      constructAbstractRegionFeatures(ALL_OBJ_NAME, test_img);
      test_scores = getFeatureVectors(test_img, ALL_OBJ_NAME, null);
    }

    NeuralNet nn = trainNN(scores_concat, test_scores);
    m_nnModels.put(ALL_OBJ_NAME, nn);
    m_nnVerifyErrors.putScore(ALL_OBJ_NAME, nn.getVerifyError());
  }


  /**
   * Train Neural Net.  The input are the features from all the training images
   * and those from all the test images.  (the latter one is just for experiment,
   * and can be null if it is not interested)
   * @param features
   * @param obj_list
   * @param test
   * @return
   */
  public NeuralNet trainNN(double[][] d, double[][] test){

    System.out.println();
    System.out.println("train NN");
    System.out.println("data # = " + d.length);

    // construct the target data
    double[][] target =  prepareNNTargetData(m_imageList);

    System.out.println("data dim = " + d[0].length);

    // construct the test data
    double[][] test_target = null;
    if (test != null) {
      System.out.println("test data # = " + test.length);
      LinkedList test_img = StatModelFoldExperiment.testing;
      test_target = prepareNNTargetData(test_img);
    }

    NeuralNet nn = new NeuralNet(d[0].length, NN_HIDDEN_VAR_NUM, target[0].length, NN_RATE);
    nn.trainWithVerification(d, target, NN_ITERATIONS, MAX_NN_ITERATION, VERIFICATION_PORTION, test, test_target);

    return nn;
  }

  double[][] prepareNNTargetData(LinkedList images) {
    double[][] target = new double[images.size()][m_objectList.size()];
    // set up a object -> int map to speed up
    HashMap obj2ind = new HashMap();
    int i_index = 0;
    for (Iterator i=m_objectList.iterator(); i.hasNext(); i_index++) {
      obj2ind.put(i.next(), new Integer(i_index));
    }

    // construct the data
    int img_index = 0;
    for (Iterator i = images.iterator(); i.hasNext(); img_index++) {
      String img = (String) i.next();
      LinkedList objs = m_imgDataCenter.getImgDescriptions(img);
      for (Iterator j = objs.iterator(); j.hasNext(); ) {
        String obj = (String) j.next();
        if (obj2ind.containsKey(obj)) {
          target[img_index][ ( (Integer) obj2ind.get(obj)).intValue()] = 1;
        }
      }
    }
    return target;
  }

  /**
   * Given score for each images in the given image list for each object.
   * The score here is after log.
   * @param imgs
   * @return StatTestRst
   */
  public StatTestRst test(java.util.LinkedList imgs) {
    StatTestRst rst = new StatTestRst();
    String object = ALL_OBJ_NAME;
    System.out.println(object);

    // set up a object -> int map to speed up
    HashMap obj2ind = new HashMap();
    int i_index = 0;
    for (Iterator i = m_objectList.iterator(); i.hasNext(); i_index++) {
      obj2ind.put(i.next(), new Integer(i_index));
    }

//      constructAbstractRegionFeatures(object, imgs);

    double[][] fv_concat = getFeatureVectors(imgs, object, null);

    NeuralNet nn = (NeuralNet) m_nnModels.get(object);

    i_index = 0;
    double[][] scores = new double[imgs.size()][];
    // for return the normalized data in NN
    double[][] normalized_data = new double[imgs.size()][];
    double[][] hidden_nodes = new double[imgs.size()][];
    double[][] final_scores = new double[imgs.size()][];
    int[][] cat = new int[imgs.size()][m_objectList.size()];
    for (Iterator j = imgs.iterator(); j.hasNext(); i_index++) {
      String img = (String) j.next();
      scores[i_index] = fv_concat[i_index];
      normalized_data[i_index] = new double[fv_concat[0].length];
      double[] nn_scores = nn.test(scores[i_index], false,
                                   normalized_data[i_index]);
      final_scores[i_index] = nn_scores;
      int o_index = 0;
      for (Iterator o = m_objectList.iterator(); o.hasNext(); o_index++) {
        String obj = (String) o.next();
        rst.putScore(obj, img, nn_scores[o_index]);
        cat[i_index][o_index] = m_imgDataCenter.getImgDescriptions(img).
            contains(obj) ? 1 : 0;
      }
      hidden_nodes[i_index] = (double[]) nn.getHiddenNodes().clone();
    }

    if (VERBOSE_IN_TEST) {
      String img_path = Configuration.getParameter("IMAGE_PATH", "");
      String ft_path = Configuration.getParameter("FEATURE_PATH", "");
      String rep_path = Configuration.getParameter("REPRESENTATION_PATH", "");
      String learn_path = Configuration.getParameter("LEARNING_PATH", "");

      int o_index = 0;
      for (Iterator o = m_objectList.iterator(); o.hasNext(); o_index++) {
        double[] final_score_obj = new double[final_scores.length];
        for (int k = 0; k < final_score_obj.length; k++) {
          final_score_obj[k] = final_scores[k][o_index];
        }

        int[] sorted_index = yi.ds.sort.SortUtil.sort(final_score_obj);
        double[][] nn_beta = nn.getBeta();
        String generator = (String) m_featureList.getFirst();
        FeatureDescriptor ft_desc = FeatureDescriptor.
            getDefaultFeatureDescriptor(generator);
        RepresentationDescriptor rpt_desc = RepresentationDescriptor.
            getDefaultRepresentationDescriptor(generator);
        for (int i = 0; i < sorted_index.length; i++) {
          int sorted_i = sorted_index[i];
          String image = (String) imgs.get(sorted_i);
          int disp_width = 0, disp_height = 0;
          try {
            String rpt_file_name = RepresentationDescriptor.
                getRepresentationOutputName(rpt_desc, image, true);
            FeatureRepresentationList fr = (FeatureRepresentationList) DIU_Ser.
                input(
                rpt_file_name);
            if (fr instanceof RgnFeatureRepresentationList) {
              RgnFeatureRepresentationList rft = (RgnFeatureRepresentationList)
                  fr;
              disp_width = rft.imgWidth;
              disp_height = rft.imgHeight;
            }
          }
          catch (Exception e) {
            e.printStackTrace();
          }
          String ft_gdemo_name = FeatureDescriptor.getFeatureGraphicDemoName(
              ft_desc, image, true);
          String ft_cdemo_name = FeatureDescriptor.getFeatureColorDemoName(
              ft_desc, image, true);
          String rpt_gdemo_name = RepresentationDescriptor.
              getRepresentationGraphicDemoName(rpt_desc, image, true);
          String rpt_cdemo_name = RepresentationDescriptor.
              getRepresentationColorDemoName(rpt_desc, image, true);

          System.out.print("<img src=\"" +
                           img_path + imgs.get(sorted_i) + "\"");
          if (disp_width != 0) {
            System.out.print(" width=" + disp_width + " height=" + disp_height);
          }
          System.out.print(">\t");
          if (ft_gdemo_name != null) {
            System.out.print("<img src=\"" + ft_gdemo_name + "\">\t");
          }
          if (ft_cdemo_name != null) {
            System.out.print("<img src=\"" + ft_cdemo_name + "\">\t");
          }
          if (rpt_gdemo_name != null) {
            System.out.print("<img src=\"" + rpt_gdemo_name + "\">\t");
          }
          if (rpt_cdemo_name != null) {
            System.out.print("<img src=\"" + rpt_cdemo_name + "\">\t");
          }

          System.out.print("<br>");
          if (cat[sorted_i][o_index] == 1) {
            System.out.print("<font color=red>");
          }
          else {
            System.out.print("<font color=black>");
          }
          System.out.print("\t");
          for (int k = 0; k < scores[sorted_i].length; k++) {
            System.out.print( (int) (scores[sorted_i][k] * 100) / 100.0);
            System.out.print("\t");
          }
          System.out.print("<br>");
          for (int k = 0; k < normalized_data[sorted_i].length; k++) {
            System.out.print( (int) (normalized_data[sorted_i][k] * 100) /
                             100.0);
            System.out.print("\t");
          }
          System.out.print("<br><i>");
          for (int k = 0; k < hidden_nodes[sorted_i].length; k++) {
            System.out.print( (int) (hidden_nodes[sorted_i][k] * 100) / 100.0);
            System.out.print("\t");
          }
          System.out.print("<br>");
          for (int k = 0; k < hidden_nodes[sorted_i].length; k++) {
            System.out.print( (int) (hidden_nodes[sorted_i][k] * nn_beta[k][0] *
                                     100) / 100.0);
            System.out.print("\t");
          }
          System.out.print("</i><b>");
          System.out.print(final_scores[sorted_i]);
          System.out.println("</b></font><br>");
        }
      }
    }

    if (OUTPUT_NN_DATA) {
      System.out.println("scores on test writen to " + object + ".test.dat");
      yi.util.io.OutputResult.print(object.toString() + ".test.dat", scores,
                                    scores[0].length, scores.length);

      System.out.println("categories on test writen to " + object +
                         ".test.cat");
      yi.util.io.OutputResult.print(object.toString() + ".test.cat", cat, 1,
                                    cat.length);

    }
    return rst;
  }

}