AnimTestbed/src/AnimGraph/AnimGraphBlendTree.h

//
// Created by martin on 17.03.24.
//

#ifndef ANIMTESTBED_ANIMGRAPHBLENDTREE_H
#define ANIMTESTBED_ANIMGRAPHBLENDTREE_H

#include <algorithm>
#include <iostream>

#include "AnimNode.h"

//
// AnimGraph (Runtime)
//
struct AnimGraphBlendTree : public AnimNode {
  std::vector<AnimNode*> m_nodes;
  std::vector<AnimNode*> m_eval_ordered_nodes;
  std::vector<std::vector<AnimGraphConnection> > m_node_input_connections;
  std::vector<std::vector<AnimGraphConnection> > m_node_output_connections;

  [[nodiscard]] const std::vector<AnimGraphConnection>&
  GetGraphInputConnections() const {
    return m_node_output_connections[1];
  }

  [[nodiscard]] const std::vector<AnimGraphConnection>&
  GetGraphOutputConnections() const {
    return m_node_input_connections[0];
  }

  std::vector<Pose*> m_pose_blocks;
  NodeDescriptorBase* m_node_descriptor = nullptr;
  char* m_input_buffer = nullptr;
  char* m_output_buffer = nullptr;
  char* m_connection_data_storage = nullptr;
  char* m_const_node_inputs = nullptr;

  std::vector<Socket>& GetGraphOutputs() {
    return m_node_descriptor->m_outputs;
  }
  std::vector<Socket>& GetGraphInputs() { return m_node_descriptor->m_inputs; }

  AnimDataAllocator m_anim_data_allocator;

  ~AnimGraphBlendTree() override { dealloc(); }

  void StartUpdateTick() { m_tick_number++; }

  // AnimNode overrides
  bool Init(AnimGraphContext& context) override;

  /// Determines which nodes in the BlendTree are active.
  ///
  /// Note: this does not use the provided input_connections, instead it marks
  /// all nodes directly connected to the BlendTree outputs as active and then
  /// propagates the node state throught the tree. For this each active node's
  /// AnimNode::MarkActiveInputs() gets called.
  void MarkActiveInputs(
      const std::vector<AnimGraphConnection>& input_connections) override;
  void CalcSyncTrack(
      const std::vector<AnimGraphConnection>& input_connections) override;
  void UpdateTime(float time_last, float time_now) override;
  void Evaluate(AnimGraphContext& context) override;

  void PropagateTimeToNodeInputs(const AnimNode* node);

  void dealloc() {
    for (size_t i = 0; i < m_pose_blocks.size(); i++) {
      m_pose_blocks[i]->m_local_matrices.vector::~vector();
    }
    m_pose_blocks.clear();

    m_node_input_connections.clear();
    m_node_output_connections.clear();

    delete[] m_input_buffer;
    delete[] m_output_buffer;
    delete[] m_connection_data_storage;
    delete[] m_const_node_inputs;

    for (int i = 0; i < m_nodes.size(); i++) {
      delete m_nodes[i];
    }
    m_nodes.clear();

    delete m_node_descriptor;
  }

  void UpdateOrderedNodes();
  void UpdateOrderedNodesRecursive(int node_index);
  bool CheckIsNodeActive(AnimNode* node) {
    return node->m_tick_number == m_tick_number;
  }

  void ResetNodeStates() {
    for (size_t i = 0, n = m_nodes.size(); i < n; i++) {
      m_nodes[i]->m_time_now = 0.f;
      m_nodes[i]->m_time_last = 0.f;
      m_nodes[i]->m_state = AnimNodeEvalState::Undefined;
    }
  }

  /** Sets the address that is used for the specified BlendTree input Socket.
   *
   * @tparam T Type of the Socket.
   * @param name Name of the Socket.
   * @param value_ptr Pointer where the input is fetched during evaluation.
   */
  template <typename T>
  void SetInput(const char* name, T* value_ptr) {
    m_node_descriptor->SetInput(name, value_ptr);

    std::vector<size_t> connected_node_indices;

    for (int i = 0; i < m_node_output_connections[1].size(); i++) {
      const AnimGraphConnection& graph_input_connection =
          m_node_output_connections[1][i];

      if (graph_input_connection.m_source_socket_name == name) {
        *graph_input_connection.m_socket.m_reference.ptr_ptr = value_ptr;
      }
    }
  }

  /** Sets the address that is used for the specified BlendTree output Socket.
   *
   * We update the pointer of the outputting node. We also have to ensure that
   * all usages of that output use the same pointer.
   *
   * @tparam T Type of the Socket.
   * @param name Name of the Socket.
   * @param value_ptr Pointer where the graph output is written to at the end of evaluation.
   */
  template <typename T>
  void SetOutput(const char* name, T* value_ptr) {
    const auto& graph_output_connection = std::find_if(
        GetGraphOutputConnections().begin(),
        GetGraphOutputConnections().end(),
        [&name](const AnimGraphConnection& connection) {
          return connection.m_target_socket_name == name;
        });

    if (graph_output_connection == GetGraphOutputConnections().end()) {
      std::cerr << "Error: could not find graph connection to output socket "
                << name << "." << std::endl;
      return;
    }

    if (graph_output_connection->m_source_node == m_nodes[1]
        && graph_output_connection->m_target_node == m_nodes[0]) {
      std::cerr << "Error: cannot set output for direct graph input to graph "
                   "output connections. Use GetOutptPtr for output instead!"
                << std::endl;

      return;
    }

    // Ensure that all connections that consume the output use the updated address.
    size_t output_node_index =
        GetAnimNodeIndex(graph_output_connection->m_source_node);

    const std::vector<AnimGraphConnection>& node_output_connections =
        m_node_output_connections[output_node_index];
    for (const AnimGraphConnection& connection : node_output_connections) {
      if (connection.m_source_socket_name
          == graph_output_connection->m_source_socket_name) {
        *connection.m_socket.m_reference.ptr_ptr = value_ptr;
      }
    }

    *graph_output_connection->m_socket.m_reference.ptr_ptr = value_ptr;

    // And additionally update the BlendTree's node descriptor:
    Socket* blend_tree_output_socket = m_node_descriptor->GetOutputSocket(name);
    assert(blend_tree_output_socket != nullptr);

    *blend_tree_output_socket->m_reference.ptr_ptr = value_ptr;
  }

  /** Returns the address that is used for the specified AnimGraph output Socket.
   *
   * This function is needed for connections that directly connect an AnimGraph
   * input Socket to an output Socket of the same AnimGraph.
   *
   * @tparam T Type of the Socket.
   * @param name Name of the Socket.
   * @return Address that is used for the specified AnimGraph output Socket.
   */
  template <typename T>
  T* GetOutputPtr(const char* name) {
    for (int i = 0; i < m_node_input_connections[0].size(); i++) {
      const AnimGraphConnection& graph_output_connection =
          m_node_input_connections[0][i];
      if (graph_output_connection.m_target_socket_name == name) {
        return static_cast<T*>(
            *graph_output_connection.m_socket.m_reference.ptr_ptr);
      }
    }

    return nullptr;
  }

  int GetAnimNodeIndex(const AnimNode* node) const {
    for (int i = 0; i < m_nodes.size(); i++) {
      if (m_nodes[i] == node) {
        return i;
      }
    }

    return -1;
  }
};

//
// BlendTreeSocketNode
//
struct BlendTreeSocketNode : public AnimNode {};

template <>
struct NodeDescriptor<BlendTreeSocketNode> : public NodeDescriptorBase {
  explicit NodeDescriptor(BlendTreeSocketNode* node_) {}
};

#endif  //ANIMTESTBED_ANIMGRAPHBLENDTREE_H