package jkind.solvers.yices;

import java.math.BigInteger;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

import jkind.lustre.NamedType;
import jkind.sexp.Sexp;
import jkind.sexp.Symbol;
import jkind.solvers.BoolValue;
import jkind.solvers.Eval;
import jkind.solvers.Model;
import jkind.solvers.NumericValue;
import jkind.solvers.Value;

public class YicesModel extends Model {
	private Map<String, String> aliases;
	private Map<String, Value> valueAssignments;
	private Map<String, Map<BigInteger, Value>> functionAssignments;

	public YicesModel() {
		this.aliases = new HashMap<>();
		this.valueAssignments = new HashMap<>();
		this.functionAssignments = new HashMap<>();
	}
	
	public void addAlias(String from, String to) {
		aliases.put(from, to);
	}
	
	public void addValue(String id, Value v) {
		valueAssignments.put(id, v);
	}
	
	public void addFunctionValue(String fn, BigInteger arg, Value v) {
		Map<BigInteger, Value> fnMap = functionAssignments.get(fn);
		if (fnMap == null) {
			fnMap = new HashMap<>();
			functionAssignments.put(fn, fnMap);
		}
		
		fnMap.put(arg, v);
	}
	
	private String getAlias(String id) {
		String result = id;
		while (aliases.containsKey(result)) {
			result = aliases.get(result);
		}
		return result;
	}
	
	@Override
	public Value getValue(Symbol sym) {
		return valueAssignments.get(getAlias(sym.toString()));
	}

	public Map<BigInteger, Value> getFunction(String fn) {
		return functionAssignments.get(getAlias(fn));
	}
	
	@Override
	public Value getFunctionValue(String fn, BigInteger index) {
		fn = getAlias(fn);
		if (functionAssignments.containsKey(fn)) {
			return functionAssignments.get(fn).get(index);
		} else if (definitions.containsKey(fn)) {
			Sexp s = definitions.get(fn).getLambda().instantiate(new Symbol(index.toString()));
			return new Eval(this).eval(s);
		} else {
			Value value = getDefaultValue(fn);
			addFunctionValue(fn, index, value);
			return value;
		}
	}
	
	private Value getDefaultValue(String fn) {
		if (declarations.get(fn).getType() == NamedType.BOOL) {
			return BoolValue.TRUE;
		} else {
			return new NumericValue("0");
		}
	}

	@Override
	public Set<String> getFunctions() {
		Set<String> fns = new HashSet<>(functionAssignments.keySet());
		for (String alias : aliases.keySet()) {
			if (getFunction(alias) != null) {
				fns.add(alias);
			}
		}
		return fns;
	}
}

@@ -717,26 +717,25 @@
 Map%3C%3E();%0A%09%7D%0A
-%09
 %0A
+
 %09public void
@@ -794,26 +794,25 @@
 om, to);%0A%09%7D%0A
-%09
 %0A
+
 %09public void
@@ -877,18 +877,17 @@
  v);%0A%09%7D%0A
-%09
 %0A
+
 %09public 
@@ -1106,11 +1106,9 @@
 %09%09%7D%0A
-%09%09
 %0A
+
 %09%09fn
@@ -1123,26 +1123,25 @@
 arg, v);%0A%09%7D%0A
-%09
 %0A
+
 %09private Str
@@ -1281,26 +1281,25 @@
  result;%0A%09%7D%0A
-%09
 %0A
+
 %09@Override%0A%09
@@ -1502,18 +1502,17 @@
 n));%0A%09%7D%0A
-%09
 %0A
+
 %09@Overri
@@ -1634,24 +1634,74 @@
 tainsKey(fn)
+ && functionAssignments.get(fn).containsKey(index)
 ) %7B%0A%09%09%09retur
@@ -2019,17 +2019,16 @@
 %0A%09%09%7D%0A%09%7D%0A
-%09
 %0A%09privat

Fix lookup of partial function values in YicesModel (bug discovered by Lucas Wagner)