"""Mass properties and rigid-body rotation, using row-major 3x3 matrices.

Point coordinates and the supplied reference point use the same frame.
Euler dynamics use a body-fixed frame about the center of mass or a fixed pivot.
"""
import math
from rotation import apply, mul, transpose
from quaternion import derivative, normalize


def dot(a, b):
    return sum(x*y for x, y in zip(a, b))


def cross(a, b):
    return [a[1]*b[2]-a[2]*b[1], a[2]*b[0]-a[0]*b[2],
            a[0]*b[1]-a[1]*b[0]]


def center_of_mass(masses, points):
    total = sum(masses)
    if total <= 0 or any(m < 0 for m in masses) or len(masses) != len(points):
        raise ValueError("Masses must be nonnegative and have a positive sum")
    return [sum(m*p[j] for m, p in zip(masses, points))/total for j in range(3)]


def point_inertia(masses, points, origin=(0, 0, 0)):
    result = [0.0]*9
    for m, p in zip(masses, points):
        r = [p[j]-origin[j] for j in range(3)]
        r2 = dot(r, r)
        for i in range(3):
            for j in range(3):
                result[3*i+j] += m*((r2 if i == j else 0)-r[i]*r[j])
    return result


def box_inertia(mass, sides):
    if mass <= 0 or any(s <= 0 for s in sides):
        raise ValueError("A solid box has positive mass and side lengths")
    a, b, c = sides
    return [mass*(b*b+c*c)/12, 0, 0, 0, mass*(a*a+c*c)/12,
            0, 0, 0, mass*(a*a+b*b)/12]


def rotate_inertia(inertia, rotation):
    return mul(mul(rotation, inertia), transpose(rotation))


def shift_inertia(inertia_com, mass, from_origin_to_com):
    d = from_origin_to_com
    d2 = dot(d, d)
    return [inertia_com[3*i+j] + mass*((d2 if i == j else 0)-d[i]*d[j])
            for i in range(3) for j in range(3)]


def rotational_energy(inertia, omega):
    return dot(omega, apply(inertia, omega))/2


def solve3(matrix, rhs):
    """Solve by partial-pivot elimination; reject singular inertias."""
    a = [list(matrix[3*i:3*i+3])+[rhs[i]] for i in range(3)]
    size = max(abs(x) for x in matrix)
    for k in range(3):
        pivot = max(range(k, 3), key=lambda i: abs(a[i][k]))
        if abs(a[pivot][k]) <= 1e-14*size:
            raise ValueError("The inertia matrix is singular")
        a[k], a[pivot] = a[pivot], a[k]
        for i in range(k+1, 3):
            factor = a[i][k]/a[k][k]
            for j in range(k, 4):
                a[i][j] -= factor*a[k][j]
    x = [0.0]*3
    for i in range(2, -1, -1):
        x[i] = (a[i][3]-sum(a[i][j]*x[j] for j in range(i+1, 3)))/a[i][i]
    return x


def angular_acceleration(inertia, omega, torque):
    gyro = cross(omega, apply(inertia, omega))
    return solve3(inertia, [torque[j]-gyro[j] for j in range(3)])


def step_free(q, omega_body, principal, h):
    """Coupled RK4 step; principal entries are positive body-frame inertias.

    Normalize the final quaternion. This is a numerical approximation,
    so energy and world angular momentum are checked separately.
    """
    if any(i <= 0 for i in principal):
        raise ValueError("Free rotation needs positive principal inertias")

    def rhs(y):
        w = y[4:]
        momentum = [i*v for i, v in zip(principal, w)]
        gyro = cross(w, momentum)
        return derivative(y[:4], w)+[-gyro[j]/principal[j] for j in range(3)]

    def shifted(y, k, factor):
        return [v+factor*dv for v, dv in zip(y, k)]

    y = list(q)+list(omega_body)
    k1 = rhs(y)
    k2 = rhs(shifted(y, k1, h/2))
    k3 = rhs(shifted(y, k2, h/2))
    k4 = rhs(shifted(y, k3, h))
    out = [v+h*(a+2*b+2*c+d)/6 for v, a, b, c, d in zip(y, k1, k2, k3, k4)]
    return normalize(out[:4]), out[4:]