package de.tu_dresden.lat.proofs.tools;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import org.semanticweb.owlapi.model.OWLAxiom;

import de.tu_dresden.lat.proofs.data.Proof;
import de.tu_dresden.lat.proofs.interfaces.IInference;
import de.tu_dresden.lat.proofs.interfaces.IProof;

public class MinimalTreeProofExtractor {

	public static IProof makeUnique(IProof proof) {
		Map<OWLAxiom, Integer> distances = new HashMap<>();
		Map<OWLAxiom, IInference> bestInferences = new HashMap<>();

		for (IInference inf : proof.getInferences().stream().filter(inf -> inf.getPremises().isEmpty())
				.collect(Collectors.toList())) {
			OWLAxiom ax = inf.getConclusion();
			distances.put(ax, 1);
			bestInferences.put(ax, inf);
		}

		bestInferences = dijkstra(proof, proof.getFinalConclusion(), distances, bestInferences);

		List<IInference> treeProof = new ArrayList<>();
		collectInferences(treeProof, proof.getFinalConclusion(), bestInferences);

		return new Proof(proof.getFinalConclusion(), treeProof);
	}

	private static Map<OWLAxiom, IInference> dijkstra(IProof proof, OWLAxiom conclusion,
			Map<OWLAxiom, Integer> distances, Map<OWLAxiom, IInference> bestInferences) {
		Map<OWLAxiom, Integer> finalDistances = new HashMap<>();

		while (!distances.isEmpty()) {
			OWLAxiom current = getMinimal(distances);
			if (current.equals(conclusion)) {
				break;
			}
			finalDistances.put(current, distances.get(current));
			distances.remove(current);
			Stream<IInference> nextInferences = getNextInferences(proof, current, distances, finalDistances);

			nextInferences.forEach(inf -> {

				Integer newDistance = 1 + inf.getPremises().stream().mapToInt(ax -> {
					if (distances.containsKey(ax)) {
						return distances.get(ax);
					} else {
						return finalDistances.get(ax);
					}
				}).sum();

				OWLAxiom ax = inf.getConclusion();
				if ((!distances.containsKey(ax)) || (newDistance < distances.get(ax))) {
					distances.put(ax, newDistance);
					bestInferences.put(ax, inf);
				}

			});

		}

		return bestInferences;
	}

	private static OWLAxiom getMinimal(Map<OWLAxiom, Integer> distances) {
		return distances.entrySet().stream().sorted((e1, e2) -> e1.getValue().compareTo(e2.getValue())).findFirst()
				.get().getKey();
	}

	private static Stream<IInference> getNextInferences(IProof proof, OWLAxiom current,
			Map<OWLAxiom, Integer> distances, Map<OWLAxiom, Integer> finalDistances) {
		return proof.getInferences().stream().filter(inf -> inf.getPremises().contains(current))
				.filter(inf -> !finalDistances.containsKey(inf.getConclusion())).filter(inf -> inf.getPremises()
						.stream().allMatch(ax -> distances.containsKey(ax) || finalDistances.containsKey(ax)));
	}

	private static void collectInferences(List<IInference> list, OWLAxiom conclusion,
			Map<OWLAxiom, IInference> bestInferences) {
		IInference current = bestInferences.get(conclusion);
		list.add(current);
		for (OWLAxiom premise : current.getPremises()) {
			collectInferences(list, premise, bestInferences);
		}
	}

}
