"""Fixed-base tree kinematics; each body has a hinge, slide or fixed joint.

Body 0 is the world. Other bodies follow their parents in storage order.
Body frames have their origins at their joint anchors. R0 maps the child's
zero-angle axes to its parent's axes; offset uses the parent's axes.
Scalar coordinates/velocities follow the movable bodies in storage order.
Spatial vectors have angular components first and use each body's frame.
"""
from dataclasses import dataclass
import json
import math
from pathlib import Path
from rotation import identity, exp_so3, mul, apply, transpose
from quaternion import step_body, conjugate, mul as quat_mul, to_rotvec
from inertia import cross, dot, box_inertia
from spatial import motion_transform, matvec

E = tuple(identity())
ZERO = (0.0, 0.0, 0.0)


@dataclass(frozen=True)
class Body:
    parent: int
    joint: str = "hinge"
    axis: tuple = (0.0, -1.0, 0.0)
    offset: tuple = ZERO
    rotation0: tuple = E
    mass: float = 1.0
    com: tuple = ZERO
    inertia_com: tuple = E


def model_data():
    """Read the same literal loaded by the offline browser; no code evaluation."""
    source = (Path(__file__).resolve().parents[1] /
              "assets/two-link-model.js").read_text()
    return json.loads(source.split("=", 1)[1].strip().removesuffix(";"))


def two_link_model():
    data = model_data()
    lengths, masses = data["lengths"], data["masses"]
    section, axis = data["cross_section"], tuple(data["axis"])
    bodies = [Body(-1, joint="fixed", mass=0)]
    for i, (length, mass) in enumerate(zip(lengths, masses), 1):
        offset = ZERO if i == 1 else (0, 0, -lengths[i-2])
        bodies.append(Body(i-1, axis=axis, offset=offset, mass=mass,
                           com=(0, 0, -length/2),
                           inertia_com=tuple(box_inertia(mass, [*section, length]))))
    return bodies


def dof_indices(model):
    if not model or model[0].parent != -1 or model[0].joint != "fixed":
        raise ValueError("Body 0 must be the fixed world")
    indices, count = [None], 0
    for i, body in enumerate(model[1:], 1):
        if not 0 <= body.parent < i:
            raise ValueError("Parents must precede their children")
        if body.joint not in ("hinge", "slide", "fixed"):
            raise ValueError("Joint must be hinge, slide or fixed")
        if body.joint != "fixed" and abs(dot(body.axis, body.axis)-1) > 1e-12:
            raise ValueError("Joint axes must have unit length")
        indices.append(None if body.joint == "fixed" else count)
        count += body.joint != "fixed"
    return indices, count


def forward(model, q, velocity=None):
    dof, nv = dof_indices(model)
    velocity = [0.0]*nv if velocity is None else velocity
    if len(q) != nv or len(velocity) != nv:
        raise ValueError("Configuration and velocity lengths must equal nv")
    n = len(model)
    state = {
        "rotation": [None]*n, "position": [None]*n, "X": [None]*n,
        "S": [None]*n, "velocity": [None]*n, "axis": [None]*n,
        "dof": dof, "nv": nv,
    }
    state["rotation"][0] = identity()
    state["position"][0] = [0.0]*3
    state["velocity"][0] = [0.0]*6
    state["X"][0] = motion_transform(identity(), ZERO)
    state["S"][0] = [0.0]*6
    state["axis"][0] = [0.0]*3
    for i, body in enumerate(model[1:], 1):
        p, j = body.parent, dof[i]
        rotation, displacement = list(body.rotation0), list(body.offset)
        s = [0.0]*6
        if body.joint == "hinge":
            rotation = mul(rotation, exp_so3([x*q[j] for x in body.axis]))
            s[:3] = body.axis
        elif body.joint == "slide":
            direction = apply(rotation, body.axis)
            displacement = [x+q[j]*y for x, y in zip(displacement, direction)]
            s[3:] = body.axis
        rp, xp = state["rotation"][p], state["position"][p]
        state["rotation"][i] = mul(rp, rotation)
        state["position"][i] = [x+y for x, y in zip(xp, apply(rp, displacement))]
        state["X"][i] = motion_transform(rotation, displacement)
        state["S"][i] = s
        state["axis"][i] = apply(mul(rp, body.rotation0), body.axis)
        inherited = matvec(state["X"][i], state["velocity"][p])
        state["velocity"][i] = [x+(velocity[j]*y if j is not None else 0)
                                for x, y in zip(inherited, s)]
    return state


def world_point(state, body, local_point):
    return [x+y for x, y in zip(state["position"][body],
                               apply(state["rotation"][body], local_point))]


def point_velocity(state, body, local_point):
    v = state["velocity"][body]
    local = [x+y for x, y in zip(v[3:], cross(v[:3], local_point))]
    return apply(state["rotation"][body], local)


def ancestors(model, body):
    result = []
    while body > 0:
        result.append(body)
        body = model[body].parent
    return result


def geometric_jacobian(model, state, body, local_point):
    """Flat row-major 6-by-nv; world axes, reference at the queried point."""
    nv, point = state["nv"], world_point(state, body, local_point)
    jac = [0.0]*(6*nv)
    for i in ancestors(model, body):
        j = state["dof"][i]
        if j is None:
            continue
        axis = state["axis"][i]
        if model[i].joint == "hinge":
            delta = [x-y for x, y in zip(point, state["position"][i])]
            column = list(axis)+cross(axis, delta)
        else:
            column = [0.0]*3+list(axis)
        for row, value in enumerate(column):
            jac[nv*row+j] = value
    return jac


def jacobian_velocity(jac, velocity):
    nv = len(velocity)
    return [sum(jac[nv*i+j]*velocity[j] for j in range(nv)) for i in range(6)]


def generalized_force(jac, wrench, nv):
    return [sum(jac[nv*i+j]*wrench[i] for i in range(6)) for j in range(nv)]


def two_link_geometry(q):
    """Independent trigonometric positions and planar x-z Jacobian."""
    l1, l2 = model_data()["lengths"]
    a, b = q[0], q[0]+q[1]
    p1 = [l1*math.sin(a), 0, -l1*math.cos(a)]
    p2 = [p1[0]+l2*math.sin(b), 0, p1[2]-l2*math.cos(b)]
    j = [l1*math.cos(a)+l2*math.cos(b), l2*math.cos(b),
         l1*math.sin(a)+l2*math.sin(b), l2*math.sin(b)]
    return p1, p2, j


def two_link_point_acceleration(q, velocity, acceleration):
    l1, l2 = model_data()["lengths"]
    a, b = q[0], q[0]+q[1]
    v1, vsum = velocity[0], velocity[0]+velocity[1]
    a1, asum = acceleration[0], acceleration[0]+acceleration[1]
    x = l1*(math.cos(a)*a1-math.sin(a)*v1*v1)
    x += l2*(math.cos(b)*asum-math.sin(b)*vsum*vsum)
    z = l1*(math.sin(a)*a1+math.cos(a)*v1*v1)
    z += l2*(math.sin(b)*asum+math.cos(b)*vsum*vsum)
    return [x, 0, z]


def integrate_configuration(joint_types, configuration, velocity, dt):
    """MuJoCo interface convention: free linear world velocity precedes body omega."""
    out, qp, vp = list(configuration), 0, 0
    for joint in joint_types:
        if joint in ("hinge", "slide"):
            out[qp] += dt*velocity[vp]
            qp, vp = qp+1, vp+1
        elif joint == "ball":
            out[qp:qp+4] = step_body(out[qp:qp+4], velocity[vp:vp+3], dt)
            qp, vp = qp+4, vp+3
        elif joint == "free":
            for i in range(3):
                out[qp+i] += dt*velocity[vp+i]
            out[qp+3:qp+7] = step_body(out[qp+3:qp+7], velocity[vp+3:vp+6], dt)
            qp, vp = qp+7, vp+6
        else:
            raise ValueError("Unsupported configuration joint")
    if qp != len(configuration) or vp != len(velocity):
        raise ValueError("Configuration or velocity length does not match joint types")
    return out


def difference_configuration(joint_types, first, second, dt):
    if dt == 0:
        raise ValueError("dt must be nonzero")
    out, p = [], 0
    for joint in joint_types:
        if joint in ("hinge", "slide"):
            out.append((second[p]-first[p])/dt)
            p += 1
        elif joint == "ball":
            relative = quat_mul(conjugate(first[p:p+4]), second[p:p+4])
            out.extend(x/dt for x in to_rotvec(relative))
            p += 4
        elif joint == "free":
            out.extend((second[p+i]-first[p+i])/dt for i in range(3))
            relative = quat_mul(conjugate(first[p+3:p+7]), second[p+3:p+7])
            out.extend(x/dt for x in to_rotvec(relative))
            p += 7
        else:
            raise ValueError("Unsupported configuration joint")
    if p != len(first) or p != len(second):
        raise ValueError("Configuration lengths do not match joint types")
    return out