#include <wfr/scheduling/scheduler/sufferage/SufferageScheduler.h>
#include <wfr/common/wep/WorkflowExecutionPlan.h>
#include <wfr/common/wep/AbstractWepNode.h>
#include <wfr/common/wep/activities/AtomicWepNode.h>
#include <wfr/common/wfs/Activity.h>
#include <wfr/common/wfs/ActivityInstance.h>
#include <wfr/common/wfs/AtomicActivityImplementation.h>
#include <wfr/common/wfs/instances/AtomicActivityInstance.h>
#include <wfr/common/wfs/templates/AtomicActivityTemplate.h>
#include <wfr/common/resource/ResourceManager.h>
#include <wfr/enactment/EnactmentEngine.h>
#include <wfr/common/scheduling/SchedulingDecision.h>
#include <wfr/execution/ExecutionManager.h>
#include <wfr/WorkflowRuntime.h>
#include <wfr/wep/utils/WepSearch.h>
#include <iostream>
#include <set>
#include <queue>


namespace wfr {
namespace scheduling {




double SufferageScheduler::_getResourceEstimatedAvailableTime(resource::ComputingResource* resource) {
	if (resource->getFreeSlotCount() > 0) {
		return 0.0;
	} else {
		// Calculating the estimated available time for resources with more than 1 slot is not trivial
		// Here we are employing a heuristic that partitions the scheduled tasks into chunks the size of the maximal slot count,
		// takes the minimum time of each chunk and adds them together.
		unsigned i = 0;
		double eat = 0;
		double minTime = -1;
		auto& tasks = scheduledTasks[resource];
		for (auto t : tasks) {
			i++;
			if (minTime < 0) {
				minTime = t->getPrediction().getExecutionTime().getValue();
			} else {
				if (i > resource->getMaxSlotCount()) {
					eat += minTime;
					i = 0;
					minTime = t->getPrediction().getExecutionTime().getValue();
				} else {
					if (t->getPrediction().getExecutionTime().getValue() < minTime) {
						minTime = t->getPrediction().getExecutionTime().getValue();
					}
				}
			}
		}
		return eat;
	}
}


std::pair<wep::AtomicWepNode*, SufferageScheduler::ECTEntry> SufferageScheduler::_getNextTaskToSchedule(std::map<wep::AtomicWepNode*, ECTEntry>& estimatedTaskCompletionTimes) {
	wep::AtomicWepNode* selectedNode = nullptr;
	ECTEntry selectedEntry;
	bool firstIteration = true;
	for (auto n : estimatedTaskCompletionTimes) {
		if (firstIteration) {
			selectedNode = n.first;
			selectedEntry = n.second;
			firstIteration = false;
		} else {
			if (std::get<2>(n.second) > std::get<2>(selectedEntry)) {
				selectedNode = n.first;
				selectedEntry = n.second;
			}
		}
	}
	return std::make_pair(selectedNode, selectedEntry);
}


void SufferageScheduler::_notifyWEPUpdate(SchedulerBase::WepMapEntry& wepe, std::set<wep::AtomicWepNode*> readyToExecute, std::unique_lock<std::mutex>& schedulerLock) {
	while (!readyToExecute.empty()) {
		std::map<wep::AtomicWepNode*, ECTEntry> estimatedTaskCompletionTimes;
		auto resourceTypes = resourceManager.getComputingResourceTypes();
		// Iterate over all tasks
		for (auto n : readyToExecute) {
			resource::ComputingResource* selectedResource = nullptr;
			double selectedEarliestEstimatedCompletionTime = -1;
			wfs::AtomicActivityImplementation* selectedImpl = nullptr;
			prediction::ExecutionTimeResult selectedPred;
			double selectedEarliestEstimatedCompletionTime2 = -1;
			std::map<resource::ComputingResource*, double> estimatedCompletionTimes;
			// Iterate over all resource types
			for (auto restype : resourceTypes) {
				// Get fastest implementation and predicted time for given resource type
				wfs::AtomicActivityImplementation* fastestImpl = nullptr;
				prediction::ExecutionTimeResult fastestPred;
				// Check cache first
				bool cacheHit = false;
				auto aaType = n->activity.baseActivity.baseTemplate.as<wfs::AtomicActivityTemplate>().atomicActivityType.get();
				auto& innerMap = implementationCache[aaType];
				auto cacheEntry = innerMap.find(restype);
				if (cacheEntry != innerMap.end()) {
					fastestImpl = cacheEntry->second.first;
					fastestPred = cacheEntry->second.second;
					cacheHit = true;
				}
				if (!cacheHit) {
					// Cache miss
					auto& aat = n->activity.baseActivity.baseTemplate.as<wfs::AtomicActivityTemplate>();
					for (auto i : aat.atomicActivityType->getImplementations()) {
						auto execPred = predictionManager.predictExecutionTime(n->activity, *i, *restype);
						if (execPred.isValid() && (!fastestPred.isValid() || execPred < fastestPred)) {
							fastestPred = execPred;
							fastestImpl = i;
						}
					}
					// store into cache
					if (fastestPred.isValid()) {
						innerMap.insert({restype, {fastestImpl, fastestPred}});
					}
				}
				if (fastestPred.isValid()) {
					// Iterate over all resources and calculate earliest completion time
					for (auto res : restype->getResources()) {
							auto ect = _getResourceEstimatedAvailableTime(res) + fastestPred.executionTime.getValue();
							if (selectedEarliestEstimatedCompletionTime < 0 || ect < selectedEarliestEstimatedCompletionTime) {
								selectedEarliestEstimatedCompletionTime2 = selectedEarliestEstimatedCompletionTime;
								selectedEarliestEstimatedCompletionTime = ect;
								selectedResource = res;
								selectedImpl = fastestImpl;
								selectedPred = fastestPred;
							} else if (selectedEarliestEstimatedCompletionTime2 < 0 || ect < selectedEarliestEstimatedCompletionTime2) {
								selectedEarliestEstimatedCompletionTime2 = ect;
							}
					} // End: Iterate over all resources
				}
			} // End: Iterate over all resource types
			// Calculate sufferage
			double sufferage = 0;
			if (selectedEarliestEstimatedCompletionTime > 0 && selectedEarliestEstimatedCompletionTime2 > 0) {
				sufferage = selectedEarliestEstimatedCompletionTime2 - selectedEarliestEstimatedCompletionTime;
			}
			//std::cout << "selectedEarliestEstimatedCompletionTime " << n->activity.baseActivity.baseTemplate.activityName << ": " << selectedEarliestEstimatedCompletionTime << std::endl;
			//std::cout << "selectedEarliestEstimatedCompletionTime2 " << n->activity.baseActivity.baseTemplate.activityName << ": " << selectedEarliestEstimatedCompletionTime2 << std::endl;
			//std::cout << "Sufferage " << n->activity.baseActivity.baseTemplate.activityName << ": " << sufferage << std::endl;
			if (selectedResource) {
				estimatedTaskCompletionTimes.insert({n, std::make_tuple(selectedResource, selectedImpl, sufferage, selectedPred)});
			}
		} // End: Iterate over all tasks
		// Select next task to schedule
		auto selectedTask = _getNextTaskToSchedule(estimatedTaskCompletionTimes);
		// Schedule task
		if (selectedTask.first) {
			auto reservation = resourceManager.reserveComputingResource(std::get<0>(selectedTask.second), 1, false);
			auto sd = std::make_shared<SchedulingDecision>(&selectedTask.first->activity, std::get<1>(selectedTask.second), wepe.wep, reservation);
			sd->setPrediction(std::get<3>(selectedTask.second));
			selectedTask.first->activity.setInstanceStatus(wfs::ActivityInstanceStatus::Scheduled);
			scheduledTasks[sd->resource].push_back(sd);
			executeSchedulingDecision(wepe, sd, schedulerLock);
			readyToExecute.erase(selectedTask.first);
		}
	}
}

void SufferageScheduler::_notifySchedulingDecisionStatusUpdate(WepMapEntry& wepe, std::shared_ptr<SchedulingDecision> sd, std::unique_lock<std::mutex>& schedulerLock) {
	auto& innerMap = scheduledTasks[sd->resource];
	for (auto i = innerMap.begin(); i != innerMap.end(); i++) {
		if (*i == sd) {
			innerMap.erase(i);
			break;
		}
	}
}

} /* namespace scheduling */
} /* namespace wfr */