"""World-space joint -> 2D pixel projection for a given camera + resolution.

Pure math, no drawing. Shared by the viewport draw override (which draws in
3D world space directly and does not strictly need this) and the render
override (which needs actual 2D pixel coordinates to rasterize the
OpenPose-style image). Kept independent of "which camera" / "which
character" decisions -- callers supply an explicit camera path and mapping.
"""

import maya.api.OpenMaya as om

from . import constants

DEFAULT_RESOLUTION_NODE = "defaultResolution"


def get_render_resolution():
    """Return (width, height) in pixels from the scene's defaultResolution node."""
    sel = om.MSelectionList()
    sel.add(DEFAULT_RESOLUTION_NODE)
    fn = om.MFnDependencyNode(sel.getDependNode(0))
    width = fn.findPlug("width", False).asInt()
    height = fn.findPlug("height", False).asInt()
    return max(1, width), max(1, height)


def world_position(dag_path):
    """World-space MPoint for any DAG transform (including joints)."""
    world_matrix = dag_path.inclusiveMatrix()
    return om.MPoint(om.MTransformationMatrix(world_matrix).translation(om.MSpace.kWorld))


def project_point(world_point, camera_path, width, height):
    """Project a world-space MPoint through ``camera_path`` to pixel coordinates.

    Returns (pixel_x, pixel_y, in_front_of_camera).
    """
    camera_space = om.MPoint(world_point) * camera_path.inclusiveMatrixInverse()

    camera_fn = om.MFnCamera(camera_path)
    # projectionMatrix() returns an MFloatMatrix; MPoint's matrix multiply
    # operator only supports MMatrix, hence the explicit conversion.
    projection_matrix = om.MMatrix(camera_fn.projectionMatrix())
    clip = om.MPoint(camera_space) * projection_matrix
    if abs(clip.w) < 1e-9:
        return 0.0, 0.0, False

    ndc_x = clip.x / clip.w
    ndc_y = clip.y / clip.w
    pixel_x = (ndc_x * 0.5 + 0.5) * width
    pixel_y = (1.0 - (ndc_y * 0.5 + 0.5)) * height
    # The camera looks down its own -Z axis, so a point in front of it has
    # negative camera-space Z.
    in_front = camera_space.z < 0.0
    return pixel_x, pixel_y, in_front


def local_axes_world(dag_path):
    """World-space (right, up) unit vectors from a DAG node's local X/Y axes."""
    world_matrix = dag_path.inclusiveMatrix()
    right = (om.MVector(1.0, 0.0, 0.0) * world_matrix).normal()
    up = (om.MVector(0.0, 1.0, 0.0) * world_matrix).normal()
    return right, up


def plane_world_points(anchor_world_point, right, up, local_points_xy, scale):
    """Place 2D local (x, y) points on a plane centered at ``anchor_world_point``
    using the given world-space ``right``/``up`` axes, scaled by ``scale``.
    """
    anchor = om.MPoint(anchor_world_point)
    return [
        om.MPoint(anchor + right * (x * scale) + up * (y * scale))
        for x, y in local_points_xy
    ]


def billboard_world_points(anchor_world_point, camera_path, local_points_xy, scale):
    """``plane_world_points`` using ``camera_path``'s own local X/Y as
    right/up, so the plane always faces whichever camera is drawing.

    Used to draw/render the procedural ARKit-blendshape face without needing
    any facial joint orientation -- it billboards toward whichever camera is
    currently drawing (the active viewport camera for the live overlay, the
    render camera for output), sidestepping the rig-specific, unanswerable
    question of "which way does this head joint's local frame face". This is
    the "Camera" faceFollow mode.
    """
    right, up = local_axes_world(camera_path)
    return plane_world_points(anchor_world_point, right, up, local_points_xy, scale)


def oriented_world_points(anchor_world_point, orientation_dag_path, local_points_xy, scale):
    """``plane_world_points`` using ``orientation_dag_path``'s own local X/Y
    as right/up, so the plane rotates *with* that joint (e.g. the mapped
    Neck) instead of always facing the camera. This is the "Neck" faceFollow
    mode -- unlike "Camera", it is rig-dependent: it assumes the joint's
    local X is the character's right and local Y is up, which matches
    common convention but is not guaranteed for every rig. If the face
    appears rotated relative to the head, adjust the neck joint's rotate
    axis/orient, or use "Camera" mode instead.
    """
    right, up = local_axes_world(orientation_dag_path)
    return plane_world_points(anchor_world_point, right, up, local_points_xy, scale)


def default_head_relative_points(anchor_world_point, orientation_dag_path, scale, local_offsets):
    """Compute head-relative fallback world points (e.g. the default
    REye/LEye/REar/LEar positions in constants.DEFAULT_EYE_EAR_LOCAL_OFFSETS)
    anchored at ``anchor_world_point`` (the mapped Nose), oriented by
    ``orientation_dag_path``'s own local X/Y axes -- always the joint
    mapped to Nose itself (the head), so these points turn with the head's
    actual rotation rather than billboarding toward the camera.

    ``local_offsets``: {index: (x, y)}. Returns {index: MPoint}.
    """
    indices = list(local_offsets.keys())
    offsets = [local_offsets[i] for i in indices]
    world_points = oriented_world_points(anchor_world_point, orientation_dag_path, offsets, scale)
    return dict(zip(indices, world_points))


def project_character(mapped_joint_paths, camera_path, width, height):
    """Project one character's mapping to OpenPose_full pixel keypoints.

    ``mapped_joint_paths``: {keypoint_index: MDagPath}, as returned by
    ``mapping_node.get_mapped_joint_paths``.

    Returns a list of length ``constants.TOTAL_KEYPOINTS`` of
    (pixel_x, pixel_y, confidence) tuples. Unmapped or off-camera keypoints
    get confidence 0.0 at (0.0, 0.0), matching OpenPose's own convention for
    undetected keypoints.
    """
    keypoints = [(0.0, 0.0, 0.0)] * constants.TOTAL_KEYPOINTS
    for index, joint_path in mapped_joint_paths.items():
        world_point = world_position(joint_path)
        pixel_x, pixel_y, in_front = project_point(world_point, camera_path, width, height)
        confidence = 1.0 if in_front else 0.0
        keypoints[index] = (pixel_x, pixel_y, confidence)
    return keypoints