package pl.edu.pw.mini.jk.sg.gteval;

import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.stream.Collectors;

import org.apache.commons.math3.util.Pair;

import com.google.common.collect.ImmutableList;

import pl.edu.pw.mini.jk.lib.CollectionHelpers;
import pl.edu.pw.mini.jk.sg.game.common.DefenderMove;
import pl.edu.pw.mini.jk.sg.mixeduct.attackerresponse.OptimalAttacker;
import pl.edu.pw.mini.jk.sg.opt.util.GamePayoff;

public final class PayoffUtils {

	public static <AM, DM extends DefenderMove> double getDefenderPayoff(GameMatrices<AM, DM, ?, ?, ?> gameMatrices,
			Map<ImmutableList<DM>, Double> dmp, Map<ImmutableList<AM>, Double> amp) {
		final double normCoeff = dmp.values().stream().mapToDouble(i -> i).sum();
		Map<ImmutableList<DM>, Double> dmp2 = dmp.entrySet().stream()
				.collect(Collectors.toMap(e -> e.getKey(), e -> e.getValue() / normCoeff));
	
		return getExpectedPayoff(gameMatrices, amp, dmp2, Side.DEFENDER);
	}

	public static <AM, DM extends DefenderMove> double getDefenderPayoff(GameMatrices<AM, DM, ?, ?, ?> gameMatrices,
			Map<ImmutableList<DM>, Double> dmp, List<Pair<ImmutableList<AM>, Double>> amp) {
		final double normCoeff = dmp.values().stream().mapToDouble(i -> i).sum();
		Map<ImmutableList<DM>, Double> dmp2 = dmp.entrySet().stream()
				.collect(Collectors.toMap(e -> e.getKey(), e -> e.getValue() / normCoeff));
	
		return getExpectedPayoff(gameMatrices, CollectionHelpers.pairsToMap(amp), dmp2, Side.DEFENDER);
	}
	
	public static <AM, DM extends DefenderMove> double getAttackerPayoff(GameMatrices<AM, DM, ?, ?, ?> gameMatrices,
			Map<ImmutableList<DM>, Double> dmp, List<Pair<ImmutableList<AM>, Double>> amp) {
		final double normCoeff = dmp.values().stream().mapToDouble(i -> i).sum();
		Map<ImmutableList<DM>, Double> dmp2 = dmp.entrySet().stream()
				.collect(Collectors.toMap(e -> e.getKey(), e -> e.getValue() / normCoeff));
	
		return getExpectedPayoff(gameMatrices, CollectionHelpers.pairsToMap(amp), dmp2, Side.ATTACKER);
	}

	public static <AM, DM extends DefenderMove> double getDefenderPayoff(GameMatrices<AM, DM, ?, ?, ?> gameMatrices,
			Map<ImmutableList<DM>, Double> dmp) {
		Map<ImmutableList<AM>, Double> amp = CollectionHelpers.pairsToMap(
				OptimalAttacker.solveAttackerOptimalStrategyPlus(gameMatrices, dmp));
		return getDefenderPayoff(gameMatrices, dmp, amp);
	}

	public static <DM extends DefenderMove, AM> double getExpectedPayoff(
			final GameMatrices<AM, DM, ?, ?, ?> gameMatrices, Map<ImmutableList<AM>, Double> amp,
			Map<ImmutableList<DM>, Double> dmp, Side side) {
		double exp = 0;
	
		Function<GamePayoff, Double> getPayoffFun = side.getPayoffExtractor();
	
		exp = amp.entrySet().stream().mapToDouble(ae -> ae.getValue() * dmp.entrySet().stream()
				.mapToDouble(
						de -> de.getValue() * getPayoffFun.apply(gameMatrices.getGamePayoff(de.getKey(), ae.getKey())))
				.sum()).sum();
	
		return exp / dmp.values().stream().mapToDouble(i -> i).sum();
	}

}
