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

#include "AnimGraph/AnimGraph.h"
#include "AnimGraph/AnimGraphEditor.h"
#include "AnimGraph/AnimGraphResource.h"
#include "catch.hpp"
#include "ozz/base/io/archive.h"
#include "ozz/base/io/stream.h"
#include "ozz/base/log.h"

bool load_skeleton(ozz::animation::Skeleton& skeleton, const char* filename) {
  assert(filename);
  ozz::io::File file(filename, "rb");
  if (!file.opened()) {
    ozz::log::Err() << "Failed to open skeleton file " << filename << "."
                    << std::endl;
    return false;
  }
  ozz::io::IArchive archive(&file);
  if (!archive.TestTag<ozz::animation::Skeleton>()) {
    ozz::log::Err() << "Failed to load skeleton instance from file " << filename
                    << "." << std::endl;
    return false;
  }

  // Once the tag is validated, reading cannot fail.
  archive >> skeleton;

  return true;
}

TEST_CASE("AnimSamplerGraph", "[AnimGraphResource]") {
  AnimGraphResource graph_resource;

  graph_resource.clear();
  graph_resource.m_name = "AnimSamplerGraph";

  // Prepare graph inputs and outputs
  size_t walk_node_index =
      graph_resource.addNode(AnimNodeResourceFactory("AnimSampler"));

  AnimNodeResource& walk_node = graph_resource.m_nodes[walk_node_index];
  walk_node.m_name = "WalkAnim";
  walk_node.m_socket_accessor->SetPropertyValue(
      "Filename",
      std::string("media/Walking-loop.ozz"));

  AnimNodeResource& graph_node = graph_resource.m_nodes[0];
  graph_node.m_socket_accessor->RegisterInput<AnimData>("GraphOutput", nullptr);

  graph_resource.connectSockets(
      walk_node,
      "Output",
      graph_resource.getGraphOutputNode(),
      "GraphOutput");

  graph_resource.saveToFile("AnimSamplerGraph.animgraph.json");
  AnimGraphResource graph_resource_loaded;
  graph_resource_loaded.loadFromFile("AnimSamplerGraph.animgraph.json");

  AnimGraph graph;
  graph_resource_loaded.createInstance(graph);
  AnimGraphContext graph_context;

  ozz::animation::Skeleton skeleton;
  REQUIRE(load_skeleton(skeleton, "media/skeleton.ozz"));
  graph_context.m_skeleton = &skeleton;

  REQUIRE(graph.init(graph_context));

  REQUIRE(graph.m_nodes.size() == 3);
  REQUIRE(graph.m_nodes[0]->m_node_type_name == "BlendTree");
  REQUIRE(graph.m_nodes[1]->m_node_type_name == "BlendTree");
  REQUIRE(graph.m_nodes[2]->m_node_type_name == "AnimSampler");

  // connections within the graph
  AnimSamplerNode* anim_sampler_walk =
      dynamic_cast<AnimSamplerNode*>(graph.m_nodes[2]);

  BlendTreeNode* graph_output_node =
      dynamic_cast<BlendTreeNode*>(graph.m_nodes[0]);

  // check node input dependencies
  size_t anim_sampler_index = anim_sampler_walk->m_index;

  REQUIRE(graph.m_node_output_connections[anim_sampler_index].size() == 1);
  CHECK(
      graph.m_node_output_connections[anim_sampler_index][0].m_target_node
      == graph_output_node);

  // Ensure animation sampler nodes use the correct files
  REQUIRE(anim_sampler_walk->m_filename == "media/Walking-loop.ozz");
  REQUIRE(anim_sampler_walk->m_animation != nullptr);

  // Ensure that outputs are properly propagated.
  AnimData output;
  output.m_local_matrices.resize(skeleton.num_soa_joints());
  graph.SetOutput("GraphOutput", &output);
  REQUIRE(anim_sampler_walk->o_output == &output);

  WHEN("Emulating Graph Evaluation") {
    CHECK(graph.m_anim_data_allocator.size() == 0);
    anim_sampler_walk->Evaluate(graph_context);
  }

  graph_context.freeAnimations();
}

/*
 * Checks that node const inputs are properly set.
 */
TEST_CASE("AnimSamplerSpeedScaleGraph", "[AnimGraphResource]") {
  AnimGraphResource graph_resource;

  graph_resource.clear();
  graph_resource.m_name = "AnimSamplerSpeedScaleGraph";

  // Prepare graph inputs and outputs
  size_t walk_node_index =
      graph_resource.addNode(AnimNodeResourceFactory("AnimSampler"));

  size_t speed_scale_node_index =
      graph_resource.addNode(AnimNodeResourceFactory("SpeedScale"));

  AnimNodeResource& walk_node = graph_resource.m_nodes[walk_node_index];
  walk_node.m_name = "WalkAnim";
  walk_node.m_socket_accessor->SetPropertyValue(
      "Filename",
      std::string("media/Walking-loop.ozz"));

  AnimNodeResource& speed_scale_node =
      graph_resource.m_nodes[speed_scale_node_index];
  speed_scale_node.m_name = "SpeedScale";
  float speed_scale_value = 1.35f;
  speed_scale_node.m_socket_accessor->SetInputValue(
      "SpeedScale",
      speed_scale_value);

  AnimNodeResource& graph_node = graph_resource.m_nodes[0];
  graph_node.m_socket_accessor->RegisterInput<AnimData>("GraphOutput", nullptr);

  graph_resource.connectSockets(walk_node, "Output", speed_scale_node, "Input");

  graph_resource.connectSockets(
      speed_scale_node,
      "Output",
      graph_resource.getGraphOutputNode(),
      "GraphOutput");

  graph_resource.saveToFile("AnimSamplerSpeedScaleGraph.animgraph.json");
  AnimGraphResource graph_resource_loaded;
  graph_resource_loaded.loadFromFile(
      "AnimSamplerSpeedScaleGraph.animgraph.json");

  Socket* speed_scale_resource_loaded_input =
      graph_resource_loaded.m_nodes[speed_scale_node_index]
          .m_socket_accessor->GetInputSocket("SpeedScale");
  REQUIRE(speed_scale_resource_loaded_input != nullptr);

  REQUIRE_THAT(
      speed_scale_resource_loaded_input->m_value.float_value,
      Catch::Matchers::WithinAbs(speed_scale_value, 0.1));

  AnimGraph graph;
  graph_resource_loaded.createInstance(graph);

  REQUIRE_THAT(*dynamic_cast<SpeedScaleNode*>(graph.m_nodes[speed_scale_node_index])->i_speed_scale,
               Catch::Matchers::WithinAbs(speed_scale_value, 0.1));
}


TEST_CASE("Blend2Graph", "[AnimGraphResource]") {
  AnimGraphResource graph_resource;

  graph_resource.clear();
  graph_resource.m_name = "WalkRunBlendGraph";

  // Prepare graph inputs and outputs
  size_t walk_node_index =
      graph_resource.addNode(AnimNodeResourceFactory("AnimSampler"));
  size_t run_node_index =
      graph_resource.addNode(AnimNodeResourceFactory("AnimSampler"));
  size_t blend_node_index =
      graph_resource.addNode(AnimNodeResourceFactory("Blend2"));

  AnimNodeResource& walk_node = graph_resource.m_nodes[walk_node_index];
  walk_node.m_name = "WalkAnim";
  walk_node.m_socket_accessor->SetPropertyValue(
      "Filename",
      std::string("media/Walking-loop.ozz"));
  AnimNodeResource& run_node = graph_resource.m_nodes[run_node_index];
  run_node.m_socket_accessor->SetPropertyValue(
      "Filename",
      std::string("media/Running0-loop.ozz"));
  run_node.m_name = "RunAnim";
  AnimNodeResource& blend_node = graph_resource.m_nodes[blend_node_index];
  blend_node.m_name = "BlendWalkRun";

  AnimNodeResource& graph_node = graph_resource.m_nodes[0];
  graph_node.m_socket_accessor->RegisterInput<AnimData>("GraphOutput", nullptr);

  REQUIRE(graph_node.m_socket_accessor->m_inputs.size() == 1);
  REQUIRE(blend_node.m_socket_accessor->GetInputIndex("Input0") == 0);
  REQUIRE(blend_node.m_socket_accessor->GetInputIndex("Input1") == 1);

  graph_resource.connectSockets(walk_node, "Output", blend_node, "Input0");
  graph_resource.connectSockets(run_node, "Output", blend_node, "Input1");
  graph_resource.connectSockets(
      blend_node,
      "Output",
      graph_resource.getGraphOutputNode(),
      "GraphOutput");

  graph_resource.saveToFile("Blend2Graph.animgraph.json");
  AnimGraphResource graph_resource_loaded;
  graph_resource_loaded.loadFromFile("Blend2Graph.animgraph.json");

  AnimGraph graph;
  graph_resource_loaded.createInstance(graph);
  AnimGraphContext graph_context;

  ozz::animation::Skeleton skeleton;
  REQUIRE(load_skeleton(skeleton, "media/skeleton.ozz"));
  graph_context.m_skeleton = &skeleton;

  REQUIRE(graph.init(graph_context));

  REQUIRE(graph.m_nodes.size() == 5);
  REQUIRE(graph.m_nodes[0]->m_node_type_name == "BlendTree");
  REQUIRE(graph.m_nodes[1]->m_node_type_name == "BlendTree");
  REQUIRE(graph.m_nodes[2]->m_node_type_name == "AnimSampler");
  REQUIRE(graph.m_nodes[3]->m_node_type_name == "AnimSampler");
  REQUIRE(graph.m_nodes[4]->m_node_type_name == "Blend2");

  // connections within the graph
  AnimSamplerNode* anim_sampler_walk =
      dynamic_cast<AnimSamplerNode*>(graph.m_nodes[2]);
  AnimSamplerNode* anim_sampler_run =
      dynamic_cast<AnimSamplerNode*>(graph.m_nodes[3]);
  Blend2Node* blend2_instance = dynamic_cast<Blend2Node*>(graph.m_nodes[4]);

  // check node input dependencies
  size_t anim_sampler_index0 = anim_sampler_walk->m_index;
  size_t anim_sampler_index1 = anim_sampler_run->m_index;
  size_t blend_index = blend2_instance->m_index;

  REQUIRE(graph.m_node_input_connections[blend_index].size() == 2);
  CHECK(
      graph.m_node_input_connections[blend_index][0].m_source_node
      == anim_sampler_walk);
  CHECK(
      graph.m_node_input_connections[blend_index][1].m_source_node
      == anim_sampler_run);

  REQUIRE(graph.m_node_output_connections[anim_sampler_index0].size() == 1);
  CHECK(
      graph.m_node_output_connections[anim_sampler_index0][0].m_target_node
      == blend2_instance);

  REQUIRE(graph.m_node_output_connections[anim_sampler_index1].size() == 1);
  CHECK(
      graph.m_node_output_connections[anim_sampler_index1][0].m_target_node
      == blend2_instance);

  // Ensure animation sampler nodes use the correct files
  REQUIRE(anim_sampler_walk->m_filename == "media/Walking-loop.ozz");
  REQUIRE(anim_sampler_walk->m_animation != nullptr);

  REQUIRE(anim_sampler_run->m_filename == "media/Running0-loop.ozz");
  REQUIRE(anim_sampler_run->m_animation != nullptr);

  WHEN("Emulating Graph Evaluation") {
    CHECK(graph.m_anim_data_allocator.size() == 0);
    CHECK(blend2_instance->i_input0 == anim_sampler_walk->o_output);
    CHECK(blend2_instance->i_input1 == anim_sampler_run->o_output);

    const Socket* graph_output_socket = graph.getOutputSocket("GraphOutput");
    AnimData* graph_output =
        static_cast<AnimData*>(*graph_output_socket->m_reference.ptr_ptr);

    CHECK(
        graph_output->m_local_matrices.size()
        == graph_context.m_skeleton->num_soa_joints());

    CHECK(
        blend2_instance->o_output == *graph_output_socket->m_reference.ptr_ptr);
  }

  graph_context.freeAnimations();
}

TEST_CASE("InputAttributeConversion", "[AnimGraphResource]") {
  int node_id = 3321;
  int input_index = 221;
  int output_index = 125;
  int parsed_node_id;
  int parsed_input_index;
  int parsed_output_index;

  int attribute_id = GenerateInputAttributeId(node_id, input_index);
  SplitInputAttributeId(attribute_id, &parsed_node_id, &parsed_input_index);
  CHECK(node_id == parsed_node_id);
  CHECK(input_index == parsed_input_index);

  attribute_id = GenerateOutputAttributeId(node_id, output_index);
  SplitOutputAttributeId(attribute_id, &parsed_node_id, &parsed_output_index);
  CHECK(node_id == parsed_node_id);
  CHECK(output_index == parsed_output_index);
}

TEST_CASE("ResourceSaveLoadMathGraphInputs", "[AnimGraphResource]") {
  AnimGraphResource graph_resource_origin;

  graph_resource_origin.clear();
  graph_resource_origin.m_name = "TestInputOutputGraph";

  size_t float_to_vec3_node_index = graph_resource_origin.addNode(
      AnimNodeResourceFactory("MathFloatToVec3Node"));

  AnimNodeResource& graph_output_node =
      graph_resource_origin.getGraphOutputNode();
  graph_output_node.m_socket_accessor->RegisterInput<float>(
      "GraphFloatOutput",
      nullptr);
  graph_output_node.m_socket_accessor->RegisterInput<Vec3>(
      "GraphVec3Output",
      nullptr);

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

  // Prepare graph inputs and outputs
  AnimNodeResource& float_to_vec3_node =
      graph_resource_origin.m_nodes[float_to_vec3_node_index];

  REQUIRE(graph_resource_origin.connectSockets(
      graph_input_node,
      "GraphFloatInput",
      graph_output_node,
      "GraphFloatOutput"));

  REQUIRE(graph_resource_origin.connectSockets(
      graph_input_node,
      "GraphFloatInput",
      float_to_vec3_node,
      "Input0"));
  REQUIRE(graph_resource_origin.connectSockets(
      graph_input_node,
      "GraphFloatInput",
      float_to_vec3_node,
      "Input1"));
  REQUIRE(graph_resource_origin.connectSockets(
      graph_input_node,
      "GraphFloatInput",
      float_to_vec3_node,
      "Input2"));

  REQUIRE(graph_resource_origin.connectSockets(
      float_to_vec3_node,
      "Output",
      graph_output_node,
      "GraphVec3Output"));

  WHEN("Saving and loading graph resource") {
    const char* filename = "ResourceSaveLoadGraphInputs.json";
    graph_resource_origin.saveToFile(filename);

    AnimGraphResource graph_resource_loaded;
    graph_resource_loaded.loadFromFile(filename);

    const AnimNodeResource& graph_loaded_output_node =
        graph_resource_loaded.m_nodes[0];
    const AnimNodeResource& graph_loaded_input_node =
        graph_resource_loaded.m_nodes[1];

    THEN("Graph inputs and outputs must be in loaded resource as well.") {
      REQUIRE(
          graph_output_node.m_socket_accessor->m_inputs.size()
          == graph_loaded_output_node.m_socket_accessor->m_inputs.size());

      REQUIRE(
          graph_input_node.m_socket_accessor->m_outputs.size()
          == graph_loaded_input_node.m_socket_accessor->m_outputs.size());

      REQUIRE(
          graph_loaded_input_node.m_socket_accessor->GetOutputSocket(
              "GraphFloatInput")
          != nullptr);

      REQUIRE(
          graph_loaded_output_node.m_socket_accessor->GetInputSocket(
              "GraphFloatOutput")
          != nullptr);
      REQUIRE(
          graph_loaded_output_node.m_socket_accessor->GetInputSocket(
              "GraphVec3Output")
          != nullptr);

      WHEN("Instantiating an AnimGraph") {
        AnimGraph anim_graph;
        graph_resource_loaded.createInstance(anim_graph);

        REQUIRE(anim_graph.getInputSocket("GraphFloatInput") != nullptr);
        REQUIRE(
            anim_graph.getInputPtr("GraphFloatInput")
            == anim_graph.m_input_buffer);

        float graph_float_input = 123.456f;
        anim_graph.SetInput("GraphFloatInput", &graph_float_input);

        AND_WHEN("Evaluating Graph") {
          AnimGraphContext context;
          context.m_graph = &anim_graph;

          anim_graph.init(context);

          // GraphFloatOutput is directly connected to GraphFloatInput therefore
          // we need to get the pointer here.
          float* graph_float_ptr = nullptr;
          graph_float_ptr = anim_graph.GetOutputPtr<float>("GraphFloatOutput");
          Vec3 graph_vec3_output;
          anim_graph.SetOutput("GraphVec3Output", &graph_vec3_output);

          anim_graph.updateTime(0.f);
          anim_graph.evaluate(context);

          THEN("output vector components equal the graph input vaulues") {
            CHECK(graph_float_ptr == &graph_float_input);
            CHECK(graph_vec3_output.v[0] == graph_float_input);
            CHECK(graph_vec3_output.v[1] == graph_float_input);
            CHECK(graph_vec3_output.v[2] == graph_float_input);
          }

          context.freeAnimations();
        }
      }
    }
  }
}

TEST_CASE("SimpleMathEvaluations", "[AnimGraphResource]") {
  AnimGraphResource graph_resource_origin;

  graph_resource_origin.clear();
  graph_resource_origin.m_name = "TestInputOutputGraph";

  size_t math_add0_node_index =
      graph_resource_origin.addNode(AnimNodeResourceFactory("MathAddNode"));

  size_t math_add1_node_index =
      graph_resource_origin.addNode(AnimNodeResourceFactory("MathAddNode"));

  AnimNodeResource& graph_output_node =
      graph_resource_origin.getGraphOutputNode();

  graph_output_node.m_socket_accessor->RegisterInput<float>(
      "GraphFloat0Output",
      nullptr);
  graph_output_node.m_socket_accessor->RegisterInput<float>(
      "GraphFloat1Output",
      nullptr);
  graph_output_node.m_socket_accessor->RegisterInput<float>(
      "GraphFloat2Output",
      nullptr);

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

  // Prepare graph inputs and outputs
  AnimNodeResource& math_add0_node =
      graph_resource_origin.m_nodes[math_add0_node_index];
  AnimNodeResource& math_add1_node =
      graph_resource_origin.m_nodes[math_add1_node_index];

  // direct output
  REQUIRE(graph_resource_origin.connectSockets(
      graph_input_node,
      "GraphFloatInput",
      graph_output_node,
      "GraphFloat0Output"));

  // add0 node
  REQUIRE(graph_resource_origin.connectSockets(
      graph_input_node,
      "GraphFloatInput",
      math_add0_node,
      "Input0"));

  REQUIRE(graph_resource_origin.connectSockets(
      graph_input_node,
      "GraphFloatInput",
      math_add0_node,
      "Input1"));

  REQUIRE(graph_resource_origin.connectSockets(
      math_add0_node,
      "Output",
      graph_output_node,
      "GraphFloat1Output"));

  // add1 node
  REQUIRE(graph_resource_origin.connectSockets(
      math_add0_node,
      "Output",
      math_add1_node,
      "Input0"));

  REQUIRE(graph_resource_origin.connectSockets(
      graph_input_node,
      "GraphFloatInput",
      math_add1_node,
      "Input1"));

  REQUIRE(graph_resource_origin.connectSockets(
      math_add1_node,
      "Output",
      graph_output_node,
      "GraphFloat2Output"));

  WHEN("Saving and loading graph resource") {
    const char* filename = "ResourceSaveLoadGraphInputs.json";
    graph_resource_origin.saveToFile(filename);

    AnimGraphResource graph_resource_loaded;
    graph_resource_loaded.loadFromFile(filename);

    const AnimNodeResource& graph_loaded_output_node =
        graph_resource_loaded.m_nodes[0];
    const AnimNodeResource& graph_loaded_input_node =
        graph_resource_loaded.m_nodes[1];

    WHEN("Instantiating an AnimGraph") {
      AnimGraph anim_graph;
      graph_resource_loaded.createInstance(anim_graph);

      REQUIRE(anim_graph.getInputSocket("GraphFloatInput") != nullptr);
      REQUIRE(
          anim_graph.getInputPtr("GraphFloatInput")
          == anim_graph.m_input_buffer);

      float graph_float_input = 123.456f;
      anim_graph.SetInput("GraphFloatInput", &graph_float_input);

      AND_WHEN("Evaluating Graph") {
        AnimGraphContext context;
        context.m_graph = &anim_graph;

        // float0 output is directly connected to the graph input, therefore
        // we have to get a ptr to the input data here.
        float* float0_output_ptr = nullptr;
        float float1_output = -1.f;
        float float2_output = -1.f;

        float0_output_ptr = anim_graph.GetOutputPtr<float>("GraphFloat0Output");

        anim_graph.SetOutput("GraphFloat1Output", &float1_output);
        anim_graph.SetOutput("GraphFloat2Output", &float2_output);

        anim_graph.updateTime(0.f);
        anim_graph.evaluate(context);

        THEN("output vector components equal the graph input vaulues") {
          CHECK(*float0_output_ptr == Approx(graph_float_input));
          CHECK(float1_output == Approx(graph_float_input * 2.f));
          REQUIRE_THAT(
              float2_output,
              Catch::Matchers::WithinAbs(graph_float_input * 3.f, 10));
        }

        context.freeAnimations();
      }
    }
  }
}