from __future__ import annotations

import numpy as np

from eval.dual_push_retarget_utils import (
    backproject_depth_to_world,
    delta_action_from_absolute_targets,
    estimate_dual_push_target_positions_from_front_rgbd,
    extract_dual_push_button_positions,
    make_bimanual_absolute_action,
    parse_dual_push_target_colors,
    retarget_demo_pose_to_live_button,
)


def test_extract_dual_push_button_positions_reads_expected_slices() -> None:
    state = np.arange(136, dtype=np.float32)

    button0, button1, button2 = extract_dual_push_button_positions(state)

    assert button0.tolist() == [7.0, 8.0, 9.0]
    assert button1.tolist() == [14.0, 15.0, 16.0]
    assert button2.tolist() == [21.0, 22.0, 23.0]


def test_retarget_demo_pose_to_live_button_translates_position_only() -> None:
    demo_pose = np.array([0.4, 0.5, 0.6, 0.0, 0.0, 0.0, 1.0], dtype=np.float32)
    demo_button = np.array([0.1, 0.2, 0.3], dtype=np.float32)
    live_button = np.array([0.3, -0.1, 0.7], dtype=np.float32)

    retargeted = retarget_demo_pose_to_live_button(demo_pose, demo_button, live_button)

    np.testing.assert_allclose(retargeted[:3], np.array([0.6, 0.2, 1.0], dtype=np.float32))
    np.testing.assert_allclose(retargeted[3:], demo_pose[3:])


def test_make_bimanual_absolute_action_packs_expected_layout() -> None:
    right_pose = np.array([1, 2, 3, 0, 0, 0, 1], dtype=np.float32)
    left_pose = np.array([4, 5, 6, 0, 0, 1, 0], dtype=np.float32)

    action = make_bimanual_absolute_action(right_pose, left_pose, 1.0, 0.0)

    assert action.shape == (18,)
    np.testing.assert_allclose(action[:7], right_pose)
    np.testing.assert_allclose(action[9:16], left_pose)
    assert action[7] == 1.0
    assert action[8] == 1.0
    assert action[16] == 0.0
    assert action[17] == 1.0


def test_delta_action_from_absolute_targets_matches_translation_and_gripper() -> None:
    current_right = np.array([0.1, 0.2, 0.3, 0.0, 0.0, 0.0, 1.0], dtype=np.float32)
    current_left = np.array([-0.1, 0.4, 0.2, 0.0, 0.0, 0.0, 1.0], dtype=np.float32)
    target_right = np.array([0.15, 0.1, 0.35, 0.0, 0.0, 0.0, 1.0], dtype=np.float32)
    target_left = np.array([-0.05, 0.5, 0.25, 0.0, 0.0, 0.0, 1.0], dtype=np.float32)

    delta = delta_action_from_absolute_targets(
        current_right,
        current_left,
        target_right,
        target_left,
        right_gripper_open=1.0,
        left_gripper_open=0.0,
    )

    np.testing.assert_allclose(delta[:3], np.array([0.05, -0.1, 0.05], dtype=np.float32), atol=1e-6)
    np.testing.assert_allclose(delta[7:10], np.array([0.05, 0.1, 0.05], dtype=np.float32), atol=1e-6)
    np.testing.assert_allclose(delta[3:6], np.zeros(3, dtype=np.float32), atol=1e-6)
    np.testing.assert_allclose(delta[10:13], np.zeros(3, dtype=np.float32), atol=1e-6)
    assert delta[6] == 1.0
    assert delta[13] == 0.0


def test_delta_action_from_absolute_targets_captures_rotation_delta() -> None:
    current_right = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0], dtype=np.float32)
    current_left = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0], dtype=np.float32)
    half_turn_z_xyzw = np.array([0.0, 0.0, np.sin(np.pi / 4.0), np.cos(np.pi / 4.0)], dtype=np.float32)
    target_right = np.concatenate([np.zeros(3, dtype=np.float32), half_turn_z_xyzw], axis=0)
    target_left = current_left.copy()

    delta = delta_action_from_absolute_targets(
        current_right,
        current_left,
        target_right,
        target_left,
        right_gripper_open=1.0,
        left_gripper_open=1.0,
    )

    np.testing.assert_allclose(delta[:3], np.zeros(3, dtype=np.float32), atol=1e-6)
    np.testing.assert_allclose(delta[7:13], np.zeros(6, dtype=np.float32), atol=1e-6)
    assert abs(float(delta[5])) > 1.4


def test_parse_dual_push_target_colors_reads_both_targets() -> None:
    assert parse_dual_push_target_colors("push the olive and the orange buttons") == ("olive", "orange")


def test_backproject_depth_to_world_handles_signed_focal_lengths() -> None:
    depth = np.array([[0.75]], dtype=np.float32)
    intrinsics = np.array([[-100.0, 0.0, 0.0], [0.0, -100.0, 0.0], [0.0, 0.0, 1.0]], dtype=np.float32)
    extrinsics = np.eye(4, dtype=np.float32)

    points = backproject_depth_to_world(depth, intrinsics, extrinsics)

    np.testing.assert_allclose(points[0, 0], np.array([0.0, 0.0, 0.75], dtype=np.float32))


def test_estimate_dual_push_target_positions_from_front_rgbd_finds_color_centers() -> None:
    rgb = np.zeros((3, 8, 8), dtype=np.float32)
    depth = np.full((8, 8), 0.75, dtype=np.float32)
    rgb[:, 2:4, 1:3] = np.array([128.0, 128.0, 0.0], dtype=np.float32)[:, None, None] / 255.0
    rgb[:, 2:4, 5:7] = np.array([255.0, 128.0, 0.0], dtype=np.float32)[:, None, None] / 255.0
    intrinsics = np.array([[-100.0, 0.0, 3.5], [0.0, -100.0, 3.5], [0.0, 0.0, 1.0]], dtype=np.float32)
    extrinsics = np.eye(4, dtype=np.float32)

    right, left = estimate_dual_push_target_positions_from_front_rgbd(
        rgb,
        depth,
        intrinsics,
        extrinsics,
        "push the olive and the orange buttons",
    )

    assert right is not None
    assert left is not None
    assert right[0] > left[0]
    np.testing.assert_allclose(right[2], 0.75, atol=1e-4)
    np.testing.assert_allclose(left[2], 0.75, atol=1e-4)