/**
 * code based on
 * https://github.com/benikabocha/saba/blob/master/src/Saba/Model/MMD/MMDIkSolver.cpp
 */
import { Matrix, Quaternion, Vector3 } from "@babylonjs/core/Maths/math.vector";
class IkChain {
    bone;
    minimumAngle;
    maximumAngle;
    prevAngle;
    savedIkRotation;
    planeModeAngle;
    constructor(bone, minimumAngle, maximumAngle) {
        this.bone = bone;
        this.minimumAngle = minimumAngle;
        this.maximumAngle = maximumAngle;
        this.prevAngle = Vector3.Zero();
        this.savedIkRotation = Quaternion.Identity();
        this.planeModeAngle = 0;
    }
}
/**
 * IK solver
 */
export class IkSolver {
    /**
     * Ik solver index
     */
    index;
    /**
     * Iteration count
     *
     * The higher the value, the more accurate the IK solver will be, but the more expensive it will be
     */
    iteration;
    /**
     * Limit angle
     */
    limitAngle;
    /**
     * The bone to which the IK solver is attached
     */
    ikBone;
    /**
     * Ik target bone
     */
    targetBone;
    _ikChains;
    /**
     * Create a new IK solver
     * @param index Ik solver index
     * @param ikBone Attach bone
     * @param targetBone Ik target bone
     */
    constructor(index, ikBone, targetBone) {
        this.index = index;
        this.iteration = 0;
        this.limitAngle = 0;
        this.ikBone = ikBone;
        this.targetBone = targetBone;
        this._ikChains = [];
    }
    /**
     * Add an IK chain
     *
     * The angle constraint must be either both min max or neither
     *
     * For better performance, we do not constrain this to a type
     * @param bone Bone to add
     * @param minimumAngle minimum angle
     * @param maximumAngle maximum angle
     */
    addIkChain(bone, minimumAngle, maximumAngle) {
        bone.ikRotation = Quaternion.Identity();
        const ikChain = new IkChain(bone, minimumAngle, maximumAngle);
        this._ikChains.push(ikChain);
    }
    static _TargetPosition = new Vector3();
    static _IkPosition = new Vector3();
    /**
     * Solve IK
     */
    solve() {
        const ikBone = this.ikBone;
        const targetBone = this.targetBone;
        const chains = this._ikChains;
        for (let i = 0; i < chains.length; ++i) {
            const chain = chains[i];
            const chainbone = chain.bone;
            chain.prevAngle.setAll(0);
            chainbone.ikRotation.set(0, 0, 0, 1);
            chain.planeModeAngle = 0;
            chainbone.updateLocalMatrix();
            chainbone.updateWorldMatrix();
        }
        let maxDistance = Number.MAX_VALUE;
        for (let i = 0; i < this.iteration; ++i) {
            this._solveCore(i);
            const targetPosition = targetBone.getWorldTranslationToRef(IkSolver._TargetPosition);
            const ikPosition = ikBone.getWorldTranslationToRef(IkSolver._IkPosition);
            const distance = Vector3.DistanceSquared(targetPosition, ikPosition);
            if (distance < maxDistance) {
                maxDistance = distance;
                for (let j = 0; j < chains.length; ++j) {
                    const chain = chains[j];
                    chain.savedIkRotation.copyFrom(chain.bone.ikRotation);
                }
            }
            else {
                for (let j = 0; j < chains.length; ++j) {
                    const chain = chains[j];
                    chain.bone.ikRotation.copyFrom(chain.savedIkRotation);
                    chain.bone.updateLocalMatrix();
                    chain.bone.updateWorldMatrix();
                }
                break;
            }
        }
    }
    static _TargetPosition2 = new Vector3();
    static _IkPosition2 = new Vector3();
    static _InversedChain = new Matrix();
    static _ChainCross = new Vector3();
    static _Rotation = new Quaternion();
    static _ChainRotation = new Quaternion();
    static _AnimatedRotation = new Quaternion();
    static _ChainRotationMatrix = new Matrix();
    static _DecomposedRotation = new Vector3();
    static _ClampedRotation = new Vector3();
    static _FinalRotationA = new Quaternion();
    static _FinalRotationB = new Quaternion();
    static _Right = Vector3.Right();
    static _Up = Vector3.Up();
    static _Forward = Vector3.Forward();
    static _InversedAnimatedRotation = new Quaternion();
    static _RadToDeg = 180 / Math.PI;
    _solveCore(iteration) {
        const ikPosition = this.ikBone.getWorldTranslationToRef(IkSolver._IkPosition2);
        const chains = this._ikChains;
        for (let i = 0; i < chains.length; ++i) {
            const chain = chains[i];
            const chainBone = chain.bone;
            if (chainBone === this.targetBone)
                continue;
            if (chain.minimumAngle !== null /* && chain.minimumAngle !== null */) {
                if ((chain.minimumAngle.x !== 0 || chain.maximumAngle.x !== 0) &&
                    (chain.minimumAngle.y === 0 || chain.maximumAngle.y === 0) &&
                    (chain.minimumAngle.z === 0 || chain.maximumAngle.z === 0)) {
                    this._solvePlane(iteration, chain, 0 /* SolveAxis.X */);
                    continue;
                }
                else if ((chain.minimumAngle.y !== 0 || chain.maximumAngle.y !== 0) &&
                    (chain.minimumAngle.x === 0 || chain.maximumAngle.x === 0) &&
                    (chain.minimumAngle.z === 0 || chain.maximumAngle.z === 0)) {
                    this._solvePlane(iteration, chain, 1 /* SolveAxis.Y */);
                    continue;
                }
                else if ((chain.minimumAngle.z !== 0 || chain.maximumAngle.z !== 0) &&
                    (chain.minimumAngle.x === 0 || chain.maximumAngle.x === 0) &&
                    (chain.minimumAngle.y === 0 || chain.maximumAngle.y === 0)) {
                    this._solvePlane(iteration, chain, 2 /* SolveAxis.Z */);
                    continue;
                }
            }
            const targetPosition = this.targetBone.getWorldTranslationToRef(IkSolver._TargetPosition2);
            const inverseChain = chainBone.getWorldMatrixToRef(IkSolver._InversedChain).invert();
            const chainIkPosition = Vector3.TransformCoordinatesToRef(ikPosition, inverseChain, IkSolver._IkPosition);
            const chainTargetPosition = Vector3.TransformCoordinatesToRef(targetPosition, inverseChain, IkSolver._TargetPosition);
            const chainIkVector = chainIkPosition.normalize();
            const chainTargetVector = chainTargetPosition.normalize();
            let dot = Vector3.Dot(chainTargetVector, chainIkVector);
            dot = Math.max(-1.0, Math.min(1.0, dot));
            let angle = Math.acos(dot);
            const angleDeg = angle * IkSolver._RadToDeg;
            if (angleDeg < 1.0e-3)
                continue;
            angle = Math.max(-this.limitAngle, Math.min(this.limitAngle, angle));
            const cross = Vector3.CrossToRef(chainTargetVector, chainIkVector, IkSolver._ChainCross).normalize();
            const rotation = Quaternion.RotationAxisToRef(cross, angle, IkSolver._Rotation);
            const chainRotation = IkSolver._ChainRotation.copyFrom(chainBone.ikRotation);
            const animatedRotation = chainBone.getAnimatedRotationToRef(IkSolver._AnimatedRotation);
            chainRotation.multiplyInPlace(animatedRotation).multiplyInPlace(rotation);
            if (chain.minimumAngle !== null /* && chain.minimumAngle !== null */) {
                const chainRotationMatrix = chainRotation.toRotationMatrix(IkSolver._ChainRotationMatrix);
                const rotXYZ = this._decomposeToRef(chainRotationMatrix, chain.prevAngle, IkSolver._DecomposedRotation);
                const clampXYZ = Vector3.ClampToRef(rotXYZ, chain.minimumAngle, chain.maximumAngle, IkSolver._ClampedRotation).subtractInPlace(chain.prevAngle);
                clampXYZ.set(Math.max(-this.limitAngle, Math.min(this.limitAngle, clampXYZ.x)) + chain.prevAngle.x, Math.max(-this.limitAngle, Math.min(this.limitAngle, clampXYZ.y)) + chain.prevAngle.y, Math.max(-this.limitAngle, Math.min(this.limitAngle, clampXYZ.z)) + chain.prevAngle.z);
                const rA = Quaternion.RotationAxisToRef(IkSolver._Right, clampXYZ.x, IkSolver._FinalRotationA);
                const rB = IkSolver._FinalRotationB;
                Quaternion.RotationAxisToRef(IkSolver._Up, clampXYZ.y, rB);
                rA.multiplyInPlace(rB);
                Quaternion.RotationAxisToRef(IkSolver._Forward, clampXYZ.z, rB);
                rA.multiplyInPlace(rB);
                Matrix.FromQuaternionToRef(rA, chainRotationMatrix);
                chain.prevAngle.copyFrom(clampXYZ);
                Quaternion.FromRotationMatrixToRef(chainRotationMatrix, chainRotation);
            }
            chainRotation.multiplyToRef(Quaternion.InverseToRef(animatedRotation, IkSolver._InversedAnimatedRotation), chainBone.ikRotation);
            chainBone.updateLocalMatrix();
            chainBone.updateWorldMatrix();
        }
    }
    static _IkPosition3 = new Vector3();
    static _TargetPosition3 = new Vector3();
    static _InversedChain2 = new Matrix();
    static _ChainIkPosition = new Vector3();
    static _ChainTargetPosition = new Vector3();
    static _Rotation2 = new Quaternion();
    static _TargetVector = new Vector3();
    static _InversedAnimatedRotation2 = new Quaternion();
    _solvePlane(iteration, chain, solveAxis) {
        let minimumAngle;
        let maximumAngle;
        let rotateAxis;
        switch (solveAxis) {
            case 0 /* SolveAxis.X */:
                minimumAngle = chain.minimumAngle.x;
                maximumAngle = chain.maximumAngle.x;
                rotateAxis = IkSolver._Right;
                break;
            case 1 /* SolveAxis.Y */:
                minimumAngle = chain.minimumAngle.y;
                maximumAngle = chain.maximumAngle.y;
                rotateAxis = IkSolver._Up;
                break;
            case 2 /* SolveAxis.Z */:
                minimumAngle = chain.minimumAngle.z;
                maximumAngle = chain.maximumAngle.z;
                rotateAxis = IkSolver._Forward;
                break;
            default:
                throw new Error("Invalid solve axis");
        }
        const ikPosition = this.ikBone.getWorldTranslationToRef(IkSolver._IkPosition3);
        const targetPosition = this.targetBone.getWorldTranslationToRef(IkSolver._TargetPosition3);
        const inverseChain = chain.bone.getWorldMatrixToRef(IkSolver._InversedChain2).invert();
        const chainIkPosition = Vector3.TransformCoordinatesToRef(ikPosition, inverseChain, IkSolver._ChainIkPosition);
        const chainTargetPosition = Vector3.TransformCoordinatesToRef(targetPosition, inverseChain, IkSolver._ChainTargetPosition);
        const chainIkVector = chainIkPosition.normalize();
        const chainTargetVector = chainTargetPosition.normalize();
        let dot = Vector3.Dot(chainTargetVector, chainIkVector);
        dot = Math.max(-1.0, Math.min(1.0, dot));
        let angle = Math.acos(dot);
        angle = Math.max(-this.limitAngle, Math.min(this.limitAngle, angle));
        const rot1 = Quaternion.RotationAxisToRef(rotateAxis, angle, IkSolver._Rotation2);
        const targetVec1 = chainTargetVector.applyRotationQuaternionToRef(rot1, IkSolver._TargetVector);
        const dot1 = Vector3.Dot(targetVec1, chainIkVector);
        const rot2 = Quaternion.RotationAxisToRef(rotateAxis, -angle, IkSolver._Rotation2);
        const targetVec2 = chainTargetVector.applyRotationQuaternionToRef(rot2, IkSolver._TargetVector);
        const dot2 = Vector3.Dot(targetVec2, chainIkVector);
        let newAngle = chain.planeModeAngle;
        if (dot1 > dot2)
            newAngle += angle;
        else
            newAngle -= angle;
        if (iteration === 0) {
            if (newAngle < minimumAngle || newAngle > maximumAngle) {
                if (-newAngle > minimumAngle && -newAngle < maximumAngle)
                    newAngle = -newAngle;
                else {
                    const halfRad = (minimumAngle + maximumAngle) * 0.5;
                    if (Math.abs(halfRad - newAngle) > Math.abs(halfRad + newAngle))
                        newAngle = -newAngle;
                }
            }
        }
        newAngle = Math.max(minimumAngle, Math.min(maximumAngle, newAngle));
        chain.planeModeAngle = newAngle;
        const inversedAnimatedRotation = Quaternion.InverseToRef(chain.bone.getAnimatedRotationToRef(IkSolver._InversedAnimatedRotation2), IkSolver._InversedAnimatedRotation2);
        const ikRotation = Quaternion.RotationAxisToRef(rotateAxis, newAngle, chain.bone.ikRotation);
        ikRotation.multiplyInPlace(inversedAnimatedRotation);
        chain.bone.updateLocalMatrix();
        chain.bone.updateWorldMatrix();
    }
    static _TwoPi = Math.PI * 2;
    _normalizeAngle(angle) {
        while (angle >= IkSolver._TwoPi)
            angle -= IkSolver._TwoPi;
        while (angle < 0)
            angle += IkSolver._TwoPi;
        return angle;
    }
    _diffAngle(a, b) {
        const diff = this._normalizeAngle(a) - this._normalizeAngle(b);
        if (diff > Math.PI)
            return diff - IkSolver._TwoPi;
        else if (diff < -Math.PI)
            return diff + IkSolver._TwoPi;
        return diff;
    }
    static _Tests = [
        new Vector3(),
        new Vector3(),
        new Vector3(),
        new Vector3(),
        new Vector3(),
        new Vector3(),
        new Vector3(),
        new Vector3()
    ];
    _decomposeToRef(m, before, result) {
        const r = result;
        const sy = -m.m[2];
        const e = 1.0e-6;
        if (Math.abs(1.0 - Math.abs(sy)) < e) {
            r.y = Math.asin(sy);
            const sx = Math.sin(before.x);
            const sz = Math.sin(before.z);
            if (Math.abs(sx) < Math.abs(sz)) {
                const cx = Math.cos(before.x);
                if (cx > 0) {
                    r.x = 0;
                    r.z = Math.asin(-m.m[4]);
                }
                else {
                    r.x = Math.PI;
                    r.z = Math.asin(m.m[4]);
                }
            }
            else {
                const cz = Math.cos(before.z);
                if (cz > 0) {
                    r.z = 0;
                    r.x = Math.asin(-m.m[9]);
                }
                else {
                    r.z = Math.PI;
                    r.x = Math.asin(m.m[9]);
                }
            }
        }
        else {
            r.x = Math.atan2(m.m[6], m.m[10]);
            r.y = Math.asin(-m.m[2]);
            r.z = Math.atan2(m.m[1], m.m[0]);
        }
        const pi = Math.PI;
        const tests = IkSolver._Tests;
        tests[0].set(r.x + pi, pi - r.y, r.z + pi);
        tests[1].set(r.x + pi, pi - r.y, r.z - pi);
        tests[2].set(r.x + pi, -pi - r.y, r.z + pi);
        tests[3].set(r.x + pi, -pi - r.y, r.z - pi);
        tests[4].set(r.x - pi, pi - r.y, r.z + pi);
        tests[5].set(r.x - pi, pi - r.y, r.z - pi);
        tests[6].set(r.x - pi, -pi - r.y, r.z + pi);
        tests[7].set(r.x - pi, -pi - r.y, r.z - pi);
        const errX = Math.abs(this._diffAngle(r.x, before.x));
        const errY = Math.abs(this._diffAngle(r.y, before.y));
        const errZ = Math.abs(this._diffAngle(r.z, before.z));
        let minErr = errX + errY + errZ;
        for (let i = 0; i < tests.length; ++i) {
            const test = tests[i];
            const err = Math.abs(this._diffAngle(test.x, before.x))
                + Math.abs(this._diffAngle(test.y, before.y))
                + Math.abs(this._diffAngle(test.z, before.z));
            if (err < minErr) {
                minErr = err;
                r.copyFrom(test);
            }
        }
        return r;
    }
}
