code: package edu.berkeley.nlp.syntax;
  
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.HashMap;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.Set;
  
  import edu.berkeley.nlp.util.CollectionUtils;
  import edu.berkeley.nlp.util.Counter;
  import edu.berkeley.nlp.util.Factory;
  
  /**
   * Assumes the type V is hashable
   * 
   * @author adampauls
   * 
   * @param <V>
   */
  public class UnaryClosureComputer<V> {
  
  	public static class Edge<V> {
  
  		@Override
  		public int hashCode() {
  			final int prime = 31;
  			int result = 1;
  
  			result = prime * result + ((child == null) ? 0 : child.hashCode());
  			result = prime * result
  					+ ((parent == null) ? 0 : parent.hashCode());
  			return result;
  		}
  
  		@Override
  		public boolean equals(Object obj) {
  			if (this == obj)
  				return true;
  			if (obj == null)
  				return false;
  			if (getClass() != obj.getClass())
  				return false;
  			Edge other = (Edge) obj;
  
  			if (child == null) {
  				if (other.child != null)
  					return false;
  			} else if (!child.equals(other.child))
  				return false;
  			if (parent == null) {
  				if (other.parent != null)
  					return false;
  			} else if (!parent.equals(other.parent))
  				return false;
  			return true;
  		}
  
  		public void setParent(V parent) {
  			this.parent = parent;
  		}
  
  		public void setChild(V child) {
  			this.child = child;
  		}
  
  		private V parent;
  
  		private V child;
  
  		private double score;
  
  		private Edge(V parent, V child) {
  			this.parent = parent;
  			this.child = child;
  		}
  
  		public V getParent() {
  			return parent;
  		}
  
  		public V getChild() {
  			return child;
  		}
  
  		public double getScore() {
  			return score;
  		}
  
  		public void setScore(double d) {
  			score = d;
  		}
  
  	}
  
  	private Factory<Edge> unaryRuleFactory = new Factory<Edge>() {
  
  		public Edge newInstance(Object... args) {
  			return new Edge(args[0], args[1]);
  		}
  	};
  
  	Map<V, List<Edge<V>>> closedUnaryRulesByChild = new HashMap<V, List<Edge<V>>>();
  
  	Map<V, List<Edge<V>>> closedUnaryRulesByParent = new HashMap<V, List<Edge<V>>>();
  
  	Map<Edge<V>, List<V>> pathMap = new HashMap<Edge<V>, List<V>>();
  
  	Set<Edge<V>> unaryRules = new HashSet<Edge<V>>();
  
  	private boolean sumInsteadOfMultipy;
  
  	/**
  	 * First is parent, second is child;
  	 * 
  	 * @return
  	 */
  	public Map<V, List<Edge<V>>> getAllClosedRulesByChildren() {
  		return closedUnaryRulesByChild;
  	}
  
  	public List<Edge<V>> getClosedUnaryRulesByChild(V child) {
  		return CollectionUtils.getValueList(closedUnaryRulesByChild, child);
  	}
  
  	public List<Edge<V>> getClosedUnaryRulesByParent(V parent) {
  		return CollectionUtils.getValueList(closedUnaryRulesByParent, parent);
  	}
  
  	public List<V> getPath(Edge unaryRule) {
  		return pathMap.get(unaryRule);
  	}
  
  	@Override
  	public String toString() {
  		StringBuilder sb = new StringBuilder();
  		for (V parent : closedUnaryRulesByParent.keySet()) {
  			for (Edge unaryRule : getClosedUnaryRulesByParent(parent)) {
  				List<V> path = getPath(unaryRule);
  				// if (path.size() == 2) continue;
  				sb.append(unaryRule);
  				sb.append("  ");
  				sb.append(path);
  				sb.append("\n");
  			}
  		}
  		return sb.toString();
  	}
  
  	public UnaryClosureComputer(boolean sumInsteadOfMultiply) {
  		this.sumInsteadOfMultipy = sumInsteadOfMultiply;
  	}
  
  	public void add(V parent, V child, double score) {
  		final Edge edge = new Edge(parent, child);
  		edge.setScore(score);
  		unaryRules.add(edge);
  	}
  
  	public void solve() {
  		Map<Edge<V>, List<V>> closureMap = computeUnaryClosure(unaryRules);
  		for (Edge<V> unaryRule : closureMap.keySet()) {
  			addUnary(unaryRule, closureMap.get(unaryRule));
  		}
  	}
  
  	private void addUnary(Edge<V> unaryRule, List<V> path) {
  		CollectionUtils.addToValueList(closedUnaryRulesByChild,
  				unaryRule.getChild(), unaryRule);
  		CollectionUtils.addToValueList(closedUnaryRulesByParent,
  				unaryRule.getParent(), unaryRule);
  		pathMap.put(unaryRule, path);
  	}
  
  	private Map<Edge<V>, List<V>> computeUnaryClosure(
  			Collection<Edge<V>> unaryRules) {
  
  		Map<Edge<V>, V> intermediateStates = new HashMap<Edge<V>, V>();
  		Counter<Edge<V>> pathCosts = new Counter<Edge<V>>();
  		Map<V, List<Edge<V>>> closedUnaryRulesByChild = new HashMap<V, List<Edge<V>>>();
  		Map<V, List<Edge<V>>> closedUnaryRulesByParent = new HashMap<V, List<Edge<V>>>();
  
  		Set<V> states = new HashSet<V>();
  
  		for (Edge<V> unaryRule : unaryRules) {
  			relax(pathCosts, intermediateStates, closedUnaryRulesByChild,
  					closedUnaryRulesByParent, unaryRule, null,
  					unaryRule.getScore());
  			states.add(unaryRule.getParent());
  			states.add(unaryRule.getChild());
  		}
  
  		for (V intermediateState : states) {
  			List<Edge<V>> incomingRules = closedUnaryRulesByChild
  					.get(intermediateState);
  			List<Edge<V>> outgoingRules = closedUnaryRulesByParent
  					.get(intermediateState);
  			if (incomingRules == null || outgoingRules == null)
  				continue;
  			for (Edge<V> incomingRule : incomingRules) {
  				for (Edge<V> outgoingRule : outgoingRules) {
  					Edge<V> rule = unaryRuleFactory.newInstance(
  							incomingRule.getParent(), outgoingRule.getChild());
  					double newScore = combinePathCosts(pathCosts, incomingRule,
  							outgoingRule);
  					relax(pathCosts, intermediateStates,
  							closedUnaryRulesByChild, closedUnaryRulesByParent,
  							rule, intermediateState, newScore);
  				}
  			}
  		}
  
  		for (V state : states)
  
  		{
  			Edge<V> selfLoopRule = unaryRuleFactory.newInstance(state, state);
  			relax(pathCosts, intermediateStates, closedUnaryRulesByChild,
  					closedUnaryRulesByParent, selfLoopRule, null, 0.0);
  		}
  
  		Map<Edge<V>, List<V>> closureMap = new HashMap<Edge<V>, List<V>>();
  
  		for (Edge<V> unaryRule : pathCosts.keySet()) {
  			unaryRule.setScore(pathCosts.getCount(unaryRule));
  			List<V> path = extractPath(unaryRule, intermediateStates);
  			closureMap.put(unaryRule, path);
  		}
  
  		return closureMap;
  
  	}
  
  	/**
  	 * @param pathCosts
  	 * @param incomingRule
  	 * @param outgoingRule
  	 * @return
  	 */
  	private double combinePathCosts(Counter<Edge<V>> pathCosts,
  			Edge<V> incomingRule, Edge<V> outgoingRule) {
  		return this.sumInsteadOfMultipy ? (pathCosts.getCount(incomingRule) + pathCosts
  				.getCount(outgoingRule))
  				: (pathCosts.getCount(incomingRule) * pathCosts
  						.getCount(outgoingRule));
  	}
  
  	private List<V> extractPath(Edge<V> unaryRule,
  			Map<Edge<V>, V> intermediateStates) {
  		List<V> path = new ArrayList<V>();
  		path.add(unaryRule.getParent());
  		V intermediateState = intermediateStates.get(unaryRule);
  		if (intermediateState != null) {
  			List<V> parentPath = extractPath(unaryRuleFactory.newInstance(
  					unaryRule.getParent(), intermediateState),
  					intermediateStates);
  			for (int i = 1; i < parentPath.size() - 1; i++) {
  				V state = parentPath.get(i);
  				path.add(state);
  			}
  			path.add(intermediateState);
  			List<V> childPath = extractPath(
  					unaryRuleFactory.newInstance(intermediateState,
  							unaryRule.getChild()), intermediateStates);
  			for (int i = 1; i < childPath.size() - 1; i++) {
  				V state = childPath.get(i);
  				path.add(state);
  			}
  		}
  		if (path.size() == 1 && unaryRule.getParent() == unaryRule.getChild())
  			return path;
  		path.add(unaryRule.getChild());
  		return path;
  	}
  
  	private void relax(Counter<Edge<V>> pathCosts,
  			Map<Edge<V>, V> intermediateStates,
  			Map<V, List<Edge<V>>> closedUnaryRulesByChild,
  			Map<V, List<Edge<V>>> closedUnaryRulesByParent, Edge<V> unaryRule,
  			V intermediateState, double newScore) {
  		if (intermediateState != null
  				&& (intermediateState.equals(unaryRule.getParent()) || intermediateState
  						.equals(unaryRule.getChild())))
  			return;
  		boolean isNewRule = !pathCosts.containsKey(unaryRule);
  		double oldScore = (isNewRule ? Double.NEGATIVE_INFINITY : pathCosts
  				.getCount(unaryRule));
  		if (oldScore > newScore)
  			return;
  		if (isNewRule) {
  			CollectionUtils.addToValueList(closedUnaryRulesByChild,
  					unaryRule.getChild(), unaryRule);
  			CollectionUtils.addToValueList(closedUnaryRulesByParent,
  					unaryRule.getParent(), unaryRule);
  		}
  		pathCosts.setCount(unaryRule, newScore);
  		intermediateStates.put(unaryRule, intermediateState);
  	}
  
  	public double getProb(V parent, V child) {
  		if (parent == child)
  			return 0.0;
  		final List<Edge<V>> byParent = closedUnaryRulesByParent.get(parent);
  		if (byParent == null)
  			return Double.POSITIVE_INFINITY;
  		int childIndex = byParent.indexOf(unaryRuleFactory.newInstance(parent,
  				child));
  		if (childIndex < 0)
  			return Double.POSITIVE_INFINITY;
  		final Edge<V> unaryRule = byParent.get(childIndex);
  
  		return unaryRule.getScore();
  	}
  
  }
  