//
// 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 "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_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();
}

void AnimationController::Update(float dt) {
  m_current_time += dt;
  if (!m_override_ratio_A) {
    m_anim_time_A += dt * m_time_scale_A;
    const float duration_A = m_blend_anim_A->duration();
    m_anim_ratio_A = fmodf((float)m_anim_time_A / duration_A, 1.0f);
  }

  if (!m_override_ratio_B) {
    m_anim_time_B += dt * m_time_scale_B;
    const float duration_B = m_blend_anim_B->duration();
    m_anim_ratio_B = fmodf((float)m_anim_time_B / duration_B, 1.0f);
  }
}

void AnimationController::Evaluate() {
  // sample animation A
  ozz::animation::SamplingJob sampling_job;
  sampling_job.animation = m_blend_anim_A;
  sampling_job.cache = &m_sampling_cache_A;
  sampling_job.ratio = m_anim_ratio_A;
  sampling_job.output = make_span(m_local_matrices_A);
  if (!sampling_job.Run()) {
    ozz::log::Err() << "Error sampling animation." << std::endl;
  }

  // sample animation B
  sampling_job.animation = m_blend_anim_B;
  sampling_job.cache = &m_sampling_cache_B;
  sampling_job.ratio = m_anim_ratio_B;
  sampling_job.output = make_span(m_local_matrices_B);
  if (!sampling_job.Run()) {
    ozz::log::Err() << "Error sampling animation." << std::endl;
  }

  // perform blend
  ozz::animation::BlendingJob::Layer layers[2];
  layers[0].transform = make_span(m_local_matrices_A);
  layers[0].weight = (1.0f - m_blend_weight);

  layers[1].transform = make_span(m_local_matrices_B);
  layers[1].weight = (m_blend_weight);

  ozz::animation::BlendingJob blend_job;
  blend_job.threshold = ozz::animation::BlendingJob().threshold;
  blend_job.layers = layers;
  blend_job.bind_pose = m_skinned_mesh->m_skeleton.joint_bind_poses();
  blend_job.output = make_span(m_local_matrices_blended);

  if (!blend_job.Run()) {
    ozz::log::Err() << "Error blending animations." << std::endl;
  }

  // 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_blended);
  ltm_job.output = make_span(m_skinned_mesh->m_model_matrices);
  ltm_job.Run();
};

void AnimationController::DrawDebugUi() {
  ImGui::Begin("AnimationController");

  int anim_count = m_skinned_mesh->m_animation_names.size();
  const char* items[255] = {0};
  int item_current_A = 0;
  int item_current_B = 0;
  for (int i = 0; i < anim_count; i++) {
    items[i] = m_skinned_mesh->m_animation_names[i].c_str();
    if (m_skinned_mesh->m_animations[i] == m_blend_anim_A) {
      item_current_A = i;
    }
    if (m_skinned_mesh->m_animations[i] == m_blend_anim_B) {
      item_current_B = i;
    }
  }

  bool anim_changed = false;
  anim_changed =
      ImGui::Combo("Animation A", &item_current_A, items, anim_count);
  ImGui::SliderFloat("Time Scale", &m_time_scale_A, 0.01f, 5.f);
  ImGui::Checkbox("Override", &m_override_ratio_A);
  ImGui::SameLine();
  ImGui::SliderFloat("Ratio", &m_anim_ratio_A, 0.f, 1.f);
  anim_changed = ImGui::Combo("Animation B", &item_current_B, items, anim_count)
                 || anim_changed;
  ImGui::SliderFloat("Time Scale##", &m_time_scale_B, 0.01f, 5.f);
  ImGui::Checkbox("Override##", &m_override_ratio_B);
  ImGui::SameLine();
  ImGui::SliderFloat("Ratio##", &m_anim_ratio_B, 0.f, 1.f);

  SetBlendAnims(
      m_skinned_mesh->m_animations[item_current_A],
      m_skinned_mesh->m_animations[item_current_B]);
  ImGui::SliderFloat("Weight", &m_blend_weight, 0.f, 1.f);
  ImGui::End();
}