#include "rive/constraints/ik_constraint.hpp"
#include "rive/bones/bone.hpp"
#include "rive/math/math_types.hpp"
#include <algorithm>

using namespace rive;

static float atan2(Vec2D v) { return std::atan2(v.y, v.x); }

void IKConstraint::buildDependencies()
{
    Super::buildDependencies();

    // IK Constraint needs to depend on the target so that world transform
    // changes can propagate to the bones (and they can be reset before IK
    // runs).
    if (auto* t = target())
    {
        t->addDependent(this);
    }
}

StatusCode IKConstraint::onAddedClean(CoreContext* context)
{
    if (!parent()->is<Bone>())
    {
        return StatusCode::InvalidObject;
    }

    auto boneCount = parentBoneCount();
    auto bone = parent()->as<Bone>();
    std::vector<Bone*> bones;
    bones.push_back(bone);
    // Get the reverse FK chain of bones (from tip up).
    while (bone->parent()->is<Bone>() && boneCount > 0)
    {
        boneCount--;
        bone = bone->parent()->as<Bone>();
        bone->addPeerConstraint(this);
        bones.push_back(bone);
    }
    int numBones = (int)bones.size();
    m_FkChain.resize(numBones);
    // Now put them in FK order (top to bottom).
    int idx = 0;
    for (auto boneItr = bones.rbegin(); boneItr != bones.rend(); ++boneItr)
    {
        BoneChainLink& link = m_FkChain[idx];
        link.index = idx++;
        link.bone = *boneItr;
        link.angle = 0.0f;
    }

    // Make sure all of the first level children of each bone (other than the
    // chain bones themselves) depend on the tip (constrainedComponent). The
    // chain is built by walking up parents, so the only chain bone that can
    // appear among an ancestor's direct children is the next bone down the
    // chain — bones[i - 1].
    auto tip = parent()->as<Bone>();
    for (int i = 1; i < numBones; i++)
    {
        Bone* ancestor = bones[i];
        Bone* chainChild = bones[i - 1];
        for (Component* child : ancestor->children())
        {
            if (!child->is<TransformComponent>() || child == chainChild)
            {
                continue;
            }
            tip->addDependent(child->as<TransformComponent>());
        }
    }
    return Super::onAddedClean(context);
}

void IKConstraint::markConstraintDirty()
{
    Super::markConstraintDirty();
    // We automatically propagate dirt to the parent constrained bone, but we
    // also need to make sure the other bones we influence above it rebuild
    // their transforms.
    for (int i = 0, length = (int)m_FkChain.size() - 1; i < length; i++)
    {
        m_FkChain[i].bone->markTransformDirty();
    }
}

void IKConstraint::solve1(BoneChainLink* fk1,
                          const Vec2D& worldTargetTranslation)
{
    Mat2D iworld = fk1->parentWorldInverse;
    Vec2D pA = fk1->bone->worldTranslation();
    Vec2D pBT(worldTargetTranslation);

    // To target in worldspace
    const Vec2D toTarget = pBT - pA;

    // Note this is directional, hence not transformMat2d
    Vec2D toTargetLocal = Vec2D::transformDir(toTarget, iworld);
    float r = atan2(toTargetLocal);

    constrainRotation(*fk1, r);
    fk1->angle = r;
}

void IKConstraint::solve2(BoneChainLink* fk1,
                          BoneChainLink* fk2,
                          const Vec2D& worldTargetTranslation)
{
    Bone* b1 = fk1->bone;
    Bone* b2 = fk2->bone;
    BoneChainLink* firstChild = &(m_FkChain[fk1->index + 1]);

    const Mat2D& iworld = fk1->parentWorldInverse;

    Vec2D pA = b1->worldTranslation();
    Vec2D pC = firstChild->bone->worldTranslation();
    Vec2D pB = b2->tipWorldTranslation();
    Vec2D pBT(worldTargetTranslation);

    pA = iworld * pA;
    pC = iworld * pC;
    pB = iworld * pB;
    pBT = iworld * pBT;

    // http://mathworld.wolfram.com/LawofCosines.html
    Vec2D av = pB - pC, bv = pC - pA, cv = pBT - pA;
    float a = av.length(), b = bv.length(), c = cv.length();

    float A = std::acos(
        std::max(-1.0f,
                 std::min(1.0f, (-a * a + b * b + c * c) / (2.0f * b * c))));
    float C = std::acos(
        std::max(-1.0f,
                 std::min(1.0f, (a * a + b * b - c * c) / (2.0f * a * b))));

    float r1, r2;
    if (b2->parent() != b1)
    {
        BoneChainLink& secondChild = m_FkChain[fk1->index + 2];

        const Mat2D& secondChildWorldInverse = secondChild.parentWorldInverse;

        pC = firstChild->bone->worldTranslation();
        pB = b2->tipWorldTranslation();

        Vec2D avLocal = Vec2D::transformDir(pB - pC, secondChildWorldInverse);
        float angleCorrection = -atan2(avLocal);

        if (invertDirection())
        {
            r1 = atan2(cv) - A;
            r2 = -C + math::PI + angleCorrection;
        }
        else
        {
            r1 = A + atan2(cv);
            r2 = C - math::PI + angleCorrection;
        }
    }
    else if (invertDirection())
    {
        r1 = atan2(cv) - A;
        r2 = -C + math::PI;
    }
    else
    {
        r1 = A + atan2(cv);
        r2 = C - math::PI;
    }

    constrainRotation(*fk1, r1);
    constrainRotation(*firstChild, r2);
    if (firstChild != fk2)
    {
        fk2->bone->composeWorldTransform();
    }

    // Simple storage, need this for interpolation.
    fk1->angle = r1;
    firstChild->angle = r2;
}

void IKConstraint::invertDirectionChanged() { markConstraintDirty(); }

void IKConstraint::constrainRotation(BoneChainLink& fk, float rotation)
{
    Bone* bone = fk.bone;
    Mat2D& transform = bone->mutableTransform();
    TransformComponents& c = fk.transformComponents;

    transform = Mat2D::fromRotation(rotation);

    // Translate
    transform[4] = c.x();
    transform[5] = c.y();
    // Scale
    float scaleX = c.scaleX();
    float scaleY = c.scaleY();
    transform[0] *= scaleX;
    transform[1] *= scaleX;
    transform[2] *= scaleY;
    transform[3] *= scaleY;

    // Skew
    const float skew = c.skew();
    if (skew != 0.0f)
    {
        transform[2] = transform[0] * skew + transform[2];
        transform[3] = transform[1] * skew + transform[3];
    }

    bone->composeWorldTransform();
}

void IKConstraint::constrain(TransformComponent* component)
{
    auto* tgt = target();
    if (tgt == nullptr || tgt->isCollapsed())
    {
        return;
    }

    Vec2D worldTargetTranslation = tgt->worldTranslation();

    // Decompose the chain where it currently stands, before rebuilding any of
    // it, so no bone is measured against a parent we already rebuilt. A
    // constraint that ran before us left its work here and nowhere else.
    for (BoneChainLink& item : m_FkChain)
    {
        item.parentWorldInverse = getParentWorld(*item.bone).invertOrIdentity();
        item.transformComponents =
            (item.parentWorldInverse * item.bone->worldTransform()).decompose();
    }

    // Take the angle back from the FK base. Rotation is the only component we
    // write, so reading ours back would blend from the pose we solved last
    // pass; translation, scale and skew we merely carry, so they stay as
    // whatever constrained the bone left them.
    for (BoneChainLink& item : m_FkChain)
    {
        Bone* bone = item.bone;
        bone->updateTransform();
        item.transformComponents.rotation(
            bone->transform().decompose().rotation());

        item.parentWorldInverse = getParentWorld(*bone).invertOrIdentity();
        constrainRotation(item, item.transformComponents.rotation());
    }

    int count = (int)m_FkChain.size();
    switch (count)
    {
        case 1:
            solve1(&m_FkChain[0], worldTargetTranslation);
            break;
        case 2:
            solve2(&m_FkChain[0], &m_FkChain[1], worldTargetTranslation);
            break;
        default:
        {
            auto last = count - 1;
            BoneChainLink* tip = &m_FkChain[last];
            for (int i = 0; i < last; i++)
            {
                BoneChainLink* item = &m_FkChain[i];
                solve2(item, tip, worldTargetTranslation);
                for (int j = item->index + 1, end = (int)m_FkChain.size() - 1;
                     j < end;
                     j++)
                {
                    BoneChainLink& fk = m_FkChain[j];
                    fk.parentWorldInverse =
                        getParentWorld(*fk.bone).invertOrIdentity();
                }
            }
            break;
        }
    }
    // At the end, mix the FK angle with the IK angle by strength
    if (strength() != 1.0f)
    {
        for (BoneChainLink& fk : m_FkChain)
        {
            float fromAngle =
                std::fmod(fk.transformComponents.rotation(), math::PI * 2);
            float toAngle = std::fmod(fk.angle, math::PI * 2);
            float diff = toAngle - fromAngle;
            if (diff > math::PI)
            {
                diff -= math::PI * 2;
            }
            else if (diff < -math::PI)
            {
                diff += math::PI * 2;
            }
            float angle = fromAngle + diff * strength();
            constrainRotation(fk, angle);
        }
    }
}