"""Angular-first spatial algebra with flat, row-major matrices.

R maps new-frame coordinates to old-frame coordinates. d is the displacement
from the old reference point to the new one, expressed in the old frame.
Both reference points are attached to the same rigid body for acceleration
transforms. A = [alpha, a_point - omega x v_point].
"""
from rotation import identity, skew, mul, transpose
from inertia import cross, dot, shift_inertia


def matvec(a, x):
    n = len(x)
    return [sum(a[n*i+j]*x[j] for j in range(n)) for i in range(n)]


def matmul(a, b, n=6):
    return [sum(a[n*i+k]*b[n*k+j] for k in range(n))
            for i in range(n) for j in range(n)]


def transposed(a, n=6):
    return [a[n*j+i] for i in range(n) for j in range(n)]


def blocks(a, b, c, d):
    """Assemble four row-major 3x3 blocks into a row-major 6x6 matrix."""
    return [((a if i < 3 else c) if j < 3 else
             (b if i < 3 else d))[3*(i % 3)+j % 3]
            for i in range(6) for j in range(6)]


def motion_transform(rotation, displacement):
    rt = transpose(rotation)
    lower = [-x for x in mul(rt, skew(displacement))]
    return blocks(rt, [0.0]*9, lower, rt)


def inverse_motion_transform(rotation, displacement):
    return blocks(rotation, [0.0]*9, mul(skew(displacement), rotation), rotation)


def force_transform(rotation, displacement):
    return transposed(inverse_motion_transform(rotation, displacement))


def spatial_inertia(mass, origin_to_com, inertia_com):
    """COM inertia and displacement are expressed in the reference frame."""
    if mass <= 0:
        raise ValueError("Mass must be positive")
    md = [mass*x for x in skew(origin_to_com)]
    return blocks(shift_inertia(inertia_com, mass, origin_to_com), md,
                  [-x for x in md], [mass*x for x in identity()])


def cross_motion(v, u):
    return cross(v[:3], u[:3])+[
        a+b for a, b in zip(cross(v[3:], u[:3]), cross(v[:3], u[3:]))]


def cross_force(v, f):
    return [a+b for a, b in zip(cross(v[:3], f[:3]),
                                cross(v[3:], f[3:]))]+cross(v[:3], f[3:])


def motion_cross_matrix(v):
    w, linear = skew(v[:3]), skew(v[3:])
    return blocks(w, [0.0]*9, linear, w)


def force_cross_matrix(v):
    return [-x for x in transposed(motion_cross_matrix(v))]


def spatial_acceleration(omega, alpha, velocity, point_acceleration):
    correction = cross(omega, velocity)
    return list(alpha)+[a-c for a, c in zip(point_acceleration, correction)]


def point_acceleration(v, a):
    correction = cross(v[:3], v[3:])
    return [x+y for x, y in zip(a[3:], correction)]


def required_wrench(inertia, velocity, acceleration):
    linear = matvec(inertia, acceleration)
    bias = cross_force(velocity, matvec(inertia, velocity))
    return [x+y for x, y in zip(linear, bias)]


def screw_axis(velocity):
    """Return (displacement to closest axis point, pitch); None for translation."""
    omega, linear = velocity[:3], velocity[3:]
    magnitude2 = dot(omega, omega)
    if magnitude2 == 0:
        return None
    return ([x/magnitude2 for x in cross(omega, linear)],
            dot(omega, linear)/magnitude2)


DOOR = {
    "mass": 2.0, "width": 0.8, "height": 1.2,
    "com": [0.4, 0.0, 0.0],
    # Uniform thin plate in the local x-z plane.
    "inertia_com": [0.24, 0, 0, 0, 2*(0.8**2+1.2**2)/12,
                    0, 0, 0, 2*0.8**2/12],
}