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

#include "AnimGraphBlendTree.h"

#include <algorithm>
#include <cstring>

bool AnimGraphBlendTree::Init(AnimGraphContext& context) {
  for (size_t i = 2; i < m_nodes.size(); i++) {
    if (!m_nodes[i]->Init(context)) {
      return false;
    }
  }

  for (size_t i = 0; i < m_animdata_blocks.size(); i++) {
    int num_soa_joints = context.m_skeleton->num_soa_joints();
    m_animdata_blocks[i]->m_local_matrices.resize(num_soa_joints);
  }

  return true;
}

void AnimGraphBlendTree::UpdateOrderedNodes() {
  m_eval_ordered_nodes.clear();
  UpdateOrderedNodesRecursive(0);
}

void AnimGraphBlendTree::UpdateOrderedNodesRecursive(int node_index) {
  AnimNode* node = m_nodes[node_index];
  const std::vector<AnimGraphConnection>& node_input_connections =
      m_node_input_connections[node_index];
  for (size_t i = 0, n = node_input_connections.size(); i < n; i++) {
    if (node_input_connections[i].m_crosses_hierarchy) {
      continue;
    }

    int input_node_index =
        GetAnimNodeIndex(node_input_connections[i].m_source_node);

    if (input_node_index == 1) {
      continue;
    }

    UpdateOrderedNodesRecursive(input_node_index);
  }

  if (node_index != 0) {
    // In case we have multiple output connections from the node we here
    // ensure that use the node evaluation that is the furthest away from
    // the output.
    std::vector<AnimNode*>::iterator find_iter = std::find(
        m_eval_ordered_nodes.begin(),
        m_eval_ordered_nodes.end(),
        node);
    if (find_iter != m_eval_ordered_nodes.end()) {
      m_eval_ordered_nodes.erase(find_iter);
    }

    m_eval_ordered_nodes.push_back(node);
  }
}

void AnimGraphBlendTree::MarkActiveInputs(const std::vector<AnimGraphConnection>& input_connections) {
  for (size_t i = 0, n = m_nodes.size(); i < n; i++) {
    if (m_nodes[i]->m_tick_number != m_tick_number) {
      m_nodes[i]->m_state = AnimNodeEvalState::Deactivated;
    }
  }

  const std::vector<AnimGraphConnection>& graph_output_inputs =
      m_node_input_connections[0];
  for (size_t i = 0, n = graph_output_inputs.size(); i < n; i++) {
    const AnimGraphConnection& graph_input = graph_output_inputs[i];
    AnimNode* node = graph_input.m_source_node;
    if (node != nullptr) {
      node->m_tick_number = m_tick_number;
      node->m_state = AnimNodeEvalState::Activated;
    }
  }

  for (int i = m_eval_ordered_nodes.size() - 1; i >= 0; i--) {
    AnimNode* node = m_eval_ordered_nodes[i];
    if (CheckIsNodeActive(node)) {
      size_t node_index = GetAnimNodeIndex(node);
      node->MarkActiveInputs(m_node_input_connections[node_index]);

      // Non-animation data inputs are always active.
      for (size_t j = 0, nj = m_node_input_connections[node_index].size();
           j < nj;
           j++) {
        const AnimGraphConnection& input =
            m_node_input_connections[node_index][j];
        if (input.m_source_node != nullptr
            && input.m_target_socket.m_type
                   != SocketType::SocketTypeAnimation) {
          input.m_source_node->m_state = AnimNodeEvalState::Activated;
        }
      }
    }
  }
}

void AnimGraphBlendTree::CalcSyncTrack(const std::vector<AnimGraphConnection>& input_connections) {
  for (size_t i = m_eval_ordered_nodes.size() - 1; i >= 0; i--) {
    AnimNode* node = m_eval_ordered_nodes[i];
    if (node->m_state == AnimNodeEvalState::Deactivated) {
      continue;
    }

    node->CalcSyncTrack(m_node_input_connections[GetAnimNodeIndex(node)]);
  }
}

void AnimGraphBlendTree::UpdateTime(float time_last, float time_now) {
  float dt = time_now - time_last;

  const std::vector<AnimGraphConnection>& graph_output_inputs =
      m_node_input_connections[0];
  for (size_t i = 0, n = graph_output_inputs.size(); i < n; i++) {
    AnimNode* node = graph_output_inputs[i].m_source_node;
    if (node != nullptr && node->m_state != AnimNodeEvalState::TimeUpdated) {
      node->UpdateTime(time_last, time_now);
    }
  }

  for (int i = m_eval_ordered_nodes.size() - 1; i >= 0; --i) {
    AnimNode* node = m_eval_ordered_nodes[i];
    if (node->m_state != AnimNodeEvalState::TimeUpdated) {
      continue;
    }

    PropagateTimeToNodeInputs(node);
  }

  m_state = AnimNodeEvalState::TimeUpdated;
}

void AnimGraphBlendTree::Evaluate(AnimGraphContext& context) {
  for (int i = 0, n = m_eval_ordered_nodes.size(); i < n; i++) {
    AnimNode* node = m_eval_ordered_nodes[i];

    if (node->m_state == AnimNodeEvalState::Deactivated) {
      continue;
    }

    node->Evaluate(context);
  }
}

void AnimGraphBlendTree::PropagateTimeToNodeInputs(const AnimNode* node) {
  size_t node_index = GetAnimNodeIndex(node);

  float node_time_now = node->m_time_now;
  float node_time_last = node->m_time_last;

  const std::vector<AnimGraphConnection>& node_input_connections =
      m_node_input_connections[node_index];
  for (size_t i = 0, n = node_input_connections.size(); i < n; i++) {
    AnimNode* input_node = node_input_connections[i].m_source_node;

    // Only propagate time updates via animation sockets.
    if (input_node != nullptr
        && node_input_connections[i].m_target_socket.m_type
               == SocketType::SocketTypeAnimation
        && input_node->m_state == AnimNodeEvalState::Activated) {
      input_node->UpdateTime(node_time_last, node_time_now);
    }
  }
}

Socket* AnimGraphBlendTree::GetInputSocket(const std::string& name) {
  for (size_t i = 0, n = m_node_output_connections[1].size(); i < n; i++) {
    AnimGraphConnection& connection = m_node_output_connections[1][i];
    if (connection.m_source_socket.m_name == name) {
      return &connection.m_source_socket;
    }
  }

  return nullptr;
}

Socket* AnimGraphBlendTree::GetOutputSocket(const std::string& name) {
  for (size_t i = 0, n = m_node_input_connections[0].size(); i < n; i++) {
    AnimGraphConnection& connection = m_node_input_connections[0][i];
    if (connection.m_target_socket.m_name == name) {
      return &connection.m_target_socket;
    }
  }

  return nullptr;
}

const Socket* AnimGraphBlendTree::GetInputSocket(const std::string& name) const {
  for (size_t i = 0, n = m_node_output_connections[1].size(); i < n; i++) {
    const AnimGraphConnection& connection = m_node_output_connections[1][i];
    if (connection.m_source_socket.m_name == name) {
      return &connection.m_source_socket;
    }
  }

  return nullptr;
}

const Socket* AnimGraphBlendTree::GetOutputSocket(const std::string& name) const {
  for (size_t i = 0, n = m_node_input_connections[0].size(); i < n; i++) {
    const AnimGraphConnection& connection = m_node_input_connections[0][i];
    if (connection.m_target_socket.m_name == name) {
      return &connection.m_target_socket;
    }
  }

  return nullptr;
}