"""Drives the 70-point OpenPose_full face part from ARKit's 52 facial
blendshape weights, for rigs that animate the face with blendShape targets
(or live ARKit/iPhone TrueDepth mocap) instead of facial joints.

``ARKIT_BLENDSHAPE_NAMES`` (order, spelling) is verbatim from Apple's own
``ARFaceAnchor.BlendShapeLocation`` enumeration.

There is no published, authoritative "which ARKit weight moves which
OpenPose/iBUG landmark, by how much" table anywhere -- unlike BODY_25/COCO,
which OpenPose's own source defines precisely, this is a genuinely
unstandardized problem. ``NEUTRAL_FACE_TEMPLATE`` and ``BLENDSHAPE_OFFSETS``
below are therefore a hand-authored, geometrically-reasoned *approximation*
(each shape nudges only the anatomically obvious nearby landmarks, in an
intuitive direction) -- good enough to produce a recognizable, responsive
OpenPose-style face for ControlNet conditioning, but not a biomechanically
accurate face simulation. A few shapes with no representable effect on a
frontal 2D landmark set (``jawForward``, ``tongueOut``, ``mouthFunnel`` --
visually redundant with ``mouthPucker`` on this coarse a template) are
intentionally no-ops.
"""

import math

from . import constants

# Verbatim from Apple's ARFaceAnchor.BlendShapeLocation.
ARKIT_BLENDSHAPE_NAMES = [
    # Left eye
    "eyeBlinkLeft", "eyeLookDownLeft", "eyeLookInLeft", "eyeLookOutLeft",
    "eyeLookUpLeft", "eyeSquintLeft", "eyeWideLeft",
    # Right eye
    "eyeBlinkRight", "eyeLookDownRight", "eyeLookInRight", "eyeLookOutRight",
    "eyeLookUpRight", "eyeSquintRight", "eyeWideRight",
    # Jaw
    "jawForward", "jawLeft", "jawRight", "jawOpen",
    # Mouth
    "mouthClose", "mouthFunnel", "mouthPucker", "mouthLeft", "mouthRight",
    "mouthSmileLeft", "mouthSmileRight", "mouthFrownLeft", "mouthFrownRight",
    "mouthDimpleLeft", "mouthDimpleRight", "mouthStretchLeft", "mouthStretchRight",
    "mouthRollLower", "mouthRollUpper", "mouthShrugLower", "mouthShrugUpper",
    "mouthPressLeft", "mouthPressRight", "mouthLowerDownLeft", "mouthLowerDownRight",
    "mouthUpperUpLeft", "mouthUpperUpRight",
    # Brow
    "browDownLeft", "browDownRight", "browInnerUp", "browOuterUpLeft", "browOuterUpRight",
    # Cheek
    "cheekPuff", "cheekSquintLeft", "cheekSquintRight",
    # Nose
    "noseSneerLeft", "noseSneerRight",
    # Tongue
    "tongueOut",
]
assert len(ARKIT_BLENDSHAPE_NAMES) == 52
BLENDSHAPE_COUNT = 52
BLENDSHAPE_NAME_TO_INDEX = {name: i for i, name in enumerate(ARKIT_BLENDSHAPE_NAMES)}

_FACE_INDEX = {name: i for i, name in enumerate(constants.FACE_NAMES)}

# Mouth outer contour = Mouth0..11 (12 pts), inner contour = Mouth12..19 (8
# pts). Within the outer contour: 0 = left corner, 1-5 = upper lip (left to
# right), 6 = right corner, 7-11 = lower lip (right to left).
_MOUTH_UPPER_OUTER = [0, 1, 2, 3, 4, 5, 6]
_MOUTH_LOWER_OUTER = [6, 7, 8, 9, 10, 11, 0]


def _mouth_outer(i):
    return "Mouth{0}".format(i)


def _mouth_inner(i):
    return "Mouth{0}".format(12 + i)


def _arc(cx, cy, rx, ry, start_deg, end_deg, count):
    points = []
    for i in range(count):
        t = i / (count - 1) if count > 1 else 0.0
        angle = math.radians(start_deg + (end_deg - start_deg) * t)
        points.append((cx + rx * math.cos(angle), cy + ry * math.sin(angle)))
    return points


def _build_neutral_template():
    """70 (x, y) points in a local, camera-facing "face plane" (roughly a
    unit square, +X = screen-right, +Y = up, origin ~ between the eyes), in
    the same order as constants.FACE_NAMES. Proportions are a plausible
    generic face, not a specific published dataset's exact mean shape.
    """
    template = [None] * constants.FACE_COUNT

    # Jaw0..16: jawline, ear to ear around the chin.
    for i in range(17):
        angle = math.radians(180 + 180 * i / 16.0)
        x = 1.05 * math.cos(angle)
        y = -0.15 + 0.85 * math.sin(angle)
        template[_FACE_INDEX["Jaw{0}".format(i)]] = (x, y)

    # Eyebrows: shallow arcs above each eye. Right eyebrow sits at negative
    # local X, left eyebrow mirrors at positive X (an arbitrary but
    # internally consistent left/right convention, same idea as BODY_25).
    for i, (x, y) in enumerate(_arc(-0.5, 0.62, 0.32, 0.10, 200, 340, 5)):
        template[_FACE_INDEX["REyebrow{0}".format(i)]] = (x, y)
    for i, (x, y) in enumerate(_arc(0.5, 0.62, 0.32, 0.10, -20, 140, 5)):
        template[_FACE_INDEX["LEyebrow{0}".format(i)]] = (x, y)

    # Nose: bridge (top to tip) then the nostril-base arc.
    for i in range(4):
        t = i / 3.0
        template[_FACE_INDEX["Nose{0}".format(i)]] = (0.0, 0.42 - 0.38 * t)
    for i, (x, y) in enumerate(_arc(0.0, -0.02, 0.20, 0.06, 200, 340, 5)):
        template[_FACE_INDEX["Nose{0}".format(4 + i)]] = (x, y)

    # Eyes: 6 points each (outer corner, 2x upper lid, inner corner, 2x lower lid).
    for i, (x, y) in enumerate(_arc(-0.5, 0.32, 0.30, 0.14, 0, 300, 6)):
        template[_FACE_INDEX["REye{0}".format(i)]] = (x, y)
    for i, (x, y) in enumerate(_arc(0.5, 0.32, 0.30, 0.14, 180, -120, 6)):
        template[_FACE_INDEX["LEye{0}".format(i)]] = (x, y)

    # Mouth: 12-point outer contour, then 8-point inner contour.
    for i, (x, y) in enumerate(_arc(0.0, -0.55, 0.42, 0.20, 0, 300, 12)):
        template[_FACE_INDEX["Mouth{0}".format(i)]] = (x, y)
    for i, (x, y) in enumerate(_arc(0.0, -0.55, 0.28, 0.11, 0, 315, 8)):
        template[_FACE_INDEX["Mouth{0}".format(12 + i)]] = (x, y)

    template[_FACE_INDEX["RPupil"]] = (-0.5, 0.32)
    template[_FACE_INDEX["LPupil"]] = (0.5, 0.32)

    assert all(p is not None for p in template)
    return template


NEUTRAL_FACE_TEMPLATE = _build_neutral_template()


def _sign(value):
    return 1.0 if value >= 0 else -1.0


# {blendshape_name: [(face_point_name, dx, dy), ...]} -- displacement applied
# at weight 1.0, linearly scaled by the live weight. See module docstring for
# the honesty caveat on precision.
BLENDSHAPE_OFFSETS = {
    # --- Eyes: eyelid closure / widening (upper lid moves toward or away
    # from the lower lid; lower lid nudges slightly the opposite way).
    "eyeBlinkLeft": [("LEye1", 0.0, -0.10), ("LEye2", 0.0, -0.10)],
    "eyeBlinkRight": [("REye1", 0.0, -0.10), ("REye2", 0.0, -0.10)],
    "eyeWideLeft": [("LEye1", 0.0, 0.05), ("LEye2", 0.0, 0.05), ("LEye4", 0.0, -0.03), ("LEye5", 0.0, -0.03)],
    "eyeWideRight": [("REye1", 0.0, 0.05), ("REye2", 0.0, 0.05), ("REye4", 0.0, -0.03), ("REye5", 0.0, -0.03)],
    "eyeSquintLeft": [("LEye1", 0.0, -0.04), ("LEye2", 0.0, -0.04), ("LEye4", 0.0, 0.02), ("LEye5", 0.0, 0.02)],
    "eyeSquintRight": [("REye1", 0.0, -0.04), ("REye2", 0.0, -0.04), ("REye4", 0.0, 0.02), ("REye5", 0.0, 0.02)],
    # Gaze direction moves the pupil point within the eye socket.
    "eyeLookDownLeft": [("LPupil", 0.0, -0.06)],
    "eyeLookUpLeft": [("LPupil", 0.0, 0.06)],
    "eyeLookInLeft": [("LPupil", -0.08, 0.0)],
    "eyeLookOutLeft": [("LPupil", 0.08, 0.0)],
    "eyeLookDownRight": [("RPupil", 0.0, -0.06)],
    "eyeLookUpRight": [("RPupil", 0.0, 0.06)],
    "eyeLookInRight": [("RPupil", 0.08, 0.0)],
    "eyeLookOutRight": [("RPupil", -0.08, 0.0)],

    # --- Brows.
    "browDownLeft": [("LEyebrow{0}".format(i), 0.0, -0.08) for i in range(5)],
    "browDownRight": [("REyebrow{0}".format(i), 0.0, -0.08) for i in range(5)],
    "browInnerUp": [("LEyebrow0", 0.0, 0.09), ("REyebrow4", 0.0, 0.09)],
    "browOuterUpLeft": [("LEyebrow3", 0.0, 0.08), ("LEyebrow4", 0.0, 0.08)],
    "browOuterUpRight": [("REyebrow0", 0.0, 0.08), ("REyebrow1", 0.0, 0.08)],

    # --- Jaw. jawForward has no representable 2D effect (frontal view) --
    # intentionally omitted below.
    "jawOpen": (
        [("Jaw{0}".format(i), 0.0, -0.30) for i in range(6, 11)]
        + [(_mouth_outer(i), 0.0, -0.22) for i in _MOUTH_LOWER_OUTER]
        + [(_mouth_inner(i), 0.0, -0.16) for i in range(4, 7)]
    ),
    "jawLeft": [("Jaw{0}".format(i), 0.06, 0.0) for i in range(6, 11)],
    "jawRight": [("Jaw{0}".format(i), -0.06, 0.0) for i in range(6, 11)],

    # --- Mouth shape.
    "mouthSmileLeft": [(_mouth_outer(0), 0.05, 0.10), (_mouth_outer(1), 0.03, 0.06), (_mouth_outer(11), 0.03, 0.06)],
    "mouthSmileRight": [(_mouth_outer(6), -0.05, 0.10), (_mouth_outer(5), -0.03, 0.06), (_mouth_outer(7), -0.03, 0.06)],
    "mouthFrownLeft": [(_mouth_outer(0), 0.02, -0.08)],
    "mouthFrownRight": [(_mouth_outer(6), -0.02, -0.08)],
    "mouthLeft": [(_mouth_outer(i), 0.08, 0.0) for i in range(12)],
    "mouthRight": [(_mouth_outer(i), -0.08, 0.0) for i in range(12)],
    "mouthDimpleLeft": [(_mouth_outer(0), 0.02, 0.02)],
    "mouthDimpleRight": [(_mouth_outer(6), -0.02, 0.02)],
    "mouthStretchLeft": [(_mouth_outer(0), 0.10, 0.0)],
    "mouthStretchRight": [(_mouth_outer(6), -0.10, 0.0)],
    "mouthPucker": [
        (_mouth_outer(i), -0.10 * _sign(NEUTRAL_FACE_TEMPLATE[_FACE_INDEX[_mouth_outer(i)]][0]), 0.0)
        for i in range(12)
    ],
    "mouthFunnel": [],  # visually redundant with mouthPucker on this coarse template
    "mouthClose": (
        [(_mouth_outer(i), 0.0, -0.02) for i in _MOUTH_UPPER_OUTER]
        + [(_mouth_outer(i), 0.0, 0.02) for i in _MOUTH_LOWER_OUTER]
    ),
    "mouthRollLower": [(_mouth_outer(i), 0.0, 0.03) for i in _MOUTH_LOWER_OUTER],
    "mouthRollUpper": [(_mouth_outer(i), 0.0, -0.03) for i in _MOUTH_UPPER_OUTER],
    "mouthShrugLower": [(_mouth_outer(i), 0.0, -0.04) for i in _MOUTH_LOWER_OUTER],
    "mouthShrugUpper": [(_mouth_outer(i), 0.0, 0.04) for i in _MOUTH_UPPER_OUTER],
    "mouthPressLeft": [(_mouth_outer(0), 0.02, -0.01)],
    "mouthPressRight": [(_mouth_outer(6), -0.02, -0.01)],
    "mouthLowerDownLeft": [(_mouth_outer(7), 0.0, -0.06), (_mouth_outer(8), 0.0, -0.06)],
    "mouthLowerDownRight": [(_mouth_outer(10), 0.0, -0.06), (_mouth_outer(11), 0.0, -0.06)],
    "mouthUpperUpLeft": [(_mouth_outer(1), 0.0, 0.05), (_mouth_outer(2), 0.0, 0.05)],
    "mouthUpperUpRight": [(_mouth_outer(4), 0.0, 0.05), (_mouth_outer(5), 0.0, 0.05)],

    # --- Cheeks / nose: subtle.
    "cheekPuff": [("Jaw{0}".format(i), 0.03, 0.0) for i in (3, 4)] + [("Jaw{0}".format(i), -0.03, 0.0) for i in (12, 13)],
    "cheekSquintLeft": [("LEye4", 0.0, 0.02), ("LEye5", 0.0, 0.02)],
    "cheekSquintRight": [("REye4", 0.0, 0.02), ("REye5", 0.0, 0.02)],
    "noseSneerLeft": [("Nose7", 0.0, 0.03), ("Nose8", 0.0, 0.03)],
    "noseSneerRight": [("Nose4", 0.0, 0.03), ("Nose5", 0.0, 0.03)],

    # --- No 2D-representable effect on a frontal landmark set.
    "jawForward": [],
    "tongueOut": [],
}
assert set(BLENDSHAPE_OFFSETS.keys()) == set(ARKIT_BLENDSHAPE_NAMES)


def compute_face_local_points(weights):
    """``weights``: {blendshape_name: 0..1}. Returns 70 (x, y) local points
    (constants.FACE_NAMES order), the neutral template plus every active
    blendshape's weighted displacement.
    """
    points = list(NEUTRAL_FACE_TEMPLATE)
    for name, weight in weights.items():
        if not weight:
            continue
        for point_name, dx, dy in BLENDSHAPE_OFFSETS.get(name, ()):
            idx = _FACE_INDEX[point_name]
            x, y = points[idx]
            points[idx] = (x + dx * weight, y + dy * weight)
    return points
