#!/usr/bin/env python3
"""Compatibility adapter for the original MuJoCo controller import path."""

from time import sleep

import mujoco.viewer
import numpy as np

from rm75_ik.mujoco_backend import DualArmMuJoCo


class MuJoCoPositionController:
    """Legacy facade backed by the threadless normalized dual-arm scene."""

    def __init__(
        self,
        urdf_path=None,
        smoothness=0.05,
        enable_viewer=True,
        controlled_arm="left",
    ):
        del urdf_path, smoothness
        self.backend = DualArmMuJoCo(controlled_arm=controlled_arm)
        self.controlled_arm = controlled_arm
        self.viewer = (
            mujoco.viewer.launch_passive(self.backend.model, self.backend.data)
            if enable_viewer
            else None
        )

    def start(self):
        return None

    def stop(self):
        if self.viewer is not None:
            self.viewer.close()
            sleep(0.2)
            self.viewer = None

    def send_command(self, joint_positions):
        self.backend.set_arm_configuration(
            self.controlled_arm, np.asarray(joint_positions, dtype=float)
        )
        if self.viewer is not None:
            self.viewer.sync()

    def get_feedback(self):
        return self.backend.get_arm_configuration(self.controlled_arm)

    def get_target(self):
        return self.get_feedback()

    def move_to_joints(self, target, duration=1.0):
        start = self.get_feedback()
        target_q = np.asarray(target, dtype=float)
        points = max(2, int(round(duration * 90.0)))
        blend = 0.5 - 0.5 * np.cos(np.linspace(0.0, np.pi, points))
        trajectory = start[None, :] + blend[:, None] * (target_q - start)[None, :]
        self.backend.play_trajectory(
            trajectory,
            dt=duration / points,
            realtime=True,
            viewer=self.viewer,
        )

    def wait_until_reached(self, tolerance=0.01, timeout=10.0):
        del tolerance, timeout
        return True

    def print_state(self):
        print("Current joints (rad):", self.get_feedback().tolist())


def demo_position_control():
    from rm75_ik.stage2_demo import main

    return main([])


if __name__ == "__main__":
    raise SystemExit(demo_position_control())