"""Tests for gripper detection in array_intelligence.py.

Validates that the bimodality threshold of 0.70 correctly distinguishes
real gripper dimensions (bimodal, few transitions) from joint angles
that happen to be slightly bimodal (many transitions).
"""

import numpy as np
import pytest

from orbit.analyzer.array_intelligence import profile_array


def _make_gripper_episodes(n_episodes: int = 20, length: int = 200) -> list[np.ndarray]:
    """Create synthetic episodes with 7 joint dims + 1 gripper dim.

    The gripper dim has binary open/close with 1-3 transitions per episode.
    """
    rng = np.random.RandomState(42)
    episodes = []
    for _ in range(n_episodes):
        # 7 joint angles: smooth sinusoidal motion
        t = np.linspace(0, 2 * np.pi, length)
        joints = np.column_stack([
            np.sin(t + rng.uniform(0, np.pi)) * rng.uniform(0.5, 1.5)
            for _ in range(7)
        ])
        # 1 gripper: binary (0 or 1) with 1-3 transitions
        gripper = np.zeros(length)
        n_transitions = rng.randint(1, 4)
        transition_points = sorted(rng.choice(range(20, length - 20), n_transitions, replace=False))
        state = 0.0
        for tp in transition_points:
            state = 1.0 - state
            gripper[tp:] = state
        actions = np.column_stack([joints, gripper.reshape(-1, 1)])
        episodes.append(actions)
    return episodes


def _make_bimodal_joint_episodes(n_episodes: int = 20, length: int = 200) -> list[np.ndarray]:
    """Create episodes where a joint is slightly bimodal but NOT a gripper.

    The joint oscillates between two regions with many transitions — this is
    a joint that moves between start and end positions, not a discrete gripper.
    """
    rng = np.random.RandomState(42)
    episodes = []
    for _ in range(n_episodes):
        t = np.linspace(0, 2 * np.pi, length)
        joints = np.column_stack([
            np.sin(t + rng.uniform(0, np.pi)) * rng.uniform(0.5, 1.5)
            for _ in range(7)
        ])
        # 8th dim: bimodal but many transitions (oscillating, not discrete)
        bimodal_joint = np.sin(t * 5) * 0.8  # many zero crossings
        bimodal_joint = np.where(bimodal_joint > 0, 1.2, -0.8)  # make bimodal
        bimodal_joint += rng.normal(0, 0.05, length)  # add noise
        actions = np.column_stack([joints, bimodal_joint.reshape(-1, 1)])
        episodes.append(actions)
    return episodes


def test_gripper_detected():
    """A real gripper (bimodal, few transitions) should be detected."""
    episodes = _make_gripper_episodes()
    profile = profile_array(episodes)
    assert 7 in profile.gripper_dimensions, (
        f"Expected dim 7 to be detected as gripper, got gripper_dims={profile.gripper_dimensions}"
    )


def test_bimodal_joint_not_gripper():
    """A bimodal joint with many transitions should NOT be flagged as gripper."""
    episodes = _make_bimodal_joint_episodes()
    profile = profile_array(episodes)
    assert 7 not in profile.gripper_dimensions, (
        f"Dim 7 is a bimodal joint, not a gripper, but was detected as gripper. "
        f"gripper_dims={profile.gripper_dimensions}"
    )


def test_dead_dimension_not_gripper():
    """A constant (dead) dimension should not be flagged as gripper."""
    rng = np.random.RandomState(42)
    episodes = []
    for _ in range(10):
        data = rng.randn(100, 4)
        data[:, 3] = 0.0  # dead dim
        episodes.append(data)
    profile = profile_array(episodes)
    assert 3 not in profile.gripper_dimensions
    assert 3 in profile.dead_dimensions
