#include "RigInverter.hpp"
#include "DataSpec/DNAMP1/CINF.hpp"
#include "DataSpec/DNAMP2/CINF.hpp"
#include "DataSpec/DNAMP3/CINF.hpp"
#include "hecl/Blender/Connection.hpp"

namespace DataSpec::DNAANIM {

template <class CINFType>
RigInverter<CINFType>::Bone::Bone(const CINFType& cinf, const typename CINFType::Bone& origBone)
: m_origBone(origBone) {
  atUint32 parentIdx = cinf.getInternalBoneIdxFromId(origBone.parentId);
  zeus::CVector3f boneOrigin(origBone.origin);
  zeus::CVector3f naturalTail = boneOrigin + zeus::CVector3f{0.f, 0.5f, 0.f};
  if (parentIdx != -1) {
    const typename CINFType::Bone& pBone = cinf.bones[parentIdx];
    m_parentDelta = boneOrigin - zeus::CVector3f(pBone.origin);
  }

  size_t actualChildren = 0;
  for (atUint32 chId : origBone.linked) {
    if (chId == origBone.parentId)
      continue;
    atUint32 chIdx = cinf.getInternalBoneIdxFromId(chId);
    if (chIdx != -1)
      ++actualChildren;
  }

  const std::string* bName = cinf.getBoneNameFromId(origBone.id);
  bool isLCTR = false;
  if (bName)
    isLCTR = bName->find("_LCTR") != std::string::npos;

  if (parentIdx == -1) {
    /* Root will always use +Y tail */
    m_tail = naturalTail;
  } else if (actualChildren) {
    /* Position tail to average of children */
    for (atUint32 chId : origBone.linked) {
      if (chId == origBone.parentId)
        continue;
      atUint32 chIdx = cinf.getInternalBoneIdxFromId(chId);
      if (chIdx != -1) {
        const typename CINFType::Bone& chBone = cinf.bones[chIdx];
        m_tail += chBone.origin;
      }
    }
    m_tail /= zeus::CVector3f(float(actualChildren));
    if ((m_tail - boneOrigin).magSquared() < 0.001f)
      m_tail = naturalTail;
    else if (isLCTR)
      m_tail = boneOrigin + zeus::CVector3f{0.f, 1.0f, 0.f} * (m_tail - boneOrigin).magnitude();
  } else if (parentIdx != -1) {
    /* Extrapolate by delta with parent */
    m_tail = boneOrigin + m_parentDelta;
    float deltaMag = m_parentDelta.magnitude();
    if (deltaMag < 0.001f) {
      deltaMag = 0.5f;
      m_tail = naturalTail;
    } else if (deltaMag > 0.5f) {
      /* Extreme bones capped to +0.5 value */
      deltaMag = 0.5f;
      m_tail = boneOrigin + m_parentDelta.normalized() * 0.5f;
    }

    if (isLCTR)
      m_tail = boneOrigin + zeus::CVector3f{0.f, 1.0f, 0.f} * deltaMag;
  } else {
    /* Fallback to +Y tail */
    m_tail = naturalTail;
  }
}

template <class CINFType>
RigInverter<CINFType>::RigInverter(const CINFType& cinf) : m_cinf(cinf) {
  m_bones.reserve(cinf.bones.size());
  for (const typename CINFType::Bone& b : cinf.bones)
    m_bones.emplace_back(cinf, b);
}

template <class CINFType>
RigInverter<CINFType>::RigInverter(const CINFType& cinf,
                                   const std::unordered_map<std::string, hecl::blender::Matrix3f>& matrices)
: m_cinf(cinf) {
  m_bones.reserve(cinf.bones.size());
  for (const typename CINFType::Bone& b : cinf.bones) {
    m_bones.emplace_back(cinf, b);

    const std::string* name = cinf.getBoneNameFromId(b.id);
    if (name) {
      auto search = matrices.find(*name);
      if (search != matrices.cend()) {
        zeus::CMatrix3f boneMtx(search->second[0], search->second[1], search->second[2]);
        m_bones.back().m_inverter = boneMtx.transposed();
        m_bones.back().m_restorer = boneMtx;
      }
    }
  }
}

template <class CINFType>
zeus::CQuaternion RigInverter<CINFType>::invertRotation(atUint32 boneId, const zeus::CQuaternion& origRot) const {
  for (const Bone& b : m_bones)
    if (b.m_origBone.id == boneId)
      return b.m_restorer * zeus::CMatrix3f(origRot) * b.m_inverter;
  return origRot;
}

template <class CINFType>
zeus::CVector3f RigInverter<CINFType>::invertPosition(atUint32 boneId, const zeus::CVector3f& origPos,
                                                      bool subDelta) const {
  for (const Bone& b : m_bones)
    if (b.m_origBone.id == boneId) {
      zeus::CVector3f localPos = origPos;
      if (subDelta)
        localPos -= b.m_parentDelta;
      return b.m_restorer * localPos;
    }
  return origPos;
}

template <class CINFType>
zeus::CQuaternion RigInverter<CINFType>::restoreRotation(atUint32 boneId, const zeus::CQuaternion& origRot) const {
  for (const Bone& b : m_bones)
    if (b.m_origBone.id == boneId)
      return b.m_inverter * zeus::CMatrix3f(origRot) * b.m_restorer;
  return origRot;
}

template <class CINFType>
zeus::CVector3f RigInverter<CINFType>::restorePosition(atUint32 boneId, const zeus::CVector3f& origPos,
                                                       bool subDelta) const {
  for (const Bone& b : m_bones)
    if (b.m_origBone.id == boneId) {
      zeus::CVector3f localPos = b.m_inverter * origPos;
      if (subDelta)
        localPos += b.m_parentDelta;
      return localPos;
    }
  return origPos;
}

template class RigInverter<DNAMP1::CINF>;
template class RigInverter<DNAMP2::CINF>;
template class RigInverter<DNAMP3::CINF>;

} // namespace DataSpec::DNAANIM