AnimTestbed/src/AnimGraph/AnimGraphBlendTree.h

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

#ifndef ANIMTESTBED_ANIMGRAPHBLENDTREE_H
#define ANIMTESTBED_ANIMGRAPHBLENDTREE_H

#include "AnimNode.h"

//
// AnimGraph (Runtime)
//
struct AnimGraphBlendTree : public AnimNode {
  AnimData m_local_transforms;

  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;
  std::vector<AnimData*> m_animdata_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_inputs; }
  std::vector<Socket>& getGraphInputs() { return m_node_descriptor->m_outputs; }

  AnimDataAllocator m_anim_data_allocator;

  ~AnimGraphBlendTree() override { dealloc(); }

  void StartUpdateTick() { m_tick_number++; }

  // AnimNode overrides
  bool Init(AnimGraphContext& context) override;
  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_animdata_blocks.size(); i++) {
      m_animdata_blocks[i]->m_local_matrices.vector::~vector();
    }
    m_animdata_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;
    }
  }

  Socket* GetInputSocket(const std::string& name);
  Socket* GetOutputSocket(const std::string& name);

  const Socket* GetInputSocket(const std::string& name) const;
  const Socket* GetOutputSocket(const std::string& name) const;

  /** Sets the address that is used for the specified AnimGraph 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->SetOutput(name, value_ptr);

    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.m_name == name) {
        *graph_input_connection.m_target_socket.m_reference.ptr_ptr = value_ptr;
      }
    }
  }

  /** Sets the address that is used for the specified AnimGraph output Socket.
   *
   * @tparam T Type of the Socket.
   * @param name Name of the Socket.
   * @param value_ptr Pointer where the graph output output is written to at the end of evaluation.
   */
  template <typename T>
  void SetOutput(const char* name, T* value_ptr) {
    m_node_descriptor->SetInput(name, value_ptr);

    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.m_name == name) {
        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;
        }

        *graph_output_connection.m_source_socket.m_reference.ptr_ptr =
            value_ptr;

        // Make sure all other output connections of this pin use the same output pointer
        int source_node_index =
            GetAnimNodeIndex(graph_output_connection.m_source_node);
        for (int j = 0; j < m_node_output_connections[source_node_index].size();
             j++) {
          const AnimGraphConnection& source_output_connection =
              m_node_output_connections[source_node_index][j];
          if (source_output_connection.m_target_node == m_nodes[0]) {
            continue;
          }

          if (source_output_connection.m_source_socket.m_name
              == graph_output_connection.m_source_socket.m_name) {
            *source_output_connection.m_target_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.m_name == name) {
        return static_cast<float*>(
            *graph_output_connection.m_source_socket.m_reference.ptr_ptr);
      }
    }

    return nullptr;
  }

  void* GetInputPtr(const std::string& name) const {
    const Socket* input_socket = GetInputSocket(name);
    if (input_socket != nullptr) {
      return input_socket->m_reference.ptr;
    }

    return nullptr;
  }

  void* GetOutputPtr(const std::string& name) const {
    const Socket* input_socket = GetOutputSocket(name);
    if (input_socket != nullptr) {
      return input_socket->m_reference.ptr;
    }

    return nullptr;
  }

  int getNodeEvalOrderIndex(const AnimNode* node) {
    for (size_t i = 0, n = m_eval_ordered_nodes.size(); i < n; i++) {
      if (m_eval_ordered_nodes[i] == node) {
        return i;
      }
    }

    return -1;
  }
  const AnimNode* getAnimNodeForInput(
      size_t node_index,
      const std::string& input_name) const {
    assert(node_index < m_nodes.size());

    const std::vector<AnimGraphConnection>& input_connection =
        m_node_input_connections[node_index];
    for (size_t i = 0, n = input_connection.size(); i < n; i++) {
      if (input_connection[i].m_target_socket.m_name == input_name) {
        return input_connection[i].m_source_node;
      }
    }

    return nullptr;
  }

  AnimNode* getAnimNode(const char* name) {
    for (size_t i = 0; i < m_nodes.size(); i++) {
      if (m_nodes[i]->m_name == name) {
        return m_nodes[i];
      }
    }

    return nullptr;
  }

  size_t GetAnimNodeIndex(const AnimNode* node) const {
    for (size_t 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 {
  NodeDescriptor(BlendTreeSocketNode* node_) {}
};

#endif  //ANIMTESTBED_ANIMGRAPHBLENDTREE_H