package org.albaross.agents4j.learning;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Locale;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;

/**
 * 
 * @author Manuel Barbi
 *
 * @param <S> state
 * @param <A> action
 */
public class QTable<S, A> implements ValueFunction<S, A> {

	protected Map<S, Map<A, Double>> map;
	protected int size;

	public QTable() {
		this.map = new HashMap<>();
		this.size = 0;
	}

	public int size() {
		return size;
	}

	public boolean isEmpty() {
		return size() == 0;
	}

	@Override
	public void update(S state, A action, double reward, S nextState) {
		update(state, action, reward, nextState, 0.1, 0.9);
	}

	/**
	 * Calculates and updates the weight of an executed action.
	 * 
	 * @param state
	 * @param action
	 * @param reward
	 * @param nextState
	 * @param learningRate
	 * @param discountRate
	 */
	public void update(S state, A action, double reward, S nextState, double learningRate, double discountRate) {
		double old = get(state, action);
		double max = getMaxWeight(nextState);
		double weight = (1 - learningRate) * old + learningRate * (reward + discountRate * max);
		put(state, action, weight);
	}

	public boolean containsKey(S state, A action) {
		Map<A, Double> tmp = map.get(state);
		return (tmp != null) ? tmp.containsKey(action) : false;
	}

	public double get(S state, A action) {
		Map<A, Double> tmp = map.get(state);
		if (tmp == null)
			return 0;

		Double weight = tmp.get(action);
		if (weight == null)
			return 0;

		return weight;
	}

	public void put(S state, A action, double weight) {
		Map<A, Double> tmp = map.get(state);

		if (tmp == null)
			map.put(state, (tmp = new HashMap<>()));

		if (!tmp.containsKey(action))
			size++;

		tmp.put(action, weight);
	}

	public void remove(S state, A action) {
		Map<A, Double> tmp = map.get(state);
		if (tmp != null) {
			tmp.remove(action);
			size--;

			if (tmp.isEmpty())
				map.remove(state);
		}
	}

	public void clear() {
		map.clear();
	}

	public Collection<Double> values(S state) {
		Map<A, Double> tmp = map.get(state);
		return (tmp != null) ? tmp.values() : new ArrayList<>();
	}

	public Set<Entry<A, Double>> entrySet(S state) {
		Map<A, Double> tmp = map.get(state);
		return (tmp != null) ? tmp.entrySet() : new HashSet<>();
	}

	public double getMaxWeight(S state) {
		if (!map.containsKey(state))
			return 0;

		double maxWeight = -Double.MAX_VALUE;
		for (double weight : values(state)) {
			if (weight > maxWeight)
				maxWeight = weight;
		}

		return maxWeight;
	}

	public A getBestAction(S state) {
		if (!map.containsKey(state))
			return null;

		double maxWeight = -Double.MAX_VALUE;
		A bestAction = null;

		for (Entry<A, Double> e : entrySet(state)) {
			if (e.getValue() > maxWeight) {
				bestAction = e.getKey();
				maxWeight = e.getValue();
			}
		}

		return bestAction;
	}

	@Override
	public String toString() {
		StringBuilder sb = new StringBuilder();
		sb.append('{');
		String state;
		int i = 0;

		for (Entry<S, Map<A, Double>> s : map.entrySet()) {
			state = s.getKey().toString();

			for (Entry<A, Double> a : s.getValue().entrySet()) {
				if (i > 0)
					sb.append(',').append(' ');

				sb.append(state);
				sb.append(a.getKey());
				sb.append('=');
				sb.append(String.format(Locale.ENGLISH, "%.2f", a.getValue()));

				i++;
			}
		}

		sb.append('}');
		return sb.toString();
	}

}

@@ -91,16 +91,46 @@
 ection;%0A
+import java.util.Collections;%0A
 import j
@@ -164,22 +164,19 @@
 va.util.
-HashSe
+Lis
 t;%0Aimpor
@@ -245,16 +245,41 @@
 .Entry;%0A
+import java.util.Random;%0A
 import j
@@ -433,16 +433,54 @@
 , A%3E %7B%0A%0A
+%09protected Random rnd = new Random();%0A
 %09protect
@@ -2152,32 +2152,85 @@
 ap.clear();%0A%09%7D%0A%0A
+%09public Set%3CS%3E states() %7B%0A%09%09return map.keySet();%0A%09%7D%0A%0A
 %09public Collecti
@@ -2341,23 +2341,29 @@
 ) : 
-new Arra
+Collections.empt
 yList
-%3C%3E
 ();%0A
@@ -2501,21 +2501,28 @@
 ) : 
-new Hash
+Collections.empty
 Set
-%3C%3E
 ();%0A
@@ -2877,30 +2877,38 @@
 t = 
--Double.MAX_VALUE;%0A%09%09A
+getMaxWeight(state);%0A%09%09List%3CA%3E
  bes
@@ -2918,15 +2918,29 @@
 tion
+s
  = n
-ull
+ew ArrayList%3C%3E()
 ;%0A%0A%09
@@ -3005,17 +3005,18 @@
 Value() 
-%3E
+==
  maxWeig
@@ -3018,18 +3018,16 @@
 xWeight)
- %7B
 %0A%09%09%09%09bes
@@ -3033,19 +3033,22 @@
 stAction
- = 
+s.add(
 e.getKey
@@ -3053,49 +3053,78 @@
 ey()
+)
 ;%0A%09%09
-%09%09maxWeight = e.getValue(
+%7D%0A%0A%09%09if(bestActions.size() == 1)%0A%09%09%09return bestActions.get(0
 );%0A%09%09
-%09
-%7D%0A%09%09%7D%0A
 %0A%09%09r
@@ -3139,16 +3139,54 @@
 stAction
+s.get(rnd.nextInt(bestActions.size()))
 ;%0A%09%7D%0A%0A%09@

return random action at multiple best actions