//
// Created by martin on 12.11.21.
//

#include "AnimationController.h"

#include <imgui.h>
#include <ozz/animation/runtime/local_to_model_job.h>
#include <ozz/animation/runtime/sampling_job.h>

#include <queue>

#include "AnimNodes/AnimSamplerNode.h"
#include "AnimNodes/BlendNode.h"
#include "AnimNodes/BlendSpace1D.h"
#include "AnimNodes/LockTranslationNode.h"
#include "AnimNodes/SpeedScaleNode.h"
#include "SkinnedMesh.h"

AnimationController::AnimationController(SkinnedMesh* skinned_mesh)
    : m_current_time(0.f), m_paused(true), m_skinned_mesh(skinned_mesh) {
  const int num_soa_joints = skinned_mesh->m_skeleton.num_soa_joints();
  const int num_joints = skinned_mesh->m_skeleton.num_joints();
  skinned_mesh->m_local_matrices.resize(num_soa_joints);
  skinned_mesh->m_model_matrices.resize(num_joints);

  ResetAnims();

  AnimSamplerNode* sampler_node0 = new AnimSamplerNode(this);
  sampler_node0->m_name = "AnimSampler0";
  sampler_node0->SetAnimation(
      skinned_mesh->m_animations[1],
      skinned_mesh->m_animation_sync_track[1]);
  m_anim_nodes.push_back(sampler_node0);

  AnimSamplerNode* sampler_node1 = new AnimSamplerNode(this);
  sampler_node1->m_name = "AnimSampler1";
  sampler_node1->SetAnimation(
      skinned_mesh->m_animations[2],
      skinned_mesh->m_animation_sync_track[2]);
  m_anim_nodes.push_back(sampler_node1);

  AnimSamplerNode* sampler_node2 = new AnimSamplerNode(this);
  sampler_node2->m_name = "AnimSampler2";
  sampler_node2->SetAnimation(
      skinned_mesh->m_animations[3],
      skinned_mesh->m_animation_sync_track[3]);
  m_anim_nodes.push_back(sampler_node2);

  BlendSpace1D* blend_space = new BlendSpace1D(this);
  blend_space->m_name = "BlendSpace0";
  blend_space->m_weight = 0.f;
  blend_space->m_inputs.push_back(sampler_node0);
  blend_space->m_input_weights.push_back(-1.f);
  blend_space->m_inputs.push_back(sampler_node1);
  blend_space->m_input_weights.push_back(0.f);
  blend_space->m_inputs.push_back(sampler_node2);
  blend_space->m_input_weights.push_back(1.f);
  m_anim_nodes.push_back(blend_space);

  SpeedScaleNode* speed_node = new SpeedScaleNode(this);
  speed_node->m_name = "SpeedNode0";
  speed_node->m_input_node = sampler_node0;
  m_anim_nodes.push_back(speed_node);

  BlendNode* blend_node = new BlendNode(this);
  blend_node->m_name = "Blend0";
  blend_node->m_input_A = speed_node;
  blend_node->m_input_B = sampler_node1;
  blend_node->m_sync_inputs = true;
  m_anim_nodes.push_back(blend_node);

  SpeedScaleNode* speed_node1 = new SpeedScaleNode(this);
  speed_node1->m_name = "SpeedNode1";
  speed_node1->m_input_node = blend_space;
  m_anim_nodes.push_back(speed_node1);

  LockTranslationNode* lock_node = new LockTranslationNode(this);
  lock_node->m_name = "LockNode0";
  lock_node->m_locked_bone_index = 0;
  lock_node->m_input = speed_node1;
  m_anim_nodes.push_back(lock_node);

  m_output_node = m_anim_nodes.back();

  UpdateOrderedNodes();

  m_output_node->Reset();
}

AnimationController::~AnimationController() {
  while (m_anim_nodes.size() > 0) {
    delete m_anim_nodes[m_anim_nodes.size() - 1];
    m_anim_nodes.pop_back();
  }

  m_output_node = nullptr;
}

void AnimationController::ResetAnims() {
  for (int i = 0; i < m_ordered_nodes.size(); i++) {
    m_ordered_nodes[i]->Reset();
  }
}

void AnimationController::UpdateOrderedNodes() {
  std::vector<AnimNode*> node_stack;
  node_stack.push_back(m_output_node);

  m_ordered_nodes.clear();

  while (node_stack.size() > 0) {
    AnimNode* node = node_stack.back();
    m_ordered_nodes.push_back(node);
    node_stack.pop_back();

    std::vector<AnimNode*> node_inputs;
    node->GetInputNodes(node_inputs);
    for (int i = node_inputs.size() - 1; i >= 0; i--) {
      node_stack.push_back(node_inputs[i]);
    }
  }
}

void AnimationController::UpdateTime(float dt) {
  if (m_output_node == nullptr) {
    return;
  }

  // Mark all nodes that evaluate time using sync tracks.
  m_output_node->UpdateIsSynced(false);

  if (m_paused) {
    return;
  }

  // For all synced nodes calculate their current sync track durations
  for (int i = m_ordered_nodes.size() - 1; i >= 0; i--) {
    AnimNode* node = m_ordered_nodes[i];
    if (node->m_is_time_synced) {
      node->UpdateSyncTrack();
    }
  }

  // Update the time of all nodes.
  m_output_node->UpdateTime(dt);
}

void AnimationController::Evaluate() {
  if (m_output_node == nullptr) {
    return;
  }

  m_output_node->Evaluate(
      &m_skinned_mesh->m_local_matrices,
      m_calc_root_transform ? &m_root_transform : nullptr);
};

void AnimationController::DrawDebugUi() {
  ImGui::SetNextWindowSize(ImVec2(500, 300), ImGuiCond_FirstUseEver);
  ImGui::Begin("AnimationController");

  if (ImGui::Button("Reset")) {
    ResetAnims();
  }
  ImGui::SameLine();
  if (m_paused) {
    if (ImGui::Button("Play")) {
      m_paused = false;
    }
  } else {
    if (ImGui::Button("Pause")) {
      m_paused = true;
    }
  }
  ImGui::SameLine();
  if (ImGui::Button("Step")) {
    bool was_paused = m_paused;
    m_paused = false;
    UpdateTime(0.1);
    Evaluate();
    m_paused = was_paused;
  }

  ImVec2 node_size(200, 100);
  for (int i = 0; i < m_ordered_nodes.size(); i++) {
    AnimNode* node = m_ordered_nodes[i];

    ImGui::SetNextWindowSize(node_size, ImGuiCond_FirstUseEver);
    ImGui::SetNextWindowPos(
        ImVec2((m_ordered_nodes.size() - 1 - i) * node_size.x - i * 10, 300),
        ImGuiCond_FirstUseEver);
    ImGui::Begin(node->m_name.c_str());
    node->DrawDebugUi();
    ImGui::End();
  }

  ImGui::Text("Node States");

  ImGui::Columns(4, "Node States");  // 4-ways, with border
  ImGui::Separator();
  ImGui::Text("Name");
  ImGui::NextColumn();
  ImGui::Text("Synced");
  ImGui::NextColumn();
  ImGui::Text("Duration");
  ImGui::NextColumn();
  ImGui::Text("Time");
  ImGui::NextColumn();

  ImGui::Separator();
  static int selected = -1;

  for (int i = 0; i < m_ordered_nodes.size(); i++) {
    AnimNode* node = m_ordered_nodes[i];
    if (ImGui::Selectable(
            node->m_name.c_str(),
            selected == i,
            ImGuiSelectableFlags_SpanAllColumns))
      selected = i;
    bool hovered = ImGui::IsItemHovered();
    ImGui::NextColumn();

    ImGui::Text(node->m_is_time_synced ? "X" : "-");
    ImGui::NextColumn();

    ImGui::PushID((void*)&node->m_sync_track.m_duration);
    ImGui::Text("%2.3f", node->m_sync_track.m_duration);
    ImGui::NextColumn();
    ImGui::PopID();

    ImGui::PushID((void*)&node->m_time_current);
    ImGui::Text("%2.3f", node->m_time_current);
    ImGui::NextColumn();
    ImGui::PopID();
  }

  ImGui::Columns(1);
  ImGui::Separator();

  ImGui::End();
}