//=============================================================================

#include <graphene/character/data_structure/Skeleton.h>
#include <graphene/geometry/Matrix3x3.h>
#include <Filename_utilities.h>

//=============================================================================

namespace graphene {
namespace character {

//=============================================================================


Skeleton::
Skeleton()
  : root_(NULL)
{
}


//-----------------------------------------------------------------------------

Skeleton
::Skeleton(const Skeleton &rhs)
{
    operator=(rhs);
}

//-----------------------------------------------------------------------------


Skeleton::
~Skeleton()
{
    size_t i;
    for (i=0; i < joints_.size(); ++i)
    {
        delete joints_[i];
    }
}


//-----------------------------------------------------------------------------

Skeleton& Skeleton::operator =(const Skeleton& rhs)
{
    if (this != &rhs)
    {
        for (size_t i=0; i < joints_.size(); ++i)
        {
            delete joints_[i];
        }
        root_ = NULL;
        skinning_matrices_.clear();

        size_t i;
        Joint* joint,*new_joint;

        joints_.resize(rhs.joints_.size());
        for (i=0; i < rhs.joints_.size(); ++i)
        {
            joint = rhs.joints_[i];
            joints_[i] = new Joint(joint->get_name().c_str(), joint->local_, joint->bind_pose_local_, joint->initial_,
                                   joint->global_, joint->global_inv_, joint->final_);
        }

        for (i=0; i < rhs.joints_.size(); ++i)
        {
            joint = rhs.joints_[i];
            new_joint = joints_[i];

            if (joint->parent_)
            {
                new_joint->parent_ = get_joint(joint->parent_->get_name().c_str());
                new_joint->parent_->children_.push_back(new_joint);
            }
            else
            {
                root_ = new_joint;
            }
        }

        skinning_matrices_ = rhs.skinning_matrices_;
    }

    return *this;
}

//-----------------------------------------------------------------------------

void
Skeleton::
init()
{
    skinning_matrices_.clear();
    skinning_matrices_.resize(joints_.size(), Mat4f::identity());

    if (root_)
        root_->init();

    set_current_pose_as_bindpose();
    update();
}


//-----------------------------------------------------------------------------


void
Skeleton::
update()
{
    if (root_ == NULL)
        return;

    root_->update();

    Joint* j;
    for (unsigned int i = 0; i < joints_.size(); ++i)
    {
        j = joints_[i];

        skinning_matrices_[i] = j->final_;
    }
}

//-----------------------------------------------------------------------------

void
Skeleton::
update_global_bindpose()
{
    for (auto* j : joints_)
    {
        j->global_inv_ = inverse(j->global_);
    }
}

//-----------------------------------------------------------------------------


void
Skeleton::set_current_pose_as_bindpose()
{
    size_t i;
    Joint* joint;

    for (i = 0; i < joints_.size(); ++i)
    {
        joint = joints_[i];
        joint->bind_pose_local_ = joint->local_;
    }

    update();
    update_global_bindpose();
    update();
}


//-----------------------------------------------------------------------------


void
Skeleton::reset_to_initial(bool keep_translation)
{
    Joint* joint;
    size_t i;

    if (keep_translation)
    {
        int m, n;
        for (i = 0; i < joints_.size(); ++i)
        {
            joint = joints_[i];
            for (m = 0; m < 3; ++m)
                for (n = 0; n < 3; ++n)
                    joint->local_(m, n) = joint->initial_(m, n);
        }
    }
    else
    {
        for (i = 0; i < joints_.size(); ++i)
        {
            joint = joints_[i];
            joint->local_ = joint->initial_;
        }
    }

    update();
}


//-----------------------------------------------------------------------------

void
Skeleton::
reset_to_bindpose()
{
    Joint* joint;
    for (size_t i = 0; i < joints_.size(); ++i)
    {
        joint = joints_[i];
        joint->local_ = joint->bind_pose_local_;
    }

    update();
}

//-----------------------------------------------------------------------------

void
Skeleton::
apply_transform(const Mat4f& transform)
{
    root_->local_ = transform * root_->local_;

    update();
}

//-----------------------------------------------------------------------------

Joint*
Skeleton::
get_joint(const char *joint_name) const
{
    size_t i;
    Joint* j;
    for (i=0; i < joints_.size(); ++i)
    {
        j = joints_[i];
        if (j->name_ == joint_name)
        {
            return j;
        }
    }
    return NULL;
}

//-----------------------------------------------------------------------------


int
Skeleton::
get_joint_idx(const char *joint_name) const
{
    size_t i;
    Joint* j;
    for (i=0; i < joints_.size(); ++i)
    {
        j = joints_[i];
        if (j->name_ == joint_name)
        {
            return (int)i;
        }
    }
    return -1;
}


//-----------------------------------------------------------------------------


void
Skeleton::
align_axes()
{
    reset_to_bindpose();
    root_->init_aligned_axes();
}


//-----------------------------------------------------------------------------


bool
Skeleton::
match_angles_between_joints(const Skeleton& desired_skeleton, bool only_hand_joints)
{
    if (only_hand_joints)
    {
        // compensate difference in angles between joints
        for (size_t joint_index = 0; joint_index < joints_.size(); joint_index++)
        {
            auto* desired_joint = desired_skeleton.joints_[joint_index];
            auto* current_joint = joints_[joint_index];

            if (current_joint->children_.size() == 0)
                continue;

            if (current_joint->children_.size() > 1)
                continue;

            Mat3f global_inv = transpose(Mat3f(current_joint->global_));

            auto initial_bone_direction = normalize(global_inv * (desired_joint->children_[0]->get_global_translation() -
                                                                  desired_joint->get_global_translation()));
            auto current_bone_direction = normalize(global_inv * (current_joint->children_[0]->get_global_translation() -
                                                                  current_joint->get_global_translation()));

            double cos_angle = dot(current_bone_direction, initial_bone_direction);

            // NOTE(swenninger): Do nothing, if directions are already parallel
            if (cos_angle >= 0.9999f)
                continue;

            const std::string& joint_name = current_joint->name_;

            // TODO(swenninger): @HARDCODED

            bool is_hand_joint =
                filename::starts_with(joint_name, "LeftFinger")  ||
                filename::starts_with(joint_name, "LeftHand")    ||
                filename::starts_with(joint_name, "RightFinger") ||
                filename::starts_with(joint_name, "RightHand");

            bool is_hand_base = joint_name == "RightHand" ||
                                joint_name == "LeftHand";

            if (!is_hand_joint || is_hand_base)
            {
                continue;
            }

            Vec3f axis = normalize(cross(current_bone_direction, initial_bone_direction));
            Mat4f R = Mat4f::rotate(axis, acos(cos_angle) * 180 / M_PI);

            current_joint->local_ = current_joint->local_ * R;
            update();
        }
    }
    else
    {
        // Set same rotation
        for (size_t joint_index = 0; joint_index < joints_.size(); joint_index++)
        {
            auto* desired_joint = desired_skeleton.joints_[joint_index];
            auto* current_joint = joints_[joint_index];

            for (int i = 0; i < 3; ++i)
                for (int j = 0; j < 3; ++j)
                    current_joint->local_(i,j) = desired_joint->local_(i,j);
        }

        update();


        {
            // NOTE(swenninger): Treat spine seperately

            auto* current_spine_hip  = get_joint("Spine");
            auto* current_spine_head = get_joint("Neck");
            auto* desired_spine_hip  = desired_skeleton.get_joint("Spine");
            auto* desired_spine_head = desired_skeleton.get_joint("Neck");

            if (!current_spine_hip  ||
                !current_spine_head ||
                !desired_spine_hip  ||
                !desired_spine_head)
            {
                std::cerr << "[ERROR] Set TPose: Cannot get joints. Joint names are hardcoded?" << std::endl;
                return false;
            }


            Mat3f global_inv = transpose(Mat3f(current_spine_hip->global_));

            auto desired_bone_direction = normalize(global_inv * (desired_spine_head->get_global_translation() -
                                                                  desired_spine_hip->get_global_translation()));
            auto current_bone_direction = normalize(global_inv * (current_spine_head->get_global_translation() -
                                                                  current_spine_hip->get_global_translation()));

            double cos_angle = dot(current_bone_direction, desired_bone_direction);

            if (cos_angle < 0.9999f)
            {
                Vec3f axis = normalize(cross(current_bone_direction, desired_bone_direction));
                Mat4f R = Mat4f::rotate(axis, acos(cos_angle) * 180 / M_PI);

                current_spine_hip->local_ = current_spine_hip->local_ * R;
                update();
            }
        }

        // compensate difference in angles between joints
        for (size_t joint_index = 0; joint_index < joints_.size(); joint_index++)
        {
            auto* desired_joint = desired_skeleton.joints_[joint_index];
            auto* current_joint = joints_[joint_index];

            if (current_joint->children_.size() == 0)
                continue;

            if (current_joint->children_.size() > 1)
                continue;

            Mat3f global_inv = transpose(Mat3f(current_joint->global_));

            auto initial_bone_direction = normalize(global_inv * (desired_joint->children_[0]->get_global_translation() -
                                                                  desired_joint->get_global_translation()));
            auto current_bone_direction = normalize(global_inv * (current_joint->children_[0]->get_global_translation() -
                                                                  current_joint->get_global_translation()));

            double cos_angle = dot(current_bone_direction, initial_bone_direction);

            // NOTE(swenninger): Do nothing, if directions are already parallel
            if (cos_angle >= 0.9999f)
                continue;

            // TODO(swenninger): @HARDCODED
            if (filename::starts_with(current_joint->name_, "LeftFinger") ||
                filename::starts_with(current_joint->name_, "LeftHand")   ||
                filename::starts_with(current_joint->name_, "RightFinger") ||
                filename::starts_with(current_joint->name_, "RightHand") ||
                filename::starts_with(current_joint->name_, "Spine")     ||
                filename::starts_with(current_joint->name_, "Hips"))
            {
                continue;
            }

            Vec3f axis = normalize(cross(current_bone_direction, initial_bone_direction));
            Mat4f R = Mat4f::rotate(axis, acos(cos_angle) * 180 / M_PI);

            current_joint->local_ = current_joint->local_ * R;
            update();
        }

        {
            // Left Hand
            auto* current_joint = get_joint("LeftFingerBase");
            auto* current_child = current_joint->children_[0];
            auto* desired_joint = desired_skeleton.get_joint("LeftFingerBase");
            auto* desired_child = desired_joint->children_[0];

            Mat3f global_inv = transpose(Mat3f(current_joint->global_));

            auto desired_bone_direction = normalize(global_inv * (desired_child->get_global_translation() -
                                                                  desired_joint->get_global_translation()));
            auto current_bone_direction = normalize(global_inv * (current_child->get_global_translation() -
                                                                  current_joint->get_global_translation()));

            double cos_angle = dot(current_bone_direction, desired_bone_direction);

            if (cos_angle < 0.9999f)
            {
                Vec3f axis = normalize(cross(current_bone_direction, desired_bone_direction));
                Mat4f R = Mat4f::rotate(axis, acos(cos_angle) * 180 / M_PI);

                current_joint->local_ = current_joint->local_ * R;
                update();
            }

            current_child = get_joint("LeftHandMiddle4");
            desired_child = desired_skeleton.get_joint("LeftHandMiddle4");

            desired_bone_direction = normalize(global_inv * (desired_child->get_global_translation() -
                                                             desired_joint->get_global_translation()));
            current_bone_direction = normalize(global_inv * (current_child->get_global_translation() -
                                                             current_joint->get_global_translation()));

            cos_angle = dot(current_bone_direction, desired_bone_direction);

            if (cos_angle < 0.9999f)
            {
                Vec3f axis = normalize(cross(current_bone_direction, desired_bone_direction));
                Mat4f R = Mat4f::rotate(axis, acos(cos_angle) * 180 / M_PI);

                current_joint->local_ = current_joint->local_ * R;
                update();
            }
        }

        {
            // Right Hand
            auto* current_joint = get_joint("RightFingerBase");
            auto* desired_joint = desired_skeleton.get_joint("RightFingerBase");

            if (!current_joint ||
                !desired_joint)
            {
                std::cerr << "[ERROR] Set TPose: Cannot get joints. Joint names are hardcoded?" << std::endl;
                return false;
            }

            auto* current_child = current_joint->children_[0];
            auto* desired_child = desired_joint->children_[0];

            Mat3f global_inv = transpose(Mat3f(current_joint->global_));

            auto desired_bone_direction = normalize(global_inv * (desired_child->get_global_translation() -
                                                                  desired_joint->get_global_translation()));
            auto current_bone_direction = normalize(global_inv * (current_child->get_global_translation() -
                                                                  current_joint->get_global_translation()));

            double cos_angle = dot(current_bone_direction, desired_bone_direction);

            if (cos_angle < 0.9999f)
            {
                Vec3f axis = normalize(cross(current_bone_direction, desired_bone_direction));
                Mat4f R = Mat4f::rotate(axis, acos(cos_angle) * 180 / M_PI);

                current_joint->local_ = current_joint->local_ * R;
                update();
            }

            current_child = get_joint("RightHandMiddle4");
            desired_child = desired_skeleton.get_joint("RightHandMiddle4");

            desired_bone_direction = normalize(global_inv * (desired_child->get_global_translation() -
                                                             desired_joint->get_global_translation()));
            current_bone_direction = normalize(global_inv * (current_child->get_global_translation() -
                                                             current_joint->get_global_translation()));

            cos_angle = dot(current_bone_direction, desired_bone_direction);

            if (cos_angle < 0.9999f)
            {
                Vec3f axis = normalize(cross(current_bone_direction, desired_bone_direction));
                Mat4f R = Mat4f::rotate(axis, acos(cos_angle) * 180 / M_PI);

                current_joint->local_ = current_joint->local_ * R;
                update();
            }
        }

        update();
    }

    return true;
}


//-----------------------------------------------------------------------------


bool
Skeleton::
match_angles_between_joints_with_single_child(const Skeleton& desired_skeleton)
{
    // compensate difference in angles between joints

    // Set same rotation
    for (size_t joint_index = 0; joint_index < joints_.size(); joint_index++)
    {
        auto* desired_joint = desired_skeleton.joints_[joint_index];
        auto* current_joint = joints_[joint_index];

        if (filename::starts_with(desired_joint->name_, "clavicle"))
            continue;

        for (int i = 0; i < 3; ++i)
            for (int j = 0; j < 3; ++j)
                current_joint->local_(i,j) = desired_joint->local_(i,j);
    }

    update();

    for (size_t joint_index = 0; joint_index < joints_.size(); joint_index++)
    {
        auto* desired_joint = desired_skeleton.joints_[joint_index];
        auto* current_joint = joints_[joint_index];

        if (filename::starts_with(desired_joint->name_, "clavicle"))
            continue;


        if (current_joint->children_.size() == 0)
            continue;

        if (current_joint->children_.size() > 1)
            continue;

        Mat3f global_inv = transpose(Mat3f(current_joint->global_));

        auto initial_bone_direction = normalize(global_inv * (desired_joint->children_[0]->get_global_translation() -
                                                              desired_joint->get_global_translation()));
        auto current_bone_direction = normalize(global_inv * (current_joint->children_[0]->get_global_translation() -
                                                              current_joint->get_global_translation()));

        double cos_angle = dot(current_bone_direction, initial_bone_direction);

        // Debug
        //std::cout << "Angle between " << current_joint->name_ << " and " << current_joint->children_[0]->name_ << ": " << acos(cos_angle);

        // NOTE(swenninger): Do nothing, if directions are already parallel
        if (cos_angle >= 0.9999f)
            continue;

        Vec3f axis = normalize(cross(current_bone_direction, initial_bone_direction));
        Mat4f R = Mat4f::rotate(axis, acos(cos_angle) * 180 / M_PI);

        current_joint->local_ = current_joint->local_ * R;
        update();
    }

    return true;
}

//-----------------------------------------------------------------------------

bool Skeleton::match_hand_joint_angles(const Skeleton& desired_skeleton)
{
    // TODO(swenninger) hardcoded joint names

    const static std::vector<std::string> joint_names_to_adapt = {
//        "LeftHandThumb1",
        "LeftHandThumb2",
        "LeftHandThumb3",
        "LeftHandIndex1",
        "LeftHandIndex2",
        "LeftHandIndex3",
        "LeftHandMiddle1",
        "LeftHandMiddle2",
        "LeftHandMiddle3",
        "LeftHandRing1",
        "LeftHandRing2",
        "LeftHandRing3",
        "LeftHandPinky1",
        "LeftHandPinky2",
        "LeftHandPinky3",

//        "RightHandThumb1",
        "RightHandThumb2",
        "RightHandThumb3",
        "RightHandIndex1",
        "RightHandIndex2",
        "RightHandIndex3",
        "RightHandMiddle1",
        "RightHandMiddle2",
        "RightHandMiddle3",
        "RightHandRing1",
        "RightHandRing2",
        "RightHandRing3",
        "RightHandPinky1",
        "RightHandPinky2",
        "RightHandPinky3",
    };


    for (size_t i = 0; i < joint_names_to_adapt.size(); ++i)
    {
        Joint* current_joint = get_joint(joint_names_to_adapt[i].c_str());
        Joint* desired_joint = desired_skeleton.get_joint(joint_names_to_adapt[i].c_str());

        if (!current_joint || !desired_joint)
        {
            std::cerr << "[WARNING] Cannot find " << joint_names_to_adapt[i]
                      << ". Maybe joint names are hardcoded?" << std::endl;
            return false;
        }

        Mat3f global_inv = transpose(Mat3f(current_joint->global_));

        auto initial_bone_direction = normalize(global_inv * (desired_joint->children_[0]->get_global_translation() -
                                                              desired_joint->get_global_translation()));
        auto current_bone_direction = normalize(global_inv * (current_joint->children_[0]->get_global_translation() -
                                                              current_joint->get_global_translation()));

        double cos_angle = dot(current_bone_direction, initial_bone_direction);

        // Debug
        //std::cout << "Angle between " << current_joint->name_ << " and " << current_joint->children_[0]->name_ << ": " << acos(cos_angle);

        // NOTE(swenninger): Do nothing, if directions are already parallel
        if (cos_angle >= 0.9999f)
            continue;

        Vec3f axis = normalize(cross(current_bone_direction, initial_bone_direction));
        Mat4f R = Mat4f::rotate(axis, acos(cos_angle) * 180 / M_PI);

        current_joint->local_ = current_joint->local_ * R;
        update();

    }

    return true;
}

//-----------------------------------------------------------------------------

void Skeleton::print_angles()
{
    for (size_t i = 0; i < joints_.size(); ++i)
    {
        auto* j = joints_[i];

        Mat4f bindpose = j->bind_pose_local_;
        Mat4f local    = j->local_;

        bindpose[12] = bindpose[13] = bindpose[14] = 0.0f;
        local[12] = local[13] = local[14] = 0.0f;

        Mat4f m = inverse(bindpose) * local;

        Vec3f euler = Mat4f::to_euler_XYZ(m);
        euler *= 180.0f / M_PI;

        //if (euler[0] > 0.01f ||
//            euler[1] > 0.01f ||
//            euler[2] > 0.01f)
        {
            std::cout << j->name_ << std::endl;
            printf("(%f, %f, %f)\n", euler[0], euler[1], euler[2]);
        }
    }
}

//=============================================================================
} // namespace character
} // namespace graphene
//=============================================================================
