import numpy as np

from pose_tracking_exp.common.joints import BODY20_INDEX_BY_NAME
from pose_tracking_exp.tracking.kinematics import seed_state_from_pose3d


def _sample_pose3d() -> np.ndarray:
    pose = np.zeros((20, 4), dtype=np.float64)
    joint_positions = {
        "hip_middle": [0.0, 1.0, 3.0],
        "hip_left": [0.12, 1.0, 3.0],
        "hip_right": [-0.12, 1.0, 3.0],
        "shoulder_middle": [0.0, 1.52, 3.0],
        "shoulder_left": [0.18, 1.52, 3.0],
        "shoulder_right": [-0.18, 1.52, 3.0],
        "elbow_left": [0.42, 1.48, 3.02],
        "elbow_right": [-0.42, 1.48, 3.02],
        "wrist_left": [0.64, 1.45, 3.04],
        "wrist_right": [-0.64, 1.45, 3.04],
        "knee_left": [0.1, 0.58, 3.0],
        "knee_right": [-0.1, 0.58, 3.0],
        "ankle_left": [0.1, 0.15, 3.02],
        "ankle_right": [-0.1, 0.15, 3.02],
        "head": [0.0, 1.82, 3.02],
        "nose": [0.0, 1.8, 3.06],
        "eye_left": [0.03, 1.81, 3.05],
        "eye_right": [-0.03, 1.81, 3.05],
        "ear_left": [0.06, 1.81, 3.02],
        "ear_right": [-0.06, 1.81, 3.02],
    }
    for name, position in joint_positions.items():
        pose[BODY20_INDEX_BY_NAME[name], :3] = position
        pose[BODY20_INDEX_BY_NAME[name], 3] = 1.0
    return pose


def test_seed_state_from_pose3d_does_not_call_least_squares(monkeypatch) -> None:
    def fail_least_squares(*args: object, **kwargs: object) -> object:
        raise AssertionError("seed_state_from_pose3d should not call scipy.optimize.least_squares")

    monkeypatch.setattr("pose_tracking_exp.tracking.kinematics.least_squares", fail_least_squares)
    state = seed_state_from_pose3d(_sample_pose3d())

    assert state.parameters.shape == (31,)
    assert state.beta.shape == (8,)
    np.testing.assert_allclose(
        state.parameters[:3],
        _sample_pose3d()[BODY20_INDEX_BY_NAME["hip_middle"], :3],
    )