package at.ac.tuwien.kr.alpha.core.programs.transformation.aggregates;

import static at.ac.tuwien.kr.alpha.commons.util.Util.oops;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import org.apache.commons.lang3.tuple.ImmutablePair;
import org.stringtemplate.v4.ST;

import at.ac.tuwien.kr.alpha.api.ComparisonOperator;
import at.ac.tuwien.kr.alpha.api.programs.atoms.AggregateAtom;
import at.ac.tuwien.kr.alpha.api.programs.atoms.AggregateAtom.AggregateFunctionSymbol;
import at.ac.tuwien.kr.alpha.api.programs.atoms.Atom;
import at.ac.tuwien.kr.alpha.api.programs.atoms.BasicAtom;
import at.ac.tuwien.kr.alpha.api.programs.literals.AggregateLiteral;
import at.ac.tuwien.kr.alpha.api.programs.literals.Literal;
import at.ac.tuwien.kr.alpha.api.programs.rules.Rule;
import at.ac.tuwien.kr.alpha.api.programs.rules.heads.Head;
import at.ac.tuwien.kr.alpha.api.programs.terms.Term;
import at.ac.tuwien.kr.alpha.api.programs.terms.VariableTerm;
import at.ac.tuwien.kr.alpha.commons.Predicates;
import at.ac.tuwien.kr.alpha.commons.programs.atoms.Atoms;
import at.ac.tuwien.kr.alpha.commons.programs.terms.Terms;

/**
 * Holds all information about aggregate literals that need to be rewritten within a program.
 * Specifically, a rewriting context holds a list of rules with aggregates, as well as detailed analysis results like
 * global variables and dependencies per aggregate (see AggregateInfo).
 * 
 * Copyright (c) 2021, the Alpha Team.
 */
public final class AggregateRewritingContext {

	private static final ST AGGREGATE_RESULT_TEMPLATE = new ST("<id>_result");
	private static final ST AGGREGATE_ARGS_FUNCTION_SYMBOL = new ST("<id>_args");
	private static final ST AGGREGATE_ARGS_NOARGS_CONST = new ST("<id>_no_args");

	private int idCounter;
	private Map<AggregateLiteral, AggregateInfo> aggregateInfos = new HashMap<>(); // Maps aggregate literals to their respective AggregateInfo.
	private Map<ImmutablePair<AggregateFunctionSymbol, ComparisonOperator>, Set<AggregateInfo>> aggregateFunctionsToRewrite = new LinkedHashMap<>();
	private List<Rule<Head>> rulesWithAggregates = new ArrayList<>();

	public AggregateRewritingContext() {
		idCounter = 0;
	}

	/**
	 * Registers a rule that potentially contains one or more {@link AggregateLiteral}s with this context.
	 * In case aggregates are found in the rule, global variables and dependencies of the aggregate are calculated and the
	 * aggregate literal is stored in the rewriting context along with a reference to the rule it occurs in
	 * 
	 * @param rule
	 * @return true if the given rule contains one or more aggregate literals, false otherwise
	 */
	public boolean registerRule(Rule<Head> rule) {
		AggregateRewritingRuleAnalysis ruleAnalysis = AggregateRewritingRuleAnalysis.analyzeRuleDependencies(rule);
		if (ruleAnalysis.aggregatesInRule.isEmpty()) {
			// Rule has no aggregates.
			return false;
		}
		// Do initial registration of each aggregate literal and keep the ids.
		for (Map.Entry<AggregateLiteral, Set<VariableTerm>> entry : ruleAnalysis.globalVariablesPerAggregate.entrySet()) {
			registerAggregateLiteral(entry.getKey(), entry.getValue());
		}
		// Now go through dependencies and replace the actual aggregate literals with their rewritten versions
		for (Map.Entry<AggregateLiteral, Set<Literal>> entry : ruleAnalysis.dependenciesPerAggregate.entrySet()) {
			AggregateInfo aggregateInfo = getAggregateInfo(entry.getKey());
			for (Literal dependency : entry.getValue()) {
				if (dependency instanceof AggregateLiteral) {
					AggregateInfo dependencyInfo = getAggregateInfo((AggregateLiteral) dependency);
					aggregateInfo.addDependency(dependencyInfo.getOutputAtom().toLiteral(!dependency.isNegated()));
				} else {
					aggregateInfo.addDependency(dependency);
				}
			}
		}
		rulesWithAggregates.add(rule);
		return true;
	}

	private void registerAggregateLiteral(AggregateLiteral lit, Set<VariableTerm> globalVariables) {
		AggregateAtom atom = lit.getAtom();
		String id = atom.getAggregateFunction().toString().toLowerCase() + "_" + (++idCounter);
		AggregateInfo info = new AggregateInfo(id, lit, globalVariables);
		if (aggregateInfos.containsKey(lit)) {
			throw oops("AggregateInfo for AggregateLiteral already existing.");
		}
		aggregateInfos.put(lit, info);
		aggregateFunctionsToRewrite.putIfAbsent(new ImmutablePair<>(atom.getAggregateFunction(), atom.getLowerBoundOperator()), new LinkedHashSet<>());
		aggregateFunctionsToRewrite.get(new ImmutablePair<>(atom.getAggregateFunction(), atom.getLowerBoundOperator())).add(info);
	}

	AggregateInfo getAggregateInfo(AggregateLiteral aggregateLiteral) {
		return aggregateInfos.get(aggregateLiteral);
	}

	List<Rule<Head>> getRulesWithAggregates() {
		return rulesWithAggregates;
	}

	public Map<ImmutablePair<AggregateFunctionSymbol, ComparisonOperator>, Set<AggregateInfo>> getAggregateFunctionsToRewrite() {
		return Collections.unmodifiableMap(aggregateFunctionsToRewrite);
	}

	public static class AggregateInfo {
		private final String id;
		private final AggregateLiteral literal;
		private final BasicAtom outputAtom;
		private final Set<VariableTerm> globalVariables;
		private final Term aggregateArguments;
		private final Set<Literal> dependencies = new LinkedHashSet<>();

		AggregateInfo(String id, AggregateLiteral literal, Set<VariableTerm> globalVariables) {
			this.id = id;
			this.literal = literal;
			this.globalVariables = globalVariables;
			this.aggregateArguments = buildArguments();
			this.outputAtom = buildOutputAtom();
		}

		/**
		 * Builds a term representing the global variables of the {@link AggregateLiteral} represented by this
		 * {@link AggregateInfo}. For a literal with global variables "A, B, C", the argument term is "$aggregate_id_$args(A, B,
		 * C)". For a
		 * literal without global variables, the constant "$aggregate_id$_no_args" is returned.
		 */
		private Term buildArguments() {
			Term argumentTerm;
			if (globalVariables.isEmpty()) {
				argumentTerm = Terms.newSymbolicConstant(new ST(AGGREGATE_ARGS_NOARGS_CONST).add("id", id).render());
			} else {
				argumentTerm = Terms.newFunctionTerm(new ST(AGGREGATE_ARGS_FUNCTION_SYMBOL).add("id", id).render(),
						new ArrayList<>(globalVariables));
			}
			return argumentTerm;
		}

		/**
		 * Builds an {@link Atom} which is inserted into rule bodies instead of the respective {@link AggregateLiteral}.
		 * The atom is of structure "$aggregate_id$_result($argTerm$, $resultTerm$)", where aggregate_id is the unique id of the
		 * aggregate literal as generated by {@link AggregateRewritingContext}, arg_term contains the global variables of the
		 * aggregate literal in question (or a constant "$aggregate_id$_no_args" for a literal without global variables), and
		 * resultTerm is the
		 * "lowerBoundTerm" of the aggregate literal.
		 * 
		 * @return
		 */
		private BasicAtom buildOutputAtom() {
			String outputPredicateName = new ST(AGGREGATE_RESULT_TEMPLATE).add("id", id).render();
			Term argumentTerm = aggregateArguments;
			Term resultTerm = literal.getAtom().getLowerBoundTerm();
			return Atoms.newBasicAtom(Predicates.getPredicate(outputPredicateName, 2, true), argumentTerm, resultTerm);
		}

		private void addDependency(Literal dependency) {
			dependencies.add(dependency);
		}

		public String getId() {
			return id;
		}

		public AggregateLiteral getLiteral() {
			return literal;
		}

		public BasicAtom getOutputAtom() {
			return outputAtom;
		}

		public Set<VariableTerm> getGlobalVariables() {
			return globalVariables;
		}

		public Term getAggregateArguments() {
			return aggregateArguments;
		}

		public Set<Literal> getDependencies() {
			return dependencies;
		}

	}

}
