AnimTestbed/src/AnimGraph/AnimGraphResource.cc

//
// Created by martin on 04.02.22.
//

#include "AnimGraphResource.h"

#include <fstream>

#include "3rdparty/json/json.hpp"

using json = nlohmann::json;

//
// Socket <-> json
//
std::string sSocketTypeToStr(SocketType pin_type) {
  if (pin_type < SocketType::SocketTypeUndefined
      || pin_type >= SocketType::SocketTypeLast) {
    return "Unknown";
  }

  return SocketTypeNames[static_cast<int>(pin_type)];
}

json sSocketToJson(const Socket& socket) {
  json result;
  result["name"] = socket.m_name;
  result["type"] = sSocketTypeToStr(socket.m_type);

  if (socket.m_reference.ptr != nullptr) {
    if (socket.m_type == SocketType::SocketTypeBool) {
      result["value"] = socket.m_value.flag;
    } else if (socket.m_type == SocketType::SocketTypeAnimation) {
    } else if (socket.m_type == SocketType::SocketTypeFloat) {
      result["value"] = socket.m_value.float_value;
    } else if (socket.m_type == SocketType::SocketTypeVec3) {
      result["value"][0] = socket.m_value.vec3.v[0];
      result["value"][1] = socket.m_value.vec3.v[1];
      result["value"][2] = socket.m_value.vec3.v[2];
    } else if (socket.m_type == SocketType::SocketTypeQuat) {
      result["value"][0] = socket.m_value.quat.v[0];
      result["value"][1] = socket.m_value.quat.v[1];
      result["value"][2] = socket.m_value.quat.v[2];
      result["value"][3] = socket.m_value.quat.v[3];
    } else if (socket.m_type == SocketType::SocketTypeString) {
      result["value"] = *static_cast<std::string*>(socket.m_reference.ptr);
    } else {
      std::cerr << "Invalid socket type '" << static_cast<int>(socket.m_type)
                << "'." << std::endl;
    }
  }
  return result;
}

Socket sJsonToSocket(const json& json_data) {
  Socket result;
  result.m_type = SocketType::SocketTypeUndefined;
  result.m_reference.ptr = nullptr;
  result.m_name = json_data["name"];

  std::string type_string = json_data["type"];

  if (type_string == "Bool") {
    result.m_type = SocketType::SocketTypeBool;
    result.m_type_size = sizeof(bool);
  } else if (type_string == "Animation") {
    result.m_type = SocketType::SocketTypeAnimation;
    result.m_type_size = sizeof(AnimData);
  } else if (type_string == "Float") {
    result.m_type = SocketType::SocketTypeFloat;
    result.m_type_size = sizeof(float);
  } else if (type_string == "Vec3") {
    result.m_type = SocketType::SocketTypeVec3;
    result.m_type_size = sizeof(Vec3);
  } else if (type_string == "Quat") {
    result.m_type = SocketType::SocketTypeQuat;
    result.m_type_size = sizeof(Quat);
  } else if (type_string == "String") {
    result.m_type = SocketType::SocketTypeString;
    result.m_type_size = sizeof(std::string);
  } else {
    std::cerr << "Invalid socket type '" << type_string << "'." << std::endl;
  }

  return result;
}

//
// AnimGraphNode <-> json
//
json sAnimGraphNodeToJson(const AnimNodeResource& node) {
  json result;

  result["name"] = node.m_name;
  result["type"] = "AnimNodeResource";
  result["node_type"] = node.m_type_name;

  for (size_t j = 0; j < 2; j++) {
    result["position"][j] = node.m_position[j];
  }

  for (size_t j = 0, n = node.m_socket_accessor->m_properties.size(); j < n;
       j++) {
    Socket& property = node.m_socket_accessor->m_properties[j];
    result["properties"][property.m_name] = sSocketToJson(property);
  }

  return result;
}

AnimNodeResource sAnimGraphNodeFromJson(const json& json_node) {
  AnimNodeResource result;

  result.m_name = json_node["name"];
  result.m_type_name = json_node["node_type"];
  result.m_position[0] = json_node["position"][0];
  result.m_position[1] = json_node["position"][1];

  result.m_anim_node = AnimNodeFactory(result.m_type_name);
  result.m_socket_accessor =
      AnimNodeDescriptorFactory(result.m_type_name, result.m_anim_node);

  for (size_t j = 0, n = result.m_socket_accessor->m_properties.size(); j < n;
       j++) {
    Socket& property = result.m_socket_accessor->m_properties[j];
    json json_property = json_node["properties"][property.m_name];

    if (sSocketTypeToStr(property.m_type) == json_property["type"]) {
      if (property.m_type == SocketType::SocketTypeBool) {
        property.m_value.flag = json_property["value"];
      } else if (property.m_type == SocketType::SocketTypeAnimation) {
      } else if (property.m_type == SocketType::SocketTypeFloat) {
        property.m_value.float_value = json_property["value"];
      } else if (property.m_type == SocketType::SocketTypeVec3) {
        property.m_value.vec3.v[0] = json_property["value"][0];
        property.m_value.vec3.v[1] = json_property["value"][1];
        property.m_value.vec3.v[2] = json_property["value"][2];
      } else if (property.m_type == SocketType::SocketTypeQuat) {
        Quat* property_quat = reinterpret_cast<Quat*>(property.m_reference.ptr);
        property.m_value.quat.v[0] = json_property["value"][0];
        property.m_value.quat.v[1] = json_property["value"][1];
        property.m_value.quat.v[2] = json_property["value"][2];
        property.m_value.quat.v[3] = json_property["value"][3];
      } else if (property.m_type == SocketType::SocketTypeString) {
        property.m_value_string = json_property["value"].get<std::string>();
      } else {
        std::cerr << "Invalid type for property '" << property.m_name
                  << "'. Cannot parse json to type '"
                  << static_cast<int>(property.m_type) << std::endl;
        break;
      }
    } else {
      std::cerr << "Invalid type for property '" << property.m_name
                << "': expected " << sSocketTypeToStr(property.m_type)
                << " but got " << json_property["type"] << std::endl;
    }
  }

  return result;
}

//
// AnimGraphConnectionResource <-> Json
//
json sAnimGraphConnectionToJson(
    const AnimGraphResource& graph_resource,
    const AnimGraphConnectionResource& connection) {
  json result;

  result["type"] = "AnimGraphConnectionResource";

  result["source_node_index"] = connection.source_node_index;
  result["source_socket_name"] = connection.source_socket_name;

  result["target_node_index"] = connection.target_node_index;
  result["target_socket_name"] = connection.target_socket_name;

  return result;
}

AnimGraphConnectionResource sAnimGraphConnectionFromJson(
    const AnimGraphResource& graph_resource,
    const json& json_node) {
  AnimGraphConnectionResource connection;

  connection.source_node_index = json_node["source_node_index"];
  connection.source_socket_name = json_node["source_socket_name"];

  connection.target_node_index = json_node["target_node_index"];
  connection.target_socket_name = json_node["target_socket_name"];

  return connection;
}

void AnimGraphResource::clear() {
  m_name = "";

  clearNodes();
  m_connections.clear();

  initGraphConnectors();
}

void AnimGraphResource::clearNodes() {
  for (size_t i = 0; i < m_nodes.size(); i++) {
    delete m_nodes[i].m_socket_accessor;
    m_nodes[i].m_socket_accessor = nullptr;
    delete m_nodes[i].m_anim_node;
    m_nodes[i].m_anim_node = nullptr;
  }
  m_nodes.clear();
}

void AnimGraphResource::initGraphConnectors() {
  m_nodes.push_back(AnimNodeResourceFactory("BlendTree"));
  m_nodes[0].m_name = "Outputs";
  m_nodes.push_back(AnimNodeResourceFactory("BlendTree"));
  m_nodes[1].m_name = "Inputs";
}

bool AnimGraphResource::saveToFile(const char* filename) const {
  json result;

  result["name"] = m_name;
  result["type"] = "AnimGraphResource";

  for (size_t i = 0; i < m_nodes.size(); i++) {
    const AnimNodeResource& node = m_nodes[i];
    result["nodes"][i] = sAnimGraphNodeToJson(node);
  }

  for (size_t i = 0; i < m_connections.size(); i++) {
    const AnimGraphConnectionResource& connection = m_connections[i];
    result["connections"][i] = sAnimGraphConnectionToJson(*this, connection);
  }

  // Graph inputs and outputs
  {
    const AnimNodeResource& graph_output_node = m_nodes[0];
    const std::vector<Socket> graph_inputs =
        graph_output_node.m_socket_accessor->m_inputs;
    for (size_t i = 0; i < graph_inputs.size(); i++) {
      result["nodes"][0]["inputs"][i] = sSocketToJson(graph_inputs[i]);
    }

    const AnimNodeResource& graph_input_node = m_nodes[1];
    const std::vector<Socket> graph_outputs =
        graph_input_node.m_socket_accessor->m_outputs;
    for (size_t i = 0; i < graph_outputs.size(); i++) {
      result["nodes"][1]["outputs"][i] = sSocketToJson(graph_outputs[i]);
    }
  }

  std::ofstream output_file;
  output_file.open(filename);
  output_file << result.dump(4, ' ') << std::endl;
  output_file.close();

  return true;
}

bool AnimGraphResource::loadFromFile(const char* filename) {
  std::ifstream input_file;
  input_file.open(filename);
  std::stringstream buffer;
  buffer << input_file.rdbuf();

  json json_data = json::parse(buffer.str(), nullptr, false);
  if (json_data.is_discarded()) {
    std::cerr << "Error parsing json of file '" << filename << "'."
              << std::endl;
  }

  if (json_data["type"] != "AnimGraphResource") {
    std::cerr
        << "Invalid json object. Expected type 'AnimGraphResource' but got '"
        << json_data["type"] << "'." << std::endl;
  }

  clear();
  clearNodes();

  m_name = json_data["name"];

  // Load nodes
  for (size_t i = 0; i < json_data["nodes"].size(); i++) {
    const json& json_node = json_data["nodes"][i];
    if (json_node["type"] != "AnimNodeResource") {
      std::cerr
          << "Invalid json object. Expected type 'AnimNodeResource' but got '"
          << json_node["type"] << "'." << std::endl;
      return false;
    }

    AnimNodeResource node = sAnimGraphNodeFromJson(json_node);
    m_nodes.push_back(node);
  }

  // Setup graph inputs and outputs
  const json& graph_outputs = json_data["nodes"][0]["inputs"];
  for (size_t i = 0; i < graph_outputs.size(); i++) {
    AnimNodeResource& graph_node = m_nodes[0];
    graph_node.m_socket_accessor->m_inputs.push_back(
        sJsonToSocket(graph_outputs[i]));
  }

  const json& graph_inputs = json_data["nodes"][1]["outputs"];
  for (size_t i = 0; i < graph_inputs.size(); i++) {
    AnimNodeResource& graph_node = m_nodes[1];
    graph_node.m_socket_accessor->m_outputs.push_back(
        sJsonToSocket(graph_inputs[i]));
  }

  // Load connections
  for (size_t i = 0; i < json_data["connections"].size(); i++) {
    const json& json_connection = json_data["connections"][i];
    if (json_connection["type"] != "AnimGraphConnectionResource") {
      std::cerr
          << "Invalid json object. Expected type 'AnimGraphConnectionResource' "
             "but got '"
          << json_connection["type"] << "'." << std::endl;
      return false;
    }

    AnimGraphConnectionResource connection =
        sAnimGraphConnectionFromJson(*this, json_connection);
    m_connections.push_back(connection);
  }

  return true;
}

void AnimGraphResource::createInstance(AnimGraph& result) const {
  createRuntimeNodeInstances(result);
  prepareGraphIOData(result);
  setRuntimeNodeProperties(result);

  result.updateOrderedNodes();
  result.resetNodeStates();
}

void AnimGraphResource::createRuntimeNodeInstances(AnimGraph& instance) const {
  for (int i = 0; i < m_nodes.size(); i++) {
    const AnimNodeResource& node_resource = m_nodes[i];
    AnimNode* node = AnimNodeFactory(node_resource.m_type_name.c_str());
    node->m_name = node_resource.m_name;
    node->m_node_type_name = node_resource.m_type_name;
    node->m_index = i;
    instance.m_nodes.push_back(node);

    // runtime node connections
    instance.m_node_input_connections.push_back(
        std::vector<AnimGraphConnection>());
    instance.m_node_output_connections.push_back(
        std::vector<AnimGraphConnection>());
  }
}

void AnimGraphResource::prepareGraphIOData(AnimGraph& instance) const {
  instance.m_socket_accessor =
      AnimNodeDescriptorFactory("BlendTree", instance.m_nodes[0]);
  instance.m_socket_accessor->m_outputs =
      m_nodes[1].m_socket_accessor->m_outputs;
  instance.m_socket_accessor->m_inputs = m_nodes[0].m_socket_accessor->m_inputs;
  instance.m_socket_accessor->m_outputs =
      m_nodes[1].m_socket_accessor->m_outputs;

  // inputs
  int input_block_size = 0;
  std::vector<Socket>& graph_inputs = instance.getGraphInputs();
  for (int i = 0; i < graph_inputs.size(); i++) {
    input_block_size += sizeof(void*);
  }

  if (input_block_size > 0) {
    instance.m_input_buffer = new char[input_block_size];
    memset(instance.m_input_buffer, 0, input_block_size);
  }

  int input_block_offset = 0;
  for (int i = 0; i < graph_inputs.size(); i++) {
    graph_inputs[i].m_reference.ptr =
        (void*)&instance.m_input_buffer[input_block_offset];
    input_block_offset += sizeof(void*);
  }

  // outputs
  int output_block_size = 0;
  std::vector<Socket>& graph_outputs = instance.getGraphOutputs();
  for (int i = 0; i < graph_outputs.size(); i++) {
    output_block_size += graph_outputs[i].m_type_size;
  }

  if (output_block_size > 0) {
    instance.m_output_buffer = new char[output_block_size];
    memset(instance.m_output_buffer, 0, output_block_size);
  }

  int output_block_offset = 0;
  for (int i = 0; i < graph_outputs.size(); i++) {
    graph_outputs[i].m_reference.ptr =
        &instance.m_output_buffer[output_block_offset];
    output_block_offset += graph_outputs[i].m_type_size;
  }

  // connections: make source and target sockets point to the same address in the connection data storage.
  int connection_data_storage_size = 0;
  for (int i = 0; i < m_connections.size(); i++) {
    const AnimGraphConnectionResource& connection = m_connections[i];
    const AnimNodeResource& source_node = m_nodes[connection.source_node_index];
    Socket* source_socket = source_node.m_socket_accessor->GetOutputSocket(
        connection.source_socket_name.c_str());
    connection_data_storage_size += source_socket->m_type_size;
  }

  if (connection_data_storage_size > 0) {
    instance.m_connection_data_storage = new char[connection_data_storage_size];
    memset(instance.m_connection_data_storage, 0, connection_data_storage_size);
  }

  std::vector<NodeDescriptorBase*> instance_node_descriptors(
      m_nodes.size(),
      nullptr);
  for (int i = 0; i < m_nodes.size(); i++) {
    instance_node_descriptors[i] = AnimNodeDescriptorFactory(
        m_nodes[i].m_type_name.c_str(),
        instance.m_nodes[i]);
  }

  instance_node_descriptors[0]->m_inputs = instance.getGraphOutputs();
  instance_node_descriptors[1]->m_outputs = instance.getGraphInputs();

  int connection_data_offset = 0;
  for (int i = 0; i < m_connections.size(); i++) {
    const AnimGraphConnectionResource& connection = m_connections[i];
    const AnimNodeResource& source_node_resource =
        m_nodes[connection.source_node_index];
    const AnimNodeResource& target_node_resource =
        m_nodes[connection.target_node_index];

    NodeDescriptorBase* source_node_descriptor =
        instance_node_descriptors[connection.source_node_index];
    NodeDescriptorBase* target_node_descriptor =
        instance_node_descriptors[connection.target_node_index];

    AnimNode* source_node = instance.m_nodes[connection.source_node_index];
    AnimNode* target_node = instance.m_nodes[connection.target_node_index];

    Socket* source_socket = source_node_descriptor->GetOutputSocket(
        connection.source_socket_name.c_str());
    Socket* target_socket = target_node_descriptor->GetInputSocket(
        connection.target_socket_name.c_str());

    AnimGraphConnection instance_connection;
    instance_connection.m_source_node = source_node;
    instance_connection.m_source_socket = *source_socket;
    instance_connection.m_target_node = target_node;
    instance_connection.m_target_socket = *target_socket;
    instance.m_node_input_connections[connection.target_node_index].push_back(
        instance_connection);
    instance.m_node_output_connections[connection.source_node_index].push_back(
        instance_connection);

    source_node_descriptor->SetOutputUnchecked(
        connection.source_socket_name.c_str(),
        &instance.m_connection_data_storage[connection_data_offset]);

    target_node_descriptor->SetInputUnchecked(
        connection.target_socket_name.c_str(),
        &instance.m_connection_data_storage[connection_data_offset]);

    if (source_socket->m_type == SocketType::SocketTypeAnimation) {
      instance.m_animdata_blocks.push_back((AnimData*)(
          &instance.m_connection_data_storage[connection_data_offset]));
    }

    connection_data_offset += source_socket->m_type_size;
  }

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

/*
void AnimGraphResource::connectRuntimeNodes(AnimGraph& instance) const {
  for (int i = 0; i < m_connections.size(); i++) {
    const AnimGraphConnectionResource& connection = m_connections[i];
    std::string source_node_type = "";
    std::string target_node_type = "";
    AnimNode* source_node = nullptr;
    AnimNode* target_node = nullptr;
    NodeSocketAccessorBase* source_node_accessor = nullptr;
    NodeSocketAccessorBase* target_node_accessor = nullptr;
    SocketType source_type;
    SocketType target_type;
    size_t source_socket_index = -1;
    size_t target_socket_index = -1;

    if (connection.source_node_index < 0
        || connection.source_node_index >= m_nodes.size()) {
      std::cerr << "Could not find source node index." << std::endl;
      continue;
    }

    source_node = instance.m_nodes[connection.source_node_index];
    source_node_type = source_node->m_node_type_name;
    if (connection.source_node_index == 1) {
      source_node_accessor = instance.m_socket_accessor;
    } else {
      source_node_accessor =
          AnimNodeAccessorFactory(source_node_type, source_node);
    }

    if (connection.target_node_index < 0
        || connection.target_node_index >= m_nodes.size()) {
      std::cerr << "Could not find source node index." << std::endl;
      continue;
    }

    target_node = instance.m_nodes[connection.target_node_index];
    target_node_type = target_node->m_node_type_name;
    if (connection.target_node_index == 0) {
      target_node_accessor = instance.m_socket_accessor;
    } else {
      target_node_accessor =
          AnimNodeAccessorFactory(target_node_type, target_node);
    }

    assert(source_node != nullptr);
    assert(target_node != nullptr);

    //
    // Map resource node sockets to graph instance node sockets
    //
    source_socket_index =
        source_node_accessor->GetOutputIndex(connection.source_socket_name);
    if (source_socket_index == -1) {
      std::cerr << "Invalid source socket " << connection.source_socket_name
                << " for node " << source_node->m_name << "." << std::endl;
      continue;
    }
    Socket* source_socket =
        &source_node_accessor->m_outputs[source_socket_index];

    target_socket_index =
        target_node_accessor->GetInputIndex(connection.target_socket_name);
    if (target_socket_index == -1) {
      std::cerr << "Invalid target socket " << connection.target_socket_name
                << " for node " << target_node->m_name << "." << std::endl;
      continue;
    }
    Socket* target_socket =
        &target_node_accessor->m_inputs[target_socket_index];

    if (source_socket->m_type != target_socket->m_type) {
      std::cerr << "Cannot connect sockets: invalid types!" << std::endl;
    }

    //
    // Wire up outputs to inputs.
    //
    // Skip animation connections and connections to the output node as the
    // pointers are already set up in AnimGraphResource::prepareGraphIOData().
    if (target_socket->m_type != SocketType::SocketTypeAnimation
        && connection.target_node_index != 0) {
      (*target_socket->m_reference.ptr_ptr) = source_socket->m_reference.ptr;
    }

    size_t target_node_index = target_node->m_index;

    // Register the runtime connection
    AnimGraphConnection runtime_connection = {
        source_node,
        *source_socket,
        target_node,
        *target_socket};

    std::vector<AnimGraphConnection>& target_input_connections =
        instance.m_node_input_connections[target_node_index];
    target_input_connections.push_back(runtime_connection);

    std::vector<AnimGraphConnection>& source_output_connections =
        instance.m_node_output_connections[source_node->m_index];
    source_output_connections.push_back(runtime_connection);

    if (target_node_accessor != instance.m_socket_accessor) {
      delete target_node_accessor;
    }

    if (source_node_accessor != instance.m_socket_accessor) {
      delete source_node_accessor;
    }
  }
}
*/

void AnimGraphResource::setRuntimeNodeProperties(AnimGraph& instance) const {
  for (int i = 2; i < m_nodes.size(); i++) {
    const AnimNodeResource& node_resource = m_nodes[i];

    NodeDescriptorBase* node_instance_accessor = AnimNodeDescriptorFactory(
        node_resource.m_type_name,
        instance.m_nodes[i]);

    std::vector<Socket>& resource_properties =
        node_resource.m_socket_accessor->m_properties;
    for (size_t j = 0, n = resource_properties.size(); j < n; j++) {
      const Socket& property = resource_properties[j];
      const std::string& name = property.m_name;

      switch (property.m_type) {
        case SocketType::SocketTypeBool:
          node_instance_accessor->SetProperty(
              name.c_str(),
              property.m_value.flag);
          break;
        case SocketType::SocketTypeFloat:
          node_instance_accessor->SetProperty(
              name.c_str(),
              property.m_value.float_value);
          break;
        case SocketType::SocketTypeVec3:
          node_instance_accessor->SetProperty<Vec3>(
              name.c_str(),
              property.m_value.vec3);
          break;
        case SocketType::SocketTypeQuat:
          node_instance_accessor->SetProperty(
              name.c_str(),
              property.m_value.quat);
          break;
        case SocketType::SocketTypeString:
          node_instance_accessor->SetProperty(
              name.c_str(),
              property.m_value_string);
          break;
        default:
          std::cerr << "Invalid socket type "
                    << static_cast<int>(property.m_type) << std::endl;
      }
    }

    delete node_instance_accessor;
  }
}