import java.util.Map;
import java.util.HashMap;
import java.util.List;
import java.util.Iterator;
import java.util.ArrayList;

import edu.mit.six825.bn.functiontable.*;
import edu.mit.six825.bn.bayesnet.*;
import edu.mit.six825.bn.inputs.*;
import techniques.utils.*;



/**
 * Some utilities for implementing sampling-based solvers
 *
 * @author drkp
 */
public class SampleUtils {
    private static final boolean DEBUG = false;
    
    private SampleUtils() {

    }

    public static Map buildChildMap(BayesNetNodeSet bnset) {
        Map childMap = new HashMap();
        for (Iterator it = bnset.iterator(); it.hasNext();) {
            BayesNetNode node = (BayesNetNode) it.next();
            childMap.put(node, new ArrayList());
        }

        for (Iterator it = bnset.iterator(); it.hasNext();) {
            BayesNetNode parent = (BayesNetNode) it.next();
            for (Iterator it2 = bnset.iterator(); it2.hasNext();) {
                BayesNetNode child = (BayesNetNode) it2.next();
                if (child.parents.contains(parent)) {
                    List lst = (List) childMap.get(parent);
                    lst.add(child);
                }

            }
        }
        
        return childMap;
    }

    
    /**
     * Return the probability of each value in the node's variable's domain,
     * given an assignment that includes its parents values.
     */
    public static double [] evaluateNodeGivenParents(BayesNetNode node,
                                                     Assignment evidence) {
        double [] entries = new double[node.var.domain.size()];

        for (int i = 0; i < node.var.domain.size(); i++) {
            Comparable val = node.var.domain.getValue(i);
            Assignment fullAss = new Assignment(evidence, node.var, val);
            entries[i] = node.cpt.evaluate(fullAss);
        }

	return entries;
    }

    /**
     * Return the probability of each value in the node's variable's domain,
     * consistent with the current assignments of all of its children. This
     * requires that the assignment provided includes the full Markov blanket
     * (parents, children, and children's parents) of the given node.
     *
     * Implemented per formula (14.11) of AIMA.
     */
    public static double [] evaluateNodeGivenMarkovBlanket(BayesNetNode node,
                                                           Assignment as,
                                                           Map childMap) {

        Function cpt = node.cpt;
        double sum = 0;
        
	if (DEBUG) {
	    System.out.println("Evaluating node " + node.var);
            System.out.println("Assignment is " + as);
	}
	
	double [] entries = new double[node.var.domain.size()];

        for (int i = 0; i < node.var.domain.size(); i++) {
            Comparable val = node.var.domain.getValue(i);
            Assignment fullAss = new Assignment(as, node.var, val);
            if (DEBUG) {
                System.out.println(" Considering assignment " + val);
            }
            
	    entries[i] = cpt.evaluate(fullAss);
	    if (DEBUG) {
		System.out.println("  CPT evaluates to " + entries[i]);
	    }
            
	    // Multiply by the probabilities of each child
	    for (Iterator it2 = ((List)childMap.get(node)).iterator();
                 it2.hasNext();) {
		BayesNetNode child = (BayesNetNode) it2.next();
                
                double childProb = child.cpt.evaluate(fullAss);
                entries[i] *= childProb;
                if (DEBUG) {
                    System.out.println("  Child " + child.var +
                                       " gives prob " +
                                       childProb);
                }
            }
            sum += entries[i];
        }
        
        for (int i = 0; i < node.var.domain.size(); i++) {
            entries[i] /= sum;
            if (DEBUG) {
                System.out.println("Final: " + entries[i]);
            }
        }

        return entries;
    }
    
    /**
     * Select a value of a node's variable according to the probability
     * distribution specified by the given assignment.
     *
     * If useOnlyParents is specified, then the sampling is performed based
     * solely on the distribution given by the parents' values; otherwise it
     * is performed based on the full Markov blanket. So setting this false is
     * useful when generating an initial sample, and true useful when sampling
     * given a full assignment.
     */
    public static Comparable sampleNode(BayesNetNode node,
                                        Assignment evidence,
                                        boolean useOnlyParents,
                                        Map childMap) {
	double[] dist;
	if (useOnlyParents) {
	    dist = evaluateNodeGivenParents(node, evidence);
	} else {
	    dist = evaluateNodeGivenMarkovBlanket(node, evidence, childMap);
	}
        
	int r = Random.sampleCompleteProbDist(dist);
	return node.var.domain.getValue(r);
    }
}