#include "Ragdoll.h"

#include <cfloat>
#include <iostream>
#include <string>

#include "Animator.h"
#include "Character.h"
#include "CollisionMesh.h"
#include "Physics.h"

struct PartDefinition {
    const char* bone;
    const char* endBone;   // set for hinges (elbows, knees): the far end of the part, used to find the bend
    float mass;
    float maxSwing;        // degrees the part can swing around the joint, in any direction
    float coneTilt;        // degrees the swing cone leans forward, so hips reach further forward than back
    float maxTwist;        // degrees the part can twist along its own length, either way
    float minBend;         // hinge range in degrees, 0 = straight
    float maxBend;
    float bendDirection;   // 1 = folds forward, -1 = folds backward. Only used if the pose starts straight
};

// The hand, foot, finger and neck bones ride along with the nearest part above them.
// The root part has to come first, and parents have to come before their children.
static const PartDefinition partDefinitions[] = {
    //  bone          endBone    mass  swing   tilt  twist  bend range     direction
    { "Hip",        nullptr,    14.0f,   0.0f,  0.0f,  0.0f, 0.0f,   0.0f,  0.0f },
    { "UpperTorso", nullptr,    22.0f,  30.0f,  0.0f, 25.0f, 0.0f,   0.0f,  0.0f },
    { "Head",       nullptr,     5.0f,  45.0f,  0.0f, 60.0f, 0.0f,   0.0f,  0.0f },
    { "L_UpperArm", nullptr,     3.0f, 110.0f,  0.0f, 60.0f, 0.0f,   0.0f,  0.0f },
    { "L_LowerArm", "L_Hand",    2.5f,   0.0f,  0.0f,  0.0f, 0.0f, 150.0f,  1.0f },
    { "R_UpperArm", nullptr,     3.0f, 110.0f,  0.0f, 60.0f, 0.0f,   0.0f,  0.0f },
    { "R_LowerArm", "R_Hand",    2.5f,   0.0f,  0.0f,  0.0f, 0.0f, 150.0f,  1.0f },
    { "L_UpperLeg", nullptr,     9.0f,  65.0f, 35.0f, 30.0f, 0.0f,   0.0f,  0.0f },
    { "L_LowerLeg", "L_Foot",    6.0f,   0.0f,  0.0f,  0.0f, 0.0f, 140.0f, -1.0f },
    { "R_UpperLeg", nullptr,     9.0f,  65.0f, 35.0f, 30.0f, 0.0f,   0.0f,  0.0f },
    { "R_LowerLeg", "R_Foot",    6.0f,   0.0f,  0.0f,  0.0f, 0.0f, 140.0f, -1.0f },
};
static constexpr int PART_COUNT = sizeof(partDefinitions) / sizeof(partDefinitions[0]);

static constexpr float PART_MIN_SIZE = 0.05f;
static constexpr float PART_MAX_SPEED = 30.0f;
static constexpr float PART_MAX_ANGULAR_SPEED = 30.0f;
static constexpr float SETTLE_SPEED = 0.2f;
static constexpr float SETTLE_ANGULAR_SPEED = 0.4f;
static constexpr float SETTLE_TIME = 0.6f;
static constexpr float FALL_LIMIT = -20.0f;
static constexpr float HINGE_MIN_POSE_BEND = 0.17f;   // sin of 10 degrees, below this the pose counts as straight

static const glm::quat IDENTITY_ROTATION(1.0f, 0.0f, 0.0f, 0.0f);

// Walks the skeleton, handing every bone to the part it hangs from and recording each part's parent
static void AssignBones(const AssimpNodeData& node, int currentPart, const Rig& rig,
    std::vector<int>& boneParts, int (&partBones)[PART_COUNT], int (&partParents)[PART_COUNT])
{
    const auto info = rig.m_BoneInfoMap.find(node.name);

    for (int i = 0; i < PART_COUNT; i++) {
        if (node.name != partDefinitions[i].bone || info == rig.m_BoneInfoMap.end())
            continue;
        partBones[i] = info->second.id;
        partParents[i] = i == 0 ? -1 : currentPart;
        currentPart = i;
        break;
    }

    if (info != rig.m_BoneInfoMap.end() && info->second.id >= 0 && info->second.id < (int)boneParts.size())
        boneParts[info->second.id] = currentPart;

    for (const AssimpNodeData& child : node.children)
        AssignBones(child, currentPart, rig, boneParts, partBones, partParents);
}

static glm::vec3 SkinVertex(const RigVertex& vertex, const std::vector<glm::mat4>& bones)
{
    glm::vec4 position(0.0f);
    float totalWeight = 0.0f;
    for (int i = 0; i < MAX_BONE_INFLUENCE; i++) {
        const int bone = vertex.m_BoneIDs[i];
        if (bone < 0 || bone >= (int)bones.size())
            continue;
        position += bones[bone] * glm::vec4(vertex.Position, 1.0f) * vertex.m_Weights[i];
        totalWeight += vertex.m_Weights[i];
    }
    return totalWeight > 0.0f ? glm::vec3(position) / totalWeight : vertex.Position;
}

Ragdoll::Ragdoll(CharacterBody& characterBody, const glm::mat4& rigMatrix, Physics& physics)
{
    static int nextGroup = 1;
    group = nextGroup++;

    Rig* rig = characterBody.rig;
    Animator* animator = characterBody.animator;
    if (!rig || !animator || !animator->m_CurrentAnimation) {
        std::cout << "RAGDOLL NEEDS A RIG AND AN ANIMATION\n";
        return;
    }

    // Bone matrices aren't filled in until the animator has stepped once
    characterBody.Step(0.001f);
    const std::vector<glm::mat4> poseBones = characterBody.GetBlendedBoneMatrices();

    int partBones[PART_COUNT];
    int partParents[PART_COUNT];
    for (int i = 0; i < PART_COUNT; i++) {
        partBones[i] = -1;
        partParents[i] = -1;
    }
    boneParts.assign(poseBones.size(), 0);
    AssignBones(animator->m_CurrentAnimation->GetRootNode(), 0, *rig, boneParts, partBones, partParents);

    restBones.resize(poseBones.size());
    for (int i = 0; i < (int)poseBones.size(); i++)
        restBones[i] = rigMatrix * poseBones[i];

    // Each triangle goes to whichever part most of its skin weight belongs to
    std::vector<Triangle> partTriangles[PART_COUNT];
    for (RigMesh* mesh : rig->meshes) {
        for (int i = 0; i + 2 < (int)mesh->indices.size(); i += 3) {
            const RigVertex* corners[3] = {
                &mesh->vertices[mesh->indices[i]],
                &mesh->vertices[mesh->indices[i + 1]],
                &mesh->vertices[mesh->indices[i + 2]]
            };

            float partWeights[PART_COUNT] = {};
            for (const RigVertex* corner : corners) {
                for (int j = 0; j < MAX_BONE_INFLUENCE; j++) {
                    const int bone = corner->m_BoneIDs[j];
                    if (bone >= 0 && bone < (int)boneParts.size())
                        partWeights[boneParts[bone]] += corner->m_Weights[j];
                }
            }
            int bestPart = 0;
            for (int j = 1; j < PART_COUNT; j++) {
                if (partWeights[j] > partWeights[bestPart])
                    bestPart = j;
            }

            partTriangles[bestPart].push_back({
                glm::vec3(rigMatrix * glm::vec4(SkinVertex(*corners[0], poseBones), 1.0f)),
                glm::vec3(rigMatrix * glm::vec4(SkinVertex(*corners[1], poseBones), 1.0f)),
                glm::vec3(rigMatrix * glm::vec4(SkinVertex(*corners[2], poseBones), 1.0f))
            });
        }
    }

    // Parts start axis aligned around their triangles, in unit space like every other body
    for (int i = 0; i < PART_COUNT; i++) {
        glm::vec3 boundsMin(FLT_MAX);
        glm::vec3 boundsMax(-FLT_MAX);
        for (const Triangle& tri : partTriangles[i]) {
            boundsMin = glm::min(boundsMin, glm::min(tri.a, glm::min(tri.b, tri.c)));
            boundsMax = glm::max(boundsMax, glm::max(tri.a, glm::max(tri.b, tri.c)));
        }
        if (partTriangles[i].empty()) {
            const int bone = partBones[i] >= 0 ? partBones[i] : 0;
            boundsMin = boundsMax = glm::vec3(restBones[bone][3]);
        }

        const glm::vec3 center = (boundsMin + boundsMax) * 0.5f;
        const glm::vec3 size = glm::max(boundsMax - boundsMin, glm::vec3(PART_MIN_SIZE));

        std::vector<Triangle> unitTriangles;
        unitTriangles.reserve(partTriangles[i].size());
        for (const Triangle& tri : partTriangles[i])
            unitTriangles.push_back({ (tri.a - center) / size, (tri.b - center) / size, (tri.c - center) / size });

        Part* part = new Part();
        part->restPosition = center;
        part->transform.t.TranslateTo(center);
        part->transform.t.ScaleTo(size);
        part->transform.t.RotateToQuaternion(IDENTITY_ROTATION);
        part->mesh = new CollisionMesh(unitTriangles);

        part->body = new p(&part->transform);
        part->body->mass = partDefinitions[i].mass;
        part->body->restitution = 0.0f;
        part->body->friction = 0.6f;
        part->body->collisionMesh = part->mesh;
        part->body->collisionGroup = group;
        part->body->maxSpeed = PART_MAX_SPEED;
        part->body->maxAngularSpeed = PART_MAX_ANGULAR_SPEED;
        physics.AddObject(part->body);

        parts.push_back(part);
        restInverses.push_back(glm::inverse(part->transform.t.GetTranslationMatrix()));
    }

    auto BonePosition = [&](const char* name, const glm::vec3& fallback) {
        const auto info = rig->m_BoneInfoMap.find(name);
        if (info == rig->m_BoneInfoMap.end())
            return fallback;
        return glm::vec3(restBones[info->second.id] * glm::inverse(info->second.offset) * glm::vec4(0.0f, 0.0f, 0.0f, 1.0f));
    };

    // The rig faces its local -Z, same as Character::GetForwardDirection
    const glm::vec3 forward = glm::normalize(glm::vec3(rigMatrix * glm::vec4(0.0f, 0.0f, -1.0f, 0.0f)));

    // Joints sit on the child bone's origin. Everything starts unrotated, so local = world here.
    glm::vec3 jointAnchors[PART_COUNT];
    jointAnchors[0] = parts[0]->restPosition;

    for (int i = 1; i < PART_COUNT; i++) {
        const PartDefinition& definition = partDefinitions[i];
        if (partBones[i] < 0 || partParents[i] < 0) {
            std::cout << "RAGDOLL BONE NOT FOUND: " << definition.bone << "\n";
            jointAnchors[i] = parts[i]->restPosition;
            continue;
        }

        const glm::vec3 anchor = BonePosition(definition.bone, parts[i]->restPosition);
        jointAnchors[i] = anchor;

        Part* parent = parts[partParents[i]];
        Part* child = parts[i];
        glm::vec3 axis = child->restPosition - anchor;
        axis = glm::length(axis) > 0.0001f ? glm::normalize(axis) : glm::vec3(0.0f, -1.0f, 0.0f);

        // Leaning only the parent's axis moves the cone without moving the part
        glm::vec3 coneAxis = axis;
        const glm::vec3 forwardOfAxis = forward - axis * glm::dot(forward, axis);
        if (glm::length(forwardOfAxis) > 0.001f) {
            const float tilt = glm::radians(definition.coneTilt);
            coneAxis = glm::normalize(axis * std::cos(tilt) + glm::normalize(forwardOfAxis) * std::sin(tilt));
        }

        BallJoint* limits = new BallJoint();
        limits->parent = parent->body;
        limits->child = child->body;
        limits->anchorParent = anchor - parent->restPosition;
        limits->anchorChild = anchor - child->restPosition;
        limits->axisParent = coneAxis;
        limits->axisChild = axis;
        limits->maxAngle = glm::radians(definition.maxSwing);
        limits->maxTwist = glm::radians(definition.maxTwist);

        if (definition.endBone) {
            const glm::vec3 upperDirection = glm::normalize(anchor - jointAnchors[partParents[i]]);
            const glm::vec3 lowerDirection = glm::normalize(BonePosition(definition.endBone, child->restPosition) - anchor);

            // A bent pose shows the real hinge. A straight one has to go by which way the limb should fold.
            glm::vec3 hinge = glm::cross(upperDirection, lowerDirection);
            if (glm::length(hinge) < HINGE_MIN_POSE_BEND)
                hinge = glm::cross(upperDirection, forward * definition.bendDirection);
            hinge = glm::length(hinge) > 0.001f ? glm::normalize(hinge) : glm::vec3(1.0f, 0.0f, 0.0f);

            limits->isHinge = true;
            limits->hingeAxisParent = hinge;
            limits->hingeAxisChild = hinge;
            limits->hingeRefParent = glm::normalize(upperDirection - hinge * glm::dot(upperDirection, hinge));
            limits->hingeRefChild = glm::normalize(lowerDirection - hinge * glm::dot(lowerDirection, hinge));
            limits->minBend = glm::radians(definition.minBend);
            limits->maxBend = glm::radians(definition.maxBend);
        }

        physics.AddJoint(limits);
    }

    // A nudge so it topples instead of balancing perfectly on its feet
    parts[1]->body->linearvelocity = glm::vec3(0.8f, 0.0f, 0.4f);
}

void Ragdoll::Update(float dt)
{
    if (parts.empty())
        return;

    if (parts[0]->transform.t.GetTranslationY() < FALL_LIMIT) {
        Reset();
        return;
    }

    if (parts[0]->body->frozen)
        return;

    bool settled = true;
    for (Part* part : parts) {
        if (glm::length(part->body->linearvelocity) > SETTLE_SPEED ||
            glm::length(part->body->angularvelocity) > SETTLE_ANGULAR_SPEED) {
            settled = false;
            break;
        }
    }

    settleTimer = settled ? settleTimer + dt : 0.0f;
    if (settleTimer < SETTLE_TIME)
        return;

    // Freezing saves the triangle tests and stops resting contacts jittering. Contact wakes it again.
    settleTimer = 0.0f;
    for (Part* part : parts) {
        part->body->frozen = true;
        part->body->linearvelocity = glm::vec3(0.0f);
        part->body->angularvelocity = glm::vec3(0.0f);
    }
}

void Ragdoll::Reset()
{
    settleTimer = 0.0f;
    for (Part* part : parts) {
        part->transform.t.TranslateTo(part->restPosition);
        part->transform.t.RotateToQuaternion(IDENTITY_ROTATION);
        part->body->linearvelocity = glm::vec3(0.0f);
        part->body->angularvelocity = glm::vec3(0.0f);
        part->body->frozen = false;
    }
}

std::vector<glm::mat4> Ragdoll::GetBoneMatrices() const
{
    std::vector<glm::mat4> bones = restBones;
    if (parts.empty())
        return bones;

    std::vector<glm::mat4> partMoves;
    partMoves.reserve(parts.size());
    for (int i = 0; i < (int)parts.size(); i++) {
        const t& transform = parts[i]->transform.t;
        partMoves.push_back(transform.GetTranslationMatrix() * transform.GetRotationMatrix() * restInverses[i]);
    }

    for (int i = 0; i < (int)bones.size(); i++)
        bones[i] = partMoves[boneParts[i]] * restBones[i];
    return bones;
}