content: package at.ac.tuwien.kr.alpha.core.programs.transformation;
  
  import java.util.ArrayList;
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.HashSet;
  import java.util.LinkedHashMap;
  import java.util.LinkedHashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.Set;
  import java.util.Stack;
  
  import org.apache.commons.collections4.SetUtils;
  import org.apache.commons.lang3.tuple.ImmutablePair;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import at.ac.tuwien.kr.alpha.api.grounder.Substitution;
  import at.ac.tuwien.kr.alpha.api.programs.Predicate;
  import at.ac.tuwien.kr.alpha.api.programs.analysis.ComponentGraph;
  import at.ac.tuwien.kr.alpha.api.programs.analysis.DependencyGraph;
  import at.ac.tuwien.kr.alpha.api.programs.atoms.Atom;
  import at.ac.tuwien.kr.alpha.api.programs.literals.Literal;
  import at.ac.tuwien.kr.alpha.commons.programs.atoms.Atoms;
  import at.ac.tuwien.kr.alpha.commons.substitutions.BasicSubstitution;
  import at.ac.tuwien.kr.alpha.commons.substitutions.Instance;
  import at.ac.tuwien.kr.alpha.core.depgraph.StratificationAlgorithm;
  import at.ac.tuwien.kr.alpha.core.grounder.IndexedInstanceStorage;
  import at.ac.tuwien.kr.alpha.core.grounder.RuleGroundingInfo;
  import at.ac.tuwien.kr.alpha.core.grounder.RuleGroundingOrder;
  import at.ac.tuwien.kr.alpha.core.grounder.WorkingMemory;
  import at.ac.tuwien.kr.alpha.core.grounder.instantiation.AssignmentStatus;
  import at.ac.tuwien.kr.alpha.core.grounder.instantiation.LiteralInstantiationResult;
  import at.ac.tuwien.kr.alpha.core.grounder.instantiation.LiteralInstantiator;
  import at.ac.tuwien.kr.alpha.core.grounder.instantiation.WorkingMemoryBasedInstantiationStrategy;
  import at.ac.tuwien.kr.alpha.core.programs.AnalyzedProgram;
  import at.ac.tuwien.kr.alpha.core.programs.InternalProgram;
  import at.ac.tuwien.kr.alpha.core.programs.rules.CompiledRule;
  
  /**
   * Evaluates the stratifiable part of a given (analyzed) ASP program.
   * 
   * Copyright (c) 2019-2020, the Alpha Team.
   */
  public class StratifiedEvaluation extends ProgramTransformation<AnalyzedProgram, InternalProgram> {
  
  	private static final Logger LOGGER = LoggerFactory.getLogger(StratifiedEvaluation.class);
  
  	private WorkingMemory workingMemory = new WorkingMemory();
  	private Map<Predicate, LinkedHashSet<CompiledRule>> predicateDefiningRules;
  
  	private Map<Predicate, Set<Instance>> modifiedInLastEvaluationRun = new HashMap<>();
  
  	private List<Atom> additionalFacts = new ArrayList<>(); // The additional facts derived by stratified evaluation. Note that it may contain duplicates.
  	private Set<Integer> solvedRuleIds = new HashSet<>(); // Set of rules that have been completely evaluated.
  
  	private LiteralInstantiator literalInstantiator;
  
  	@Override
  	// Note: ideally this returns a "PartiallyEvaluatedProgram" such that the grounder can directly use the working
  	// memories created here rather than re-initialize everything.
  	public InternalProgram apply(AnalyzedProgram inputProgram) {
  		// Calculate a stratification and initialize the working memory.
  		ComponentGraph componentGraph = inputProgram.getComponentGraph();
  		List<ComponentGraph.SCComponent> strata = StratificationAlgorithm.calculateStratification(componentGraph);
  		predicateDefiningRules = inputProgram.getPredicateDefiningRules();
  
  		// Set up list of atoms which are known to be true - these will be expand by the evaluation.
  		Map<Predicate, Set<Instance>> knownFacts = new LinkedHashMap<>(inputProgram.getFactsByPredicate());
  		for (Map.Entry<Predicate, Set<Instance>> entry : knownFacts.entrySet()) {
  			workingMemory.initialize(entry.getKey());
  			workingMemory.addInstances(entry.getKey(), true, entry.getValue());
  		}
  
  		// Create working memories for all predicates occurring in each rule.
  		for (CompiledRule nonGroundRule : inputProgram.getRulesById().values()) {
  			for (Predicate predicate : nonGroundRule.getOccurringPredicates()) {
  				workingMemory.initialize(predicate);
  			}
  		}
  
  		workingMemory.reset();
  
  		// Set up literal instantiator.
  		literalInstantiator = new LiteralInstantiator(new WorkingMemoryBasedInstantiationStrategy(workingMemory));
  
  		// Evaluate the program part covered by the calculated stratification.
  		for (ComponentGraph.SCComponent currComponent : strata) {
  			evaluateComponent(currComponent);
  		}
  
  		// Build the program resulting from evaluating the stratified part.
  		additionalFacts.addAll(inputProgram.getFacts()); // Add original input facts to newly derived ones.
  		List<CompiledRule> outputRules = new ArrayList<>();
  		inputProgram.getRulesById().entrySet().stream().filter((entry) -> !solvedRuleIds.contains(entry.getKey()))
  				.forEach((entry) -> outputRules.add(entry.getValue()));
  
  		// NOTE: if InternalProgram requires solved rules, they should be added here.
  		return new InternalProgram(outputRules, additionalFacts);
  	}
  
  	private void evaluateComponent(ComponentGraph.SCComponent comp) {
  		LOGGER.debug("Evaluating component {}", comp);
  		ComponentEvaluationInfo evaluationInfo = getRulesToEvaluate(comp);
  		if (evaluationInfo.isEmpty()) {
  			LOGGER.debug("No rules to evaluate for component {}", comp);
  			return;
  		}
  
  		// Rules outside of dependency cycles only need to be evaluated once.
  		if (!evaluationInfo.nonRecursiveRules.isEmpty()) {
  			prepareInitialEvaluation(evaluationInfo.nonRecursiveRules);
  			evaluateRules(evaluationInfo.nonRecursiveRules, true);
  			for (IndexedInstanceStorage instanceStorage : workingMemory.modified()) {
  				// Directly record all newly derived instances as additional facts.
  				for (Instance recentlyAddedInstance : instanceStorage.getRecentlyAddedInstances()) {
  					additionalFacts.add(Atoms.newBasicAtom(instanceStorage.getPredicate(), recentlyAddedInstance.terms));
  				}
  				instanceStorage.markRecentlyAddedInstancesDone();
  			}
  		}
  		boolean isInitialRun = true;
  		if (!evaluationInfo.recursiveRules.isEmpty()) {
  			do {
  				// Now do the rules that cyclically depend on each other,
  				// evaluate these until nothing new can be derived any more.
  				if (isInitialRun) {
  					prepareInitialEvaluation(evaluationInfo.recursiveRules);
  				}
  				evaluateRules(evaluationInfo.recursiveRules, isInitialRun);
  				isInitialRun = false;
  				modifiedInLastEvaluationRun = new HashMap<>();
  				// Since we are stratified we never have to backtrack, therefore just collect the added instances.
  				for (IndexedInstanceStorage instanceStorage : workingMemory.modified()) {
  					// Directly record all newly derived instances as additional facts.
  					for (Instance recentlyAddedInstance : instanceStorage.getRecentlyAddedInstances()) {
  						additionalFacts.add(Atoms.newBasicAtom(instanceStorage.getPredicate(), recentlyAddedInstance.terms));
  					}
  					modifiedInLastEvaluationRun.putIfAbsent(instanceStorage.getPredicate(), new LinkedHashSet<>());
  					modifiedInLastEvaluationRun.get(instanceStorage.getPredicate()).addAll(instanceStorage.getRecentlyAddedInstances());
  					instanceStorage.markRecentlyAddedInstancesDone();
  				}
  				// If the evaluation of rules did not modify the working memory we have a fixed-point.
  			} while (!workingMemory.modified().isEmpty());
  		}
  		LOGGER.debug("Evaluation done - reached a fixed point on component {}", comp);
  		SetUtils.union(evaluationInfo.nonRecursiveRules, evaluationInfo.recursiveRules)
  				.forEach((rule) -> solvedRuleIds.add(rule.getRuleId()));
  	}
  
  	private void evaluateRules(Set<CompiledRule> rules, boolean isInitialRun) {
  		workingMemory.reset();
  		LOGGER.debug("Starting component evaluation run...");
  		for (CompiledRule r : rules) {
  			evaluateRule(r, !isInitialRun);
  		}
  	}
  
  	/**
  	 * To be called at the start of evaluateComponent. Adds all known instances of the predicates occurring in the given set
  	 * of rules to the "modifiedInLastEvaluationRun" map in order to "bootstrap" incremental grounding, i.e. making sure
  	 * that those instances are taken into account for ground substitutions by evaluateRule.
  	 */
  	private void prepareInitialEvaluation(Set<CompiledRule> rulesToEvaluate) {
  		modifiedInLastEvaluationRun = new HashMap<>();
  		for (CompiledRule rule : rulesToEvaluate) {
  			// Register rule head instances.
  			Predicate headPredicate = rule.getHeadAtom().getPredicate();
  			IndexedInstanceStorage headInstances = workingMemory.get(headPredicate, true);
  			modifiedInLastEvaluationRun.putIfAbsent(headPredicate, new LinkedHashSet<>());
  			if (headInstances != null) {
  				modifiedInLastEvaluationRun.get(headPredicate).addAll(headInstances.getAllInstances());
  			}
  			// Register positive body literal instances.
  			for (Literal lit : rule.getPositiveBody()) {
  				Predicate bodyPredicate = lit.getPredicate();
  				IndexedInstanceStorage bodyInstances = workingMemory.get(bodyPredicate, true);
  				modifiedInLastEvaluationRun.putIfAbsent(bodyPredicate, new LinkedHashSet<>());
  				if (bodyInstances != null) {
  					modifiedInLastEvaluationRun.get(bodyPredicate).addAll(bodyInstances.getAllInstances());
  				}
  			}
  		}
  	}
  
  	private void evaluateRule(CompiledRule rule, boolean checkAllStartingLiterals) {
  		LOGGER.debug("Evaluating rule {}", rule);
  		List<Substitution> satisfyingSubstitutions = calculateSatisfyingSubstitutionsForRule(rule, checkAllStartingLiterals);
  		for (Substitution subst : satisfyingSubstitutions) {
  			fireRule(rule, subst);
  		}
  	}
  
  	private List<Substitution> calculateSatisfyingSubstitutionsForRule(CompiledRule rule, boolean checkAllStartingLiterals) {
  		LOGGER.debug("Grounding rule {}", rule);
  		RuleGroundingInfo groundingOrders = rule.getGroundingInfo();
  
  		// Treat rules with fixed instantiation first.
  		LOGGER.debug("Is fixed rule? {}", rule.getGroundingInfo().hasFixedInstantiation());
  		if (groundingOrders.hasFixedInstantiation()) {
  			RuleGroundingOrder fixedGroundingOrder = groundingOrders.getFixedGroundingOrder();
  			return calcSubstitutionsWithGroundingOrder(fixedGroundingOrder, Collections.singletonList(new BasicSubstitution()));
  		}
  
  		List<Literal> startingLiterals = groundingOrders.getStartingLiterals();
  		// Check only one starting literal if indicated by the parameter.
  		if (!checkAllStartingLiterals) {
  			// If this is the first evaluation run, it suffices to start from the first starting literal only.
  			Literal lit = startingLiterals.get(0);
  			return calcSubstitutionsWithGroundingOrder(groundingOrders.orderStartingFrom(lit), substituteFromRecentlyAddedInstances(lit));
  		}
  
  		// Ground from all starting literals.
  		List<Substitution> groundSubstitutions = new ArrayList<>(); // Collection of full ground substitutions for the given rule.
  		for (Literal lit : startingLiterals) {
  			List<Substitution> substitutionsForStartingLiteral = calcSubstitutionsWithGroundingOrder(groundingOrders.orderStartingFrom(lit),
  					substituteFromRecentlyAddedInstances(lit));
  			groundSubstitutions.addAll(substitutionsForStartingLiteral);
  		}
  		return groundSubstitutions;
  	}
  
  	/**
  	 * Use this to find initial substitutions for a starting literal when grounding a rule.
  	 * In order to avoid finding the same ground instantiations of rules again, only look at
  	 * <code>modifiedInLastEvaluationRun</code> to obtain instances.
  	 * 
  	 * @param lit the literal to substitute.
  	 * @return valid ground substitutions for the literal based on the recently added instances (i.e. instances derived in
  	 *         the last evaluation run).
  	 */
  	private List<Substitution> substituteFromRecentlyAddedInstances(Literal lit) {
  		List<Substitution> retVal = new ArrayList<>();
  		Set<Instance> instances = modifiedInLastEvaluationRun.get(lit.getPredicate());
  		if (instances == null) {
  			return Collections.emptyList();
  		}
  		for (Instance instance : instances) {
  			Substitution unifyingSubstitution = BasicSubstitution.specializeSubstitution(lit, instance, BasicSubstitution.EMPTY_SUBSTITUTION);
  			if (unifyingSubstitution != null) {
  				retVal.add(unifyingSubstitution);
  			}
  		}
  		return retVal;
  	}
  
  	private List<Substitution> calcSubstitutionsWithGroundingOrder(RuleGroundingOrder groundingOrder, List<Substitution> startingSubstitutions) {
  		// Iterate through the grounding order and whenever instantiation of a Literal with a given substitution
  		// causes a result with a type other than CONTINUE, discard that substitution.
  
  		// Note that this function uses a stack of partial substitutions to simulate a recursive function.
  		Stack<ArrayList<Substitution>> substitutionStack = new Stack<>(); // For speed, we really want ArrayLists on the stack.
  		if (startingSubstitutions instanceof ArrayList) {
  			substitutionStack.push((ArrayList<Substitution>) startingSubstitutions);
  		} else {
  			substitutionStack.push(new ArrayList<>(startingSubstitutions)); // Copy startingSubstitutions into ArrayList. Note: mostly happens for empty or
  																			// singleton lists.
  		}
  		int currentOrderPosition = 0;
  		List<Substitution> fullSubstitutions = new ArrayList<>();
  		while (!substitutionStack.isEmpty()) {
  			List<Substitution> currentSubstitutions = substitutionStack.peek();
  			// If no more substitutions remain at current position, all have been processed, continue on next lower level.
  			if (currentSubstitutions.isEmpty()) {
  				substitutionStack.pop();
  				currentOrderPosition--;
  				continue;
  			}
  			// In case the full grounding order has been worked on, all current substitutions are full substitutions, add them to
  			// result.
  			Literal currentLiteral = groundingOrder.getLiteralAtOrderPosition(currentOrderPosition);
  			if (currentLiteral == null) {
  				fullSubstitutions.addAll(currentSubstitutions);
  				currentSubstitutions.clear();
  				// Continue on next lower level.
  				substitutionStack.pop();
  				currentOrderPosition--;
  				continue;
  			}
  			// Take one substitution from the top-list of the stack and try extending it.
  			Substitution currentSubstitution = currentSubstitutions.remove(currentSubstitutions.size() - 1); // Work on last element (removing last element is
  																												// O(1) for ArrayList).
  			LiteralInstantiationResult currentLiteralResult = literalInstantiator.instantiateLiteral(currentLiteral, currentSubstitution);
  			if (currentLiteralResult.getType() == LiteralInstantiationResult.Type.CONTINUE) {
  				// The currentSubstitution could be extended, push the extensions on the stack and continue working on them.
  				ArrayList<Substitution> furtheredSubstitutions = new ArrayList<>();
  				for (ImmutablePair<Substitution, AssignmentStatus> resultSubstitution : currentLiteralResult.getSubstitutions()) {
  					furtheredSubstitutions.add(resultSubstitution.left);
  				}
  				substitutionStack.push(furtheredSubstitutions);
  				// Continue work on the higher level.
  				currentOrderPosition++;
  			}
  		}
  		return fullSubstitutions;
  	}
  
  	private void fireRule(CompiledRule rule, Substitution substitution) {
  		Atom newAtom = rule.getHeadAtom().substitute(substitution);
  		if (!newAtom.isGround()) {
  			throw new IllegalStateException("Trying to fire rule " + rule.toString() + " with incompatible substitution " + substitution.toString());
  		}
  		LOGGER.debug("Firing rule - got head atom: {}", newAtom);
  		workingMemory.addInstance(newAtom, true);
  	}
  
  	private ComponentEvaluationInfo getRulesToEvaluate(ComponentGraph.SCComponent comp) {
  		Set<CompiledRule> nonRecursiveRules = new HashSet<>();
  		Set<CompiledRule> recursiveRules = new HashSet<>();
  
  		// Collect all predicates occurring in heads of rules of the given component.
  		Set<Predicate> headPredicates = new HashSet<>();
  		for (DependencyGraph.Node node : comp.getNodes()) {
  			headPredicates.add(node.getPredicate());
  		}
  		// Check each predicate whether its defining rules depend on some of the head predicates, i.e., whether there is a
  		// cycle.
  		for (Predicate headPredicate : headPredicates) {
  			HashSet<CompiledRule> definingRules = predicateDefiningRules.get(headPredicate);
  			if (definingRules == null) {
  				// Predicate only occurs in facts, skip.
  				continue;
  			}
  			// Note: here we assume that all rules defining a predicate belong to the same SC component.
  			for (CompiledRule rule : definingRules) {
  				boolean isRuleRecursive = false;
  				for (Literal lit : rule.getPositiveBody()) {
  					if (headPredicates.contains(lit.getPredicate())) {
  						// The rule body contains a predicate that is defined in the same
  						// component, the rule is therefore part of a cyclic dependency within
  						// this component.
  						isRuleRecursive = true;
  					}
  				}
  				if (isRuleRecursive) {
  					recursiveRules.add(rule);
  				} else {
  					nonRecursiveRules.add(rule);
  				}
  			}
  		}
  		return new ComponentEvaluationInfo(nonRecursiveRules, recursiveRules);
  	}
  
  	/**
  	 * Internal helper class to group rules within an {@SCComponent} into rules that are recursive, i.e. part of some cyclic
  	 * dependency chain within that component, and non-recursive rules, i.e. rules where all body predicates occur in lower
  	 * strata. The reason for this grouping is that, when evaluating rules within a component, non-recursive rules only need
  	 * to be evaluated once, while recursive rules need to be evaluated until a fixed-point has been reached.
  	 * 
  	 * Copyright (c) 2020, the Alpha Team.
  	 */
  	private class ComponentEvaluationInfo {
  		final Set<CompiledRule> nonRecursiveRules;
  		final Set<CompiledRule> recursiveRules;
  
  		ComponentEvaluationInfo(Set<CompiledRule> nonRecursive, Set<CompiledRule> recursive) {
  			nonRecursiveRules = Collections.unmodifiableSet(nonRecursive);
  			recursiveRules = Collections.unmodifiableSet(recursive);
  		}
  
  		boolean isEmpty() {
  			return nonRecursiveRules.isEmpty() && recursiveRules.isEmpty();
  		}
  
  	}
  
  }
  