package ua.ansymo.hintco

import java.io.IOException
import java.util.HashMap
import java.util.List
import java.util.Map
import java.util.Set
import org.eclipse.core.runtime.Assert
import org.slf4j.Logger
import org.slf4j.LoggerFactory

import static ua.ansymo.hintco.CosimRunUtils.*

class StrongCouplingRunner implements ICosimRunner {
	/*
	 * Runs the string coupling version of the co-simulation.
	 * Does not require FMUs to be able to rollback, but extremely inefficient
	 * 
	  Algorithm:
		1. Initialize cosim and record the first outputs.
		2. Set step=1
		3. Count total steps
		4. while step<=total_steps:
		4.1 Converge for step.
		4.2 step=step+1
		5. Output cosim results
		
	  Converge for step
		1. Given a step > 0
		2. Let iteration counter i=0
		3. Assume that all fmus are reset.
		4. Run each FMU up to step-1
		5. If i>0 set the previously recorded outputs to the rollback inputs
		6. Run the cosim step.
		7. Get new outputs, and check for convergence
		8. i = i + 1 and reset fmus
		9. If no convergence, go to 3.
		10. Otherwise return true unless max iterations has exceeded.
		11. Otherwise, issue warning and ?
	 */
	
	protected IOutputProcessor outProcessor
	protected Logger logger = LoggerFactory.getLogger(typeof(StrongCouplingRunner))
	protected IFmuLoader fmuLoader
	
	protected boolean hardReset
	
	new (IOutputProcessor o, IFmuLoader l, boolean hardReset){
		outProcessor = o
		fmuLoader = l
		this.hardReset = hardReset
	}
	
	override run(RootCandidateScenario rootScenario, Map<Scenario, List<PrecendenceNode>> scenarioNodesMap, String variantID) {
		Assert.isTrue(scenarioNodesMap.values.forall[sortedNodes | sortedNodes.forall[n1 | rootScenario.eAllContents.exists[n2 | n1===n2]]])
		
		val nodes = scenarioNodesMap.get(rootScenario)
		
		val units = nodes.filter(UnitInstance).toList
		// Takes into account hierarchical unit ports.
		val childOutputPorts = nodes.filter(OutputPortInstance).filter[!it.isInput].toList
		
		val childInputPorts = nodes.filter(InputPortInstance).filter[it.isInput].toSet
		val rollbackOutputPorts = nodes.filter(InputPortInstance).filter[it.isInput]
										.filter[p | p.adaptation instanceof RollbackInterpolationAdaptation]
										.map[p | p.valueFrom]
										.toSet
		
		val Map<UnitInstance, IFmuInstance> fmuInstanceMap = newHashMap()

		// Instantiate each unit, and record the instance reference.
		instantiateUnits(units, fmuInstanceMap, fmuLoader, scenarioNodesMap)
		
		val totalCosimSteps = Math.toIntExact(Math.round(rootScenario.stopTime/rootScenario.stepSize)+2)
		
		val HashMap<OutputPortInstance, Double>[] cosimSignals = newArrayOfSize(totalCosimSteps)
		
		var stepToConverge = 0
		
		while (stepToConverge < totalCosimSteps){
			logger.debug("String coupling iteration for step {} started.", stepToConverge)
			
			// run co-simulation up to step convergedStep and converge a new step.
			convergeStep(stepToConverge, units, fmuInstanceMap, childInputPorts, 
						childOutputPorts, rollbackOutputPorts, cosimSignals, rootScenario,
						nodes, scenarioNodesMap)
			
			logger.debug("String coupling iteration for step {} terminated.", stepToConverge)
			stepToConverge=stepToConverge+1		
		}
		
		Assert.isTrue(cosimSignals.get(totalCosimSteps-1) !== null)
		
		// Produce outputs
		printFinalOutputs(variantID, units, fmuInstanceMap, childInputPorts, 
						childOutputPorts, rollbackOutputPorts, cosimSignals, rootScenario,
						nodes)
		
		logger.debug("Outputs produced.")
		
		freeUnits(units, fmuInstanceMap)
		
		logger.debug("Cleanup complete.")
	}
	
	def convergeStep(int stepToConverge, List<UnitInstance> units, 
						Map<UnitInstance, IFmuInstance> fmuInstanceMap,
						Set<InputPortInstance> childInputPorts,
						List<OutputPortInstance> childOutputPorts, 
						Set<OutputPortInstance> rollbackOutputPorts, 
						HashMap<OutputPortInstance, Double>[] cosimSignals,
						RootCandidateScenario rootScenario,
						List<PrecendenceNode> nodes,
						Map<Scenario, List<PrecendenceNode>> scenarioNodesMap
	) {
		if (stepToConverge == 0){
			// Initialize cosim and record first outputs.
			
			// setup experiment
			setupExperiment(units, fmuInstanceMap)
			
			// Go into init mode
			intoInitMode(units, fmuInstanceMap)
			
			fixPointOutputs(childOutputPorts, fmuInstanceMap, rootScenario.maxInitIterations)
			
			// record first line for outputs that are to be used in the convergence
			val outVals = newHashMap() 
			getOutputSnapshot(units, fmuInstanceMap, outVals, rollbackOutputPorts)
			cosimSignals.set(0,outVals as HashMap<OutputPortInstance, Double>)
			logger.debug("Output line recorded for time {} and step {}.", 0.0, 0.0)
			
			exitInitMode(units, fmuInstanceMap)
			resetUnits(units, fmuInstanceMap, fmuLoader, scenarioNodesMap, hardReset)
		} else {
			val convergedStep = stepToConverge-1
			
			val maxCosims = rootScenario.maxInitIterations
			var cosimCounter = 0
			var hasConverged = false
			
			while (!hasConverged && cosimCounter<maxCosims){
				logger.debug("Cosim {} started.", cosimCounter)
				
				val time = initializeFMUsToStep(convergedStep, units, fmuInstanceMap, childInputPorts, childOutputPorts, 
									cosimSignals, rootScenario, nodes)
				
				hasConverged = runNewCosimStep(time, convergedStep, cosimCounter, units, 
												fmuInstanceMap,
												childInputPorts, childOutputPorts, rollbackOutputPorts,
												cosimSignals,
												rootScenario,
												nodes)
				
				resetUnits(units, fmuInstanceMap, fmuLoader, scenarioNodesMap, hardReset)
				
				cosimCounter = cosimCounter + 1
				
				logger.debug("Cosim {} ended.", cosimCounter)
			}
			
			if (cosimCounter >= maxCosims){
				logger.warn("Convergence of cosim step {} exceeded max iteration ({})", stepToConverge, maxCosims)
			}
		}
		logger.debug("Step {} converged.", stepToConverge)
	}
	
	def runNewCosimStep(double time, int convergedStep, int cosimCounter, 
						List<UnitInstance> units, 
						Map<UnitInstance, IFmuInstance> fmuInstanceMap,
						Set<InputPortInstance> childInputPorts,
						List<OutputPortInstance> childOutputPorts,
						Set<OutputPortInstance> rollbackOutputPorts,
						HashMap<OutputPortInstance, Double>[] cosimSignals,
						RootCandidateScenario rootScenario,
						List<PrecendenceNode> nodes) {
		val stepSize = rootScenario.stepSize
		
		val nonConvergedStep = convergedStep+1
		
		val replayValues = cosimSignals.get(nonConvergedStep)
		
		// Note that replayValues might be null.
		executeStep(rootScenario, time, stepSize, nodes, units, childInputPorts, fmuInstanceMap, replayValues)
		
		val newReplayValues = newHashMap() 
		getOutputSnapshot(units, fmuInstanceMap, newReplayValues, rollbackOutputPorts)
		
		// Store the new values for the next iteration
		cosimSignals.set(nonConvergedStep, newReplayValues)
		
		return MathUtils.converged(newReplayValues, replayValues, 1e-8)
	}
	
	def initializeFMUsToStep(int convergedStep, List<UnitInstance> units, 
						Map<UnitInstance, IFmuInstance> fmuInstanceMap,
						Set<InputPortInstance> childInputPorts,
						List<OutputPortInstance> childOutputPorts,
						HashMap<OutputPortInstance, Double>[] cosimSignals,
						RootCandidateScenario rootScenario,
						List<PrecendenceNode> nodes) {
		
		setupExperiment(units, fmuInstanceMap)
		intoInitMode(units, fmuInstanceMap)
		fixPointOutputs(childOutputPorts, fmuInstanceMap, rootScenario.maxInitIterations)
		exitInitMode(units, fmuInstanceMap)
		
		var i = 0
		var time = 0.0d
		val stepSize = rootScenario.stepSize
		
		while (i < convergedStep){
			val replayValues = cosimSignals.get(i+1)
			Assert.isTrue(replayValues !== null)
			executeStep(rootScenario, time, stepSize, nodes, units, childInputPorts, fmuInstanceMap, replayValues)
			time = time + stepSize
			i = i + 1
		}
		
		return time // Fmus are ready to run the convergedStep+1
	}
	
	def printFinalOutputs(String variantID, List<UnitInstance> units, 
						Map<UnitInstance, IFmuInstance> fmuInstanceMap,
						Set<InputPortInstance> childInputPorts,
						List<OutputPortInstance> childOutputPorts, 
						Set<OutputPortInstance> rollbackOutputPorts, 
						HashMap<OutputPortInstance, Double>[] cosimSignals,
						RootCandidateScenario rootScenario,
						List<PrecendenceNode> nodes) {
		// Since we stored only the output that needed to be replayed, in order to plot the final outputs we need to run a final co-sim.
		
		val outVals = newHashMap()
		
		setupExperiment(units, fmuInstanceMap)
		intoInitMode(units, fmuInstanceMap)
		
		fixPointOutputs(childOutputPorts, fmuInstanceMap, rootScenario.maxInitIterations)
		
		outProcessor.initialize(rootScenario, variantID)
		
		outProcessor.setOutputs(0, 0.0, getOutputSnapshot(units, fmuInstanceMap, outVals, null))
		
		exitInitMode(units, fmuInstanceMap)
		
		val stepSize = rootScenario.stepSize
		val totalCosimSteps = cosimSignals.length
		var i = 1
		var time = 0.0
		
		while (i < totalCosimSteps){
			val replayValues = cosimSignals.get(i)
			
			Assert.isTrue(replayValues !== null)
			executeStep(rootScenario, time, stepSize, nodes, units, childInputPorts, fmuInstanceMap, replayValues)
			
			outVals.clear()
			time = time + stepSize
			outProcessor.setOutputs(i, time, getOutputSnapshot(units, fmuInstanceMap, outVals, null))
			i = i + 1
		}
		
		val stopTime = rootScenario.stopTime
		Assert.isTrue(MathUtils.isApproximatelyEqual(time-stepSize, stopTime, stepSize*1e-3))
		
		outProcessor.terminate()
	}
	
	override close() throws IOException {
		fmuLoader.free()
	}
	
}