import numpy as np
import pytest
from llm_panda.task_nullspace import (
    CylinderTaskManifold,
    BoxTaskManifold,
    TaskManifoldFactory
)

def _get_spatial_velocity_fd(manifold, u, du, delta=1e-6):
    """Compute finite difference spatial twist for testing."""
    T_plus = manifold.get_transform(u + du * delta)
    T_minus = manifold.get_transform(u - du * delta)
    
    # spatial velocity [v; w] such that T_plus = T_minus * exp(twist * 2 * delta)
    # roughly, T_dot = (T_plus - T_minus) / (2 * delta)
    # Spatial velocity in base frame:
    # v = p_dot
    # [w]_x = R_dot * R^T
    
    R_plus = T_plus[:3, :3]
    R_minus = T_minus[:3, :3]
    p_plus = T_plus[:3, 3]
    p_minus = T_minus[:3, 3]
    
    v = (p_plus - p_minus) / (2 * delta)
    R_dot = (R_plus - R_minus) / (2 * delta)
    # mid R is approx R_minus
    R_mid = manifold.get_transform(u)[:3, :3]
    w_cross = R_dot @ R_mid.T
    w = np.array([w_cross[2, 1], w_cross[0, 2], w_cross[1, 0]])
    
    return np.concatenate([v, w])

def test_cylinder_manifold_jacobian():
    # Random object pose
    T_obj = np.eye(4)
    T_obj[:3, 3] = [0.1, 0.2, 0.3]
    # Rotate 90 deg around X
    T_obj[:3, :3] = np.array([
        [1, 0, 0],
        [0, 0, -1],
        [0, 1, 0]
    ])
    
    T_offset = np.eye(4)
    T_offset[:3, 3] = [0, 0, 0.05]
    
    manifold = TaskManifoldFactory.create_manifold("cylinder", T_obj, {"height": 0.2}, T_offset)
    
    # Check bounds
    bounds = manifold.get_bounds()
    assert len(bounds) == 2
    assert bounds[0] == (-0.1, 0.1)
    
    u = np.array([0.05, np.pi/4])
    J_analytical = manifold.get_transform_jacobian(u)
    
    # Finite difference for z
    v_z_fd = _get_spatial_velocity_fd(manifold, u, np.array([1.0, 0.0]))
    # Finite difference for theta
    v_theta_fd = _get_spatial_velocity_fd(manifold, u, np.array([0.0, 1.0]))
    
    J_fd = np.column_stack([v_z_fd, v_theta_fd])
    
    np.testing.assert_allclose(J_analytical, J_fd, atol=1e-5)

def test_box_manifold_jacobian():
    T_face = np.eye(4)
    T_face[:3, 3] = [0.5, -0.2, 0.1]
    
    T_offset = np.eye(4)
    T_offset[:3, 3] = [0, 0.02, 0.05]
    
    manifold = TaskManifoldFactory.create_manifold("box", T_face, {"width": 0.1, "height": 0.2}, T_offset)
    
    u = np.array([0.02, -0.03, np.pi/3])
    J_analytical = manifold.get_transform_jacobian(u)
    
    J_u_fd = _get_spatial_velocity_fd(manifold, u, np.array([1.0, 0.0, 0.0]))
    J_v_fd = _get_spatial_velocity_fd(manifold, u, np.array([0.0, 1.0, 0.0]))
    J_theta_fd = _get_spatial_velocity_fd(manifold, u, np.array([0.0, 0.0, 1.0]))
    
    J_fd = np.column_stack([J_u_fd, J_v_fd, J_theta_fd])
    
    np.testing.assert_allclose(J_analytical, J_fd, atol=1e-5)