/*
    This file is part of Mitsuba, a physically based rendering system.

    Copyright (c) 2007-2014 by Wenzel Jakob and others.

    Mitsuba is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License Version 3
    as published by the Free Software Foundation.

    Mitsuba is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program. If not, see <http://www.gnu.org/licenses/>.
*/

#pragma once
#if !defined(__MITSUBA_CORE_SCHED_H_)
#define __MITSUBA_CORE_SCHED_H_

#include <mitsuba/core/serialization.h>
#include <mitsuba/core/lock.h>
#include <deque>

/**
 * Uncomment this to enable scheduling debug messages
 */
//#define DEBUG_SCHED 1

MTS_NAMESPACE_BEGIN

/**
 * \brief Abstract work unit -- represents a small amount of information
 * that encodes part of a larger processing task.
 *
 * Instead of the usual serialization function and unserialization
 * constructor, implementations of this class supply \ref load()
 * and \ref save() methods that can be used for essentially the
 * same purpose, but without requiring any memory allocations.
 *
 * \sa WorkProcessor
 * \ingroup libcore
 * \ingroup libpython
 */
class MTS_EXPORT_CORE WorkUnit : public Object {
public:
	/// Copy the content of another work unit of the same type
	virtual void set(const WorkUnit *workUnit) = 0;

	/// Fill the work unit with content acquired from a binary data stream
	virtual void load(Stream *stream) = 0;

	/// Serialize a work unit to a binary data stream
	virtual void save(Stream *stream) const = 0;

	/// Return a string representation
	virtual std::string toString() const = 0;

	MTS_DECLARE_CLASS()
protected:
	/// Virtual destructor
	virtual ~WorkUnit() { }
};

/**
 * \brief Abstract work result -- represents the result of a
 * processed \ref WorkUnit instance.
 *
 * Instead of the usual serialization function and unserialization
 * constructor, implementations of this class supply \ref load()
 * and \ref save() methods that can be used for essentially the
 * same purpose, but without requiring any memory allocations.
 *
 * \sa WorkProcessor
 * \ingroup libcore
 * \ingroup libpython
 */
class MTS_EXPORT_CORE WorkResult : public Object {
public:
	/// Fill the work result with content acquired from a binary data stream
	virtual void load(Stream *stream) = 0;

	/// Serialize a work result to a binary data stream
	virtual void save(Stream *stream) const = 0;

	/// Return a string representation
	virtual std::string toString() const = 0;

	MTS_DECLARE_CLASS()
protected:
	/// Virtual destructor
	virtual ~WorkResult() { }
};

/**
 * \brief Abstract work processor -- takes work units and turns them into
 *\ref WorkResult instances.
 *
 * When executing a parallel task using Mitsuba's scheduling system, the
 * actual work is done in an implementation of this interface.
 *
 * The class is serializable so that it can be sent over the network if
 * required. It is possible to keep local state in \ref WorkProcessor
 * instances (e.g. scratch space for computations), though anything not
 * returned in the form of a \ref WorkResult will eventually be lost.
 * Each \ref Worker (both locally and remotely) has its own \ref WorkProcessor,
 * and therefore no form of locking is required within instances of this class.
 *
 * \sa WorkUnit
 * \sa WorkResult
 * \sa ParallelProcess
 * \sa Scheduler
 * \ingroup libcore
 * \ingroup libpython
 */
class MTS_EXPORT_CORE WorkProcessor : public SerializableObject {
	friend class Scheduler;
public:
	/// Create a work unit of the proper type and size.
	virtual ref<WorkUnit> createWorkUnit() const = 0;

	/// Create a work result of the proper type and size
	virtual ref<WorkResult> createWorkResult() const = 0;

	/**
	 * \brief Create a copy of this work processor instance.
	 *
	 * \remark In practice, before the cloned work processor
	 * is actually used, its \ref prepare() method will be called.
	 * Therefore, any state that is initialized in \ref prepeare()
	 * does not have to be copied.
	 */
	virtual ref<WorkProcessor> clone() const = 0;

	/**
	 * \brief Called once before processing starts.
	 *
	 * This is useful for allocating scratch space or resolving references
	 * to resource objects. Lengthy computations should be performed in
	 * process() instead of here, since this this method will be called
	 * while the central scheduler lock is held. A thrown exception will
	 * lead to the termination of the parallel process.
	 */
	virtual void prepare() = 0;

	/**
	 * \brief Process a work unit and store the computed results.
	 *
	 * The <tt>active</tt> parameter can be used to signal a premature
	 * stop of the execution flow. In this case, the work result is allowed
	 * to be undefined (it will simply be ignored). A thrown exception will
	 * lead to the termination of the parallel process.
	 */
	virtual void process(const WorkUnit *workUnit, WorkResult *workResult,
		const bool &stop) = 0;

	MTS_DECLARE_CLASS()
protected:
	/// Virtual destructor
	virtual ~WorkProcessor() { }
	/// Protected constructors
	inline WorkProcessor() { }
	inline WorkProcessor(Stream *stream, InstanceManager *manager)
		: SerializableObject(stream, manager) { }

	/**
	 * \brief Look up a named resource, which has been bound to
	 * the associated parallel process.
	 *
	 * Throws an exception if the resource is not known / bound.
	 */
	SerializableObject *getResource(const std::string &name);
protected:
	std::map<std::string, SerializableObject *> m_resources;
};

/**
 * \brief Abstract parallelizable task.
 *
 * Instances of this class model a larger piece of work that can be split
 * into independent `units' and subsequently farmed out over a cluster or
 * processed locally. After the work units have been completed, the results
 * are pieced back together to a solution of the original large-scale problem.
 *
 * This class implements the core logic running on the central scheduling
 * server, i.e. the part that is responsible for generating work units and
 * accepting their results. The module that performs the actual computation
 * is an instance of \ref WorkProcessor, which is also specified here.
 * Finally, the this class references `resources', which denote
 * chunks of globally shared read-only data required during execution.
 *
 * \ingroup libcore
 * \ingroup libpython
 */
class MTS_EXPORT_CORE ParallelProcess : public Object {
	friend class Scheduler;
public:
	/// Binding from local resource names to global resource IDs
	typedef std::map<std::string, int> ResourceBindings;

	/// Return codes used by generateWork() and getReturnStatus()
	enum EStatus {
		EUnknown, ///< Unknown return status
		EPause,   ///< Temporarily, no work units can be created
		ESuccess, ///< The process finished / a piece of work was generated
		EFailure  ///< The process failed / no more work is available
	};

	/**
	 * \brief Generate a piece of work.
	 *
	 * Takes a pre-allocated \ref WorkUnit instance of
	 * the appropriate sub-type and size (as specified by
	 * \ref ParallelProcess::getWorkUnitName()) and
	 * fills it with the appropriate content. Returns ESuccess
	 * on success and EFailure or EPause when no more work is
	 * left -- in that case, the work unit will be ignored and
	 * the process completed (\ref EFailure) or temporarily
	 * paused (\ref EPause). When \ref EPause was used,
	 * resubmission via \ref Scheduler::schedule() will
	 * be required once more work is available. In some cases, it
	 * is useful to distribute 'nearby' pieces of work to the same
	 * processor -- the \c worker parameter can be used to
	 * implement this.
	 * This function should run as quickly as possible, since it
	 * will be executed while the scheduler mutex is held. A
	 * thrown exception will lead to the termination of the
	 * parallel process.
	 *
	 * \param unit   Work unit data structure to be filled
	 * \param worker ID of the worker executing this function
	 */
	virtual EStatus generateWork(WorkUnit *unit, int worker) = 0;

	/**
	 * \brief Called whenever a work unit has been completed.
	 *
	 * Note that this function may concurrently be executed by
	 * multiple threads. Also, processing of work results will
	 * generally be out of order with respect to the creation
	 * in \ref generateWork().
	 *
	 * When a work unit is only partially completed due to
	 * a call to \ref Scheduler::cancel(), the second
	 * parameter is set to true.
	 * A thrown exception will lead to the termination of
	 * the parallel process.
	 *
	 * \param result Work result to be processed
	 * \param cancelled Was the associated work unit not fully completed
	 */
	virtual void processResult(const WorkResult *result,
		bool cancelled) = 0;

	/**
	 * \brief Called when the parallel process is canceled by
	 * \ref Scheduler::cancel().
	 *
	 * The default implementation does nothing.
	 */
	virtual void handleCancellation();

	/**
	 * \brief Query the return status of a process after its
	 * execution has finished.
	 *
	 * Returns one of \ref Success, \ref Failure or \ref Unknown.
	 * (\ref EUnknown means that the process is either still running
	 * or has never been scheduled).
	 */
	inline EStatus getReturnStatus() const { return m_returnStatus; }

	/**
	 * \brief Create an instance of the algorithm responsible
	 * for executing the work units of this parallel process.
	 */
	virtual ref<WorkProcessor> createWorkProcessor() const = 0;

	/**
	 * \brief Bind a resource to this parallel process.
	 *
	 * Takes a resource ID as given by the scheduler and associates it
	 * with a name. This name can later be used by the work processor
	 * to access the resource data.
	 *
	 * \param name Process-specific name of the resource
	 * \param id Resource ID as returned by \ref Scheduler::registerResource()
	 * \sa WorkProcessor::getResource
	 */
	virtual void bindResource(const std::string &name, int id);

	/**
	 * \brief Is this process strictly local?
	 *
	 * If a process is marked as local, it shouldn't be distributed
	 * to remote processing nodes. The default implementation
	 * returns false.
	 */
	virtual bool isLocal() const;

	/**
	 * \brief Return the log level for events associated with this process.
	 *
	 * By default, this is set to EDebug
	 */
	inline ELogLevel getLogLevel() const { return m_logLevel; }

	/**
	 * \brief Return a list of all bound resources
	 */
	inline const ResourceBindings &getResourceBindings() const { return m_bindings; }

	/**
	 * \brief Return a list of plugins required by this parallel process.
	 *
	 * This is required so that remote machines can load the plugins before
	 * they accept work from this process. The default implementation just
	 * returns all plugins that are loaded in the current application.
	 */
	virtual std::vector<std::string> getRequiredPlugins();

	MTS_DECLARE_CLASS()
protected:
	/// Protected constructor
	inline ParallelProcess() : m_returnStatus(EUnknown),
		m_logLevel(EDebug) { }
	/// Virtual destructor
	virtual ~ParallelProcess() { }
protected:
	ResourceBindings m_bindings;
	EStatus m_returnStatus;
	ELogLevel m_logLevel;
};

class Worker;

/**
 * \brief Centralized task scheduler implementation.
 *
 * Accepts parallelizable jobs and distributes their computational load
 * both locally and remotely. This is done by associating different types
 * of \ref Worker instances with the scheduler. These try to acquire work
 * units from the scheduler, which are then executed on the current machine
 * or sent to remote nodes over a network connection.
 *
 * \ingroup libcore
 * \ingroup libpython
 */
class MTS_EXPORT_CORE Scheduler : public Object {
	friend class Worker;
public:
	/**
	 * \brief Schedule a parallelizable process for execution.
	 *
	 * If the scheduler is currently running and idle, its execution
	 * will begin immediately. Returns \c false if the process
	 * is already scheduled and has not yet terminated and \c true
	 * in any other case.
	 */
	bool schedule(ParallelProcess *process);

	/**
	 * \brief Block until the process has successfully been completed
	 * or canceled prematurely.
	 *
	 * Returns false if the process does not exist or has already
	 * finished by the time \ref wait() is invoked.
	 */
	bool wait(const ParallelProcess *process);

	/**
	 * \brief Cancel the execution of a parallelizable process.
	 *
	 * Upon return, no more work from this process is running.
	 * Returns false if the process does not exist (anymore).
	 */
	inline bool cancel(ParallelProcess *proc) {
		return cancel(proc, false);
	}

	/**
	 * \brief Register a serializable resource with the scheduler.
	 *
	 * A resource should be thought of as a constant state that is shared
	 * amongst all processing nodes. Resources can be reused by
	 * subsequent parallel processes, and consequently do not have to be
	 * re-transmitted over the network. Returns a resource ID, which can be
	 * used to reference the associated data.
	 */
	int registerResource(SerializableObject *resource);

	/**
	 * \brief Register a \a multiple resource with the scheduler.
	 *
	 * \a Multi means that in comparison to the previous method, a separate
	 * instance is provided  for every core. An example where this is useful
	 * is to distribute random generator state when performing parallel
	 * Monte Carlo simulations. \c resources must be a vector whose
	 * length is equal to \ref getCoreCount().
	 */
	int registerMultiResource(std::vector<SerializableObject *> &resources);

	/**
	 * \brief Increase the reference count of a previously registered resource.
	 *
	 * The resource must be unregistered an additional time after calling
	 * this function.
	 *
	 * \sa unregisterResource
	 */
	void retainResource(int resourceID);

	/**
	 * \brief Unregister a resource from the scheduler
	 *
	 * Note that the resource's won't be removed until all processes using
	 * it have terminated)
	 *
	 * \return \c false if the resource could not be found
	 */
	bool unregisterResource(int id);

	/**
	 * \brief Return the ID of a registered resource
	 *
	 * Throws an exception if the resource cannot be found.
	 */
	int getResourceID(const SerializableObject *resource) const;

	/// Register a worker with the scheduler
	void registerWorker(Worker *processor);

	/// Unregister a worker from the scheduler
	void unregisterWorker(Worker *processor);

	/// Get the number of workers
	size_t getWorkerCount() const;

	/// Get the number of local workers
	size_t getLocalWorkerCount() const;

	/// Retrieve one of the workers by index
	Worker *getWorker(int index);

	/// Start all workers and begin execution of any scheduled processes
	void start();

	/**
	 * \brief Pause the distribution of work units and shut down all
	 * running workers.
	 *
	 * Any currently scheduled work units are still completed.
	 * Processing can be resumed via \ref start().
	 */
	void pause();

	/// Cancel all running processes and free memory used by resources
	void stop();

	/// Return the total number of cores exposed through this scheduler
	size_t getCoreCount() const;

	/// Does the scheduler have one or more local workers?
	bool hasLocalWorkers() const;

	/// Does the scheduler have one or more remote workers?
	bool hasRemoteWorkers() const;

	/// Return a pointer to the scheduler of this process
	inline static Scheduler *getInstance() { return m_scheduler; }

	/// Has the scheduler been started?
	inline bool isRunning() const { return m_running; }

	/// Is the scheduler currently executing work?
	bool isBusy() const;

	/// Initialize the scheduler of this process -- called once in main()
	static void staticInitialization();

	/// Free the memory taken by staticInitialization()
	static void staticShutdown();

	MTS_DECLARE_CLASS()
public:
	// Public, but shouldn't be part of the documentation
	/// \cond
	struct ProcessRecord {
		/* Unique ID value assigned to this process */
		int id;
		/* Current number of in-flight work units */
		int inflight;
		/* Is the parallel process still generating work */
		bool morework;
		/* Was the process cancelled using \c cancel()?*/
		bool cancelled;
		/* Is the process currently in the queue? */
		bool active;
		/* Signaled every time a work unit arrives */
		ref<ConditionVariable> cond;
		/* Set when the process is done/canceled */
		ref<WaitFlag> done;
		/* Log level for events associated with this process */
		ELogLevel logLevel;

		inline ProcessRecord(int id, ELogLevel logLevel, Mutex *mutex)
		 : id(id), inflight(0), morework(true), cancelled(false),
		 	active(true), logLevel(logLevel) {
			cond = new ConditionVariable(mutex);
			done = new WaitFlag();
		}
	};

	/**
	 * Data structure, which contains a piece of work
	 * as well as the information required to either execute
	 * it locally or submit it to a processing node over the
	 * network.
	 */
	struct Item {
		int id;
		int workerIndex;
		int coreOffset;
		ParallelProcess *proc;
		ProcessRecord *rec;
		ref<WorkProcessor> wp;
		ref<WorkUnit> workUnit;
		ref<WorkResult> workResult;
		bool stop;

		inline Item() : id(-1), workerIndex(-1), coreOffset(-1),
			proc(NULL), rec(NULL), stop(false) {
		}

		std::string toString() const;
	};

	struct ResourceRecord {
		std::vector<SerializableObject *> resources;
		ref<MemoryStream> stream;
		int refCount;
		bool multi;

		inline ResourceRecord(SerializableObject *resource)
		 : resources(1), refCount(1), multi(false) {
			resources[0] = resource;
		}

		inline ResourceRecord(std::vector<SerializableObject *> resources)
		 : resources(resources), refCount(1), multi(true) {
		}
	};

	/// A list of status codes returned by acquireWork()
	enum EStatus {
		/// Sucessfully acquired a work unit
		EOK,
		/// There is currently no work (and onlyTry was set to true)
		ENone,
		/// The scheduler is shutting down
		EStop
	};
	/// \endcond

	/// Look up a resource by ID & core index
	SerializableObject *getResource(int id, int coreIndex = -1);

	/// Return a resource in the form of a binary data stream
	const MemoryStream *getResourceStream(int id);

	/**
	 * \brief Test whether this is a multi-resource,
	 * i.e. different for every core.
	 */
	bool isMultiResource(int id) const;
protected:
	/// Protected constructor
	Scheduler();

	/// Virtual destructor
	virtual ~Scheduler();

	/**
	 * Acquire a piece of work from the scheduler -- internally
	 * used by the different worker implementations.
	 */
	EStatus acquireWork(Item &item, bool local, bool onlyTry, bool keepLock);

	/// Release the main scheduler lock -- internally used by the remote worker
	inline void releaseLock() { m_mutex->unlock(); }

	inline void releaseWork(Item &item) {
		ProcessRecord *rec = item.rec;
		try {
			item.proc->processResult(item.workResult, item.stop);
		} catch (const std::exception &ex) {
			Log(EWarn, "Caught an exception - canceling process %i: %s",
				item.id, ex.what());
			cancel(item.proc, true);
			return;
		}
		LockGuard lock(m_mutex);
		--rec->inflight;
		rec->cond->signal();
		if (rec->inflight == 0 && !rec->morework && !item.stop)
			signalProcessTermination(item.proc, item.rec);
	}

	/**
	 * Cancel the execution of a parallelizable process. Upon
	 * return, no more work from this process is running. When
	 * the second parameter is set to true, the number of in-flight
	 * work units for this process is reduced by one.
	 */
	bool cancel(ParallelProcess *proc, bool reduceInflight);

	/**
	 * Internally used to prepare a Scheduler::Item structure
	 * when only the process ID is known.
	 */
	inline void setProcessByID(Item &item, int id) {
		LockGuard lock(m_mutex);
		ParallelProcess *proc = m_idToProcess[id];
		if (proc == NULL) {
			Log(EError, "Process %i is not locally known!", id);
		};
		item.proc = proc;
		item.id = id;
		item.rec = m_processes[proc];
		item.wp = proc->createWorkProcessor();
		const ParallelProcess::ResourceBindings &bindings = item.proc->getResourceBindings();
		for (ParallelProcess::ResourceBindings::const_iterator it = bindings.begin();
			it != bindings.end(); ++it)
			item.wp->m_resources[(*it).first] = m_scheduler->getResource((*it).second, item.coreOffset);
		try {
			item.wp->prepare();
			item.workUnit = item.wp->createWorkUnit();
			item.workResult = item.wp->createWorkResult();
		} catch (std::exception &) {
			throw;
		}
	}

	/// Announces the termination of a process
	void signalProcessTermination(ParallelProcess *proc, ProcessRecord *rec);
private:
	/// Global scheduler instance
	static ref<Scheduler> m_scheduler;
	/// Mutex, which protects local data structures
	mutable ref<Mutex> m_mutex;
	/// CV used to signal the availability of work
	ref<ConditionVariable> m_workAvailable;
	/// Scheduled processes in FIFO order
	std::deque<int> m_localQueue, m_remoteQueue;
	/// Set of all currently scheduled processes
	std::map<const ParallelProcess *, ProcessRecord *> m_processes;
	/// Maps process IDs to processes
	std::map<int, ParallelProcess *> m_idToProcess;
	/// List of shared resources
	std::map<int, ResourceRecord *> m_resources;
	/// List of all active workers
	std::vector<Worker *> m_workers;
	int m_resourceCounter, m_processCounter;
	bool m_running;
};

/**
 * \brief Base class of all worker implementations
 * \ingroup libpython
 */
class MTS_EXPORT_CORE Worker : public Thread {
	friend class Scheduler;
public:
	/// Return the number of cores exposed by this worker
	inline size_t getCoreCount() const { return m_coreCount; }
	/// Is this a remote worker?
	inline bool isRemoteWorker() const { return m_isRemote; };

	MTS_DECLARE_CLASS()
protected:
	/// Virtual destructor
	virtual ~Worker() { }

	/// Protected constructor
	Worker(const std::string &name);

	/* Decrement reference counts to any referenced objects */
	virtual void clear();

	/// Used internally by the scheduler
	virtual void start(Scheduler *scheduler,
		int workerIndex, int coreOffset);

	/**
	 * \brief Called to inform a worker that a resource is no longer in use.
	 *
	 * The remote worker uses this to notify the machine on the other end that
	 * the memory used by this resource can now be released.
	 */
	virtual void signalResourceExpiration(int id) = 0;

	/**
	 * \brief Called to inform a worker that a process has been cancelled
	 *
	 * Guaranteed to be called while the Scheduler's main lock is held.
	 */
	virtual void signalProcessCancellation(int id) = 0;

	/**
	 * \brief Called to inform a worker that a process has successfully been
	 * completed and any associated resources can be freed.
	 */
	virtual void signalProcessTermination(int id) = 0;

	/* Inline functions to access protected members of Scheduler */
	inline Scheduler::EStatus acquireWork(bool local,
			bool onlyTry = false, bool keepLock = false) {
		return m_scheduler->acquireWork(m_schedItem, local, onlyTry, keepLock);
	}

	void releaseSchedulerLock() {
		return m_scheduler->releaseLock();
	}

	/// Release a processed work unit
	inline void releaseWork(Scheduler::Item &item) {
		m_scheduler->releaseWork(item);
	}

	/// Initialize the m_schedItem data structure when only the process ID is known
	void setProcessByID(Scheduler::Item &item, int id) {
		return m_scheduler->setProcessByID(item, id);
	}

	/**
	 * Cancel the currently scheduled parallel process and possibly
	 * reduce the number of in-flight work units
	 * Returns false if the process does not exist (anymore).
	 */
	inline void cancel(bool reduceInflight) {
		m_scheduler->cancel(m_schedItem.proc, reduceInflight);
	}
protected:
	Scheduler *m_scheduler;
	Scheduler::Item m_schedItem;
	size_t m_coreCount;
	bool m_isRemote;
};

/**
 * \brief Acquires work from the scheduler and executes
 * it locally.
 *
 * \ingroup libcore
 * \ingroup libpython
 */
class MTS_EXPORT_CORE LocalWorker : public Worker {
public:
	/**
	 * \brief Create a new local worker thread
	 *
	 * \param coreID
	 *   When an CPU core ID (>=0) is specified here, the worker
	 *   thread will attempt to register core affinity with the
	 *   operating system. Passing -1 disables this.
	 *
	 * \param name
	 *   An identifying string for this thread
	 *
	 * \param priority
	 *   The desired thread priority (not supported on some
	 *   operating systems)
	 */
	LocalWorker(int coreID, const std::string &name,
		Thread::EThreadPriority priority = Thread::ENormalPriority);

	MTS_DECLARE_CLASS()
protected:
	/// Virtual destructor
	virtual ~LocalWorker();
	/* Worker implementation */
	virtual void run();
	virtual void signalResourceExpiration(int id);
	virtual void signalProcessCancellation(int id);
	virtual void signalProcessTermination(int id);
};

/**
 * \brief Dummy work unit without contents
 * \ingroup libcore
 */
class MTS_EXPORT_CORE DummyWorkUnit : public WorkUnit {
public:
	void set(const WorkUnit *workUnit);
	void load(Stream *stream);
	void save(Stream *stream) const;
	std::string toString() const;

	MTS_DECLARE_CLASS()
protected:
	virtual ~DummyWorkUnit() { }
};

MTS_NAMESPACE_END

#endif /* __MITSUBA_CORE_SCHED_H_ */