"""Fixed-base tree dynamics; angular-first vectors in local body frames.

Body 0 is the world. Each body has one fixed-axis hinge, slide or fixed
connection. external maps body indices to local wrenches about body origins.
Gravity is a world vector; returned RNEA acceleration includes the -g shift.
All matrices are flat, row-major. No third-party dependency is required.
"""
from kinematics import Body, forward, dof_indices, two_link_model
from rotation import exp_so3, transpose, apply
from inertia import cross
from spatial import (spatial_inertia, matvec, matmul, transposed,
                     cross_motion, cross_force)
from analytic_dynamics import solve_spd


def body_inertias(model):
    return [[0.0]*36] + [spatial_inertia(b.mass, b.com, b.inertia_com)
                         for b in model[1:]]


def point_wrench(state, body, local_point, world_force, world_couple=(0, 0, 0)):
    """Point is body-local; applied force and free couple use world axes."""
    rt = transpose(state["rotation"][body])
    f, t = apply(rt, world_force), apply(rt, world_couple)
    return [x+y for x, y in zip(t, cross(local_point, f))]+f


def rnea(model, q, velocity, acceleration, gravity=(0, 0, -9.81), external=None):
    state = forward(model, q, velocity)
    nv, n = state["nv"], len(model)
    if len(acceleration) != nv or len(gravity) != 3:
        raise ValueError("Acceleration or gravity length does not match model")
    external = {} if external is None else external
    if any(not 0 < i < n or len(f) != 6 for i, f in external.items()):
        raise ValueError("External wrenches require a body index and six components")
    inertia = body_inertias(model)
    acc, own = [[0.0]*6 for _ in model], [[0.0]*6 for _ in model]
    acc[0] = [0.0]*3+[-x for x in gravity]
    for i, b in enumerate(model[1:], 1):
        j, s, x, v = state["dof"][i], state["S"][i], state["X"][i], state["velocity"][i]
        joint_v = [k*(velocity[j] if j is not None else 0) for k in s]
        transported = matvec(x, acc[b.parent])
        product = cross_motion(v, joint_v)
        acc[i] = [transported[k]+product[k]+s[k]*(acceleration[j] if j is not None else 0)
                  for k in range(6)]
        gyro = cross_force(v, matvec(inertia[i], v))
        linear = matvec(inertia[i], acc[i])
        applied = external.get(i, [0.0]*6)
        own[i] = [linear[k]+gyro[k]-applied[k] for k in range(6)]
    total, tau, steps = [list(f) for f in own], [0.0]*nv, []
    for i in range(n-1, 0, -1):
        j, p = state["dof"][i], model[i].parent
        if j is not None:
            tau[j] = sum(x*y for x, y in zip(state["S"][i], total[i]))
        transmitted = matvec(transposed(state["X"][i]), total[i])
        total[p] = [x+y for x, y in zip(total[p], transmitted)]
        steps.append({"body": i, "parent": p, "force": list(total[i]),
                      "transmitted": transmitted, "tau": tau[j] if j is not None else None})
    return {"state": state, "inertia": inertia, "acceleration": acc,
            "body_force": own, "subtree_force": total, "tau": tau, "steps": steps}


def crba(model, q):
    state = forward(model, q)
    composite, steps = body_inertias(model), []
    nv, n = state["nv"], len(model)
    for i in range(n-1, 0, -1):
        p, x = model[i].parent, state["X"][i]
        moved = matmul(transposed(x), matmul(composite[i], x))
        composite[p] = [a+b for a, b in zip(composite[p], moved)]
        steps.append({"body": i, "parent": p, "contribution": moved,
                      "parent_inertia": list(composite[p])})
    mass = [0.0]*(nv*nv)
    for i in range(1, n):
        j = state["dof"][i]
        if j is None:
            continue
        h = matvec(composite[i], state["S"][i])
        mass[nv*j+j] = sum(x*y for x, y in zip(state["S"][i], h))
        k = i
        while model[k].parent > 0:
            h = matvec(transposed(state["X"][k]), h)
            k = model[k].parent
            ancestor = state["dof"][k]
            if ancestor is not None:
                value = sum(x*y for x, y in zip(state["S"][k], h))
                mass[nv*ancestor+j] = mass[nv*j+ancestor] = value
    return {"mass": mass, "composite": composite, "steps": steps, "state": state}


def mass_by_rnea(model, q):
    _, nv = dof_indices(model)
    result = [0.0]*(nv*nv)
    for j in range(nv):
        unit = [float(k == j) for k in range(nv)]
        column = rnea(model, q, [0.0]*nv, unit, gravity=(0, 0, 0))["tau"]
        for i, value in enumerate(column):
            result[nv*i+j] = value
    return result


def forward_dynamics(model, q, velocity, force, gravity=(0, 0, -9.81), external=None):
    _, nv = dof_indices(model)
    if len(force) != nv:
        raise ValueError("Generalized force length must equal nv")
    c = rnea(model, q, velocity, [0.0]*nv, gravity, external)["tau"]
    mass = crba(model, q)["mass"]
    return solve_spd(mass, [x-y for x, y in zip(force, c)])


def branch_model():
    """Three DOFs, one fixed leaf: geometry shared with chapter 6."""
    from inertia import box_inertia
    def link(parent, joint, axis, offset, rotation0, mass, com, size):
        return Body(parent, joint=joint, axis=axis, offset=offset,
                    rotation0=tuple(rotation0), mass=mass, com=com,
                    inertia_com=tuple(box_inertia(mass, size)))
    from rotation import identity
    return [Body(-1, joint="fixed", mass=0),
            link(0, "hinge", (0, 0, 1), (.2, 0, .1), exp_so3([0, .3, 0]),
                 1.2, (0, 0, -.3), [.15, .2, .6]),
            link(1, "slide", (1, 0, 0), (0, 0, -1), identity(),
                 .8, (.2, 0, 0), [.4, .15, .15]),
            link(1, "hinge", (1, 0, 0), (.4, 0, -.5), identity(),
                 1.0, (0, 0, -.2), [.2, .2, .4]),
            link(3, "fixed", (0, 0, 1), (.1, .2, 0), exp_so3([0, 0, .2]),
                 .4, (.1, 0, 0), [.2, .1, .1])]


if __name__ == "__main__":
    from math import pi
    model, q, v, a = two_link_model(), [pi/6, pi/3], [1, -.5], [.4, -.2]
    inverse = rnea(model, q, v, a)
    print("M =", crba(model, q)["mass"])
    print("c =", rnea(model, q, v, [0, 0])["tau"])
    print("tau =", inverse["tau"])
    print("reconstructed a =", forward_dynamics(model, q, v, inverse["tau"]))