using Flyshot.Core.Domain;

namespace Flyshot.Core.Planning.Kinematics;

/// <summary>
/// 机器人正运动学计算，包含关节耦合解析、四元数/旋转矩阵转换和齐次变换链。
///
/// 为什么自己写而不是用 System.Numerics？
/// ---
/// System.Numerics.Matrix4x4 使用行向量约定（平移在 M41-M43），而 numpy 使用列向量
/// 约定（平移在 M14-M24-M34）。两者是转置关系，如果混用会导致 FK 结果完全错误。
/// 自己实现 Mat4 可以确保与 Python 参考代码的矩阵约定完全一致，便于 golden-sample 对拍。
/// </summary>
public static class RobotKinematics
{
    /// <summary>
    /// 执行正运动学，输出末端位姿 [x, y, z, qx, qy, qz, qw]。
    /// </summary>
    public static double[] ForwardKinematics(RobotKinematicsModel model, double[] jointPositions)
    {
        var kinematicAngles = ResolveKinematicAngles(model, jointPositions);
        var transform = Mat4.Identity;
        int revoluteIndex = 0;

        foreach (var joint in model.Joints)
        {
            // 先叠加固定偏移 origin，再叠加关节旋转。
            transform = transform * TransformFromPose(joint.OriginXyz, joint.OriginQuatXyzw);

            if (joint.JointType == 2)
            {
                var rot = AxisAngleToMatrix(joint.Axis, kinematicAngles[revoluteIndex]);
                var jointTransform = Mat4.FromRotation(rot);
                transform = transform * jointTransform;
                revoluteIndex++;
            }
        }

        var quat = MatrixToQuatXyzw(ExtractRotation(transform));
        return new double[]
        {
            transform.M14, transform.M24, transform.M34,
            quat[0], quat[1], quat[2], quat[3]
        };
    }

    /// <summary>
    /// 解析关节耦合，把控制器关节角转换成运动学关节角。
    ///
    /// 公式：q_kin[i] = q_raw[i] + multiplier * q_kin[master] + offset
    /// </summary>
    public static double[] ResolveKinematicAngles(RobotKinematicsModel model, double[] rawPositions)
    {
        var revoluteJoints = model.Joints.Where(j => j.JointType == 2).ToArray();
        if (rawPositions.Length != revoluteJoints.Length)
        {
            throw new ArgumentException(
                $"关节位置数量 {rawPositions.Length} 与旋转关节数量 {revoluteJoints.Length} 不匹配。");
        }

        var nameToIndex = revoluteJoints
            .Select((j, i) => (j.Name, i))
            .ToDictionary(x => x.Name, x => x.i);

        var kinematicValues = new double[rawPositions.Length];
        var resolved = new bool[rawPositions.Length];
        var resolving = new bool[rawPositions.Length];

        double Resolve(int index)
        {
            if (resolved[index])
                return kinematicValues[index];
            if (resolving[index])
                throw new InvalidOperationException($"检测到循环关节耦合：{revoluteJoints[index].Name}");

            resolving[index] = true;
            var joint = revoluteJoints[index];
            double value = rawPositions[index];

            if (!string.IsNullOrEmpty(joint.CoupleMaster))
            {
                if (!nameToIndex.TryGetValue(joint.CoupleMaster, out var masterIndex))
                    throw new InvalidOperationException($"关节 {joint.Name} 引用了未知的主关节 {joint.CoupleMaster}");
                value += joint.CoupleMultiplier * Resolve(masterIndex) + joint.CoupleOffset;
            }

            kinematicValues[index] = value;
            resolving[index] = false;
            resolved[index] = true;
            return value;
        }

        for (int i = 0; i < revoluteJoints.Length; i++)
        {
            Resolve(i);
        }

        return kinematicValues;
    }

    /// <summary>
    /// 从 Mat4 中提取 3x3 旋转矩阵。
    /// </summary>
    private static double[,] ExtractRotation(Mat4 m)
    {
        return new double[,]
        {
            { m.M11, m.M12, m.M13 },
            { m.M21, m.M22, m.M23 },
            { m.M31, m.M32, m.M33 },
        };
    }

    /// <summary>
    /// 把平移 + 四元数拼成 4x4 齐次变换矩阵。
    /// </summary>
    private static Mat4 TransformFromPose(double[] xyz, double[] quatXyzw)
    {
        var rot = QuatToMatrix(quatXyzw);
        var m = Mat4.Identity;
        m.M11 = rot[0, 0]; m.M12 = rot[0, 1]; m.M13 = rot[0, 2];
        m.M21 = rot[1, 0]; m.M22 = rot[1, 1]; m.M23 = rot[1, 2];
        m.M31 = rot[2, 0]; m.M32 = rot[2, 1]; m.M33 = rot[2, 2];
        m.M14 = xyz[0];
        m.M24 = xyz[1];
        m.M34 = xyz[2];
        return m;
    }

    /// <summary>
    /// 四元数（xyzw 顺序）转旋转矩阵。
    /// </summary>
    private static double[,] QuatToMatrix(double[] q)
    {
        double x = q[0], y = q[1], z = q[2], w = q[3];
        double xx = x * x, yy = y * y, zz = z * z;
        double xy = x * y, xz = x * z, yz = y * z;
        double wx = w * x, wy = w * y, wz = w * z;

        return new double[,]
        {
            { 1.0 - 2.0 * (yy + zz), 2.0 * (xy - wz), 2.0 * (xz + wy) },
            { 2.0 * (xy + wz), 1.0 - 2.0 * (xx + zz), 2.0 * (yz - wx) },
            { 2.0 * (xz - wy), 2.0 * (yz + wx), 1.0 - 2.0 * (xx + yy) },
        };
    }

    /// <summary>
    /// 旋转矩阵转四元数（xyzw 顺序），使用 Shepperd 分支策略保证数值稳定。
    /// </summary>
    private static double[] MatrixToQuatXyzw(double[,] r)
    {
        double trace = r[0, 0] + r[1, 1] + r[2, 2];
        double x, y, z, w;

        if (trace > 0.0)
        {
            double s = Math.Sqrt(trace + 1.0) * 2.0;
            w = 0.25 * s;
            x = (r[2, 1] - r[1, 2]) / s;
            y = (r[0, 2] - r[2, 0]) / s;
            z = (r[1, 0] - r[0, 1]) / s;
        }
        else if (r[0, 0] > r[1, 1] && r[0, 0] > r[2, 2])
        {
            double s = Math.Sqrt(1.0 + r[0, 0] - r[1, 1] - r[2, 2]) * 2.0;
            w = (r[2, 1] - r[1, 2]) / s;
            x = 0.25 * s;
            y = (r[0, 1] + r[1, 0]) / s;
            z = (r[0, 2] + r[2, 0]) / s;
        }
        else if (r[1, 1] > r[2, 2])
        {
            double s = Math.Sqrt(1.0 + r[1, 1] - r[0, 0] - r[2, 2]) * 2.0;
            w = (r[0, 2] - r[2, 0]) / s;
            x = (r[0, 1] + r[1, 0]) / s;
            y = 0.25 * s;
            z = (r[1, 2] + r[2, 1]) / s;
        }
        else
        {
            double s = Math.Sqrt(1.0 + r[2, 2] - r[0, 0] - r[1, 1]) * 2.0;
            w = (r[1, 0] - r[0, 1]) / s;
            x = (r[0, 2] + r[2, 0]) / s;
            y = (r[1, 2] + r[2, 1]) / s;
            z = 0.25 * s;
        }

        double norm = Math.Sqrt(x * x + y * y + z * z + w * w);
        return new double[] { x / norm, y / norm, z / norm, w / norm };
    }

    /// <summary>
    /// Rodrigues 公式：绕单位向量 axis 旋转 angle 弧度。
    /// </summary>
    private static double[,] AxisAngleToMatrix(double[] axis, double angle)
    {
        double norm = Math.Sqrt(axis[0] * axis[0] + axis[1] * axis[1] + axis[2] * axis[2]);
        if (norm == 0.0)
            return new double[,] { { 1, 0, 0 }, { 0, 1, 0 }, { 0, 0, 1 } };

        double ux = axis[0] / norm, uy = axis[1] / norm, uz = axis[2] / norm;
        double c = Math.Cos(angle);
        double s = Math.Sin(angle);
        double oneC = 1.0 - c;

        return new double[,]
        {
            { c + ux * ux * oneC, ux * uy * oneC - uz * s, ux * uz * oneC + uy * s },
            { uy * ux * oneC + uz * s, c + uy * uy * oneC, uy * uz * oneC - ux * s },
            { uz * ux * oneC - uy * s, uz * uy * oneC + ux * s, c + uz * uz * oneC },
        };
    }
}
