"""3D anatomical kinematics for the Exercise Library stick figure.

Model space follows the ISB convention: X anterior (the figure's facing
direction), Y superior (up), Z toward the figure's anatomical right.
All angles are degrees, measured from the neutral standing pose (upright,
facing +X, arms hanging, legs straight, toes forward).

A key frame poses joints with anatomical coordinates:
  - root: canvas anchor `pos` plus trunk orientation `yaw` (facing: 0 = +X,
    180 = -X), `pitch` (forward bow positive), `roll` (toward the right
    positive) - applied as Ry(yaw) . Rz(-pitch) . Rx(roll).
  - spine: two chained segments, each {flexion, lateral, rotation}
    (forward curl / right side-bend / turn right positive).
  - neck {flexion, rotation}, head {flexion} (extra gaze pitch).
  - shoulder/hip {flexion, abduction, rotation}: forward, away from the
    midline, external positive. elbow/knee {flexion}: bend positive.
    ankle {flexion}: dorsiflexion positive (toes up).
A bare number is shorthand for {"flexion": n}.

The camera is orthographic and rotates about the vertical axis through the
root anchor: yaw 0 is the classic side view (from +Z, anatomical right side
near), 90 views the figure face-on. Poses resolve in *view space* (x right,
y up, z toward the camera); the renderer maps to y-down canvas points.

Pins are canvas-space IK targets for hands/feet (`hand_r`, `foot_l`, ...):
the two-bone chain is solved analytically in 3D, in the plane picked by the
authored (FK) elbow/knee, then converted back to anatomical angles so key
frames always interpolate in anatomical space.

skeleton.json carries the bone-length profiles (including shoulder/pelvis
half-widths and feet) and each joint type's degrees of freedom with their
physiological range of motion (ROM), used to validate authored frames.
"""

import json
import math
from pathlib import Path

LIB = Path(__file__).parent

# ------------------------------------------------------------ vectors / mats

def _cs(deg):
    r = math.radians(deg)
    return math.cos(r), math.sin(r)


def rot_x(deg):
    c, s = _cs(deg)
    return ((1, 0, 0), (0, c, -s), (0, s, c))


def rot_y(deg):
    c, s = _cs(deg)
    return ((c, 0, s), (0, 1, 0), (-s, 0, c))


def rot_z(deg):
    c, s = _cs(deg)
    return ((c, -s, 0), (s, c, 0), (0, 0, 1))


IDENTITY = ((1, 0, 0), (0, 1, 0), (0, 0, 1))


def mmul(a, b):
    return tuple(tuple(sum(a[i][k] * b[k][j] for k in range(3)) for j in range(3))
                 for i in range(3))


def chain(*mats):
    m = mats[0]
    for n in mats[1:]:
        m = mmul(m, n)
    return m


def mvec(m, v):
    return tuple(m[i][0] * v[0] + m[i][1] * v[1] + m[i][2] * v[2] for i in range(3))


def mtrans(m):
    return tuple(tuple(m[j][i] for j in range(3)) for i in range(3))


def vadd(a, b):
    return (a[0] + b[0], a[1] + b[1], a[2] + b[2])


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


def vscale(v, s):
    return (v[0] * s, v[1] * s, v[2] * s)


def vdot(a, b):
    return a[0] * b[0] + a[1] * b[1] + a[2] * b[2]


def vcross(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 vlen(v):
    return math.sqrt(vdot(v, v))


def vnorm(v):
    d = vlen(v)
    return vscale(v, 1 / d) if d > 1e-9 else (0.0, 0.0, 0.0)


def _clamp(x, lo=-1.0, hi=1.0):
    return max(lo, min(hi, x))


# Axial rotation of a two-bone limb is recoverable only from the lower bone's
# lateral tip; below this magnitude the limb is effectively in-plane and its
# rotation is left at 0 (see invert_limb).
ROT_MIN_LATERAL = 0.08

# When a leg's authored knee sits within this fraction of the thigh length off the
# hip->ankle line (~sin 8.6 deg), the leg is treated as straight and its IK knee is
# bent anatomically forward rather than trusting the (unreliable) authored side
# (see solve_limb).
KNEE_STRAIGHT_FRAC = 0.15


# ---------------------------------------------------------------- the frame

# DoF names per joint type; the first is the shorthand a bare number sets.
JOINT_DOFS = {
    "spine": ("flexion", "lateral", "rotation"),
    "neck": ("flexion", "rotation"),
    "head": ("flexion",),
    "shoulder": ("flexion", "abduction", "rotation"),
    "elbow": ("flexion",),
    "hip": ("flexion", "abduction", "rotation"),
    "knee": ("flexion",),
    "ankle": ("flexion",),
}

# Frame keys -> joint type (spine is a two-element list handled separately).
FRAME_JOINTS = {
    "neck": "neck", "head": "head",
    "shoulder_r": "shoulder", "shoulder_l": "shoulder",
    "elbow_r": "elbow", "elbow_l": "elbow",
    "hip_r": "hip", "hip_l": "hip",
    "knee_r": "knee", "knee_l": "knee",
    "ankle_r": "ankle", "ankle_l": "ankle",
}

LIMBS = {  # limb -> (attach point key, side sign, pin key)
    "arm_r": ("shoulder_r", 1, "hand_r"),
    "arm_l": ("shoulder_l", -1, "hand_l"),
    "leg_r": ("hip_r", 1, "foot_r"),
    "leg_l": ("hip_l", -1, "foot_l"),
}


def load_skeleton():
    return json.loads((LIB / "skeleton.json").read_text())


def _full(value, joint_type):
    """Expand a joint value (number, partial dict, or None) to a full DoF dict."""
    dofs = JOINT_DOFS[joint_type]
    if value is None:
        return {d: 0.0 for d in dofs}
    if isinstance(value, (int, float)):
        out = {d: 0.0 for d in dofs}
        out[dofs[0]] = float(value)
        return out
    return {d: float(value.get(d, 0.0)) for d in dofs}


def normalize_frame(kf):
    """Expand a key frame to full anatomical dicts with defaults filled in."""
    root = kf.get("root", {})
    out = {
        "root": {"pos": [float(root["pos"][0]), float(root["pos"][1])],
                 "yaw": float(root.get("yaw", 0.0)),
                 "pitch": float(root.get("pitch", 0.0)),
                 "roll": float(root.get("roll", 0.0))},
        "spine": [_full(s, "spine") for s in (kf.get("spine") or [0, 0])],
        "pins": {k: [float(x), float(y)] for k, (x, y) in kf.get("pins", {}).items()},
        "hold": float(kf.get("hold", 0.5)),
        "tween": float(kf.get("tween", 0.6)),
    }
    for key, jt in FRAME_JOINTS.items():
        out[key] = _full(kf.get(key), jt)
    return out


def lerp_frames(a, b, t):
    """Interpolate two normalized frames in anatomical space. A pin survives
    the tween only when planted in both neighboring key frames."""
    def num(x, y):
        return x + (y - x) * t
    out = {"root": {"pos": [num(a["root"]["pos"][0], b["root"]["pos"][0]),
                    num(a["root"]["pos"][1], b["root"]["pos"][1])],
                    "yaw": num(a["root"]["yaw"], b["root"]["yaw"]),
                    "pitch": num(a["root"]["pitch"], b["root"]["pitch"]),
                    "roll": num(a["root"]["roll"], b["root"]["roll"])},
           "spine": [{d: num(sa[d], sb[d]) for d in sa}
                     for sa, sb in zip(a["spine"], b["spine"])],
           "pins": {k: [num(a["pins"][k][0], b["pins"][k][0]),
                        num(a["pins"][k][1], b["pins"][k][1])]
                    for k in a["pins"] if k in b["pins"]}}
    for key in FRAME_JOINTS:
        out[key] = {d: num(a[key][d], b[key][d]) for d in a[key]}
    return out


# ------------------------------------------------------------------ posing

def _ball(joint, sigma):
    """Local rotation of a ball joint (shoulder/hip) for side sign sigma
    (+1 right, -1 left): flexion forward, abduction away from the midline,
    rotation external."""
    return chain(rot_z(joint["flexion"]),
                 rot_x(-sigma * joint["abduction"]),
                 rot_y(-sigma * joint["rotation"]))


def fk_limb(kind, attach, joint, lower, ankle, prof, parent, sigma):
    """FK one limb from its resolved attach point. Returns the point chain
    (arm: [shoulder, elbow, hand]; leg: [hip, knee, ankle, toe])."""
    fu = mmul(parent, _ball(joint, sigma))
    if kind == "arm":
        elbow = vadd(attach, mvec(fu, (0, -prof["upperArm"], 0)))
        fl = mmul(fu, rot_z(lower["flexion"]))
        hand = vadd(elbow, mvec(fl, (0, -prof["foreArm"], 0)))
        return [attach, elbow, hand]
    knee = vadd(attach, mvec(fu, (0, -prof["thigh"], 0)))
    fl = mmul(fu, rot_z(-lower["flexion"]))
    ank = vadd(knee, mvec(fl, (0, -prof["shin"], 0)))
    toe = vadd(ank, mvec(mmul(fl, rot_z(ankle["flexion"])), (prof["foot"], 0, 0)))
    return [attach, knee, ank, toe]


def pose(nf, prof, cam_yaw, cam_pitch=0.0):
    """FK a normalized frame into view space (x right, y up, z toward the
    camera; origin at the root anchor). `cam_pitch` tilts the viewpoint down
    from slightly above (the scene rotates about the root). Returns points,
    parent frames (for IK inversion), the nose direction, and the lateral
    depth factor k."""
    r = nf["root"]
    f_root = chain(rot_x(cam_pitch), rot_y(-cam_yaw),
                   rot_y(r["yaw"]), rot_z(-r["pitch"]), rot_x(r["roll"]))
    origin = (0.0, 0.0, 0.0)

    s1, s2 = nf["spine"]
    f1 = chain(f_root, rot_z(-s1["flexion"]), rot_x(s1["lateral"]), rot_y(-s1["rotation"]))
    mid = vadd(origin, mvec(f1, (0, prof["spine1"], 0)))
    f2 = chain(f1, rot_z(-s2["flexion"]), rot_x(s2["lateral"]), rot_y(-s2["rotation"]))
    neck_base = vadd(mid, mvec(f2, (0, prof["spine2"], 0)))

    fn = chain(f2, rot_z(-nf["neck"]["flexion"]), rot_y(-nf["neck"]["rotation"]))
    head = vadd(neck_base, mvec(fn, (0, prof["neck"], 0)))
    nose_dir = mvec(mmul(fn, rot_z(-nf["head"]["flexion"])), (1, 0, 0))

    points = {
        "pelvis": origin, "mid": mid, "neckB": neck_base, "head": head,
        "shoulder_r": vadd(neck_base, mvec(f2, (0, 0, prof["shoulderHalf"]))),
        "shoulder_l": vadd(neck_base, mvec(f2, (0, 0, -prof["shoulderHalf"]))),
        "hip_r": vadd(origin, mvec(f_root, (0, 0, prof["hipHalf"]))),
        "hip_l": vadd(origin, mvec(f_root, (0, 0, -prof["hipHalf"]))),
    }
    for limb, (attach, sigma, _pin) in LIMBS.items():
        kind = "arm" if limb.startswith("arm") else "leg"
        side = limb[-1]
        parent = f2 if kind == "arm" else f_root
        upper = nf[("shoulder_" if kind == "arm" else "hip_") + side]
        lower = nf[("elbow_" if kind == "arm" else "knee_") + side]
        ankle = nf.get("ankle_" + side)
        points[limb] = fk_limb(kind, points[attach], upper, lower, ankle, prof, parent, sigma)

    # How side-on the view is: 1 when the body's lateral axis is pure depth
    # (profile view), 0 when it lies in the screen (face-on).
    k = abs(mvec(f_root, (0, 0, 1))[2])
    return {"points": points, "f2": f2, "f_root": f_root,
            "nose_dir": nose_dir, "k": k}


# ---------------------------------------------------------------------- IK

def solve_limb(kind, attach, target, guess_mid, lengths, parent, sigma):
    """Analytic two-bone IK in 3D: reach from `attach` toward `target` in the
    plane picked by the authored (FK) mid joint, then convert back to
    anatomical angles. Returns (upper joint dict, lower joint dict)."""
    a, b = lengths
    to_t = vsub(target, attach)
    d = _clamp(vlen(to_t), abs(a - b) + 0.5, a + b - 0.01)
    dir_t = vnorm(to_t) if vlen(to_t) > 1e-9 else (0.0, -1.0, 0.0)

    normal = vcross(dir_t, vsub(guess_mid, attach))
    if vlen(normal) < 1e-6:  # chain straight along the target: any plane works
        normal = vcross(dir_t, (0, 0, 1))
        if vlen(normal) < 1e-6:
            normal = vcross(dir_t, (0, 1, 0))
    normal = vnorm(normal)
    perp = vcross(normal, dir_t)

    along = (a * a + d * d - b * b) / (2 * d)
    h = math.sqrt(max(a * a - along * along, 0.0))
    if kind == "leg":
        # A knee bends one way only. Near full extension the two knee solutions
        # straddle the hip->ankle line and the authored guess (also near that line)
        # cannot reliably pick a side - the bend plane's normal is a cross of two
        # near-parallel vectors, so its sign is noise and the knee can flip behind
        # the leg, swinging the whole thigh backward. When the authored knee sits
        # within KNEE_STRAIGHT_FRAC of the line, treat the leg as straight and bend
        # the knee anatomically forward (anterior); otherwise honor the authored side.
        gm = vsub(guess_mid, attach)
        gm_perp = vsub(gm, vscale(dir_t, vdot(gm, dir_t)))
        ref = mvec(parent, (1, 0, 0)) if vlen(gm_perp) < KNEE_STRAIGHT_FRAC * a else gm_perp
        sign = 1.0 if vdot(perp, ref) >= 0 else -1.0
        mid = vadd(attach, vadd(vscale(dir_t, along), vscale(perp, sign * h)))
    else:
        best = None
        for sign in (1.0, -1.0):
            mid = vadd(attach, vadd(vscale(dir_t, along), vscale(perp, sign * h)))
            dist = vlen(vsub(mid, guess_mid))
            if best is None or dist < best[0]:
                best = (dist, mid)
        mid = best[1]
    end = vadd(mid, vscale(vnorm(vsub(target, mid)), b))
    return invert_limb(kind, attach, mid, end, lengths, parent, sigma)


def invert_limb(kind, attach, mid, end, lengths, parent, sigma):
    """Recover anatomical angles from limb joint positions (the inverse of
    fk_limb, ignoring the foot). Assumes |abduction| < 90."""
    pt = mtrans(parent)
    u = vnorm(mvec(pt, vsub(mid, attach)))
    abd = math.degrees(math.asin(_clamp(sigma * u[2])))
    flex = math.degrees(math.atan2(u[0], -u[1]))
    # Peel flexion/abduction off the lower bone to read rotation + hinge bend.
    peel = mtrans(chain(rot_z(flex), rot_x(-sigma * abd)))
    w = vnorm(mvec(peel, mvec(pt, vsub(end, mid))))
    bend = math.degrees(math.acos(_clamp(-w[1])))
    # Axial rotation is observable only through the lower bone's lateral tip
    # (w[2]); a near-sagittal limb (w[2] ~ 0) carries no recoverable rotation.
    # Gating on `bend` alone is too weak: a limb can straighten through this
    # degeneracy while still visibly bent, and atan2 then snaps to +/-180 on the
    # sign of a near-zero anterior component - twisting the limb a half-turn and
    # flipping a pinned hand/foot backward.
    twist = bend > 0.5 and abs(w[2]) > ROT_MIN_LATERAL
    if kind == "leg":  # knees hinge backward; the sign convention flips
        rot = sigma * math.degrees(math.atan2(-w[2], -w[0])) if twist else 0.0
    else:
        rot = sigma * math.degrees(math.atan2(w[2], w[0])) if twist else 0.0
    upper = {"flexion": flex, "abduction": abd, "rotation": rot}
    return upper, {"flexion": bend}


# ------------------------------------------------------------------ resolve

def view_from_canvas(pt, anchor, depth):
    return (pt[0] - anchor[0], anchor[1] - pt[1], depth)


def resolve(nf, prof, cam_yaw, cam_pitch=0.0):
    """Pose a normalized frame and apply pins: for each pinned limb, solve IK
    against the canvas target (at the limb's FK depth), write the solved
    anatomical angles back into the frame, and re-pose. Returns
    (frame with IK-resolved angles, pose dict)."""
    p = pose(nf, prof, cam_yaw, cam_pitch)
    anchor = nf["root"]["pos"]
    solved = False
    for limb, (attach_key, sigma, pin) in LIMBS.items():
        if pin not in nf["pins"]:
            continue
        kind = "arm" if limb.startswith("arm") else "leg"
        side = limb[-1]
        chain_pts = p["points"][limb]
        attach = chain_pts[0]
        end_idx = 2
        target = view_from_canvas(nf["pins"][pin], anchor, chain_pts[end_idx][2])
        lengths = ((prof["upperArm"], prof["foreArm"]) if kind == "arm"
                   else (prof["thigh"], prof["shin"]))
        parent = p["f2"] if kind == "arm" else p["f_root"]
        upper, lower = solve_limb(kind, attach, target, chain_pts[1],
                                  lengths, parent, sigma)
        if kind == "arm":
            nf["shoulder_" + side], nf["elbow_" + side] = upper, lower
        else:
            nf["hip_" + side], nf["knee_" + side] = upper, lower
        solved = True
    if solved:
        p = pose(nf, prof, cam_yaw, cam_pitch)
    return nf, p


# -------------------------------------------------------------------- ROM

def validate_rom(nf, joints, label=""):
    """Check a normalized frame's anatomical angles against each joint's
    range of motion. Returns a list of human-readable violations."""
    issues = []

    def check(joint_type, name, value):
        for dof, angle in value.items():
            lo_hi = joints.get(joint_type, {}).get(dof)
            if lo_hi and not (lo_hi[0] - 1e-6 <= angle <= lo_hi[1] + 1e-6):
                issues.append(f"{label}{name}.{dof} = {angle:.1f} outside "
                              f"[{lo_hi[0]}, {lo_hi[1]}]")

    for i, seg in enumerate(nf["spine"], start=1):
        check("spine", f"spine{i}", seg)
    for key, jt in FRAME_JOINTS.items():
        check(jt, key, nf[key])
    return issues