package jkind.results.layout;

import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import jkind.lustre.Node;
import jkind.lustre.Program;
import jkind.util.Util;

/**
 * A layout for the inputs, ouputs, and locals of a Lustre node
 */
public class NodeLayout implements Layout {
	private final Map<String, String> map;
	
	private static final String INPUTS = "Inputs";
	private static final String OUTPUTS = "Outputs";
	private static final String LOCALS = "Locals";
	private static final String INLINED = "Inlined";
	private static final String[] CATEGORIES = { INPUTS, OUTPUTS, LOCALS, INLINED };

	public NodeLayout(Node node) {
		if (node == null) {
			throw new IllegalArgumentException("Unable to create layout for null node");
		}
		
		this.map = new HashMap<>();
		for (String input : Util.getIds(node.inputs)) {
			map.put(input, INPUTS);
		}
		for (String input : Util.getIds(node.outputs)) {
			map.put(input, OUTPUTS);
		}
		for (String local : Util.getIds(node.locals)) {
			if (local.contains("~")) {
				map.put(local, INLINED);
			} else {
				map.put(local, LOCALS);
			}
		}
	}
	
	public NodeLayout(Program program) {
		this(program.getMainNode());
	}

	@Override
	public List<String> getCategories() {
		return Arrays.asList(CATEGORIES);
	}

	@Override
	public String getCategory(String signal) {
		return map.get(signal);
	}
}

@@ -70,19 +70,19 @@
 til.Hash
-Map
+Set
 ;%0Aimport
@@ -115,19 +115,19 @@
 va.util.
-Map
+Set
 ;%0A%0Aimpor
@@ -235,16 +235,17 @@
 puts, ou
+t
 puts, an
@@ -322,50 +322,8 @@
 t %7B%0A
-%09private final Map%3CString, String%3E map;%0A%09%0A
 %09pri
@@ -597,16 +597,123 @@
 NED %7D;%0A%0A
+%09private final Set%3CString%3E inputs;%0A%09private final Set%3CString%3E outputs;%0A%09private final Set%3CString%3E locals;%0A%0A
 %09public 
@@ -846,18 +846,16 @@
 %09%09%7D%0A
-%09%09
 %0A%09%09this.
 map 
@@ -854,53 +854,27 @@
 his.
-map = new HashMap%3C%3E();%0A%09%09for (String input : 
+inputs = getPrefix(
 Util
@@ -898,288 +898,122 @@
 ts))
- %7B%0A%09%09%09map.put(input, INPUTS);%0A%09%09%7D%0A%09%09for (String input : Util.getIds(node.outputs)) %7B%0A%09%09%09map.put(input, OUTPUTS);%0A%09%09%7D%0A%09%09for (String local : Util.getIds(node.locals)) %7B%0A%09%09%09if (local.contains(%22~%22)) %7B%0A%09%09%09%09map.put(local, INLINED);%0A%09%09%09%7D else %7B%0A%09%09%09%09map.put(local, LOCALS);%0A%09%09%09%7D%0A%09%09%7D
+;%0A%09%09this.outputs = getPrefix(Util.getIds(node.outputs));%0A%09%09this.locals = getPrefix(Util.getIds(node.locals));
 %0A%09%7D%0A
-%09
 %0A
+
 %09pub
@@ -1232,30 +1232,575 @@
 %7B%0A%09%09
-return map.get(signal)
+String prefix = getPrefix(signal);%0A%09%09if (prefix.contains(%22~%22)) %7B%0A%09%09%09return INLINED;%0A%09%09%7D else if (inputs.contains(prefix)) %7B%0A%09%09%09return INPUTS;%0A%09%09%7D else if (outputs.contains(prefix)) %7B%0A%09%09%09return OUTPUTS;%0A%09%09%7D else if (locals.contains(prefix)) %7B%0A%09%09%09return LOCALS;%0A%09%09%7D else %7B%0A%09%09%09return null;%0A%09%09%7D%0A%09%7D%0A%0A%09private String getPrefix(String signal) %7B%0A%09%09return signal.split(%22%5C%5C.%7C%5C%5C%5B%22)%5B0%5D;%0A%09%7D%0A%0A%09private Set%3CString%3E getPrefix(List%3CString%3E signals) %7B%0A%09%09Set%3CString%3E prefixes = new HashSet%3C%3E();%0A%09%09for (String signal : signals) %7B%0A%09%09%09prefixes.add(getPrefix(signal));%0A%09%09%7D%0A%09%09return prefixes
 ;%0A%09%7D

Handle complex types in NodeLayout