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

#include "AnimGraph/AnimGraph.h"
#include "AnimGraph/AnimGraphEditor.h"
#include "AnimGraph/AnimGraphResource.h"
#include "catch.hpp"
#include "ozz/animation/offline/animation_builder.h"
#include "ozz/animation/offline/raw_animation.h"
#include "ozz/animation/offline/raw_skeleton.h"
#include "ozz/animation/offline/skeleton_builder.h"
#include "ozz/animation/runtime/animation.h"
#include "ozz/base/io/archive.h"
#include "ozz/base/io/stream.h"
#include "ozz/base/log.h"

struct SimpleAnimFixture {
  ozz::unique_ptr<ozz::animation::Skeleton> skeleton = nullptr;
  ozz::animation::offline::RawAnimation raw_animation_translation_x;
  ozz::unique_ptr<ozz::animation::Animation> animation_translate_x = nullptr;
  ozz::animation::offline::RawAnimation raw_animation_translation_y;
  ozz::unique_ptr<ozz::animation::Animation> animation_translate_y = nullptr;
  ozz::vector<ozz::math::SoaTransform> animation_output;
  ozz::animation::SamplingCache sampling_cache;

  SimpleAnimFixture() {
    createSkeleton();
    createAnimations();

    animation_output.resize(skeleton->num_soa_joints());
    sampling_cache.Resize(skeleton->num_joints());
  }

  void createSkeleton() {
    using namespace ozz::animation::offline;

    RawSkeleton raw_skeleton;
    RawSkeleton::Joint raw_joint;

    raw_joint.name = "Bone0";
    raw_joint.transform.translation.x = 1.f;
    raw_joint.transform.translation.y = 2.f;
    raw_joint.transform.translation.z = 3.f;

    raw_skeleton.roots.push_back(raw_joint);

    SkeletonBuilder skeleton_builder;
    skeleton = skeleton_builder(raw_skeleton);
  }

  void createAnimations() {
    using namespace ozz::animation::offline;

    raw_animation_translation_x.name = "TranslationX";
    RawAnimation::JointTrack bone0_track;
    RawAnimation::JointTrack::Translations bone0_translations;

    // animation_translate_x
    RawAnimation::TranslationKey translation_key;
    translation_key.time = 0.f;
    translation_key.value = ozz::math::Float3(0.f, 0.f, 0.f);
    bone0_translations.push_back(translation_key);

    translation_key.time = 1.f;
    translation_key.value = ozz::math::Float3(1.f, 0.f, 9.f);
    bone0_translations.push_back(translation_key);

    bone0_track.translations = bone0_translations;
    raw_animation_translation_x.tracks.push_back(bone0_track);
    raw_animation_translation_x.duration = 1.f;
    REQUIRE(raw_animation_translation_x.Validate());

    AnimationBuilder animation_builder;
    animation_translate_x = animation_builder(raw_animation_translation_x);

    // animation_translate_y
    raw_animation_translation_y.name = "TranslationY";
    bone0_translations.clear();

    translation_key.time = 0.f;
    translation_key.value = ozz::math::Float3(0.f, 0.f, 0.f);
    bone0_translations.push_back(translation_key);

    translation_key.time = 1.f;
    translation_key.value = ozz::math::Float3(0.f, 1.f, 0.f);
    bone0_translations.push_back(translation_key);

    bone0_track.translations = bone0_translations;
    raw_animation_translation_y.tracks.push_back(bone0_track);
    raw_animation_translation_y.duration = 1.f;
    REQUIRE(raw_animation_translation_y.Validate());

    animation_translate_y = animation_builder(raw_animation_translation_y);
  }
};

TEST_CASE_METHOD(
    SimpleAnimFixture,
    "Create Animations",
    "[AnimGraphEvalTests]") {
  using namespace ozz::animation::offline;

  ozz::animation::SamplingJob sampling_job;
  sampling_job.animation = animation_translate_x.get();
  sampling_job.cache = &sampling_cache;
  sampling_job.ratio = 1.f;
  sampling_job.output = make_span(animation_output);
  REQUIRE(sampling_job.Run());

  RawAnimation::TranslationKey& translation_key =
      raw_animation_translation_x.tracks[0].translations.back();

  ozz::math::SoaFloat3& sampled_translation = animation_output[0].translation;
  CHECK(
      sampled_translation.x[0] == Approx(translation_key.value.x).margin(0.01));
  CHECK(
      sampled_translation.y[0] == Approx(translation_key.value.y).margin(0.01));
  CHECK(
      sampled_translation.z[0] == Approx(translation_key.value.z).margin(0.01));
}

TEST_CASE_METHOD(
    SimpleAnimFixture,
    "AnimGraphSimpleEval",
    "[AnimGraphEvalTests]") {
  AnimGraphResource graph_resource;

  // Add nodes
  size_t trans_x_node_index =
      graph_resource.addNode(AnimNodeResourceFactory("AnimSampler"));
  size_t trans_y_node_index =
      graph_resource.addNode(AnimNodeResourceFactory("AnimSampler"));
  size_t blend_node_index =
      graph_resource.addNode(AnimNodeResourceFactory("Blend2"));

  // Setup nodes
  AnimNodeResource& trans_x_node = graph_resource.m_nodes[trans_x_node_index];
  trans_x_node.m_socket_accessor->SetPropertyValue("Filename", "trans_x");
  trans_x_node.m_name = "trans_x";

  AnimNodeResource& trans_y_node = graph_resource.m_nodes[trans_y_node_index];
  trans_y_node.m_socket_accessor->SetPropertyValue("Filename", "trans_y");
  trans_y_node.m_name = "trans_y";

  AnimNodeResource& blend_node = graph_resource.m_nodes[blend_node_index];
  blend_node.m_name = "BlendWalkRun";

  // Setup graph outputs and inputs
  AnimNodeResource& graph_output_node = graph_resource.getGraphOutputNode();
  graph_output_node.m_socket_accessor->RegisterInput<AnimData>("GraphOutput", nullptr);

  AnimNodeResource& graph_input_node =
      graph_resource.getGraphInputNode();
  graph_input_node.m_socket_accessor->RegisterOutput<float>(
      "GraphFloatInput",
      nullptr);

  // Wire up nodes
  graph_resource.connectSockets(trans_x_node, "Output", blend_node, "Input0");
  graph_resource.connectSockets(trans_y_node, "Output", blend_node, "Input1");
  graph_resource.connectSockets(
      blend_node,
      "Output",
      graph_resource.getGraphOutputNode(),
      "GraphOutput");
  REQUIRE(graph_resource.connectSockets(graph_input_node, "GraphFloatInput", blend_node, "Weight"));

  // Prepare animation maps
  AnimGraphContext graph_context;
  graph_context.m_skeleton = skeleton.get();
  graph_context.m_animation_map["trans_x"] = animation_translate_x.get();
  graph_context.m_animation_map["trans_y"] = animation_translate_y.get();

  // Instantiate graph
  AnimGraph graph = graph_resource.createInstance();
  graph_context.m_graph = &graph;
  graph.init(graph_context);

  // Get runtime graph inputs and outputs
  float* graph_float_input = nullptr;
  graph_float_input =
      static_cast<float*>(graph.getInputPtr("GraphFloatInput"));

  Socket* anim_output_socket =
      graph.getOutputSocket("GraphOutput");

  AnimData* graph_anim_output = static_cast<AnimData*>(graph.getOutputPtr("GraphOutput"));

  // Evaluate graph
  *graph_float_input = 0.1f;

  graph.markActiveNodes();
  CHECK(graph.m_nodes[trans_x_node_index]->m_state == AnimNodeEvalState::Activated);
  CHECK(graph.m_nodes[trans_y_node_index]->m_state == AnimNodeEvalState::Activated);
  CHECK(graph.m_nodes[blend_node_index]->m_state == AnimNodeEvalState::Activated);

  graph.updateTime(0.5f);
  graph.evaluate(graph_context);

  CHECK(graph_anim_output->m_local_matrices[0].translation.x[0] == Approx(0.5).margin(0.1));
  CHECK(graph_anim_output->m_local_matrices[0].translation.y[0] == Approx(0.05).margin(0.01));
}