#include "pls/internal/profiling/dag_node.h"

#include <algorithm>

namespace pls::internal::profiling {

void dag_node::dag_compact() {
  while (next_node_ && next_node_->child_nodes_.empty()) {
    max_memory_ = std::max(max_memory_, next_node_->max_memory_);
    total_runtime_ += next_node_->total_runtime_;
    next_node_ = std::move(next_node_->next_node_);
  }
  if (next_node_) {
    next_node_->dag_compact();
  }

  for (auto &child : child_nodes_) {
    child.dag_compact();
  }
}

void dag_node::dag_print(std::ostream &stream, unsigned rank) {
  stream << node_print_id()
         << " [label=\"" << spawning_thread_id_ << "\n"
         << max_memory_ << " bytes\n"
         << m_to_d(total_runtime_) << " us\""
         << " ,rank=" << rank << "];" << std::endl;
  for (auto &child : child_nodes_) {
    child.dag_print(stream, rank + 1);
    stream << node_print_id() << " -> " << child.node_print_id() << ";" << std::endl;
  }
  if (next_node_) {
    next_node_->dag_print(stream, rank);
    stream << node_print_id() << " -> " << next_node_->node_print_id() << ";" << std::endl;
  }
}

unsigned dag_node::dag_max_memory() {
  unsigned max = max_memory_;
  if (next_node_) {
    max = std::max(max, next_node_->dag_max_memory());
  }
  for (auto &child : child_nodes_) {
    max = std::max(max, child.dag_max_memory());
  }
  return max;
}

unsigned long dag_node::dag_total_user_time() {
  unsigned long total_user_time = total_runtime_;
  if (next_node_) {
    total_user_time += next_node_->dag_total_user_time();
  }
  for (auto &child : child_nodes_) {
    total_user_time += child.dag_total_user_time();
  }
  return total_user_time;
}

unsigned long dag_node::dag_critical_path() {
  unsigned long critical_path = total_runtime_;
  if (next_node_) {
    critical_path += next_node_->dag_critical_path();
  }

  unsigned long slowest_child = 0;
  for (auto &child : child_nodes_) {
    slowest_child = std::max(slowest_child, child.dag_critical_path());
  }
  critical_path += slowest_child;

  return critical_path;
}

unsigned long dag_node::dag_depth() {
  unsigned long deepest_at_level = 1;
  if (next_node_) {
    deepest_at_level = std::max(deepest_at_level, next_node_->dag_depth());
  }

  unsigned long deepest_child = 0;
  for (auto &child : child_nodes_) {
    deepest_child = std::max(deepest_child, child.dag_depth());
  }

  return std::max(deepest_child + 1, deepest_at_level);
}

}