AnimTestbed/src/AnimGraph/AnimGraphData.h

//
// Created by martin on 25.03.22.
//

#ifndef ANIMTESTBED_ANIMGRAPHDATA_H
#define ANIMTESTBED_ANIMGRAPHDATA_H

#include <ozz/base/maths/soa_transform.h>

#include <cstring>
#include <iostream>
#include <list>
#include <map>
#include <string>
#include <vector>

#include "SyncTrack.h"
#include "ozz/animation/runtime/animation.h"
#include "ozz/animation/runtime/skeleton.h"
#include "ozz/base/containers/vector.h"

//
// Data types
//

struct AnimGraph;

struct AnimData {
  ozz::vector<ozz::math::SoaTransform> m_local_matrices;
};

struct AnimDataAllocator {
  struct AnimDataList {
    AnimData* m_anim_data = nullptr;
    AnimDataList* next = nullptr;
  };

  std::list<AnimData*> m_anim_data_list;
  size_t m_num_allocations = 0;

  ~AnimDataAllocator() {
    while (!m_anim_data_list.empty()) {
      AnimData* front = m_anim_data_list.front();
#ifdef ANIM_DATA_ALLOCATOR_DEBUG
      std::cout << "about to delete with size "
                << front->m_anim_data->m_local_matrices.size()
                << front->m_anim_data << std::endl;
#endif
      delete front;
      m_anim_data_list.pop_front();
    }
  }

  AnimData* allocate(ozz::animation::Skeleton* skeleton) {
    if (m_anim_data_list.empty()) {
      AnimData* result = new AnimData();
      result->m_local_matrices.resize(skeleton->num_soa_joints());
#ifdef ANIM_DATA_ALLOCATOR_DEBUG
      std::cout << "Allocated with size " << result->m_local_matrices.size()
                << " " << result << std::endl;
#endif
      m_num_allocations++;
      return result;
    }

    AnimData* result = m_anim_data_list.front();
    m_anim_data_list.pop_front();

#ifdef ANIM_DATA_ALLOCATOR_DEBUG
    std::cout << "Reusing buffer with size " << result->m_local_matrices.size()
              << " " << result << std::endl;
#endif

    return result;
  }

  void free(AnimData* anim_data) {
#ifdef ANIM_DATA_ALLOCATOR_DEBUG
    std::cout << "Storing buffer with size "
              << anim_data->m_local_matrices.size() << " " << anim_data
              << std::endl;
#endif

    m_anim_data_list.push_front(anim_data);
  }

  size_t size() { return m_anim_data_list.size(); }
};

struct AnimGraphContext {
  AnimGraph* m_graph = nullptr;
  ozz::animation::Skeleton* m_skeleton = nullptr;

  typedef std::map<std::string, ozz::animation::Animation*> AnimationFileMap;
  AnimationFileMap m_animation_map;

  void freeAnimations() {
    AnimationFileMap::iterator animation_map_iter = m_animation_map.begin();

    while (animation_map_iter != m_animation_map.end()) {
      delete animation_map_iter->second;
      animation_map_iter++;
    }
  }
};

typedef float Vec3[3];
typedef float Quat[4];

enum class SocketType {
  SocketTypeUndefined = 0,
  SocketTypeBool,
  SocketTypeAnimation,
  SocketTypeFloat,
  SocketTypeVec3,
  SocketTypeQuat,
  SocketTypeString,
  SocketTypeLast
};

constexpr size_t cSocketStringValueMaxLength = 256;

static const char* SocketTypeNames[] =
    {"", "Bool", "Animation", "Float", "Vec3", "Quat", "String"};

enum SocketFlags { SocketFlagAffectsTime = 1 };

struct Socket {
  std::string m_name;
  SocketType m_type = SocketType::SocketTypeUndefined;
  union SocketValue {
    bool flag;
    float float_value;
    float vec3[3];
    float quat[4];
    std::string* string_ptr;
  };
  SocketValue m_value = {0};
  union SocketReference {
    void* ptr;
    void** ptr_ptr;
  };
  SocketReference m_reference = {0};
  int m_flags = 0;
  size_t m_type_size = 0;
};

template <typename T>
SocketType GetSocketType() {
  if constexpr (std::is_same<T, bool>::value) {
    return SocketType::SocketTypeBool;
  }

  if constexpr (std::is_same<T, AnimData>::value) {
    return SocketType::SocketTypeAnimation;
  }

  if constexpr (std::is_same<T, float>::value) {
    return SocketType::SocketTypeFloat;
  }

  if constexpr (std::is_same<T, Vec3>::value) {
    return SocketType::SocketTypeVec3;
  }

  if constexpr (std::is_same<T, Quat>::value) {
    return SocketType::SocketTypeQuat;
  }

  if constexpr (std::is_same<T, std::string>::value) {
    return SocketType::SocketTypeString;
  }

  assert(false && "This should not be reachable");
  abort();
  return SocketType::SocketTypeUndefined;
}

struct NodeDescriptorBase {
  std::vector<Socket> m_properties;
  std::vector<Socket> m_inputs;
  std::vector<Socket> m_outputs;

  template <typename T>
  bool RegisterInput(const char* name, T** value_ptr_ptr) {
    return RegisterSocket(name, value_ptr_ptr, m_inputs);
  }

  template <typename T>
  bool RegisterOutput(const char* name, T** value_ptr_ptr) {
    return RegisterSocket(name, value_ptr_ptr, m_outputs);
  }

  template <typename T>
  bool RegisterProperty(const char* name, T* value_ptr, int flags = 0) {
    for (int i = 0; i < m_properties.size(); i++) {
      if (m_properties[i].m_name == name) {
        return false;
      }
    }

    Socket socket;
    socket.m_name = name;
    socket.m_type = GetSocketType<T>();
    socket.m_reference.ptr = static_cast<void*>(value_ptr);
    socket.m_flags = flags;
    socket.m_type_size = sizeof(T);

    m_properties.push_back(socket);

    return true;
  }

  template <typename T>
  T* GetInput(const char* name) {
    Socket* socket = FindSocket(name, m_inputs);
    assert(GetSocketType<T>() == socket->m_type);
    return *socket->m_reference.ptr_ptr;
  }

  template <typename T>
  void SetInput(const char* name, T* value_ptr) {
    Socket* socket = FindSocket(name, m_inputs);
    assert(GetSocketType<T>() == socket->m_type);
    *socket->m_reference.ptr_ptr = value_ptr;
  }

  template <typename T>
  void SetProperty(const char* name, const T& value) {
    Socket* socket = FindSocket(name, m_properties);
    assert(GetSocketType<T>() == socket->m_type);
    *static_cast<T*>(socket->m_reference.ptr) = value;
  }

  template <typename T>
  const T& GetProperty(const char* name) {
    Socket* socket = FindSocket(name, m_properties);
    assert(GetSocketType<T>() == socket->m_type);
    return *static_cast<T*>(socket->m_reference.ptr);
  }

 private:
  Socket* FindSocket(const char* name, std::vector<Socket>& sockets) {
    for (int i = 0, n = sockets.size(); i < n; i++) {
      if (sockets[i].m_name == name) {
        return &sockets[i];
      }
    }

    return nullptr;
  }

  template <typename T>
  bool RegisterSocket(const char* name, T** value_ptr_ptr, std::vector<Socket>& sockets) {
    for (int i = 0; i < sockets.size(); i++) {
      if (sockets[i].m_name == name) {
        return false;
      }
    }

    Socket socket;
    socket.m_name = name;
    socket.m_type = GetSocketType<T>();
    socket.m_reference.ptr_ptr = (void**)(value_ptr_ptr);
    socket.m_type_size = sizeof(T);

    sockets.push_back(socket);

    return true;
  }
};

template <typename T>
struct NodeDescriptor : public NodeDescriptorBase {
  virtual ~NodeDescriptor() {}
};


struct NodeSocketAccessorBase {
  std::vector<Socket> m_properties;
  std::vector<Socket> m_inputs;
  std::vector<Socket> m_outputs;

  NodeSocketAccessorBase() {}
  virtual ~NodeSocketAccessorBase() {}

  virtual void UpdateFlags(){};

  Socket* FindSocket(std::vector<Socket>& sockets, const std::string& name) {
    Socket* result = nullptr;
    for (size_t i = 0, n = sockets.size(); i < n; i++) {
      if (sockets[i].m_name == name) {
        result = &sockets[i];
        break;
      }
    }

    return result;
  }

  const Socket* FindSocket(
      const std::vector<Socket>& sockets,
      const std::string& name) const {
    const Socket* result = nullptr;
    for (size_t i = 0, n = sockets.size(); i < n; i++) {
      if (sockets[i].m_name == name) {
        result = &sockets[i];
        break;
      }
    }

    return result;
  }

  SocketType GetSocketType(
      const std::vector<Socket>& sockets,
      const std::string& name) {
    const Socket* socket = FindSocket(sockets, name);
    if (socket == nullptr) {
      return SocketType::SocketTypeUndefined;
    }
    return socket->m_type;
  }

  size_t GetSocketIndex(
      const std::vector<Socket>& sockets,
      const std::string& name) const {
    for (size_t i = 0, n = sockets.size(); i < n; i++) {
      if (sockets[i].m_name == name) {
        return i;
      }
    }

    return -1;
  }

  template <typename T>
  T GetSocketValue(
      const std::vector<Socket>& sockets,
      const std::string& name,
      T default_value) {
    const Socket* socket = FindSocket(sockets, name);
    if (socket == nullptr) {
      return default_value;
    }

    return *static_cast<T*>(socket->m_reference.ptr);
  }

  template <typename T>
  void SetSocketReferenceValue(Socket* socket, T value) {
    std::cerr << "Could not find template specialization for socket type "
              << static_cast<int>(socket->m_type) << " ("
              << SocketTypeNames[static_cast<int>(socket->m_type)] << ")."
              << std::endl;
    //    *static_cast<T*>(socket->m_value.ptr) = value;
  }

  template <typename T>
  void SetSocketValue(Socket* socket, T value) {
    std::cerr << "Could not find template specialization for socket type "
              << static_cast<int>(socket->m_type) << " ("
              << SocketTypeNames[static_cast<int>(socket->m_type)] << ")."
              << std::endl;
    //    *static_cast<T*>(socket->m_value.ptr) = value;
  }

  template <typename T>
  bool RegisterSocket(
      std::vector<Socket>& sockets,
      const std::string& name,
      T* value_ptr,
      int flags = 0) {
    Socket* socket = FindSocket(sockets, name);
    if (socket != nullptr) {
      std::cerr << "Socket " << name << " already registered." << std::endl;
      return false;
    }

    sockets.push_back(Socket());
    socket = &sockets[sockets.size() - 1];
    socket->m_name = name;
    socket->m_type_size = sizeof(T);
    socket->m_flags = flags;

    if constexpr (std::is_same<T, float>::value) {
      socket->m_type = SocketType::SocketTypeFloat;
    } else if constexpr (std::is_same<T, bool>::value) {
      socket->m_type = SocketType::SocketTypeBool;
    } else if constexpr (std::is_same<T, Vec3>::value) {
      socket->m_type = SocketType::SocketTypeVec3;
    } else if constexpr (std::is_same<T, Quat>::value) {
      socket->m_type = SocketType::SocketTypeQuat;
    } else if constexpr (std::is_same<T, AnimData>::value) {
      socket->m_type = SocketType::SocketTypeAnimation;
    } else if constexpr (std::is_same<T, std::string>::value) {
      socket->m_type = SocketType::SocketTypeString;
      socket->m_value.string_ptr = value_ptr;
      socket->m_reference.ptr = value_ptr;
      return true;
    } else if constexpr (std::is_same<T, float*>::value) {
      socket->m_type = SocketType::SocketTypeFloat;
    } else if constexpr (std::is_same<T, bool*>::value) {
      socket->m_type = SocketType::SocketTypeBool;
    } else if constexpr (std::is_same<T, Vec3*>::value) {
      socket->m_type = SocketType::SocketTypeVec3;
    } else if constexpr (std::is_same<T, Quat*>::value) {
      socket->m_type = SocketType::SocketTypeQuat;
    } else if constexpr (std::is_same<T, AnimData*>::value) {
      socket->m_type = SocketType::SocketTypeAnimation;
    } else if constexpr (std::is_same<T, std::string*>::value) {
      socket->m_type = SocketType::SocketTypeString;
    } else {
      std::cerr << "Cannot register socket, invalid type." << std::endl;
      return false;
    }

    socket->m_reference.ptr = value_ptr;
    return true;
  }

  template <typename T>
  bool RegisterProperty(const std::string& name, T* value) {
    return RegisterSocket(m_properties, name, value);
  }
  template <typename T>
  void SetPropertyReferenceValue(const std::string& name, T value) {
    Socket* socket = FindSocket(m_properties, name);
    SetSocketReferenceValue<T>(socket, value);
  }
  template <typename T>
  void SetPropertyValue(const std::string& name, T value) {
    Socket* socket = FindSocket(m_properties, name);
    SetSocketValue<T>(socket, value);
  }
  template <typename T>
  T GetProperty(const std::string& name, T default_value) {
    return GetSocketValue(m_properties, name, default_value);
  }
  SocketType GetPropertyType(const std::string& name) {
    return GetSocketType(m_properties, name);
  }

  template <typename T>
  bool RegisterInput(const std::string& name, T* value, int flags = 0) {
    return RegisterSocket(m_inputs, name, value, flags);
  }
  template <typename T>
  T* GetInput(const std::string& name, T* value) {
    return GetSocketValue(m_inputs, name, value);
  }
  Socket* FindInputSocket(const std::string& name) {
    return FindSocket(m_inputs, name);
  }
  SocketType GetInputType(const std::string& name) {
    return GetSocketType(m_inputs, name);
  }
  size_t GetInputIndex(const std::string& name) {
    return GetSocketIndex(m_inputs, name);
  }

  template <typename T>
  bool RegisterOutput(const std::string& name, T* value, int flags = 0) {
    return RegisterSocket(m_outputs, name, value, flags);
  }
  template <typename T>
  bool RegisterOutput(const std::string& name, T** value, int flags = 0) {
    return RegisterSocket(m_outputs, name, value, flags);
  }
  SocketType GetOutputType(const std::string& name) {
    return GetSocketType(m_outputs, name);
  }
  Socket* FindOutputSocket(const std::string& name) {
    return FindSocket(m_outputs, name);
  }
  size_t GetOutputIndex(const std::string& name) {
    return GetSocketIndex(m_outputs, name);
  }
};

//
// SetSocketReferenceValue<> specializations
//
template <>
inline void NodeSocketAccessorBase::SetSocketReferenceValue<const bool&>(
    Socket* socket,
    const bool& value) {
  *static_cast<bool*>(socket->m_reference.ptr) = value;
}

template <>
inline void NodeSocketAccessorBase::SetSocketReferenceValue<const float&>(
    Socket* socket,
    const float& value) {
  *static_cast<float*>(socket->m_reference.ptr) = value;
}

template <>
inline void NodeSocketAccessorBase::SetSocketReferenceValue<const Vec3&>(
    Socket* socket,
    const Vec3& value) {
  static_cast<float*>(socket->m_reference.ptr)[0] = value[0];
  static_cast<float*>(socket->m_reference.ptr)[1] = value[1];
  static_cast<float*>(socket->m_reference.ptr)[2] = value[2];
}

template <>
inline void NodeSocketAccessorBase::SetSocketReferenceValue<const Quat&>(
    Socket* socket,
    const Quat& value) {
  static_cast<float*>(socket->m_reference.ptr)[0] = value[0];
  static_cast<float*>(socket->m_reference.ptr)[1] = value[1];
  static_cast<float*>(socket->m_reference.ptr)[2] = value[2];
  static_cast<float*>(socket->m_reference.ptr)[3] = value[3];
}

template <>
inline void NodeSocketAccessorBase::SetSocketReferenceValue<const std::string*>(
    Socket* socket,
    const std::string* value) {
  socket->m_value.string_ptr = const_cast<std::string*>(value);
}

template <>
inline void NodeSocketAccessorBase::SetSocketReferenceValue<const std::string&>(
    Socket* socket,
    const std::string& value) {
  *static_cast<std::string*>(socket->m_reference.ptr) = value;
}

template <>
inline void NodeSocketAccessorBase::SetSocketReferenceValue<const char*>(
    Socket* socket,
    const char* value) {
  std::string value_string(value);
  SetSocketReferenceValue<const std::string&>(socket, value_string);
}

//
// SetSocketValue<> specializations
//
template <>
inline void NodeSocketAccessorBase::SetSocketValue<const bool&>(
    Socket* socket,
    const bool& value) {
  socket->m_value.flag = value;
}

template <>
inline void NodeSocketAccessorBase::SetSocketValue<float>(
    Socket* socket,
    float value) {
  *static_cast<float*>(socket->m_reference.ptr) = value;
}

template <>
inline void NodeSocketAccessorBase::SetSocketValue<const Vec3&>(
    Socket* socket,
    const Vec3& value) {
  socket->m_value.vec3[0] = value[0];
  socket->m_value.vec3[1] = value[1];
  socket->m_value.vec3[2] = value[2];
}

template <>
inline void NodeSocketAccessorBase::SetSocketValue<const Quat&>(
    Socket* socket,
    const Quat& value) {
  socket->m_value.quat[0] = value[0];
  socket->m_value.quat[1] = value[1];
  socket->m_value.quat[2] = value[2];
  socket->m_value.quat[3] = value[3];
}

template <>
inline void NodeSocketAccessorBase::SetSocketValue<std::string>(
    Socket* socket,
    std::string value) {
  *socket->m_value.string_ptr = value;
}

template <>
inline void NodeSocketAccessorBase::SetSocketValue<const char*>(
    Socket* socket,
    const char* value) {
  SetSocketValue<std::string>(socket, value);
}

template <typename T>
struct NodeSocketAccessor : public NodeSocketAccessorBase {
  virtual ~NodeSocketAccessor() {}
};

#endif  //ANIMTESTBED_ANIMGRAPHDATA_H