import gradio as gr
import numpy as np
from robots import LiveRobot
import example_robot_data as erd

#ROBOTS = ["erd", "panda", "ur10" , "double_pendulum", "anymal", "talos"]
ROBOTS = [r for r in erd.ROBOTS]


# keep one MeshCat session alive
robot = LiveRobot("talos")

# how many sliders we allow at most (raise if you have bigger robots)
MAX_SLIDERS = 32

def _visible_dofs():
    """Number of actuated DoFs for current robot."""
    return len(robot.idxs)

def _slider_updates_for_current_robot():
    """Return a list of gr.update() objects of length MAX_SLIDERS
    to (re)configure all sliders for the current robot."""
    updates = []
    k = _visible_dofs()
    for i in range(MAX_SLIDERS):
        if i < k:
            label = robot.labels[i]
            lo, hi = robot.bounds[i]
            if lo==hi:
                lo = -np.pi/2
                hi = np.pi/2
            qi = robot.idxs[i]
            val = float((lo+hi)/2)
            updates.append(
                gr.update(
                    minimum=float(lo),
                    maximum=float(hi),
                    value=val,
                    label=label,
                    step=0.01,
                    visible=True,
                )
            )
        else:
            updates.append(gr.update(visible=False))
    return updates

with gr.Blocks(title="Pinocchio + MeshCat — Live") as demo:
    gr.Markdown("### 🤖 Pinocchio + MeshCat — Live Viewer\n"
                "Sliders update the robot in real time.\n"
                "Select a robot from the dropdown to load it.")

    with gr.Row():
        robot_dd = gr.Dropdown(ROBOTS, value="talos", label="Robot")
        neutral_btn = gr.Button("Neutral")

    viewer = gr.HTML(robot.iframe())

    # Pre-create a fixed pool of sliders (initially hidden)
    sliders = []
    with gr.Group():
        for _ in range(MAX_SLIDERS):
            sliders.append(
                gr.Slider(
                    minimum=-3.14, maximum=3.14, value=0.0,
                    step=0.01, label="joint", visible=False
                )
            )

    # ---------- Callbacks ----------

    def on_sliders(*vals):
        """Apply first k slider values to robot joints."""
        k = _visible_dofs()
        # only use the first k values (ignore hidden/unused sliders)
        robot.set_joints(vals[:k])
        # viewer iframe persists; nothing to update there
        #return gr.update()

    # Wire all sliders to the same handler (fixed signature via *vals)
    for s in sliders:
        s.change(on_sliders, inputs=sliders, outputs=None, queue=True)

    def on_neutral():
        """Reset robot to neutral and reflect in visible sliders."""
        values = robot.neutral()  # returns values for actuated DoFs, length k
        bounds = robot.bounds
        
        for i, v in enumerate(values):
            lo, hi = bounds[i]
            if not np.isfinite(lo): lo = -3.14
            if not np.isfinite(hi): hi =  3.14
            
            # clamp neutral value to slider range
            if v < lo: v = lo
            if v > hi: v = hi
            
            values[i] = v
        k = len(values)
        # viewer stays the same + slider value updates for first k; hide the rest
        updates = [gr.update()]  # for viewer
        for i in range(MAX_SLIDERS):
            if i < k:
                updates.append(gr.update(value=float(values[i]), visible=True))
            else:
                updates.append(gr.update(visible=False))
        return updates

    neutral_btn.click(on_neutral, outputs=[viewer] + sliders, queue=False)

    def on_change_robot(name):
        """Rebuild slider labels/ranges/values/visibility to match selected robot."""
        global robot
        robot = LiveRobot(name)
        viewer_html = robot.iframe()
        # 1) update viewer HTML
        # 2) update ALL sliders to reflect new robot
        return [viewer_html] + _slider_updates_for_current_robot()

    robot_dd.change(on_change_robot, inputs=robot_dd, outputs=[viewer] + sliders, queue=True)

    # On app load, configure sliders for the initial robot
    def _initial():
        return [robot.iframe()] + _slider_updates_for_current_robot()

    demo.load(_initial, outputs=[viewer] + sliders, queue=False)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=8501)