import type { Pose } from "../kdl/types.js";

export type Mat4 = [
  number,
  number,
  number,
  number,
  number,
  number,
  number,
  number,
  number,
  number,
  number,
  number,
  number,
  number,
  number,
  number
];

export function identityMat4(): Mat4 {
  return [1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1];
}

export function multiplyMat4(a: Mat4, b: Mat4): Mat4 {
  const out = new Array<number>(16).fill(0) as Mat4;
  for (let row = 0; row < 4; row += 1) {
    for (let col = 0; col < 4; col += 1) {
      out[row * 4 + col] =
        a[row * 4 + 0]! * b[col + 0]! +
        a[row * 4 + 1]! * b[col + 4]! +
        a[row * 4 + 2]! * b[col + 8]! +
        a[row * 4 + 3]! * b[col + 12]!;
    }
  }
  return out;
}

export function translationMat4(xyz: [number, number, number]): Mat4 {
  const [x, y, z] = xyz;
  return [1, 0, 0, x, 0, 1, 0, y, 0, 0, 1, z, 0, 0, 0, 1];
}

export function rotationFromRpyMat4(rpy: [number, number, number]): Mat4 {
  const [roll, pitch, yaw] = rpy;
  const cr = Math.cos(roll);
  const sr = Math.sin(roll);
  const cp = Math.cos(pitch);
  const sp = Math.sin(pitch);
  const cy = Math.cos(yaw);
  const sy = Math.sin(yaw);

  return [
    cy * cp,
    cy * sp * sr - sy * cr,
    cy * sp * cr + sy * sr,
    0,
    sy * cp,
    sy * sp * sr + cy * cr,
    sy * sp * cr - cy * sr,
    0,
    -sp,
    cp * sr,
    cp * cr,
    0,
    0,
    0,
    0,
    1
  ];
}

export function axisAngleMat4(axis: [number, number, number], angle: number): Mat4 {
  const [nx, ny, nz] = normalizeVector(axis);
  const c = Math.cos(angle);
  const s = Math.sin(angle);
  const t = 1 - c;

  return [
    t * nx * nx + c,
    t * nx * ny - s * nz,
    t * nx * nz + s * ny,
    0,
    t * nx * ny + s * nz,
    t * ny * ny + c,
    t * ny * nz - s * nx,
    0,
    t * nx * nz - s * ny,
    t * ny * nz + s * nx,
    t * nz * nz + c,
    0,
    0,
    0,
    0,
    1
  ];
}

export function mat4ToPose(matrix: Mat4): Pose {
  return {
    position: [matrix[3], matrix[7], matrix[11]],
    quaternion: normalizeQuaternion(rotationMat4ToQuaternion(matrix))
  };
}

export function poseToMat4(pose: Pose): Mat4 {
  const [x, y, z, w] = normalizeQuaternion(pose.quaternion);
  const xx = x * x;
  const yy = y * y;
  const zz = z * z;
  const xy = x * y;
  const xz = x * z;
  const yz = y * z;
  const wx = w * x;
  const wy = w * y;
  const wz = w * z;
  const [px, py, pz] = pose.position;

  return [
    1 - 2 * (yy + zz),
    2 * (xy - wz),
    2 * (xz + wy),
    px,
    2 * (xy + wz),
    1 - 2 * (xx + zz),
    2 * (yz - wx),
    py,
    2 * (xz - wy),
    2 * (yz + wx),
    1 - 2 * (xx + yy),
    pz,
    0,
    0,
    0,
    1
  ];
}

export function composePose(a: Pose, b: Pose): Pose {
  return mat4ToPose(multiplyMat4(poseToMat4(a), poseToMat4(b)));
}

export function inversePose(pose: Pose): Pose {
  const matrix = poseToMat4(pose);
  const r00 = matrix[0];
  const r01 = matrix[1];
  const r02 = matrix[2];
  const tx = matrix[3];
  const r10 = matrix[4];
  const r11 = matrix[5];
  const r12 = matrix[6];
  const ty = matrix[7];
  const r20 = matrix[8];
  const r21 = matrix[9];
  const r22 = matrix[10];
  const tz = matrix[11];

  return mat4ToPose([
    r00,
    r10,
    r20,
    -(r00 * tx + r10 * ty + r20 * tz),
    r01,
    r11,
    r21,
    -(r01 * tx + r11 * ty + r21 * tz),
    r02,
    r12,
    r22,
    -(r02 * tx + r12 * ty + r22 * tz),
    0,
    0,
    0,
    1
  ]);
}

export function rpyToQuaternion(rpy: [number, number, number]): [number, number, number, number] {
  return normalizeQuaternion(rotationMat4ToQuaternion(rotationFromRpyMat4(rpy)));
}

export function normalizeQuaternion(input: [number, number, number, number]): [number, number, number, number] {
  const [x, y, z, w] = input;
  const length = Math.hypot(x, y, z, w);
  if (length === 0) {
    return [0, 0, 0, 1];
  }
  return [x / length, y / length, z / length, w / length];
}

export function jointMotionMat4(type: string, axis: [number, number, number], value: number): Mat4 {
  if (type === "revolute" || type === "continuous") {
    return axisAngleMat4(axis, value);
  }
  if (type === "prismatic") {
    const [x, y, z] = normalizeVector(axis);
    return translationMat4([x * value, y * value, z * value]);
  }
  return identityMat4();
}

function rotationMat4ToQuaternion(matrix: Mat4): [number, number, number, number] {
  const m00 = matrix[0];
  const m01 = matrix[1];
  const m02 = matrix[2];
  const m10 = matrix[4];
  const m11 = matrix[5];
  const m12 = matrix[6];
  const m20 = matrix[8];
  const m21 = matrix[9];
  const m22 = matrix[10];
  const trace = m00 + m11 + m22;

  if (trace > 0) {
    const s = Math.sqrt(trace + 1) * 2;
    return [(m21 - m12) / s, (m02 - m20) / s, (m10 - m01) / s, 0.25 * s];
  }
  if (m00 > m11 && m00 > m22) {
    const s = Math.sqrt(1 + m00 - m11 - m22) * 2;
    return [0.25 * s, (m01 + m10) / s, (m02 + m20) / s, (m21 - m12) / s];
  }
  if (m11 > m22) {
    const s = Math.sqrt(1 + m11 - m00 - m22) * 2;
    return [(m01 + m10) / s, 0.25 * s, (m12 + m21) / s, (m02 - m20) / s];
  }
  const s = Math.sqrt(1 + m22 - m00 - m11) * 2;
  return [(m02 + m20) / s, (m12 + m21) / s, 0.25 * s, (m10 - m01) / s];
}

function normalizeVector(axis: [number, number, number]): [number, number, number] {
  const [x, y, z] = axis;
  const length = Math.hypot(x, y, z);
  if (length === 0) {
    return [1, 0, 0];
  }
  return [x / length, y / length, z / length];
}
