abstract_task.h

// NIST-developed software is provided by NIST as a public service. You may use, copy and distribute copies of the
// software in any medium, provided that you keep intact this entire notice. You may improve, modify and create
// derivative works of the software or any portion of the software, and you may copy and distribute such modifications
// or works. Modified works should carry a notice stating that you changed the software and should note the date and
// nature of any such change. Please explicitly acknowledge the National Institute of Standards and Technology as the
// source of the software. NIST-developed software is expressly provided "AS IS." NIST MAKES NO WARRANTY OF ANY KIND,
// EXPRESS, IMPLIED, IN FACT OR ARISING BY OPERATION OF LAW, INCLUDING, WITHOUT LIMITATION, THE IMPLIED WARRANTY OF
// MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, NON-INFRINGEMENT AND DATA ACCURACY. NIST NEITHER REPRESENTS NOR
// WARRANTS THAT THE OPERATION OF THE SOFTWARE WILL BE UNINTERRUPTED OR ERROR-FREE, OR THAT ANY DEFECTS WILL BE
// CORRECTED. NIST DOES NOT WARRANT OR MAKE ANY REPRESENTATIONS REGARDING THE USE OF THE SOFTWARE OR THE RESULTS
// THEREOF, INCLUDING BUT NOT LIMITED TO THE CORRECTNESS, ACCURACY, RELIABILITY, OR USEFULNESS OF THE SOFTWARE. You
// are solely responsible for determining the appropriateness of using and distributing the software and you assume
// all risks associated with its use, including but not limited to the risks and costs of program errors, compliance
// with applicable laws, damage to or loss of data, programs or equipment, and the unavailability or interruption of 
// operation. This software is not intended to be used in any situation where a failure could cause risk of injury or
// damage to property. The software developed by NIST employees is not subject to copyright protection within the
// United States.


#ifndef HEDGEHOG_TASK_H
#define HEDGEHOG_TASK_H

#include "../behavior/io/multi_receivers.h"
#include "../behavior/io/sender.h"
#include "../behavior/execute.h"
#include "memory_manager/abstract_memory_manager.h"
#include "../core/node/core_node.h"
#include "../core/defaults/core_default_task.h"

namespace hh {
template<class TaskOutput, class ...TaskInputs>
class AbstractTask :
    public behavior::MultiReceivers<TaskInputs...>,
    public behavior::Sender<TaskOutput>,
    public virtual behavior::Node,
    public behavior::Execute<TaskInputs> ... {

  static_assert(traits::isUnique<TaskInputs...>, "A Task can't accept multiple inputs with the same type.");
  static_assert(sizeof... (TaskInputs) >= 1, "A node need to have one output type and at least one output type.");

  std::shared_ptr<core::CoreTask<TaskOutput, TaskInputs...>> taskCore_ = nullptr; 
  std::shared_ptr<AbstractMemoryManager<TaskOutput>> mm_ = nullptr; 

 public:
  AbstractTask(){
    taskCore_ = std::make_shared<core::CoreDefaultTask<TaskOutput, TaskInputs...>>("Task",
                                                                                   1,
                                                                                   core::NodeType::Task,
                                                                                   this,
                                                                                   false);
  }

  explicit AbstractTask(std::string_view const &name) {
    taskCore_ =
        std::make_shared<core::CoreDefaultTask<TaskOutput, TaskInputs...>>(name, 1, core::NodeType::Task, this, false);
  }

  explicit AbstractTask(std::string_view const &name, size_t numberThreads) {
    taskCore_ =
        std::make_shared<core::CoreDefaultTask<TaskOutput, TaskInputs...>>(name,
                                                                           numberThreads,
                                                                           core::NodeType::Task,
                                                                           this,
                                                                           false);
  }

  explicit AbstractTask(std::string_view const &name, size_t numberThreads, bool automaticStart) {
    taskCore_ =
        std::make_shared<core::CoreDefaultTask<TaskOutput, TaskInputs...>>
            (name, numberThreads, core::NodeType::Task, this, automaticStart);
  }

  AbstractTask(std::string_view const name, size_t numberThreads, core::NodeType nodeType, bool automaticStart) {
    taskCore_ = std::make_shared<core::CoreDefaultTask<TaskOutput, TaskInputs...>>
        (name, numberThreads, nodeType, this, automaticStart);
  }

  explicit AbstractTask(AbstractTask<TaskOutput, TaskInputs ...> *rhs) {
    taskCore_ = std::make_shared<core::CoreDefaultTask<TaskOutput, TaskInputs...>>
        (rhs->name(),
         rhs->numberThreads(),
         rhs->nodeType(),
         this,
         rhs->automaticStart());
  }

  virtual ~AbstractTask() = default;

  void addResult(std::shared_ptr<TaskOutput> output) { this->taskCore_->sendAndNotify(output); }

  std::string_view name() { return this->taskCore_->name(); }

  size_t numberThreads() { return this->taskCore_->numberThreads(); }

  bool automaticStart() { return this->taskCore_->automaticStart(); }

  core::NodeType nodeType() { return this->taskCore_->type(); }

  int deviceId() { return this->taskCore_->deviceId(); }

  int graphId() { return this->taskCore_->graphId(); }

  std::shared_ptr<core::CoreNode> core() final { return taskCore_; }

  std::shared_ptr<AbstractMemoryManager<TaskOutput>> const & memoryManager() const { return mm_; }

  virtual std::shared_ptr<AbstractTask<TaskOutput, TaskInputs...>> copy() { return nullptr; }

  virtual void initialize() {}

  virtual void shutdown() {}

  void connectMemoryManager(std::shared_ptr<AbstractMemoryManager<TaskOutput>> mm) {
    static_assert(
        traits::is_managed_memory_v<TaskOutput>,
        "The type given to the memory manager should inherit \"MemoryData\", and be default constructible!");
    mm_ = mm;
  }

  std::shared_ptr<TaskOutput> getManagedMemory() {
    if (mm_ == nullptr) {
      std::cerr
          << "For the task:\"" << this->name()
          << "\"To get managed memory, you need first to connect a memory manager to the task via "
             "\"connectMemoryManager()\""
          << std::endl;
      exit(42);
    }

    auto start = std::chrono::high_resolution_clock::now();
    this->taskCore_->nvtxProfiler()->startRangeWaitingForMemory();
    auto data = mm_->getManagedMemory();
    this->taskCore_->nvtxProfiler()->endRangeWaitingForMem();
    auto finish = std::chrono::high_resolution_clock::now();

    this->core()->incrementWaitForMemoryDuration(std::chrono::duration_cast<std::chrono::microseconds>(finish - start));

    return data;
  }

  bool canTerminate() override { return !this->taskCore_->hasNotifierConnected() && this->taskCore_->receiversEmpty(); }
};
}
#endif //HEDGEHOG_TASK_H