package edu.umich.soar.hpp;

import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.TreeSet;

/**
 * Core class for HPP!
 * A Hierarchical Probabilistic Plan (HPP) consists of...
 * 1. Actions
 * 2. Action-control scheme
 * 3. Variables
 * 
 * Action selection conditions upon variable values.
 * Actions either occur in sequence (seq) or randomly (set), as
 * controlled by the control scheme (i.e. the plan "type").
 * Each action can have sub-plans (i.e. arbitrary recursive hierarchy).
 */
public class HPPlan {
	
	/**
	 * Action control scheme possibilities
	 */
	public enum PlanType { seq, set }	
	
	/**
	 * Supports read-only access to 
	 * 
	 * @return action-control scheme
	 */
	public PlanType GetType() {
		return myType;
	}	
	
	/**
	 * Supports read-only access to action list
	 * 
	 * @return iterator to current action list
	 */
	public Iterator<HPPAction> GetActions() {
		return myActions.iterator();
	}
	
	/**
	 * Supports read-only access to variable set
	 * 
	 * @return iterator to current variable set
	 */
	public Iterator<HPPVariable> GetVariables() {
		return myVariables.iterator();
	}	
	
	/**
	 * Supports appending an action to the plan
	 * 
	 * @param newAction action to be added to the plan
	 */
	public void AddAction(HPPAction newAction) {
		myActions.add(newAction);
	}
	
	/**
	 * Supports appending a variable to the plan
	 * 
	 * @param newVariable variable to be added to the plan
	 */
	public void AddVariable(HPPVariable newVariable) {
		myVariables.add(newVariable);
	}
	
	/**
	 * Basic constructor - requires type and initializes empty variables/actions
	 * 
	 * @param type action-control scheme
	 */
	public HPPlan(PlanType type) {
		myType = type;
		myActions = new LinkedList<HPPAction>();
		myVariables = new TreeSet<HPPVariable>();
	}
	
	/**
	 * Constructor to support custom initialization of variables/actions via arrays
	 * 
	 * @param type action-control scheme
	 * @param variables initial variable list (supports null to ignore)
	 * @param actions initial action list (supports null to ignore)
	 */
	public HPPlan(PlanType type, HPPVariable[] variables, HPPAction[] actions) {
		this(type);
		
		if (variables != null) {
			for (HPPVariable v : variables) {
				AddVariable(v);
			}
		}
		
		if (actions != null) {
			for (HPPAction a : actions) {
				AddAction(a);
			}
		}
	}
	
	/**
	 * Utility function for plan printing (basically string concatenation with newlines)
	 * 
	 * @param before string1
	 * @param after string2
	 * @return string1 + string2 + [OS-independent] line separator
	 */
	private String StrAppend(String before, String after) {
		return StrAppend( before, after, System.getProperty("line.separator") );
	}
	
	/**
	 * Utility function for plan printing (string concatenation)
	 * 
	 * @param before string1
	 * @param after string2
	 * @param separator string separation
	 * @return string1 + string2 + separator
	 */
	private String StrAppend(String before, String after, String separator) {
		return ( before + after + separator );
	}
	
	
	/**
	 * Debugging function for inspecting plan contents
	 * 
	 * @param prefix string to indicate plan depth via recursion
	 * @return plan contents
	 */
	private String PrintPlan(String prefix) {
		String ret = "";
		
		ret = StrAppend( ret + prefix, "Type: " + GetType());
		ret = StrAppend( ret + prefix, "" );
		ret = StrAppend( ret + prefix, "==Variables==");
		Iterator<HPPVariable> vit = GetVariables();
		while (vit.hasNext()) {
			HPPVariable v = (HPPVariable) vit.next();
			
			ret = StrAppend( ret + prefix, "Variable: " + v.getName() + " (" + v.getType() + ")");
			
			if (v instanceof HPPGivenVariable) {
				ret = StrAppend( ret + prefix, "-- value=" + ((HPPGivenVariable) v).getValue());
			} else if (v instanceof HPPMultiVariable) {
				ret = StrAppend( ret + prefix, "-- default=" + ((HPPMultiVariable) v).getDefaultValue());
				ret = StrAppend( ret + prefix, "-- cpt");
				
				Iterator<HPPCPTEntry> it = ((HPPMultiVariable) v).getCPT();
				while (it.hasNext()) {
					HPPCPTEntry e = (HPPCPTEntry) it.next();
					
					String conditions = "";
					Iterator<HPPConstraint> cit = e.getConditions();
					while (cit.hasNext()) {
						HPPConstraint c = (HPPConstraint) cit.next();
						if (conditions.isEmpty()) {
							conditions = c.toString();
						} else {
							conditions = ( conditions + ", " + c );
						}
					}
					
					String values = "";
					Iterator<HPPCPTValue> vsit = e.getValues();
					while (vsit.hasNext()) {
						HPPCPTValue vs = (HPPCPTValue) vsit.next();
						if (values.isEmpty()) {
							values = vs.toString();
						} else {
							values = ( values + ", " + vs );
						}
					}
					
					ret = StrAppend( ret + prefix, "--- (" + conditions + ") -> (" + values + ")");
				}
			}
		}
		
		ret = StrAppend( ret + prefix, "==Actions==");		
		Iterator<HPPAction> it = GetActions();				
		while (it.hasNext()) {		
			HPPAction a = (HPPAction) it.next();
			
			ret = StrAppend( ret + prefix, "Action: " + a.getName());
			Iterator<HPPConstraint> cit = a.getConditions();
			while (cit.hasNext()) {
				HPPConstraint c = (HPPConstraint) cit.next();
				ret = StrAppend( ret + prefix, "--condition: " + c);
			}
			
			ret = StrAppend( ret + prefix, "--subplan:" + ( (a.getSubPlan()==null)?(" null"):("") ) );
			if (a.getSubPlan()!=null) {
				ret = StrAppend( ret + prefix, "{" );
				ret = StrAppend( ret + prefix, a.getSubPlan().PrintPlan( prefix + prefix ) );
				ret = StrAppend( ret + prefix, "}" );
			}
		}
		
		return ret;
	}
	
	
	/* (non-Javadoc)
	 * @see java.lang.Object#toString()
	 */
	public String toString() {
		return PrintPlan(" ");
	}
	
	/**
	 * Utility function for plan visualization that maintains a set of variable name-node mappings
	 * 
	 * @param variables name-node mapping
	 * @param var name
	 * @return variable node name
	 */
	private String makeVariable(Map<String, String> variables, String var) {		
		if (variables.containsKey(var)) {
			return variables.get(var);
		} else {
			String ret = ( "V" + variables.size() );
			variables.put(var, ret);
			
			return ret;
		}
	}
	
	
	/**
	 * Utility function for plan visualization that maintains a set of action name-node mappings
	 * 
	 * @param actions name-node mapping
	 * @param act name
	 * @return action node name
	 */
	private String makeAction(Map<String, String> actions, String act) {		
		if (actions.containsKey(act)) {
			return actions.get(act);
		} else {
			String ret = ( "A" + actions.size() );
			actions.put(act, ret);
			
			return ret;
		}
	}
	
	
	/**
	 * Utility class for plan visualization.
	 * Maintains two pieces of data needed for recursive visualization.
	 */
	private class gvRet {
		public String GetString() {
			return myString;
		}
		
		public List<String> GetLastNode() {
			return myLastNode;
		}
		
		public gvRet(String string, List<String> lastNode) {
			myString = string;
			myLastNode = lastNode;
		}
		
		private final String myString;
		private final List<String> myLastNode;
	}
	
	/**
	 * Main function to support recursive visualization of plans
	 * 
	 * @param variables name-node mapping
	 * @param actions name-node mapping
	 * @param cluster monotonically increasing counter for visual clustering of sub-plans
	 * @param parent node from which to connect first node(s) of current plan
	 * @return data to visualize the sub-plan
	 */
	private gvRet planGraphviz(Map<String, String> variables, Map<String, String> actions, Integer cluster, String parent) {
		String ret = "";
		LinkedList<String> retNodes = new LinkedList<String>();
		
		ret = StrAppend( ret, "subgraph cluster_" + ( ++cluster ) + " {" );
		
		ret = StrAppend( ret, "subgraph cluster_" + cluster + "V {" );
		ret = StrAppend( ret, "style=filled;" );
		ret = StrAppend( ret, "color=lightgrey;" );
		ret = StrAppend( ret, "node [ style=filled, color=white ];" );
		{
			Map<String, Set<String>> varDependencies = new HashMap<String, Set<String>>();
			
			Iterator<HPPVariable> vit = GetVariables();		
			while (vit.hasNext()) {
				HPPVariable v = (HPPVariable) vit.next();			
				
				if (v instanceof HPPGivenVariable) {
					ret = StrAppend( ret, makeVariable(variables, v.getName()) + " [ shape=\"note\"; label=\"" + v.getName() + " = " + ((HPPGivenVariable) v).getValue() + "\"; ];" );
				} else if (v instanceof HPPMultiVariable) {
					HPPMultiVariable mv = (HPPMultiVariable) v;
					
					if (!varDependencies.containsKey(mv.getName())) {
						varDependencies.put(mv.getName(), new TreeSet<String>());
					}
					
					String label = ( mv.getName() );
					if (mv.getDefaultValue()!=null) {
						label = ( label + " | default = " + mv.getDefaultValue() );
					}
					
					Iterator<HPPCPTEntry> it = ((HPPMultiVariable) v).getCPT();
					while (it.hasNext()) {
						HPPCPTEntry e = (HPPCPTEntry) it.next();
						
						String conditions = "";
						Iterator<HPPConstraint> cit = e.getConditions();
						while (cit.hasNext()) {
							HPPConstraint c = (HPPConstraint) cit.next();
							if (conditions.isEmpty()) {
								conditions = c.toString();
							} else {
								conditions = ( conditions + ", " + c );
							}
							
							varDependencies.get(mv.getName()).add(c.getName());
						}
						
						Iterator<HPPCPTValue> vsit = e.getValues();
						while (vsit.hasNext()) {
							HPPCPTValue vs = (HPPCPTValue) vsit.next();
							
							label = ( label + " | P( " + vs.getValue() + ( (conditions.isEmpty())?(""):(" \\| " + conditions) ) + " ) = " + vs.getProbability() );
						}
					}
					
					ret = StrAppend( ret, makeVariable(variables, mv.getName()) + " [ shape=\"record\", label=\"" + label + "\"; ];");
				}
			}
			
			for (Entry<String, Set<String>> e : varDependencies.entrySet() ) {
				for (String e2 : e.getValue()) {
					ret = StrAppend( ret, makeVariable(variables, e2) + " -> " + makeVariable(variables, e.getKey()) + " [style=\"dotted\"];" );
				}
			}
		}
		ret = StrAppend( ret, "}" );
		
		LinkedList<String> last = new LinkedList<String>();
		last.add(parent);
		
		Iterator<HPPAction> it = GetActions();				
		while (it.hasNext()) {		
			HPPAction a = (HPPAction) it.next();
			String thisAction = makeAction(actions, a.getName());
			
			ret = StrAppend( ret, thisAction + " [ label=\"" + a.getName() + "\" ];" );
			
			for (String l : last) {
				ret = StrAppend( ret, l + " -> " + thisAction + ";" );
			}
			
			Iterator<HPPConstraint> cit = a.getConditions();
			while (cit.hasNext()) {
				HPPConstraint c = (HPPConstraint) cit.next();
				
				ret = StrAppend( ret, makeVariable(variables, c.getName()) + " -> " + thisAction + " [ label=\"" + c.getValue() + "\", style=\"dotted\" ];" );
			}
			
			if (a.getSubPlan()==null) {
				if (GetType()==PlanType.seq) {
					last.clear();
					last.add(thisAction);
				} else {
					retNodes.add(thisAction);
				}
			} else {
				 gvRet subRet = a.getSubPlan().planGraphviz(variables, actions, cluster, thisAction);
				 
				 ret = ret + subRet.GetString();
				 
				 last.clear();
				 for (String l : subRet.GetLastNode()) {
					 last.add(l);
					 if (GetType()==PlanType.set) {
						 retNodes.add(l);
					 }
				 }
			}
		}
		
		if (GetType()==PlanType.seq) {
			for (String l : last) {
				retNodes.add(l);
			}
		}
		
		ret = StrAppend( ret, "}" );
		
		return new gvRet(ret, retNodes);
	}
	
	/**
	 * Recursively visualizes a plan
	 * 
	 * @return plan visualized in Graphviz format
	 */
	public String toGraphviz() {
		String ret = "";
		
		ret = StrAppend(ret, "digraph G {");
		ret = StrAppend(ret, "graph [ rankdir = \"LR\" ];");
		ret = StrAppend(ret, "start [shape=Mcircle];");
		ret = StrAppend(ret, "end [shape=doublecircle];");
		
		gvRet contents = planGraphviz(new HashMap<String, String>(), new HashMap<String, String>(), new Integer(0), "start");
		ret = ret + contents.GetString();
		for (String s : contents.GetLastNode()) {
			ret = StrAppend( ret, s + " -> end;" );
		}
		
		ret = StrAppend(ret, "}");
		
		return ret;
	}
	
	private final PlanType myType;
	private final List<HPPAction> myActions;
	private final Set<HPPVariable> myVariables;
}
