import numpy as np
import pytest

from camera_retarget_l20.l20_mapping import (
    FLEXION_MOTOR_INDICES,
    MOTOR_JOINT_NAMES,
    RESERVED_MOTOR_INDICES,
    L20CommandMapper,
)


JOINT_LIMITS = {
    "thumb_joint0": (-0.297, 0.683),
    "thumb_joint1": (0.122, 1.78),
    "thumb_joint2": (0.0, 0.87),
    "thumb_joint3": (0.0, 1.29),
    "thumb_joint4": (0.0, 1.29),
    "index_joint0": (-0.26, 0.26),
    "index_joint1": (0.0, 1.4),
    "index_joint2": (0.0, 1.08),
    "index_joint3": (0.0, 1.15),
    "middle_joint0": (-0.26, 0.26),
    "middle_joint1": (0.0, 1.4),
    "middle_joint2": (0.0, 1.08),
    "middle_joint3": (0.0, 1.15),
    "ring_joint0": (-0.26, 0.26),
    "ring_joint1": (0.0, 1.4),
    "ring_joint2": (0.0, 1.08),
    "ring_joint3": (0.0, 1.15),
    "little_joint0": (-0.26, 0.26),
    "little_joint1": (0.0, 1.4),
    "little_joint2": (0.0, 1.08),
    "little_joint3": (0.0, 1.15),
}


def make_mapper(**kwargs):
    names = list(JOINT_LIMITS)
    limits = [JOINT_LIMITS[name] for name in names]
    return L20CommandMapper(
        names,
        limits,
        ema_alpha=kwargs.pop("ema_alpha", 1.0),
        max_step=kwargs.pop("max_step", 255.0),
        max_closure=kwargs.pop("max_closure", 1.0),
        **kwargs,
    )


def qpos_at(limit_index):
    return np.array([limits[limit_index] for limits in JOINT_LIMITS.values()])


def test_lower_limits_map_to_open_end_of_protocol_range():
    mapper = make_mapper()
    command = mapper.raw_command(qpos_at(0))
    for motor_index, name in enumerate(MOTOR_JOINT_NAMES):
        assert command[motor_index] == 255.0
        if name is None:
            assert motor_index in RESERVED_MOTOR_INDICES


def test_upper_limits_map_to_zero_except_reserved_slots():
    mapper = make_mapper()
    command = mapper.raw_command(qpos_at(1))
    for motor_index, name in enumerate(MOTOR_JOINT_NAMES):
        expected = 255.0 if name is None else 0.0
        assert command[motor_index] == expected


def test_max_closure_only_limits_flexion_motors():
    mapper = make_mapper(max_closure=0.60)
    command = mapper.raw_command(qpos_at(1))
    for index in FLEXION_MOTOR_INDICES:
        assert command[index] == pytest.approx(102.0)
    for index in (5, 6, 7, 8, 9, 10):
        assert command[index] == 0.0


def test_filter_limits_each_frame_step():
    mapper = make_mapper(ema_alpha=1.0, max_step=7.0)
    command = mapper.map(qpos_at(1))
    for index in FLEXION_MOTOR_INDICES:
        assert command[index] == 248
    for index in RESERVED_MOTOR_INDICES:
        assert command[index] == 255


def test_missing_controlled_joint_is_rejected():
    names = [name for name in JOINT_LIMITS if name != "thumb_joint2"]
    with pytest.raises(ValueError, match="thumb_joint2"):
        L20CommandMapper(names, [JOINT_LIMITS[name] for name in names])