package yi.stat.gm.bp;

import java.util.*;

/**
 * Belief Propagation
 */

public class BeliefPropagation {

  /**
   * The BP application
   */
  BPApplication m_bpApp;

  int[] m_nodeIDs;

  int[][] m_fromNeighborIDs;

  int[][] m_toNeighborIDs;

  int[][] m_indexInNeighbor;


  int m_valueNum;

  // neighborhood term
//  double[][] m_neighborhoodTerm;

  /**
   * Two message tables are used.  One for saving the messages in the last
   * iteration; another for in the current iteration.  They are switched after
   * each iteration
   */
  double[][][] m_messages1 = null, m_messages2 = null;
  double[][][] m_lastMessages = null, m_currentMessages = null;

  public BeliefPropagation(BPApplication app) {
    m_bpApp = app;

    m_nodeIDs = app.getNodeIDs();

    m_fromNeighborIDs = app.getFromNeighborIDs();
    m_toNeighborIDs = app.getToNeighborIDs();
    m_indexInNeighbor = setupIndexInNeighbor(m_nodeIDs, m_fromNeighborIDs, m_toNeighborIDs);

//    m_neighborhoodTerm = m_bpApp.getNeighborhoodTerm();

    m_valueNum = app.getPossibleValues();

    // initialize the messages
    System.out.println("construct message structures ...");
    m_messages1 = constructMessages(m_nodeIDs, m_fromNeighborIDs, m_toNeighborIDs);
    m_messages2 = constructMessages(m_nodeIDs, m_fromNeighborIDs, m_toNeighborIDs);
    System.out.println("construct message structures end.");

    m_lastMessages = m_messages1;
    m_currentMessages = m_messages2;

    // initialize messages sent to each node i from all its neighbor node j:
    // last_messages(j,i) is a probablity table.
    initializeMessages(m_lastMessages, 1.0);
  }

  /**
   * Belief Propagation.
   *
   * Pseudo-code:
   *
   * initialize messages sent to each node i from all its neighbor node j:
   * last_messages(j,i) is a probablity table.
   *
   * iteration n=0, 1, ...
   *   // construct messages send from node i to its neighbors (node j)
   *   for each node i
   *     for each neighbor j of node i (node j must be a variable)
   *       // calculate the possibility of each possbile value of node j
   *       for each possible value v of node j
   *         message(i,j)[v] = 0;
   *         // do intergal here (may change to max)
   *         for each possible value v_i of node i
   *           p = data_term(v_i) * neighborhood_term(v|vi);
   *           for each neighbor k of node i (node k can be either a variable node or a constant node)
   *             if k == j continue;
   *             p *= last_message(k, i)[v_i];
   *           end of k
   *           message(i,j)[v] += p;
   *         end of v_i;
   *       end of v;
   *       normalize message(i,j)[v] on v
   *     end of j
   *   end of i;
   *   // replace the old messages
   *   last_messages = messages;
   *   // check the break condition
   * end of n
   *
   * What needed in the algorithm:
   * - initial guess of messages (or initial guess of the value of each nodes,
   *   which can be guessed from only the data_term.  The sent out messages can
   *   only depend on the data_term and neighbor_term (a parametered
   *   distribution), but not the messages from neighbors.  This equals to set
   *   all the messages to 1)
   * - the order of the nodes (can be sorted by the distance from the mask, the
   *   confidence from the data_term and etc)
   * - the neighbors of each nodes (this can be inferred from the x,y position)
   * - the possible values of each nodes (pre-defined)
   * - data term and neighborhood term
   *
   */
  public void bp(int[] order, int iteration) {
    int node_num = m_nodeIDs.length;

    // iteration n=0, 1, ...
    for (int n = 1; n <= iteration; n++) {
      System.out.println("bp iteration " + n + " : message passing ..." );

//      if (n < 10) { // output current estimation of hidden variable for evaluation
//        // this violate the software engineering rule, but just for debuging
//        double[] mmse_node_alpha = getMMSEEstimate();
//        yi.graphics.mat.BPMatting app =  (yi.graphics.mat.BPMatting)m_bpApp;
//        app.demoAlpha(n, mmse_node_alpha);
//      }

      // construct messages send from node i to its neighbors (node j)
      // for each node i
      for (int i = 0; i < node_num; i++) {
        int node_i = m_nodeIDs[order[i]];

        // data term
        double[] data_term = m_bpApp.getDataTerm(node_i);
        int[] index_in_neighbor = m_indexInNeighbor[order[i]];

        // for each neighbor j of node i (node j must be a variable)
        int[] to_neighbors = m_toNeighborIDs[node_i];
        for (int j = 0; j < to_neighbors.length; j++) {
          int to_node = to_neighbors[j];

          double[] neighborhood_term = m_bpApp.getNeighborhoodTerm(node_i, to_node);

          // calculate the possibility of each possbile value of node j
          // for each possible value v of node j
          for (int v = 0; v < m_valueNum; v++) {

            // message(i,j)[v] = 0;
            double prob = 0;
            // do intergal here (may change to max)
            // for each possible value v_i of node i
            for (int v_i = 0; v_i < m_valueNum; v_i++) {
              //   p = data_term(v_i) * neighborhood_term(v|vi);
              // neighborhood_term(v|vi) = neighborhood_term_f(|v-v_i|);
              int value_diff = v_i - v;
              if (value_diff < 0) value_diff = - value_diff;
              double p = data_term[v_i] * neighborhood_term[value_diff];
              //   for each neighbor k of node i (node k can be either a variable node or a constant node)
              int[] from_neighbors = m_fromNeighborIDs[node_i];
              if (from_neighbors != null) {
                for (int k = 0; k < from_neighbors.length; k++) {
                  int from_node = from_neighbors[k];
                  // if k == j continue;
                  if (from_node == to_node)
                    continue;
                  // p *= last_message(k, i)[v_i];
                  p *= m_lastMessages[from_node][index_in_neighbor[k]][v_i];
                } //   end of k
              }
              //   message(i,j)[v] += p;
              prob += p;
            } // end of v_i;
            m_currentMessages[node_i][j][v] = prob;
            if (Double.isNaN(prob)) {
              System.err.println("message from " + node_i + " to " + to_node + " is NaN");
            }
          } //   end of v;
          //   normalize message(i,j)[v] on v
          normalizeMessages(m_currentMessages[node_i][j]);
          for (int k=0; k<m_currentMessages[node_i][j].length; k++) {
            if (Double.isNaN(m_currentMessages[node_i][j][k])) {
              System.err.println("message from " + node_i + " to " + to_node + " is NaN after normalization");
            }
          }
        } // end of j
      } //   end of i;

      // replace the old messages
      double[][][] tmp = m_lastMessages;
      m_lastMessages = m_currentMessages;
      m_currentMessages = tmp;
      //   last_messages = messages;
      //       // check the break condition
      //
    } // end of n

  }

  /**
   * Compute MMSE estimation of the nodes
   *
   * Pseudo-code:
   *
   * for each node i
   *   p[v_i] = data_term[v_i];
   *   for each possible value of v_i
   *     for each node j in from_neighbors
   *       p[v_i] *= last_messages[j][i][v_i];
   *     end of j
   *   end of v_i
   *   estimation[i] = expected(p[v_i])
   * end of i
   *
   * @return
   */
  public double[] getMMSEEstimate() {
    int node_num = m_nodeIDs.length;
    double[] mmse = new double[node_num];

      // for each node i
      for (int i = 0; i < node_num; i++) {
        int node_i = m_nodeIDs[i];

        if (node_i == 38688) {
          node_i = 38688;
        }

        // data term
        double[] data_term = m_bpApp.getDataTerm(node_i);
        int[] index_in_neighbor = m_indexInNeighbor[i];


        // message(i,j)[v] = 0;
        double[] prob = new double[m_valueNum];
        // do intergal here (may change to max)
        // for each possible value v_i of node i
        for (int v_i = 0; v_i < m_valueNum; v_i++) {
          //   p = data_term(v_i) * neighborhood_term(v|vi);
          double p = data_term[v_i];
          //   for each neighbor k of node i (node k can be either a variable node or a constant node)
          int[] from_neighbors = m_fromNeighborIDs[node_i];
          if (from_neighbors != null) {
            for (int k = 0; k < from_neighbors.length; k++) {
              int from_node = from_neighbors[k];
              // p *= last_message(k, i)[v_i];
              p *= m_lastMessages[from_node][index_in_neighbor[k]][v_i];
            } //   end of k
          }
          //   message(i,j)[v] += p;
          prob[v_i] = p;
        } // end of v_i;
        mmse[i] = getExpection(prob);
      } //   end of i;
      return mmse;

  }

  public int[] getMAPEstimate() {
    int node_num = m_nodeIDs.length;
    int[] mmse = new int[node_num];

      // for each node i
      for (int i = 0; i < node_num; i++) {
        int node_i = m_nodeIDs[i];

        // data term
        double[] data_term = m_bpApp.getDataTerm(node_i);
        int[] index_in_neighbor = m_indexInNeighbor[i];


        // message(i,j)[v] = 0;
        double[] prob = new double[m_valueNum];
        // do intergal here (may change to max)
        // for each possible value v_i of node i
        for (int v_i = 0; v_i < m_valueNum; v_i++) {
          //   p = data_term(v_i) * neighborhood_term(v|vi);
          double p = data_term[v_i];
          //   for each neighbor k of node i (node k can be either a variable node or a constant node)
          int[] from_neighbors = m_fromNeighborIDs[node_i];
          if (from_neighbors != null) {
            for (int k = 0; k < from_neighbors.length; k++) {
              int from_node = from_neighbors[k];
              // p *= last_message(k, i)[v_i];
              p *= m_lastMessages[from_node][index_in_neighbor[k]][v_i];
            } //   end of k
          }
          //   message(i,j)[v] += p;
          prob[v_i] = p;
        } // end of v_i;
        mmse[i] = getMAP(prob);
      } //   end of i;
      return mmse;

  }

  double getExpection(double[] prob) {
    double sumvalue = 0;
    double sump = 0;
    for (int i=0; i<prob.length; i++) {
      sumvalue += i * prob[i];
      sump += prob[i];
    }
    return sumvalue/sump;
  }

  int getMAP(double[] prob) {
    double max_p = prob[0];
    int max_id = 0;
    for (int i=1; i<prob.length; i++) {
      if (prob[i] > max_p) {
        max_p = prob[i];
        max_id = i;
      }
    }
    return max_id;
  }

  double[][][] constructMessages(int[] nodes, int[][] from, int[][] to) {

    int node_num = nodes.length;
    int max_node_id = yi.stat.StatUtil.maximum(nodes);
    double[][][] messages = new double[max_node_id + 1][][];

    int message_space = 0;
    int space_report = 1000000;

    for (int i = 0; i < node_num; i++) {
      int node_i = nodes[i];
      // for each variable neighbor j of node i
      int to_num = (to[node_i] == null) ? 0 : to[node_i].length;
      messages[node_i] = new double[to_num][m_valueNum];
      message_space += to_num * m_valueNum;
      if (message_space >= space_report) {
        System.out.println("message space: " + message_space);
        space_report += 1000000;
      }
    }
    return messages;
  }


  /**
   * Initialize all the messages to the specified probability
   * @param app
   */
  void initializeMessages(double[][][] messages, double initial_prob) {

    int message_num = messages.length;

    for (int i = 0; i < message_num; i++) {
      if (messages[i] == null) continue;
      // for each variable neighbor j of node i
      int neighbor_num = messages[i].length;

      for (int j = 0; j < neighbor_num; j++) {
        for (int v = 0; v < m_valueNum; v++) {
          messages[i][j][v] = initial_prob;
        }
      }
    }
  }

  void normalizeMessages(double[] probs) {
    double sum = 0;
    for (int i = 0; i<probs.length; i++) {
      sum += probs[i];
    }
    for (int i = 0; i<probs.length; i++) {
      probs[i] /= sum;
    }
  }

  /**
   * In node i, we have a neighbor list array, say
   * [0]: 1600, [1]:1699, [2]: 1701, [3]:1800.  In message passing, we have
   * messages from i to its neighbors: messages[i][0..3].  For node 1600 to
   * retrieval its message from i, it must know it should get the
   * messages[i][0], since we don't save the message as messages[i][1600], so
   * for node 1600, it must have a way to find out message[i][?] is the message
   * to it.  m_indexInNeighbor is for this purpose.  Suppose i is the j-neighbor
   * of of 1600, we have m_indexInNeighbor[1600][j] = 0 so that we know you
   * should look for messages[neighor[j]][0] to get its message
   */
  int[][] setupIndexInNeighbor(int[] nodes, int[][] from, int[][] to) {
    int  node_num = nodes.length;
    int[][] indexInNeighbor = new int[node_num][];

    // for node i
    for (int i = 0; i < node_num; i++) {
      int node_i = nodes[i];
      int[] from_ids = from[node_i];
      if (from_ids == null) continue;
      int from_num = from_ids.length;
      indexInNeighbor[i] = new int[from_num];

      // for each neighbor of node i
      for (int j=0; j<from_num; j++) {
        // look through the neighbor's neighbor's list
        int neighbor_id = from_ids[j];
        int[] neighbor_to_list = to[neighbor_id];
        indexInNeighbor[i][j] = -1;
        for (int k=0; k<neighbor_to_list.length; k++) {
          if (neighbor_to_list[k] == node_i) {
            indexInNeighbor[i][j] = k;
            break;
          }
        }
        if (indexInNeighbor[i][j] == -1) {
          System.err.println("didn't find node " + node_i + " in its neighbor " +
                             neighbor_id
                             + "'s neighbor list");
        }

      }
    }
    return indexInNeighbor;

  }
}