package yi.stat.nn;

import java.io.*;
import java.util.*;

import java.awt.*;
import java.awt.geom.*;
import java.awt.image.*;

import yi.util.io.*;
import yi.util.io.output.*;

/**
 * Neural Net Classifier
 * @author Yi Li
 * @version 1.0
 */

public class NeuralNet  implements java.io.Serializable{

  /**
   * When training with verification.  Do verification once every "VERIFY_INTERVAL"
   * iterations
   */
  public final static int VERIFY_INTERVAL = 1;

  static final long serialVersionUID =  -2095806084619108814L;

  /**
   * number of hidden layer nodes
   */
  int nh;

  /**
   * dimension of output
   */
  protected int k;

  public int getOutputDim() {
    return k;
  }

  /**
   * dimension of input
   */
  int d;

  /**
   * learning rate
   */
  double rate;

  /**
   * alpha[i][j] is the weigth connecting the i-th node in x and
   * the j-th node in layer one.  size of (d+1, nh)
   */
  double alpha[][];

  /**
   * alpha[i][j] is the weigth connecting the i-th node in hidden layer and
   * the j-th node in output.  size of (nh+1, k)
   */
  double beta[][];

  /**
   * errors (output vs target)
   * size of (k);
   */
  double[] e;

  /**
   * memory allocated for hidden layer calculation
   */
  double hidden[];

  /**
   * memory allocated for output layer calculation
   */
  double out[];

  double[] slope_alpha, slope_beta;

  // means of the training data.  Save for normalize test data
  double[] means;
  // dev of the training data.  Save for normalize test data
  double[] devs;

  /**
   * Verification error
   */
  double m_verifyError = Double.NaN;


  /**
   * Initialize the neural net
   * @param d dimension of input
   * @param nh number of hidden layer nodes
   * @param k dimension of output
   * @param rate learning rate
   */
  public NeuralNet(int d, int nh, int k, double rate){

    // number of nodes on layers
    this.d = d;
    this.nh = nh;
    this.k = k;
    this.rate = rate;

    // with standardized inputs, it is typical to take random uniform weights
    // over the range of [-0.7, +0.7]
    Random random = new Random();

    // connections between input layer and hidden layer
    alpha = new double[d+1][nh];
    for (int i=0; i<d+1; i++) {
      for (int j = 0; j < nh; j++) {
        alpha[i][j] = random.nextDouble() * 1.4 - 0.7;
      }
    }
    // connections between hidden layer and output layer
    beta = new double[nh+1][k];
    for (int i=0; i<nh+1; i++) {
      for (int j = 0; j < k; j++) {
        beta[i][j] = random.nextDouble() * 1.4 - 0.7;
      }
    }

    hidden = new double[nh];
    out = new double[k];
    e = new double[k];
    slope_alpha = new double[nh];
    slope_beta = new double[k];
  }

  protected double sigmoid(double x) {
    return 1.0 / (1 + Math.exp(-x));
  }


  /**
   * @param x
   * @param alpha alpha[i][j] is the weigth connecting the i-th node in x and
   *              the j-th node in layer one
   * @param beta
   * @return
   */
  public double[] test(double[] x, boolean ignore_normalization) {
    return test(x, ignore_normalization, null);
  }

  /**
   * test on the data.  return the normalized data if xx is not null
   * @param x
   * @param alpha alpha[i][j] is the weigth connecting the i-th node in x and
   *              the j-th node in layer one
   * @param xx return the normalized data if it is not null
   * @return
   */
  public double[] test(double[] x, boolean ignore_normalization, double[] xx) {
    if (ignore_normalization) {
      xx = x;
    } else {
      if (xx == null) xx = new double[x.length];
      for (int j = 0; j < d; j++) {
        if (devs[j] != 0) {
          xx[j] = (x[j] - means[j]) / devs[j];
        }
        else {
          xx[j] = 0;
        }
      }
    }

    return forwardPass(xx);
  }

  /**
   * Test on the data, return the test result in "y".  If you don't want to
   * have the "y", just pass in "null"
   * @param x
   * @param target
   * @param y
   * @return the error sum
   */
  public double test(double[] x, double[] target, double [] y, boolean ignore_normalization) {
    double [] ytemp = test(x, ignore_normalization);
    if (y != null) {
      for (int i = 0; i < y.length; i++) {
        y[i] = ytemp[i];
      }
    }
//    return getSquaredSumErr(ytemp, target);
    return yi.math.VectorUtil.EuclideanDistance(ytemp, target);
  }


  /**
   * Test on the data (the first dimension is the different data, the second
   * dimension is the components in one data vector).  Return the test result
   * in "y".  If you don't want to have the "y", just pass in "null"
   * @param x
   * @param target
   * @param y
   * @return
   */
  public double test(double[][] x, double[][] target, double [][] y, boolean ignore_normalization) {
    double sum = 0;
    for (int i=0; i<x.length; i++) {
      sum += test(x[i], target[i], (y == null)? null: y[i], ignore_normalization);
    }
    return sum/x.length;
  }

  /**
   * Return the squared sum of two vectors (normalized by the vector length)
   * @param x
   * @param y
   * @return
   */
  public double getSquaredSumErr(double[] x, double[] y) {
    double sum = 0;
    for (int i=0; i<x.length; i++) {
      double diff = x[i] - y[i];
      sum += diff * diff;
    }
    return sum / x.length;
  }

  /**
   * Return the error between the predicted and real value.  (Euclidean distance
   * between the two vectors)
   * @return
   */
  public double getError() {
    double error = 0;
    // sum the errors
    for (int i = 0; i < k; i++) {
      error += (e[i] * e[i]);
    }
    return Math.sqrt(error);
  }


  protected double[] forwardPass(double[] x) {
    // calculate the values on the first layer
    for (int i=0; i<nh; i++) {
      double sum = 0;
      for (int j=0; j<d; j++) {
        sum += x[j] * alpha[j][i];
      }
      sum += alpha[d][i]; // x[d] = 1;
      hidden[i] = sigmoid(sum);
    }

    // calculate the values on the output
    for (int i=0; i<k; i++) {
      double sum = 0;
      for (int j=0; j<nh; j++) {
        sum += hidden[j] * beta[j][i];
      }
      sum += beta[nh][i]; // hidden[nhid] = 1;
      out[i] = sigmoid(sum);
    }
    return out;
  }

  /**
   * Train several iterations
   * @param x
   * @param categories
   * @param iterations
   */
  public void train(double[][] x, double[][] categories, int iterations) {

    // first normailize the data
    means = yi.stat.StatUtil.meanOnCol(x);
    devs = yi.stat.StatUtil.varianceOnCol(x);
    for (int i=0; i<d; i++) devs[i] = Math.sqrt(devs[i]);

    double[][] xx = new double[x.length][x[0].length];
    for (int i=0; i<x.length; i++) {
      for (int j=0; j<d; j++) {
        xx[i][j] = (x[i][j] - means[j]) / devs[j];
      }
    }

    for (int iter = 0; iter<iterations; iter++) {
      double error = train(xx, categories);
      System.out.println("iter = " + iter + ", error = " + error);
    }

//    // print the test results on the training instances
//    for (int i=0; i<x.length; i++) {
//      double[] rst = test(x[i]);
//      for (int j=0; j<x[i].length; j++) {
//        System.out.print(x[i][j] + " ");
//      }
//      System.out.print(": ");
//      for (int j=0; j<categories[i].length; j++) {
//        System.out.print(categories[i][j] + " ");
//      }
//      System.out.print("-> ");
//      for (int j=0; j<rst.length; j++) {
//        System.out.print(rst[j] + " ");
//      }
//      System.out.println();
//    }

  }

  /**
   * Train one iteration on all the data.  User shouldn't call this method
   * directly since this method skips the "normalization" step (!!!)
   * @param x
   * @param categories
   * @return
   */
  protected double train(double[][] x, double[][] categories) {
    int nsamples = categories.length; // number of training samples
    Random random = new Random();
    // pick up the samples randomly
    int[] rand_ord = yi.stat.StatUtil.randomize(random, nsamples);

    double error = 0;
    // for each sample
    for (int j = 0; j < nsamples; j++) {
      // get the input
      int index = rand_ord[j];
      double[] xx = x[index];

      train(x[index], categories[index]);

      // sum the errors
      error += getError ();
//      int i=0; i<k; i++) {
//        error += (e[i] * e[i]);
//      }
    }

    return error/x.length;
  }

  /**
   * Train one iteration on one data sample.  User shouldn't call this method
   * directly since this method skips the "normalization" step (!!!)
   * @param xx
   * @param categories
   */
  protected void train(double[] xx, double[] categories) {

      // forward pass
      forwardPass(xx);

      // see the errors
      for (int i = 0; i < k; i++) {
        e[i] = categories[i] - out[i];
      }

      // slope of beta
      for (int i = 0; i < k; i++) {
        slope_beta[i] = e[i] * out[i] * (1 - out[i]);
      }

      // slope of alpha
      for (int i = 0; i < nh; i++) {
        // slope_alpha = h .* (1-h) .* (slope_beta * beta');
        double sum = 0;
        for (int m = 0; m < k; m++) {
          sum += slope_beta[m] * beta[i][m];
        }
        slope_alpha[i] = sum * hidden[i] * (1 - hidden[i]);
      }

      // update beta and alpha
      for (int m = 0; m < nh; m++) {
        for (int n = 0; n < k; n++) {
          beta[m][n] = beta[m][n] + (rate * (hidden[m] * slope_beta[n]));
        }
      }
      for (int n = 0; n < k; n++) {
        beta[nh][n] = beta[nh][n] + (rate * slope_beta[n]);
      }

      for (int m = 0; m < d; m++) {
        for (int n = 0; n < nh; n++) {
          alpha[m][n] = alpha[m][n] + (rate * (xx[m] * slope_alpha[n]));
        }
      }
      for (int n = 0; n < nh; n++) {
        alpha[d][n] = alpha[d][n] + (rate * slope_alpha[n]);
      }
  }

  /**
   * Return the hidden nodes.  This method is particular useful when examining
   * what really happened within the NN.
   * @return
   */
  public double[] getHiddenNodes() {
    return hidden;
  }

  public double[][] getBeta() {
    return beta;
  }

  public double[][] getAlpha() {
    return alpha;
  }

  /**
   * Generate bar chart to demonstrate the alpha values.  Each weight in the
   * alpha will be mapped to a horizontal bar.  The weights to the same hidden
   * node will be mapped to a horizontal bar chart (containing several horizontal
   * bars).  In the output, there are several bar charts (same number as hidden
   * nodes)
   * @param width the width of the output chart
   * @param height the height of the output chart
   * @param max the value will be mapped to the longest point of the horizontal
   *        bar  (if Double.POSITIVE_INFINITY is passed in, max weight will be
   *        assumed)
   * @param min the value will be mapped to 0 of the horizontal bar (if
   *        Double.NEGATIVE_INFINITY is passed in, min weight will be assumed)
   * @param colors the color of different bar can be different (if null is
   *        passed in, white will be assumed
   * @return
   */
  public int[][] demoAlpha(int width, int height, double max, double min, int[] colors) {
    return demoWeightInHorizontalBarCharts(alpha, width, height, max, min, colors);
  }

  /**
   * Generate bar chart to demonstrate the beta values.  Each weight in the
   * alpha will be mapped to a horizontal bar.  The weights to the same hidden
   * node will be mapped to a horizontal bar chart (containing several horizontal
   * bars).  In the output, there are several bar charts (same number as hidden
   * nodes)
   * @param width the width of the output chart
   * @param height the height of the output chart
   * @param max the value will be mapped to the longest point of the horizontal
   *        bar  (if Double.POSITIVE_INFINITY is passed in, max weight will be
   *        assumed)
   * @param min the value will be mapped to 0 of the horizontal bar (if
   *        Double.NEGATIVE_INFINITY is passed in, min weight will be assumed)
   * @param colors the color of different bar can be different (if null is
   *        passed in, white will be assumed
   * @return
   */
  public int[][] demoBeta(int width, int height, double max, double min, int[] colors) {
    return demoWeightInHorizontalBarCharts(beta, width, height, max, min, colors);
  }


  /**
   * Draw the weight on the constant input node
   */
  boolean DRAW_ON_CONSTANT_INPUT = false;

  /**
   * Generate bar chart to demonstrate the weights(alpha/beta).  Each weight in
   * the weights will be mapped to a horizontal bar.  The weights to the same
   * hidden node will be mapped to a horizontal bar chart (containing several
   * horizontal bars).  In the output, there are several bar charts (same number
   * as hidden/output nodes)
   * @param weights
   * @param width the width of the output chart
   * @param height the height of the output chart
   * @param max the value will be mapped to the longest point of the horizontal
   *        bar  (if Double.POSITIVE_INFINITY is passed in, max weight will be
   *        assumed)
   * @param min the value will be mapped to 0 of the horizontal bar (if
   *        Double.NEGATIVE_INFINITY is passed in, min weight will be assumed)
   * @param colors the color of different bar can be different (if null is
   *        passed in, white will be assumed
   * @return
   */
  protected int[][] demoWeightInHorizontalBarCharts(double[][] weights,
      int width, int height, double max, double min, int[] colors) {
    if (max == Double.POSITIVE_INFINITY) {
      max = yi.stat.StatUtil.maximum(weights);
    }
    if (min == Double.NEGATIVE_INFINITY) {
      min = yi.stat.StatUtil.minimum(weights);
    }
    int input_len = weights.length;
    if (!DRAW_ON_CONSTANT_INPUT) {
      input_len --;
    }
    int output_len = weights[0].length;
    if (colors == null) {
      colors = new int[input_len];
      java.util.Arrays.fill(colors, 0xffffffff);
    }
    int[][] charts = new int[output_len][];
    for (int h=0; h<output_len; h++) {
      // use a buffered image to draw a chart
      BufferedImage bimage = new BufferedImage(
          width, height, BufferedImage.TYPE_INT_ARGB);
      Graphics2D g = bimage.createGraphics();
      g.setBackground(Color.WHITE);
      g.clearRect(0,0,width, height);
      // draw horizontal bars
      int zero = (int)Math.round((0 - min) / (max - min) * width);
      for (int i=0; i<input_len; i++) {
        int y = i * height / (input_len);
        int y2 = (i+1) * height / (input_len); // the position of the next bar
        if (y2 == height) y2 --; // if on the boundary already
        int bh = y2 - y;
        double v = weights[i][h];
        int x = (int)Math.round((v - min) / (max - min) * width);
        if (x == width) x--; // if on the boundary already
        int bw = x - zero;
        if (bw < 0) {
          bw = -bw;
        } else {
          x = zero;
        }
        Shape bar = new Rectangle2D.Double(x, y, bw, bh);
        g.setPaint(new Color((i>=colors.length)? 0xffffffff: colors[i]));
        g.fill(bar);
        g.setColor(Color.BLACK);
        g.draw(bar);
      }
      g.setColor(Color.BLACK);
      g.drawLine((int)Math.round(zero), 0, (int)Math.round(zero), height);
      g.dispose();
      charts[h] = bimage.getRGB(0, 0, width, height, null, 0, width);
    }
    return charts;
  }

  public static void main(String[] args) {
    try {
      // read data
      double[][] train = yi.util.io.ReadArgu.loadMatlab("train.dat");
      double[][] train_index = yi.util.io.ReadArgu.loadMatlab("train_index.dat");

      double[][] train_cat = new double[train.length][2];
      for (int i=0; i< train_index[0].length; i++) {
        if (train_index[0][i] == 1) {
          train_cat[i][0] = 1;
          train_cat[i][1] = 0;
        } else {
          train_cat[i][0] = 0;
          train_cat[i][1] = 1;
        }
      }

      NeuralNet nn = new NeuralNet(5, 3, 2, 0.1);
      nn.train(train, train_cat, 100);
      System.out.println("done");

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

  /**
   * Write the mixture model as a java object to file
   * @param file
   */
  public void writeToFile(String file) throws java.io.IOException {
    DOU_Ser.output(file, this);
  }

  /**
   * create a mixture model from a file
   *
   * @param file
   * @return ImgDataCenter
   */
  public static NeuralNet createFromFile(String file) throws Exception {
    return (NeuralNet) yi.util.io.input.DIU_Ser.input(file);
  }

  /**
   * Train the NN with verification.  At least "min_iteration" iterations will
   * be executed.  When found a best iteration so far, then at least twice of
   * the current iteration count will be executed unless the iteration number
   * is greater than "max_iteration"
   * @param x
   * @param categories
   * @param min_iteration
   * @param max_iteration
   * @param verif_portion
   */
  public double trainWithVerification(double[][] x, double[][] categories,
                                    int min_iteration, int max_iteration, double verif_portion,
                                    double[][] test, double[][] test_cat) {

    System.out.println("NN_HIDDEN_VAR_NUM = " + nh);
    System.out.println("NN_RATE = " + rate);
    System.out.println("NN_ITERATIONS = " + min_iteration);
    System.out.println("MAX_NN_ITERATION = " + max_iteration);
    System.out.println("VERIFICATION_PORTION = " + verif_portion);

    // first normailize the data
    means = yi.stat.StatUtil.meanOnCol(x);
    devs = yi.stat.StatUtil.varianceOnCol(x);
    for (int i=0; i<d; i++) devs[i] = Math.sqrt(devs[i]);

    // normalized data saved in xx
    double[][] xx = new double[x.length][x[0].length];
    for (int i=0; i<x.length; i++) {
      for (int j = 0; j < d; j++) {
        if (devs[j] != 0) {
          xx[i][j] = (x[i][j] - means[j]) / devs[j];
        } else {
          xx[i][j] = 0;
        }
      }
    }

    int nsamples = categories.length; // number of training samples
    Random random = new Random();
    // pick up the samples randomly
    int[] rand_ord = yi.stat.StatUtil.randomize(random, nsamples);

    // for verify
    int verify_num = (int)(x.length * verif_portion);
    int train_num = x.length - verify_num;

    double[][] train = new double[train_num][];
    double[][] verify = new double[verify_num][];
    double[][] train_c = new double[train_num][];
    double[][] verify_c = new double[verify_num][];


    System.out.println("train samples " + train_num);
    for (int i=0; i<train_num; i++) {
      train[i] = xx[rand_ord[i]];
      train_c[i] = categories[rand_ord[i]];
    }


    System.out.println("verify samples " + verify_num);
    for (int i=0; i<verify_num; i++) {
      verify[i] = xx[rand_ord[i+train_num]];
      verify_c[i] = categories[rand_ord[i+train_num]];
    }

    double min_err = Double.MAX_VALUE;
    double[][] best_alpha = alpha;
    double[][] best_beta = beta;
    int best_iter = 0;

//    // roc preparation for train set
//    DoubleScoreHashMap train_scores = new DoubleScoreHashMap();
//    HashSet train_pos_set = new HashSet();
//    for (int i=0; i<train_num; i++) {
//      if (train_c[i][0] >= 0.5 ) train_pos_set.add(new Integer(i));
//    }
//    System.out.println("pos instance in train = " + train_pos_set.size());
//
//    // roc preparation for verify set
//    DoubleScoreHashMap verify_scores = new DoubleScoreHashMap();
//    HashSet verify_pos_set = new HashSet();
//    for (int i=0; i<verify_num; i++) {
//      if (verify_c[i][0] >= 0.5 ) verify_pos_set.add(new Integer(i));
//    }
//    System.out.println("pos instance in verify = " + verify_pos_set.size());
//
//    // roc preparation for test set
//    DoubleScoreHashMap test_scores = null;
//    HashSet test_pos_set = null;
//    if (test != null) {
//      test_scores = new DoubleScoreHashMap();
//      test_pos_set =new HashSet();
//      for (int i = 0; i < test_cat.length; i++) {
//        if (test_cat[i][0] >= 0.5)
//          test_pos_set.add(new Integer(i));
//      }
//      System.out.println("pos instance in test = " + test_pos_set.size());
//    }

    int stat_interval = VERIFY_INTERVAL;
//    // iteration statistical information:
//    // 0: iteration index
//    // 1: nn error (target vs. output)
//    // 2: train error (with weighted pos/neg instance)
//    // 3: verify error (with weighted pos/neg instance)
//    // 4: train roc
//    // 5: verfiy roc
//    // 6: test roc
//    double[][] iter_stat = new double[7][max_iteration/stat_interval + 1];
    // iteration statistical information:
    // 0: iteration index
    // 1: nn error (target vs. output)
    // 2: train error (with weighted pos/neg instance)
    // 3: verify error (with weighted pos/neg instance)
    // 4: test error
    double[][] iter_stat = new double[5][max_iteration/stat_interval + 1];

    for (int iter = 0; iter<=min_iteration; iter++) {
      double error = train(train, train_c);
      if (iter % stat_interval == 0) {
        int iter_index = iter/stat_interval;
//        NumberFormat nf = NumberFormat.getNumberInstance();
//        nf.setMaximumFractionDigits(3);
        iter_stat[0][iter_index] = iter;
        iter_stat[1][iter_index] = error;
        double train_error = test(train, train_c, null, true);
        iter_stat[2][iter_index] = train_error;
        double verify_error = test(verify, verify_c, null, true);
        iter_stat[3][iter_index] = verify_error;
        if (verify_error < min_err) { // this weight is better
          min_err = verify_error;
          best_alpha = yi.util.data.CloneUtil.clone(alpha);
          best_beta = yi.util.data.CloneUtil.clone(beta);
          best_iter = iter;
          min_iteration = Math.min(Math.max(min_iteration, iter * 2),
                                   max_iteration);
//          System.out.println("  set min iteration to " + min_iteration);
        }

//        // use ROC
//        for (int j = 0; j < train.length; j++) {
//          double[] scores = test(train[j], true);
//          train_scores.putScore(new Integer(j), scores[0]);
//        }
//        double[][] train_rpc = StatTestRst.rpcAnalysis(train_scores, train_pos_set, true);
//        double train_roc_score =  StatTestRst.getROCScore(  StatTestRst.rpc2roc(train_rpc));
//
//        // verify ROC
//        for (int j = 0; j < verify.length; j++) {
//          double[] scores = test(verify[j], true);
//          verify_scores.putScore(new Integer(j), scores[0]);
//        }
//        double[][] verify_rpc = StatTestRst.rpcAnalysis(verify_scores, verify_pos_set, true);
//        double verify_roc_score =  StatTestRst.getROCScore(StatTestRst.rpc2roc(verify_rpc));
//
//        iter_stat[4][iter_index] = train_roc_score;
//        iter_stat[5][iter_index] = verify_roc_score;
//
//        // test ROC
//        if (test_scores != null) {
//          for (int j = 0; j < test.length; j++) {
//            double[] scores = test(test[j], false);
//            test_scores.putScore(new Integer(j), scores[0]);
//          }
//          double[][] test_rpc = StatTestRst.rpcAnalysis(test_scores, test_pos_set, true);
//          double test_roc_score =  StatTestRst.getROCScore(  StatTestRst.rpc2roc(test_rpc));
//          iter_stat[6][iter_index] = test_roc_score;
//        }
        if (test != null) {
          double test_error = test(test, test_cat, null, false);
          iter_stat[4][iter_index] = test_error;
        }

      }

    }
    alpha = best_alpha;
    beta = best_beta;
    double verify_error = test(verify, verify_c, null, true);

    System.out.println();
    System.out.println("NN iteration stat:");
    OutputResult.print(System.out, iter_stat);
    System.out.println("best iteration = " + best_iter);
    System.out.println("verify error = " + verify_error +
                       ", train error = " + iter_stat[2][best_iter/stat_interval]);
//    System.out.println("verify error = " + verify_error +
//                       ", verify roc = " + iter_stat[5][best_iter/stat_interval] +
//                       ", verify roc = " + iter_stat[6][best_iter/stat_interval]);
    if (test != null) {
      double test_error = test(test, test_cat, null, false);
      System.out.println("verify-error\t" + verify_error +
                         "\ttest-error\t" + test_error);
    }
    m_verifyError = verify_error;
    return verify_error;
  }

  public double getVerifyError() {
    return m_verifyError;
  }

  public String toString() {
    StringWriter str = new StringWriter();
    PrintWriter out = new PrintWriter(str);

    out.println("alpha");

    // on the first layer
    for (int i=0; i<nh; i++) {
      for (int j=0; j<=d; j++) {
        out.print(alpha[j][i] + "\t");
      }
      out.println();
    }

    out.println("beta");
    // second layer
    for (int i=0; i<k; i++) {
      for (int j=0; j<=nh; j++) {
        out.print(beta[j][i] + "\t");
      }
      out.println();
    }
    return str.toString();

  }

}