# Copyright 2026 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""This script runs a simulation and viewer in separate processes communicating asynchronously.

In this example, we will run the viewer in an independent process communicating
via multiprocessing queues. A slider is provided to adjust the state send rate.

You must provide a mjcf model file via the first command-line argument.
"""

import multiprocessing
import os
import sys
import time

from absl import app as absl_app
from absl import flags as absl_flags
import mujoco
from mujoco.experimental.studio import native_viewer as _viewer
from mujoco.experimental.studio import parser
from mujoco.experimental.studio import sim as _sim
from mujoco.experimental.studio import studio_app
from mujoco.experimental.studio import ux
from mujoco.experimental.studio import viewer_protocol as vp
import numpy as np

from mujoco.experimental.dear_imgui import dear_imgui as imgui

_GFX = absl_flags.DEFINE_enum(
    'gfx', None, vp.GFX_MODES, 'Rendering graphics mode.'
)
_WIDTH = absl_flags.DEFINE_integer('width', 1200, 'Width of the output image.')
_HEIGHT = absl_flags.DEFINE_integer('height', 800, 'Height of the output image')


def view(
    sim_to_view: multiprocessing.Queue,
    view_to_sim: multiprocessing.Queue,
) -> None:
  """Entry-point for process that renders the simulation."""
  # Block until the first message (containing the model) arrives.
  first_msg = sim_to_view.get()
  assert (
      first_msg.model is not None
  ), 'First message must contain the MuJoCo model.'

  title = os.path.basename(sys.argv[0])
  xfrc_sig = int(mujoco.mjtState.mjSTATE_XFRC_APPLIED)
  xfrc_size = mujoco.mj_stateSize(first_msg.model, xfrc_sig)
  xfrc_state = np.zeros(xfrc_size, np.float64)

  data = mujoco.MjData(first_msg.model)
  app = studio_app.StudioApp(first_msg.model, data)
  send_rate = 60.0
  config = vp.ViewerConfig(
      title=title,
      width=_WIDTH.value,
      height=_HEIGHT.value,
      gfx=_GFX.value,
  )
  viewer = _viewer.NativeViewer(config)

  while viewer.is_running() and app.is_running():

    # Drain the queue, keeping only the latest message.
    msg = None
    while not sim_to_view.empty():
      msg = sim_to_view.get()

    # Update the camera and compute the perturbation.
    app.handle_mouse_events(viewer.camera, viewer.vis_options, viewer.perturb)

    # Sync state from the backend if a new payload actually arrived.
    if msg is not None and msg.state is not None:
      mujoco.mj_setState(app.model, app.data, msg.state, msg.state_sig)
      mujoco.mj_forward(app.model, app.data)

    # Always apply the perturbation forces from the viewer.
    app.apply_perturb(viewer.perturb)

    # Transmit user interaction when we get a new state
    if msg is not None:
      mujoco.mj_getState(app.model, app.data, xfrc_state, xfrc_sig)
      view_to_sim.put(
          vp.ViewToSim(
              send_rate=send_rate, state=xfrc_state, state_sig=xfrc_sig
          )
      )

    # Build the UI.
    ux.setup_theme(app.theme)
    if imgui.Begin(
        'Settings',
        flags=int(imgui.WindowFlags.AlwaysAutoResize)
        | int(imgui.WindowFlags.NoTitleBar)
        | int(imgui.WindowFlags.NoCollapse),
    ):
      imgui.PushItemWidth(200.0)
      updated, send_rate = imgui.SliderFloat(
          'Send Rate (Hz)', send_rate, 1.0, 120.0
      )
      if updated:
        view_to_sim.put(vp.ViewToSim(send_rate=send_rate))

      imgui.SetNextItemWidth(-1)
      if imgui.Button('Reset Simulation'):
        view_to_sim.put(vp.ViewToSim(reset=True, send_rate=send_rate))

      imgui.PopItemWidth()
    imgui.End()

    viewer.sync(app.model, app.data)

  # Prevent multiprocessing.Queue atexit handler from blocking on exit.
  sim_to_view.cancel_join_thread()
  view_to_sim.cancel_join_thread()


def sim(
    data: mujoco.MjData,
    model: mujoco.MjModel,
    sim_to_view: multiprocessing.Queue,
    view_to_sim: multiprocessing.Queue,
    view_process: multiprocessing.Process,
) -> None:
  """Entry-point for process that runs the simulation."""

  sim_to_view.put(vp.SimToView(model=model))

  step_control = _sim.StepControl()
  integration_sig = int(mujoco.mjtState.mjSTATE_INTEGRATION)
  integration_size = mujoco.mj_stateSize(model, integration_sig)
  integration_state = np.empty(integration_size, np.float64)

  msg = vp.ViewToSim()
  last_send_time = time.time()

  while view_process.is_alive():
    while not view_to_sim.empty():
      msg = view_to_sim.get()

    if msg.reset:
      mujoco.mj_resetData(model, data)
      mujoco.mj_forward(model, data)
      msg.reset = False

    # Apply perturbation forces received from the viewer process.
    if msg.state is not None:
      mujoco.mj_setState(model, data, msg.state, msg.state_sig)

    # Advance the simulation keeping up with real-time.
    step_control.advance(model, data)

    # Send the simulation state paced by the requested send_rate.
    now = time.time()
    if now - last_send_time >= 1.0 / max(1.0, msg.send_rate):
      mujoco.mj_getState(model, data, integration_state, integration_sig)
      sim_to_view.put(
          vp.SimToView(
              state=integration_state,
              state_sig=integration_sig,
          )
      )
      last_send_time = now


def main(argv: list[str]) -> None:
  if len(argv) < 2:
    print('Usage: async <model_path.xml>')
    sys.exit(1)

  data = parser.parse(argv[1])
  model = data.model

  # Queues for communication between the simulation and viewer processes.
  sim_to_view = multiprocessing.Queue()
  view_to_sim = multiprocessing.Queue()

  # Start the viewer process.
  view_process = multiprocessing.Process(
      target=view, args=(sim_to_view, view_to_sim)
  )
  view_process.start()

  # Start the simulation in the main process.
  sim(data, model, sim_to_view, view_to_sim, view_process)

  # Prevent multiprocessing.Queue atexit handler from blocking on exit.
  sim_to_view.cancel_join_thread()
  view_to_sim.cancel_join_thread()

  view_process.join(timeout=5.0)


if __name__ == '__main__':
  absl_app.run(main)