from typing import Any

import mujoco_warp as mjwarp
import warp as wp

# Ref: https://github.com/newton-physics/newton/blob/640095cbe1914d43e9158ec71264a0eb7272fc15/newton/_src/solvers/mujoco/solver_mujoco.py#L2587-L2612


@wp.kernel(module="unique")
def repeat_array_kernel(
  src: wp.array(dtype=Any),  # type: ignore
  nelems_per_world: int,
  dst: wp.array(dtype=Any),  # type: ignore
):
  tid = wp.tid()
  src_idx = tid % nelems_per_world  # type: ignore[operator]
  dst[tid] = src[src_idx]


def expand_model_fields(
  model: mjwarp.Model,
  nworld: int,
  fields_to_expand: list[str],
) -> None:
  if nworld == 1:
    return

  def tile(x: wp.array) -> wp.array | wp.array2d | wp.array3d | wp.array4d:
    # Create new array with same shape but first dim multiplied by nworld.
    new_shape = list(x.shape)
    new_shape[0] = nworld
    wp_array = {1: wp.array, 2: wp.array2d, 3: wp.array3d, 4: wp.array4d}[
      len(new_shape)
    ]
    dst = wp_array(shape=new_shape, dtype=x.dtype, device=x.device)

    src_flat = x.flatten()
    dst_flat = dst.flatten()  # type: ignore[possibly-missing-attribute]

    # Launch kernel to repeat data, one thread per destination element.
    n_elems_per_world = dst_flat.shape[0] // nworld
    wp.launch(
      repeat_array_kernel,
      dim=dst_flat.shape[0],
      inputs=[src_flat, n_elems_per_world],
      outputs=[dst_flat],
      device=x.device,
    )
    return dst

  for field in model.__dataclass_fields__:
    if field in fields_to_expand:
      array = getattr(model, field)
      if array.shape[0] != 1:
        continue  # Already per-world (e.g., variant compilation).
      setattr(model, field, tile(array))
