package ua.ansymo.hintco

import java.util.ArrayList
import java.util.Deque
import java.util.HashMap
import java.util.List
import java.util.Set
import org.eclipse.core.runtime.Assert
import org.slf4j.Logger
import org.slf4j.LoggerFactory

class CandidatesGenerator {

	IConstraintChecker constraintChecker
	IVariantValidator variantValidator
	IVariantProcessor processor

	Logger logger = LoggerFactory.getLogger(typeof(CandidatesGenerator))

	new(IConstraintChecker c, IVariantValidator v, IVariantProcessor p) {
		constraintChecker = c
		variantValidator = v
		processor = p
	}

	/*
	 * Recursively goes through every path on the variant diagram and creates the necessary precedence constraints, cleaning up after the function call (that's why we use a stack to keep groups of constraints together).
	 * Optimization: when optimizeAlgebraicLoops==false, check if there is a cycle every time there are a new edge is created. Then abort that function call, as any children of the variant tree will not be valid. Possible paper that addresses this: https://dl.acm.org/citation.cfm?id=1033207
	 * The simple solution to find a cycle is to check if there already is a path connecting those two ports being connected.
	 * (Not done) A better way is to compute transitive closure for quicker reachable computation after adding multiple edges. This result is passed down function calls, and cleaned up as well.
	 * For now, we just check whether there is a single path passing through all nodes, in a DFS, and return the first such path (mark it on the model, and call semantic adaptation code generation).
	 * 
	 * When optimizeAlgebraicLoops==true, we have to continue the recursion, and add constraints as needed. When a variant is fully specified, we launch the algebraic loop optimizer in order to find a schedule for the cosim scenario.
	 */
	def generateVariants(HintConfiguration cs, int maxVariants, boolean optimizeAlgebraicLoops) {

		// Assumes that there is a variant diagram created.
		Assert.isTrue(cs.root !== null)

		if (optimizeAlgebraicLoops && maxVariants > 1) {
			logger.warn(
				"Multiple variants chosen with algebraic loop optimizer enabled. This will cause the optimizer to be used for every variant, and might be slow.")
		}

		val genVariants = recGenerateVariants(newLinkedList(cs.root), cs, new ConstraintsStack(), maxVariants,
			optimizeAlgebraicLoops)
		return genVariants
	}

	def processAdaptations(HintConfiguration cs) {
		// Completes adaptations such as the power bond adaptation.
		val f = HintcoFactory.eINSTANCE

		// For each power adaptation, creates an input port and an output power.
		val adaptations = cs.eAllContents.filter(PowerBondAdaptation).toList

		for (pbAdaptation : adaptations) {
			Assert.isTrue(pbAdaptation.PIn === null)
			Assert.isTrue(pbAdaptation.POut === null)
			val adaptedUnit = pbAdaptation.adapted as CosimUnitInstance

			// Create powerOut declaration
			{
				val pOutId = pbAdaptation.effort.identifier + "_" + pbAdaptation.flow.identifier + "_pOut"
				Assert.isTrue(!adaptedUnit.ports.exists[p|p.identifier == pOutId])

				val pOutPort = f.createOutputPortInstance
				pOutPort.identifier = pOutId
				adaptedUnit.ports.add(pOutPort)
				pbAdaptation.POut = pOutPort
			}

			// Create powerIn declaration
			{
				val pInId = pbAdaptation.effort.identifier + "_" + pbAdaptation.flow.identifier + "_pIn"
				Assert.isTrue(!adaptedUnit.ports.exists[p|p.identifier == pInId])

				val pInPort = f.createInputPortInstance
				pInPort.identifier = pInId
				adaptedUnit.ports.add(pInPort)
				pbAdaptation.PIn = pInPort
			}
		}

		// Connect the power out with the power Ins
		for (pbAdaptation : adaptations) {
			if (pbAdaptation.POut.valueTo.empty) {
				pbAdaptation.POut.valueTo.add(pbAdaptation.dual.PIn)
			}
		}

	}

	def createVariantTree(HintConfiguration cs) {
		// Creates a variant tree.
		// The first level of the tree is the xor with the strongest weights.
		// Clear existing information
		cs.root = null
		cs.nodes.clear()

		// Get factory for creating objects.
		val f = HintcoFactory.eINSTANCE

		val id = 0 // Generates unique ids for VariantTree objects. Allows traceability between results and variants.
		// Created sorted alternatives: highest weight first
		val sortedAlts = newArrayList()
		cs.eAllContents.filter(Alternative).forEach [ a |
			// creates. The boolean is used later to mark alternatives that have been considered.
			sortedAlts.add(new MarkedAlternative(a))
		]
		sortedAlts.sortInplace([a, b|b.alternative.weight - a.alternative.weight])

		// recursively create variant tree, starting with the highest weight.
		val rootNode = f.createVariantDiagram
		rootNode.identifier = "" + id
		cs.root = rootNode
		cs.nodes.add(rootNode)
		recCreateVariantTree(#[rootNode], sortedAlts, 0, cs, id + 1)
		println()
	}

	def private int recGenerateVariants(Deque<VariantDiagram> ns, HintConfiguration cs, ConstraintsStack constraints,
		int maxVariants, boolean optimizeAlgebraicLoops) {
		// Store size of stack for assertion
		val callDepth = constraints.size

		var numGenVariants = 0

		// Create new constraint map and put it at the head of the queue.
		val HashMap<PrecedenceNode, Set<PrecedenceNode>> localConstraints = newHashMap()
		constraints.push(localConstraints)

		val n = ns.peek() // Head of the queue is the current element being visited.
		
		val introducedCycle = addConstraintsBasedOnAlternative(n.alternative, constraints, n, cs)

		// Check invariant on localConstraints
		localConstraints.forEach [ k, v |
			Assert.isTrue(k.precedes.containsAll(v))
		]

		if (!introducedCycle || optimizeAlgebraicLoops) {
			if (n.children.empty) {
				// Base case
				if (numGenVariants < maxVariants) {
					// Propagate precedence constraints from inner scenarios to their parents.
					// If two nodes are indirectly (via a path in the children nodes) connected, then they should be connected directly.
					// This will ensure that the topological sort can safely ignore child nodes.
					propagateConstraints(n, ns, constraints, cs)

					val variantID = ns.map[n1|n1.identifier].reduce[n1, n2|n2 + "_" + n1]
					// mark alternatives as selected, so the processor knows which alternatives have been selected.
					// Exclude root because it has no alternative
					markSelected(ns, true)
					processor.process(ns, variantID, constraints, cs)
					markSelected(ns, false)
					numGenVariants++
				}
			} else {
				// Recurse
				for (c : n.children) {
					if (numGenVariants < maxVariants &&
						variantValidator.validVariant(c, n.children.filter[a|a !== c], ns)) {
						ns.push(c)
						val childGenVariants = recGenerateVariants(ns, cs, constraints, maxVariants - numGenVariants,
							optimizeAlgebraicLoops)
						numGenVariants = numGenVariants + childGenVariants
						ns.pop()
					}
				}
			}
		} else {
			logger.debug("Constraints form dependency loop and the loop optimizer is disabled. Exploring other alternatives.")
		}

		// Revert changes in model
		localConstraints.forEach [ k, v |
			Assert.isTrue(k.precedes.removeAll(v))
		]
		constraints.pop()

		Assert.isTrue(callDepth == constraints.size)
		return numGenVariants
	}

	def addConstraintsBasedOnAlternative(Alternative alternative, ConstraintsStack constraints, VariantDiagram n,
		HintConfiguration cs) {
		var introducedCycle = false
		switch alternative {
			RootCandidateScenario: {
				/*
				 * Recursively create I/O and DoStep constraints.
				 */
				val scenario = (alternative as RootCandidateScenario)
				logger.debug("Processing Alternative RootCandidateScenario {}", scenario.identifier)

				for (outP : scenario.eAllContents.filter(OutputPortInstance).toIterable) {
					// Create constraints Outports/HierarchicalPorts -> InPorts/HierarchicalPorts
					for (inP : outP.valueTo) {
						logger.debug("Adding constraint {}.{}->{}.{}.", outP.unit.identifier, outP.identifier,
							inP.unit.identifier, inP.identifier)
						val success = addConstraint(outP, inP, constraints, n, cs)
						logger.debug("Success: {}", success)
						if (success == AddConstraintResult.CYCLE) {
							introducedCycle = true
						}
					}

					// Create feedthrough constraints:
					for (inP : outP.internalValueFrom) {
						logger.debug("Attempting to add constraint {}.{}->{}.{}.", inP.unit.identifier, inP.identifier,
							outP.unit.identifier, outP.identifier)
						val success = addConstraint(inP, outP, constraints, n, cs)
						logger.debug("Success: {}", success)
						if (success == AddConstraintResult.CYCLE) {
							introducedCycle = true
						}
					}
				}

				// Create constraints between outputs and doStep.
				// We assume that at the beginning of the cosim step, every input/output is defined. 
				// Therefore, doStep must be done before getting any output.
				for (outP : scenario.eAllContents.filter(OutputPortInstance).filter[! it.isInput] // ip can be a hierarchical port, 
				.toIterable // therefore we need to make sure it is connected to an internal unit. 
				// Otherwise it is not an input port.
				) {
					logger.debug("Adding constraint {}->{}.{}.", outP.unit.identifier, outP.unit.identifier,
						outP.identifier)
					val success = addConstraint(outP.unit, outP, constraints, n, cs)
					logger.debug("Success: {}", success)
					if (success == AddConstraintResult.CYCLE) {
						introducedCycle = true
					}
				}

				// The DoStep of each child unit must conclude before the doStep of the parent unit concludes
				// Furthermore, every output operation of a child unit must come before every getout operation of the parent unit.
				// This, together with the previous set of constraints, ensures that each hierarchical unit's doStep operation is atomic.
				for (parentUnit : scenario.eAllContents.filter(HierarchicalCosimUnit).toIterable) {
					val parentOutputPorts = parentUnit.outputPorts
					for (childUnit : parentUnit.cosimunits) {
						logger.debug("Adding constraint {}->{}.", childUnit.identifier, parentUnit.identifier)
						val success = addConstraint(childUnit, parentUnit, constraints, n, cs)
						logger.debug("Success: {}", success)
						if (success == AddConstraintResult.CYCLE) {
							introducedCycle = true
						}

						for (childOutP : childUnit.outputPorts) {
							for (parentOutP : parentOutputPorts) {
								// TODO: Use toString method for printing, and move the logging operations to inside a method.
								logger.debug("Adding constraint {}.{}->{}.{}.", childOutP.unit.identifier,
									childOutP.identifier, parentOutP.unit.identifier, parentOutP.identifier)
								val succ = addConstraint(childOutP, parentOutP, constraints, n, cs)
								logger.debug("Success: {}", succ)
								if (success == AddConstraintResult.CYCLE) {
									introducedCycle = true
								}
							}
						}
					}
				}
			}
			ExtrapolationAdaptation: {
				// In case of an extrapolation on an input, the doStep must be done before setting the input,
				// because whenever the input is set, it is being set for timestamp t+H.
				val adaptation = alternative as ExtrapolationAdaptation
				val adapted = adaptation.adapted()
				logger.debug("Processing Alternative ExtrapolationAdaptation of {}", adapted.identifier)
				Assert.isNotNull(adapted)
				Assert.isTrue(adapted instanceof InputPortInstance)
				val adaptedPort = adapted as InputPortInstance
				Assert.isTrue(adaptedPort.unit instanceof CosimUnitInstance)

				logger.debug("Adding constraint {}->{}.{}.", adaptedPort.unit.identifier, adaptedPort.unit.identifier,
					adaptedPort.identifier)
				val success = addConstraint(adaptedPort.unit, adaptedPort, constraints, n, cs)
				logger.debug("Success: {}", success)
				if (success == AddConstraintResult.CYCLE) {
					introducedCycle = true
				}

			}
			InterpolationAdaptation: {
				// In case of an interpolation on an input, the input must be set before the doStep,
				// because whenever the input is set, it is being set for timestamp t+H.
				val adaptation = alternative as InterpolationAdaptation
				val adapted = adaptation.adapted()
				logger.debug("Processing Alternative InterpolationAdaptation of {}", adapted.identifier)
				Assert.isNotNull(adapted)
				Assert.isTrue(adapted instanceof InputPortInstance)
				val adaptedPort = adapted as InputPortInstance

				logger.debug("Adding constraint {}.{}->{}.{}.", adaptedPort.unit.identifier, adaptedPort.identifier,
					adaptedPort.unit.identifier)
				val success = addConstraint(adaptedPort, adaptedPort.unit, constraints, n, cs)
				logger.debug("Success: {}", success)
				if (success == AddConstraintResult.CYCLE) {
					introducedCycle = true
				}

			}
			DecompositionPortAdaptation: {
				// Nothing to do.
			}
			DecompositionUnitAdaptation: {
				// Nothing to do.
			}
			PowerBondAdaptation: {
				// This needs to be rewritten, to accomodate the hierarchical cosim case 
				// and the assumption that at the begining of the cosim step, every port is defined.
				throw new UnsupportedOperationException(
					"Powerbond adaptations need to be re-implemented for hierarchical cosim.")

			/*
			 * val adaptation = alternative as PowerBondAdaptation
			 * val adapted = adaptation.adapted()
			 * logger.debug("Processing Alternative PowerBondAdaptation of {}", adapted.identifier)
			 * Assert.isNotNull(adapted)
			 * Assert.isTrue(adapted instanceof UnitInstance)
			 * val unit = adapted as UnitInstance
			 * 
			 * Assert.isTrue(adaptation.POut.unit === unit)
			 * 
			 * logger.debug("Attempting to add constraint {}.{}->{}.", unit.identifier, adaptation.POut.identifier, unit.identifier)
			 * val success = addConstraint(adaptation.POut, unit, constraints, n, cs)
			 * logger.debug("Success: {}", success)
			 */
			}
			MultiRateAdaptation: {
				// Nothing to do
			}
			default:
				if (alternative !== null) {
					throw new UnsupportedOperationException()
				}
		}
		return introducedCycle
	}

	def propagateConstraints(VariantDiagram n, Deque<VariantDiagram> ns, ConstraintsStack constraints,
		HintConfiguration cs) {
		val rootScenario = ModelQuery.findRootScenario(ns)
		val scenarios = rootScenario.eAllContents.filter(HierarchicalCosimUnit).toList
		for (scenario : scenarios) {
			val List<PrecedenceNode> scenarioNodes = newLinkedList()
			scenarioNodes.add(scenario as PrecedenceNode)
			scenarioNodes.addAll(scenario.inputPorts)
			scenarioNodes.addAll(scenario.outputPorts)

			for (scenarioNodeSrc : scenarioNodes) {
				for (scenarioNodeTrg : scenarioNodes) {
					if (! scenarioNodeSrc.precedes.contains(scenarioNodeTrg) &&
						ModelQuery.pathExists(scenarioNodeSrc, scenarioNodeTrg, constraints, cs)) {
						Assert.isTrue(
							addConstraint(scenarioNodeSrc, scenarioNodeTrg, constraints, n, cs) == AddConstraintResult.SUCCESS)
					}
				}
			}
		}
	}

	def markSelected(Deque<VariantDiagram> ns, boolean selected) {
		for (n : ns) {
			if (n.alternative !== null) {
				n.alternative.selected = selected
			}
		}
	}

	enum AddConstraintResult {
		SUCCESS,
		EXISTS,
		CYCLE
	}

	def private addConstraint(
		PrecedenceNode src,
		PrecedenceNode trg,
		ConstraintsStack constraints,
		VariantDiagram n,
		HintConfiguration cs
	) {
		// check if constraint already exists.
		if (!constraints.stackContains(src, trg) && !src.precedes.contains(trg)) {
			// Add constraint
			
			// the following result avoids calls to children for variants that cannot be implemented.
			val result = if (constraintChecker.constraintIsFeasible(src, trg, n, constraints, cs)) { 
				AddConstraintResult.SUCCESS
			} else {
				AddConstraintResult.CYCLE
			}
			// Even if a cycle is introduced, we add the constraint.
			// This is because either the calling code tolerates cycles, or it will undo the constraints that introduced cycles.
			constraints.safeAdd(src, trg)
			src.precedes.add(trg)
			return result
		} else {
			return AddConstraintResult.EXISTS
		}
	}

	def private void recCreateVariantTree(List<VariantDiagram> roots, ArrayList<MarkedAlternative> sortedAlts,
		int highestAltIdx, HintConfiguration cs, int id) {
		// mark highest weight alternative from the list of alternatives
		// and mark siblings as well.
		// highestAltIdx stores where we can start searching
		val f = HintcoFactory.eINSTANCE

		// Base of the recursion:
		if (highestAltIdx >= sortedAlts.length)
			return

		val highestAlt = sortedAlts.get(highestAltIdx)
		highestAlt.inVariantTree = true
		// find siblings. 
		val siblings = collectMarkSiblings(highestAlt, sortedAlts, highestAltIdx)

		Assert.isTrue(siblings.length > 0)

		// Create a branch per alternative, under each of the roots
		val newNodes = newArrayList
		siblings.forEach [ a, i |
			val v = f.createVariantDiagram
			v.identifier = "" + (id + i)
			v.alternative = a.alternative
			newNodes.add(v)
		]

		val pl = cs.nodes.length
		if(! cs.nodes.addAll(newNodes)) throw new RuntimeException()
		Assert.isTrue(cs.nodes.length == pl + newNodes.length)

		roots.forEach [ r |
			Assert.isTrue(r.children !== null)
			Assert.isTrue(r.children.length == 0)
			if (! r.children.addAll(newNodes))
				throw new RuntimeException()
			Assert.isTrue(r.children.length > 0)
		]

		roots.forEach[r|Assert.isTrue(r.children.length > 0, "Problem with " + r)]

		// move highestAltIdx to the next unmarked alternative
		var newHighestIdx = highestAltIdx + 1
		while (newHighestIdx < sortedAlts.length && sortedAlts.get(newHighestIdx).inVariantTree) {
			newHighestIdx++
		}

		// Go for next recursion: take all the newly created notes into a list, and call recursively.
		recCreateVariantTree(newNodes, sortedAlts, newHighestIdx, cs, id + newNodes.length)
	}

	def private collectMarkSiblings(MarkedAlternative highestAlt, ArrayList<MarkedAlternative> sortedAlts,
		int highestAltIdx) {

		val r = newLinkedList(highestAlt)

		// First considers the different subtypes of alternative.
		switch highestAlt.alternative {
			RootCandidateScenario: {
				// Case Scenario, the siblings are the other scenarios.
				for (var i = highestAltIdx + 1; i < sortedAlts.length; i++) {
					if (!sortedAlts.get(i).inVariantTree &&
						sortedAlts.get(i).alternative instanceof RootCandidateScenario) {
						r.add(sortedAlts.get(i))
						sortedAlts.get(i).inVariantTree = true
					}
				}
			}
			UnitAdaptation: {
				// Case adaptation, it only has siblings if there is a parent, and the siblings have the same parent.
				val parent = (highestAlt.alternative as UnitAdaptation).parent
				if (parent !== null) {
					for (var i = highestAltIdx + 1; i < sortedAlts.length; i++) {
						if (!sortedAlts.get(i).inVariantTree &&
							sortedAlts.get(i).alternative instanceof UnitAdaptation &&
							(sortedAlts.get(i).alternative as UnitAdaptation).parent === parent) {
							r.add(sortedAlts.get(i))
							sortedAlts.get(i).inVariantTree = true
						}
					}
				}
			}
			PortAdaptation: {
				// Case adaptation, it only has siblings if there is a parent, and the siblings have the same parent.
				val parent = (highestAlt.alternative as PortAdaptation).parent
				if (parent !== null) {
					for (var i = highestAltIdx + 1; i < sortedAlts.length; i++) {
						if (!sortedAlts.get(i).inVariantTree &&
							sortedAlts.get(i).alternative instanceof PortAdaptation &&
							(sortedAlts.get(i).alternative as PortAdaptation).parent === parent) {
							r.add(sortedAlts.get(i))
							sortedAlts.get(i).inVariantTree = true
						}
					}
				}
			}
			default:
				throw new UnsupportedOperationException()
		}
		return r
	}

}
