"""Hardware-style seed recovery: joint-limit violations and self-collision.

Documents current EmbodiK behavior when ``q`` starts slightly out of limits or in
penetration/shallow contact. See ``docs/recovery_robustness.md``.
"""

from __future__ import annotations

import os
import re
import sys
import tempfile
import textwrap
from dataclasses import dataclass, field
from pathlib import Path
import numpy as np
import pytest

import embodik as eik

_DQ_STALL_EPS = 1e-5


@dataclass
class RecoveryLoopMetrics:
    """Lightweight metrics for velocity-integration recovery loops."""

    steps: int = 0
    infeasible_count: int = 0
    numerical_error_count: int = 0
    success_count: int = 0
    zero_velocity_steps: int = 0
    max_zero_velocity_streak: int = 0
    collision_distances: list[float] = field(default_factory=list)

    def record_step(
        self,
        status: eik.SolverStatus,
        dq: np.ndarray,
        collision_distance: float | None,
    ) -> None:
        self.steps += 1
        if status == eik.SolverStatus.INFEASIBLE:
            self.infeasible_count += 1
        elif status == eik.SolverStatus.NUMERICAL_ERROR:
            self.numerical_error_count += 1
        elif status == eik.SolverStatus.SUCCESS:
            self.success_count += 1
        dq_norm = float(np.linalg.norm(np.asarray(dq, dtype=float).ravel()))
        if dq_norm < _DQ_STALL_EPS:
            self.zero_velocity_steps += 1
            self.max_zero_velocity_streak = max(
                self.max_zero_velocity_streak, self.zero_velocity_steps
            )
        else:
            self.zero_velocity_steps = 0
        if collision_distance is not None and np.isfinite(collision_distance):
            self.collision_distances.append(float(collision_distance))


def _one_dof_limited_urdf() -> str:
    return textwrap.dedent("""<?xml version="1.0"?>
<robot name="hw_seed_1dof">
  <link name="base_link">
    <inertial><mass value="1.0"/><origin xyz="0 0 0"/>
    <inertia ixx="1" ixy="0" ixz="0" iyy="1" iyz="0" izz="1"/></inertial>
  </link>
  <link name="link1">
    <inertial><mass value="1.0"/><origin xyz="0 0 0.5"/>
    <inertia ixx="1" ixy="0" ixz="0" iyy="1" iyz="1" izz="1"/></inertial>
  </link>
  <joint name="joint1" type="revolute">
    <parent link="base_link"/><child link="link1"/>
    <origin xyz="0 0 1" rpy="0 0 0"/><axis xyz="0 0 1"/>
    <limit lower="0.0" upper="1.0" velocity="1.0" effort="10.0"/>
  </joint>
</robot>""")


def _two_link_collision_urdf() -> str:
    return textwrap.dedent("""
<robot name="two_link">
  <link name="base_link">
    <inertial><origin xyz="0 0 0" rpy="0 0 0"/><mass value="1.0"/>
    <inertia ixx="0.01" ixy="0.0" ixz="0.0" iyy="0.01" iyz="0.0" izz="0.01"/></inertial>
    <collision><origin xyz="0 0 0" rpy="0 0 0"/><geometry><box size="0.1 0.1 0.1"/></geometry></collision>
  </link>
  <link name="link1">
    <inertial><origin xyz="0 0 0" rpy="0 0 0"/><mass value="0.5"/>
    <inertia ixx="0.005" ixy="0.0" ixz="0.0" iyy="0.005" iyz="0.0" izz="0.005"/></inertial>
    <collision><origin xyz="0.08 0 0" rpy="0 0 0"/><geometry><box size="0.08 0.08 0.08"/></geometry></collision>
  </link>
  <joint name="joint1" type="revolute">
    <parent link="base_link"/><child link="link1"/>
    <origin xyz="0.05 0 0" rpy="0 0 0"/><axis xyz="0 0 1"/>
    <limit effort="10.0" lower="-1.57" upper="1.57" velocity="1.0"/>
  </joint>
</robot>""")


def _write_urdf(content: str) -> str:
    fd, path = tempfile.mkstemp(suffix=".urdf")
    with os.fdopen(fd, "w") as f:
        f.write(content)
    return path


@pytest.fixture
def one_dof_path():
    p = _write_urdf(_one_dof_limited_urdf())
    yield p
    os.remove(p)


@pytest.fixture
def two_link_path():
    p = _write_urdf(_two_link_collision_urdf())
    yield p
    os.remove(p)


_PANDA_DEFAULT_Q = np.array(
    [0.0, -0.785, 0.0, -2.356, 0.0, 1.571, 0.785], dtype=float
)
_PANDA_GRIPPER_Q = np.array([0.02, 0.02], dtype=float)
_PANDA_EE_FRAME = "panda_hand"

_LINK_INDEX_PATTERN = re.compile(r"link_?([0-9]+)")


def _extract_link_index(name: str) -> int | None:
    match = _LINK_INDEX_PATTERN.search(name)
    return int(match.group(1)) if match else None


def _dual_iiwa_exclusions(robot: eik.RobotModel) -> list[tuple[str, str]]:
    exclusions: list[tuple[str, str]] = []
    for a, b in robot.get_collision_pair_names():
        a_l, b_l = a.lower(), b.lower()
        same_arm = (("left" in a_l and "left" in b_l) or
                    ("right" in a_l and "right" in b_l))
        if same_arm:
            ia, ib = _extract_link_index(a_l), _extract_link_index(b_l)
            if ia is not None and ib is not None and abs(ia - ib) <= 3:
                exclusions.append((a, b))
        else:
            ia, ib = _extract_link_index(a_l), _extract_link_index(b_l)
            if ia is not None and ib is not None and ia <= 1 and ib <= 1:
                exclusions.append((a, b))
    return exclusions


def _setup_dual_iiwa_stall_case():
    examples_dir = Path(__file__).resolve().parent.parent / "examples"
    if str(examples_dir) not in sys.path:
        sys.path.insert(0, str(examples_dir))

    try:
        from utils.dual_iiwa_urdf import (
            build_dual_iiwa_urdf,
            get_dual_iiwa_default_configuration,
            get_dual_iiwa_frame_names,
        )
    except ImportError:
        return None

    urdf_str = build_dual_iiwa_urdf()
    with tempfile.NamedTemporaryFile(suffix=".urdf", mode="w", delete=False) as f:
        f.write(urdf_str)
        urdf_path = f.name

    try:
        robot = eik.RobotModel(urdf_path, floating_base=False)
    except Exception:
        return None
    finally:
        os.unlink(urdf_path)

    exclusions = _dual_iiwa_exclusions(robot)
    if exclusions:
        try:
            robot.apply_collision_exclusions(exclusions)
        except Exception:
            pass

    solver = eik.KinematicsSolver(robot)
    solver.dt = 0.01
    solver.enable_position_limits(True)
    solver.set_damping(0.1)

    q = np.array(get_dual_iiwa_default_configuration(), dtype=float)
    robot.update_configuration(q)
    dbg = solver.evaluate_collision_debug(q)
    if dbg is None or not np.isfinite(dbg.distance):
        return None

    min_dist = 0.05
    solver.configure_collision_constraint(
        min_distance=min_dist, include_pairs=[], exclude_pairs=list(exclusions)
    )

    left_frame, right_frame = get_dual_iiwa_frame_names()
    solver.clear_tasks()
    left_task = solver.add_frame_task("left_stall", left_frame)
    left_task.priority = 0
    left_task.weight = 1.0
    right_task = solver.add_frame_task("right_stall", right_frame)
    right_task.priority = 0
    right_task.weight = 1.0

    left_pos = np.array(robot.get_frame_pose(left_frame).translation)
    left_rot = np.array(robot.get_frame_pose(left_frame).rotation)
    right_pos = np.array(robot.get_frame_pose(right_frame).translation)
    right_rot = np.array(robot.get_frame_pose(right_frame).rotation)
    left_target = left_pos + np.array([0.15, -0.35, -0.10])
    right_target = right_pos + np.array([0.15, 0.35, -0.10])
    left_task.set_target_pose(left_target, left_rot)
    right_task.set_target_pose(right_target, right_rot)

    return robot, solver, q, left_frame, left_target, min_dist


def test_solve_velocity_recovers_from_slight_lower_limit_violation(one_dof_path):
    """Encoder slightly below ``q_min``: velocity loop should re-enter ``[0,1]``."""
    robot = eik.RobotModel(one_dof_path, floating_base=False)
    solver = eik.KinematicsSolver(robot)
    solver.dt = 0.01
    solver.enable_position_limits(True)
    solver.set_limit_recovery_gain(0.5)
    solver.set_damping(0.05)
    solver.clear_tasks()
    jt = solver.add_joint_task("j", "joint1", 0.5)
    jt.priority = 0
    jt.weight = 1.0

    q = np.array([-0.035], dtype=float)
    q_lo, q_hi = robot.get_joint_limits()
    metrics = RecoveryLoopMetrics()
    for _ in range(80):
        robot.update_configuration(q)
        res = solver.solve_velocity(q, apply_limits=True)
        dq = np.asarray(res.joint_velocities, dtype=float).ravel()
        metrics.record_step(res.status, dq, None)
        q = robot.integrate(q, dq, solver.dt)

    assert q[0] >= float(q_lo[0]) - 1e-3
    assert q[0] <= float(q_hi[0]) + 1e-3
    # Early INFEASIBLE with non-zero recovery velocity is allowed (classified outcome).
    assert metrics.infeasible_count + metrics.success_count >= 1


def test_solve_velocity_recovers_from_slight_upper_limit_violation(one_dof_path):
    """Encoder slightly above ``q_max``: should move back toward the interior."""
    robot = eik.RobotModel(one_dof_path, floating_base=False)
    solver = eik.KinematicsSolver(robot)
    solver.dt = 0.01
    solver.enable_position_limits(True)
    solver.set_limit_recovery_gain(0.5)
    solver.set_damping(0.05)
    solver.clear_tasks()
    jt = solver.add_joint_task("j", "joint1", 0.5)
    jt.priority = 0
    jt.weight = 1.0

    q = np.array([1.04], dtype=float)
    q_lo, q_hi = robot.get_joint_limits()
    metrics = RecoveryLoopMetrics()
    for _ in range(80):
        robot.update_configuration(q)
        res = solver.solve_velocity(q, apply_limits=True)
        dq = np.asarray(res.joint_velocities, dtype=float).ravel()
        metrics.record_step(res.status, dq, None)
        q = robot.integrate(q, dq, solver.dt)

    assert q[0] >= float(q_lo[0]) - 1e-3
    assert q[0] <= float(q_hi[0]) + 1e-3
    assert q[0] < 1.04 - 1e-3


def test_solve_position_step_from_limit_violation_moves_toward_feasible_region(one_dof_path):
    """``solve_position_step`` may report INFEASIBLE while still integrating toward limits."""
    robot = eik.RobotModel(one_dof_path, floating_base=False)
    solver = eik.KinematicsSolver(robot)
    solver.dt = 0.05
    solver.enable_position_limits(True)
    solver.set_damping(0.1)
    solver.clear_tasks()
    ft = solver.add_frame_task("ee", "link1")
    ft.weight = 1.0
    ft.priority = 0

    q_seed = np.array([-0.045], dtype=float)
    robot.update_configuration(q_seed)
    pose0 = robot.get_frame_pose("link1")
    T = np.eye(4, dtype=float)
    T[:3, :3] = np.asarray(pose0.rotation, dtype=float)
    T[:3, 3] = np.asarray(pose0.translation, dtype=float) + np.array(
        [0.06, 0.0, 0.0], dtype=float
    )

    opts = eik.PositionStepOptions()
    opts.dt = 0.05
    opts.max_steps = 10

    res = solver.solve_position_step(q_seed, [eik.TaskTarget("ee", T)], opts)
    q_out = np.asarray(res.q_solution, dtype=float).ravel()
    q_lo, q_hi = robot.get_joint_limits()

    assert res.status in (
        eik.SolverStatus.SUCCESS,
        eik.SolverStatus.INFEASIBLE,
        eik.SolverStatus.NUMERICAL_ERROR,
    )
    assert q_out[0] > q_seed[0] + 1e-3
    assert q_out[0] >= float(q_lo[0]) - 2e-3
    assert q_out[0] <= float(q_hi[0]) + 2e-3


@pytest.mark.parametrize("seed_side", ["lower", "upper"])
def test_panda_solve_velocity_recovers_from_slight_limit_violation(seed_side: str):
    """On a real robot model, slight limit violations should recover with SUCCESS."""
    pytest.importorskip("robot_descriptions")
    from robot_descriptions.panda_description import URDF_PATH

    robot = eik.RobotModel(URDF_PATH, floating_base=False)
    solver = eik.KinematicsSolver(robot)
    solver.dt = 0.01
    solver.enable_position_limits(True)
    solver.set_damping(0.1)
    solver.clear_tasks()
    joint_task = solver.add_joint_task("panda_joint1_center", "panda_joint1", 0.0)
    joint_task.priority = 0
    joint_task.weight = 1.0

    q_lo, q_hi = robot.get_joint_limits()
    q = np.concatenate([_PANDA_DEFAULT_Q, _PANDA_GRIPPER_Q])
    limit_value = float(q_lo[0]) - 0.01 if seed_side == "lower" else float(q_hi[0]) + 0.01
    q[0] = limit_value

    for _ in range(8):
        robot.update_configuration(q)
        result = solver.solve_velocity(q, apply_limits=True)
        dq = np.asarray(result.joint_velocities, dtype=float)
        assert result.status == eik.SolverStatus.SUCCESS
        if seed_side == "lower":
            assert dq[0] > 0.0
        else:
            assert dq[0] < 0.0
        q = robot.integrate(q, dq, solver.dt)

    if seed_side == "lower":
        assert q[0] > limit_value + 0.05
    else:
        assert q[0] < limit_value - 0.05


def test_panda_position_step_recovers_from_slight_upper_limit_violation():
    """Real-model validation: slight over-limit seed can return SUCCESS and re-enter limits."""
    pytest.importorskip("robot_descriptions")
    from robot_descriptions.panda_description import URDF_PATH

    robot = eik.RobotModel(URDF_PATH, floating_base=False)
    solver = eik.KinematicsSolver(robot)
    solver.dt = 0.05
    solver.enable_position_limits(True)
    solver.set_damping(0.1)
    solver.set_joint_limit_barrier_task(0.3, 1.0)

    q_lo, q_hi = robot.get_joint_limits()
    q = np.concatenate([_PANDA_DEFAULT_Q, _PANDA_GRIPPER_Q])
    q[0] = float(q_hi[0]) + 0.01
    robot.update_configuration(q)
    pose0 = robot.get_frame_pose(_PANDA_EE_FRAME)

    task = solver.add_frame_task("ee", _PANDA_EE_FRAME)
    task.weight = 1.0
    task.priority = 1

    target_pose = np.eye(4, dtype=float)
    target_pose[:3, :3] = np.asarray(pose0.rotation, dtype=float)
    target_pose[:3, 3] = np.asarray(pose0.translation, dtype=float) + np.array(
        [0.06, 0.0, 0.0], dtype=float
    )

    opts = eik.PositionStepOptions()
    opts.dt = 0.05
    opts.max_steps = 8
    opts.max_linear_speed = 2.0
    opts.max_angular_speed = 2.0

    result = solver.solve_position_step(q, [eik.TaskTarget("ee", target_pose)], opts)
    q_out = np.asarray(result.q_solution, dtype=float)

    assert result.status == eik.SolverStatus.SUCCESS
    assert q_out[0] < q[0] - 5e-3
    assert q_out[0] <= float(q_hi[0]) + 1e-4
    assert np.all(q_out >= q_lo - 1e-4)
    assert np.all(q_out <= q_hi + 1e-4)


def test_penetrating_seed_strict_collision_stalls_without_stall_handler(two_link_path):
    """Baseline: penetration + nominal ``min_distance`` can yield zero motion (no stall)."""
    robot = eik.RobotModel(two_link_path, floating_base=False)
    solver = eik.KinematicsSolver(robot)
    solver.dt = 0.02
    solver.enable_position_limits(True)
    solver.set_damping(0.1)
    solver.configure_collision_constraint(0.03, [], [])
    assert not solver.stall_handler_enabled()
    solver.clear_tasks()
    jt = solver.add_joint_task("j", "joint1", 0.0)
    jt.priority = 0
    jt.weight = 1.0

    q = np.array([-1.57], dtype=float)
    metrics = RecoveryLoopMetrics()
    for _ in range(12):
        robot.update_configuration(q)
        dbg = solver.evaluate_collision_debug(q)
        d = float(dbg.distance) if dbg is not None else float("nan")
        res = solver.solve_velocity(q, apply_limits=True)
        dq = np.asarray(res.joint_velocities, dtype=float).ravel()
        metrics.record_step(res.status, dq, d)
        q = robot.integrate(q, dq, solver.dt)

    assert metrics.max_zero_velocity_streak >= 10
    assert np.isfinite(metrics.collision_distances[0])
    assert metrics.collision_distances[0] < 0.0


@pytest.mark.parametrize("use_stall", [False, True])
def test_penetrating_seed_stall_handler_improves_clearance(two_link_path, use_stall: bool):
    """With stall handler, margin relaxation unlocks separation (benchmark: on vs off)."""
    robot = eik.RobotModel(two_link_path, floating_base=False)
    solver = eik.KinematicsSolver(robot)
    solver.dt = 0.02
    solver.enable_position_limits(True)
    solver.set_damping(0.1)
    solver.configure_collision_constraint(0.03, [], [])
    if use_stall:
        solver.enable_stall_handler(0.03)
        solver.configure_stall_handler(stall_threshold=3, restore_rate=0.2, floor_fraction=0.0)
    solver.clear_tasks()
    jt = solver.add_joint_task("j", "joint1", 0.0)
    jt.priority = 0
    jt.weight = 1.0

    q = np.array([-1.57], dtype=float)
    metrics = RecoveryLoopMetrics()
    for _ in range(35):
        robot.update_configuration(q)
        dbg = solver.evaluate_collision_debug(q)
        d = float(dbg.distance) if dbg is not None else float("nan")
        res = solver.solve_velocity(q, apply_limits=True)
        dq = np.asarray(res.joint_velocities, dtype=float).ravel()
        metrics.record_step(res.status, dq, d)
        q = robot.integrate(q, dq, solver.dt)
    final_dbg = solver.evaluate_collision_debug(q)
    final_d = float(final_dbg.distance) if final_dbg is not None else float("nan")

    if use_stall:
        assert final_d > metrics.collision_distances[0] + 0.002
        assert metrics.success_count >= 5
    else:
        assert metrics.max_zero_velocity_streak >= 8


def test_solve_position_step_stall_recovery_on_penetrating_seed(two_link_path):
    """``PositionStepOptions.stall_recovery`` enables the same margin relaxation path."""
    robot = eik.RobotModel(two_link_path, floating_base=False)
    solver = eik.KinematicsSolver(robot)
    solver.dt = 0.02
    solver.enable_position_limits(True)
    solver.set_damping(0.1)
    solver.configure_collision_constraint(0.03, [], [])
    solver.clear_tasks()
    ft = solver.add_frame_task("ee", "link1")
    ft.priority = 0
    ft.weight = 1.0

    robot.update_configuration(np.array([0.0], dtype=float))
    T_home = np.eye(4, dtype=float)
    ph = robot.get_frame_pose("link1")
    T_home[:3, :3] = np.asarray(ph.rotation, dtype=float)
    T_home[:3, 3] = np.asarray(ph.translation, dtype=float)

    q = np.array([-1.57], dtype=float)
    opts = eik.PositionStepOptions()
    opts.dt = 0.02
    opts.max_steps = 15
    opts.stall_recovery = True
    opts.position_gain = 1.0
    opts.orientation_gain = 1.0

    d0 = float(solver.evaluate_collision_debug(q).distance)
    res = solver.solve_position_step(q, [eik.TaskTarget("ee", T_home)], opts)
    q1 = np.asarray(res.q_solution, dtype=float).ravel()
    d1 = float(solver.evaluate_collision_debug(q1).distance)

    assert res.status in (
        eik.SolverStatus.SUCCESS,
        eik.SolverStatus.INFEASIBLE,
        eik.SolverStatus.NUMERICAL_ERROR,
    )
    assert q1[0] > q[0] + 0.05
    assert d1 > d0 + 1e-3


def test_combined_joint_at_limit_and_penetration_with_stall(two_link_path):
    """At lower limit with penetration: stall + task toward interior reduces penetration."""
    robot = eik.RobotModel(two_link_path, floating_base=False)
    solver = eik.KinematicsSolver(robot)
    solver.dt = 0.02
    solver.enable_position_limits(True)
    solver.set_damping(0.1)
    solver.configure_collision_constraint(0.03, [], [])
    solver.enable_stall_handler(0.03)
    solver.configure_stall_handler(stall_threshold=3, restore_rate=0.2, floor_fraction=0.0)
    solver.clear_tasks()
    jt = solver.add_joint_task("j", "joint1", 0.0)
    jt.priority = 0
    jt.weight = 1.0

    q = np.array([-1.57], dtype=float)
    d0 = float(solver.evaluate_collision_debug(q).distance)
    assert d0 < 0.0

    for _ in range(40):
        robot.update_configuration(q)
        res = solver.solve_velocity(q, apply_limits=True)
        dq = np.asarray(res.joint_velocities, dtype=float).ravel()
        q = robot.integrate(q, dq, solver.dt)

    d1 = float(solver.evaluate_collision_debug(q).distance)
    assert d1 > d0 + 0.005
    assert q[0] > -1.57 + 0.05


def test_dual_iiwa_stall_recovery_improves_status_counts_and_task_progress():
    """Dual iiwa validation: stall recovery reduces freezes and improves target progress."""
    setup = _setup_dual_iiwa_stall_case()
    if setup is None:
        pytest.skip("dual iiwa fixture unavailable")

    def run_case(use_stall: bool) -> dict[str, float]:
        case_setup = _setup_dual_iiwa_stall_case()
        assert case_setup is not None
        robot, solver, q, left_frame, left_target, min_dist = case_setup
        if use_stall:
            solver.enable_stall_handler(min_dist)
            solver.configure_stall_handler(
                stall_threshold=3, restore_rate=0.2, floor_fraction=0.0
            )
        else:
            solver.disable_stall_handler()

        success_count = 0
        infeasible_count = 0
        zero_velocity_steps = 0
        for _ in range(80):
            robot.update_configuration(q)
            result = solver.solve_velocity(q, apply_limits=True)
            dq = np.asarray(result.joint_velocities, dtype=float)
            if np.linalg.norm(dq) < _DQ_STALL_EPS:
                zero_velocity_steps += 1
            if result.status == eik.SolverStatus.SUCCESS:
                success_count += 1
            elif result.status == eik.SolverStatus.INFEASIBLE:
                infeasible_count += 1
            q = robot.integrate(q, dq, solver.dt)
            q_lo, q_hi = robot.get_joint_limits()
            q = np.clip(q, q_lo, q_hi)

        robot.update_configuration(q)
        left_pos = np.array(robot.get_frame_pose(left_frame).translation)
        left_err = float(np.linalg.norm(left_target - left_pos))
        return {
            "success_count": float(success_count),
            "infeasible_count": float(infeasible_count),
            "zero_velocity_steps": float(zero_velocity_steps),
            "left_err": left_err,
            "collision_distance": float(solver.evaluate_collision_debug(q).distance),
            "current_min_distance": (
                float(solver.stall_handler_current_min_distance())
                if use_stall else min_dist
            ),
            "nominal_min_distance": float(min_dist),
        }

    baseline = run_case(use_stall=False)
    recovery = run_case(use_stall=True)

    assert recovery["success_count"] > baseline["success_count"]
    assert recovery["infeasible_count"] < baseline["infeasible_count"]
    assert recovery["zero_velocity_steps"] < baseline["zero_velocity_steps"]
    assert recovery["left_err"] < baseline["left_err"]
    assert recovery["current_min_distance"] < recovery["nominal_min_distance"]
