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

#include "AnimationController.h"

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

#include "AnimSamplerNode.h"
#include "SkinnedMesh.h"

AnimationController::AnimationController(SkinnedMesh* skinned_mesh)
    : m_current_time(0.f),
      m_skinned_mesh(skinned_mesh),
      m_blend_anim_A(nullptr),
      m_blend_anim_B(nullptr),
      m_blend_weight(0.f),
      m_time_scale_A(1.f),
      m_time_scale_B(1.f),
      m_override_ratio_A(false),
      m_override_ratio_B(false) {
  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);

  m_sampling_cache_A.Resize(num_joints);
  m_sampling_cache_B.Resize(num_joints);

  m_local_matrices.resize(num_soa_joints);
  m_local_matrices_A.resize(num_soa_joints);
  m_local_matrices_B.resize(num_soa_joints);
  m_local_matrices_blended.resize(num_soa_joints);

  assert(skinned_mesh->m_animations.size() > 1);
  SetBlendAnims(skinned_mesh->m_animations[1], skinned_mesh->m_animations[0]);
  ResetAnims();

  AnimSamplerNode* sampler_node = new AnimSamplerNode(this);
  sampler_node->SetAnimation(skinned_mesh->m_animations[0]);
  m_output_node = sampler_node;
  m_output_node->m_name = "AnimSampler0";
  m_output_node->Reset();
}

AnimationController::~AnimationController() {
  if (m_output_node) {
    delete m_output_node;
    m_output_node = nullptr;
  }
}

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

  m_output_node->Update(dt); }

void AnimationController::Evaluate() {
  if (m_output_node == nullptr) {
    return;
  }
  m_output_node->Evaluate(&m_local_matrices);

  // convert joint matrices from local to model space
  ozz::animation::LocalToModelJob ltm_job;
  ltm_job.skeleton = &m_skinned_mesh->m_skeleton;
  ltm_job.input = make_span(m_local_matrices);
  ltm_job.output = make_span(m_skinned_mesh->m_model_matrices);
  ltm_job.Run();
};

void AnimationController::DrawDebugUi() {
  ImGui::Begin("AnimationController");
  if (m_output_node && ImGui::TreeNode("Output")) {

    m_output_node->DrawDebugUi();

    ImGui::TreePop();
  }

  ImGui::End();
}