"""Frozen R5 feature enrichment from the reviewed offline candidate."""
import math


def joints(body):
    return {int(k): v for k, v in body['joints'].items()}


def midpoint(j, a, b):
    return [(j[a][k] + j[b][k]) / 2 for k in (0, 1)] if a in j and b in j else None


def enrich(f, previous, current):
    a = f['actor']; other = 1 - a
    p, q = current[a], current[other]
    old = previous[a]
    j, k, oldj = joints(p), joints(q), joints(old)
    h, oh, dt = max(1, p['height']), max(1, q['height']), f['elapsed']
    shoulder = midpoint(j, 5, 6)
    hip = midpoint(j, 11, 12)
    target_shoulder = midpoint(k, 5, 6)
    end = j[f['end']]
    hand_distances = [(math.dist(end[:2], k[i][:2]) / oh, i) for i in (9, 10) if i in k]
    hand = min(hand_distances) if hand_distances else None
    target_hip = midpoint(k, 11, 12)
    root_delta = [(j[f['root']][i] - oldj[f['root']][i]) / h / dt for i in (0, 1)]
    old_shoulder = midpoint(oldj, 5, 6)
    old_hip = midpoint(oldj, 11, 12)
    angle = None
    if shoulder and hip:
        angle = abs(shoulder[0]-hip[0]) / max(1, math.dist(shoulder, hip))
    target_upright = None
    if target_shoulder and target_hip:
        target_upright = abs(target_shoulder[0]-target_hip[0]) / max(1, math.dist(target_shoulder,target_hip))
    return dict(f, shoulder_offset=None if shoulder is None else (end[1]-shoulder[1])/h,
        horizontal_reach=abs(end[0]-j[f['root']][0])/h,
        hand_contact=None if hand is None else hand[0],
        target_hand_offset=None if hand is None or target_shoulder is None else (k[hand[1]][1]-target_shoulder[1])/oh,
        root_vx=root_delta[0],root_vy=root_delta[1],
        shoulder_vy=None if shoulder is None or old_shoulder is None else (shoulder[1]-old_shoulder[1])/h/dt,
        hip_vy=None if hip is None or old_hip is None else (hip[1]-old_hip[1])/h/dt,
        torso_horizontal=angle,target_torso_horizontal=target_upright,
        target_height_ratio=oh/h,
        target_center_distance=math.dist(p['center'],q['center'])/max(1,(h+oh)/2))