# Copyright 2025 The Newton Developers
#
# 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
#
#     http://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.
# ==============================================================================

"""Tests for io functions for JAX compatibility.

NOTE: do not edit this file without consulting with the MuJoCo/MJX team.
"""

import dataclasses
import enum
import typing
from typing import Any, Dict

import numpy as np
import warp as wp
from absl.testing import absltest
from absl.testing import parameterized

import mujoco_warp as mjw
from mujoco_warp import test_data

_IO_TEST_MODELS = (
  "pendula.xml",
  "collision_sdf/tactile.xml",
  "flex/floppy.xml",
  "actuation/tendon_force_limit.xml",
  "hfield/hfield.xml",
)


def _dims_match(test_obj, d1: Any, d2: Any, prefix: str = ""):
  """Checks that two dataclasses have fields with the same leading dims."""
  fields1, fields2 = dataclasses.fields(d1), dataclasses.fields(d2)
  for f1, f2 in zip(fields1, fields2):
    full_name = prefix + f1.name
    a1, a2 = getattr(d1, f1.name), getattr(d2, f2.name)
    if dataclasses.is_dataclass(a1) or dataclasses.is_dataclass(a2):
      _dims_match(test_obj, a1, a2, prefix + f1.name + ".")
      continue

    if isinstance(f1.type, wp.array) or isinstance(f2.type, wp.array):
      s1, s2 = a1.shape, a2.shape
      test_obj.assertEqual(len(s1), len(s2), f"{full_name} dims mismatch. Got {s1} and {s2}.")
      test_obj.assertEqual(s1, s2, f"{full_name} dims mismatch. Got {s1} and {s2}.")


def _get_np_scalar_type(val: Any) -> bool | int | float | None:
  """Returns the python type from a numpy scalar."""
  is_np_scalar = list(isinstance(val, t) for t in (np.integer, np.floating, np.bool_))
  if any(is_np_scalar):
    return [int, float, bool][is_np_scalar.index(True)]


def _check_type_matches_annotation(test_obj, obj: Any, prefix: str = ""):
  """Checks that dataclass annotations match the runtime types."""
  assert dataclasses.is_dataclass(obj), prefix + " must be dataclass."
  msg = "Type of {val_type} does not match annotation {type_} for field {prefix}{field_name}"

  for field in dataclasses.fields(obj):
    field_name = field.name

    val = getattr(obj, field_name)
    val_type = type(val)
    type_ = field.type

    if dataclasses.is_dataclass(val):
      test_obj.assertTrue(dataclasses.is_dataclass(type_), msg.format(**locals()))
      _check_type_matches_annotation(test_obj, val, prefix + field_name + ".")
      continue

    np_scalar_type = _get_np_scalar_type(val)
    if np_scalar_type:
      test_obj.assertIsInstance(np_scalar_type(val), type_, msg.format(**locals()))
      continue

    if isinstance(type_, wp.array):
      test_obj.assertIsInstance(val, wp.array, msg.format(**locals()))
      continue

    origin_type = typing.get_origin(type_)
    if tuple in (val_type, origin_type):
      test_obj.assertEqual(val_type, origin_type, msg.format(**locals()))
      field_name += ".tuple[]"
      type_ = typing.get_args(type_)[0]

      items = val
      for val in items:
        val_type = type(val)
        if dataclasses.is_dataclass(val):
          _check_type_matches_annotation(test_obj, val, prefix + field_name)
          continue

        np_scalar_type = _get_np_scalar_type(val)
        if np_scalar_type:
          test_obj.assertIsInstance(np_scalar_type(val), type_, msg.format(**locals()))
          continue

        if isinstance(type_, wp.array):
          test_obj.assertIsInstance(val, wp.array, msg.format(**locals()))
          continue

        test_obj.assertEqual(type(val), type_, msg.format(**locals()))
      continue

    test_obj.assertEqual(type(val), field.type, msg.format(**locals()))


def _check_annotation_compat(
  annotations: Dict[str, Any], prefix: str = "", in_cls: bool = False, in_tuple: bool = False
) -> Dict[str, Any]:
  """Checks that dataclass annotations match criteria for JAX API compat."""
  for k, v in annotations.items():
    full_key = f"{prefix}{k}"
    info = f"Found {v} for annotation {full_key}."

    if v in (int, bool, float):
      continue

    if isinstance(v, wp.array):
      continue

    if wp.types.type_is_composite(v):
      raise AssertionError(f"Vector types are not allowed. {info}")

    if typing.get_origin(v) == tuple and (in_cls or in_tuple):
      raise AssertionError(f"Nested args in Model/Data must not be tuple. {info}")

    if typing.get_origin(v) == tuple:
      tuple_args = typing.get_args(v)
      if len(tuple_args) != 2 and tuple_args[1] != ...:
        raise AssertionError(f"Tuple args must be variadic. {info}")

      _check_annotation_compat(
        {"[]": tuple_args[0]},
        prefix=f"{full_key}.tuple",
        in_cls=in_cls,
        in_tuple=True,
      )
      continue

    if hasattr(v, "__class__") and in_cls and not issubclass(v, enum.Enum):
      raise AssertionError(f"Nested object args in Model/Data are not allowed. {info}")

    if hasattr(v, "__class__") and not dataclasses.is_dataclass(v) and not issubclass(v, enum.Enum):
      raise AssertionError(f"Args that are objects must be dataclass. {info}")

    if hasattr(v, "__class__") and not v.__module__.startswith("mujoco_warp"):
      raise AssertionError(f"dataclass args must be within the mujoco_warp module. {info}")

    if hasattr(v, "__class__"):
      _check_annotation_compat(v.__annotations__, prefix=f"{full_key}{v.__name__}.", in_cls=True, in_tuple=in_tuple)
      continue

    raise AssertionError(f"Model/Data annotation is not allowed. {info}")


def _leading_dims_scale_w_nworld(test_obj, d1: Any, d2: Any, nworld1: int, nworld2: int, prefix: str = ""):
  """Checks that dataclass fields that scale with nworld have leading dim nworld."""
  msg = "Arrays that scale with nworld should have leading dim nworld."
  fields1, fields2 = dataclasses.fields(d1), dataclasses.fields(d2)
  for f1, f2 in zip(fields1, fields2):
    full_name = prefix + f1.name
    a1, a2 = getattr(d1, f1.name), getattr(d2, f2.name)
    if dataclasses.is_dataclass(a1) or dataclasses.is_dataclass(a2):
      _leading_dims_scale_w_nworld(test_obj, a1, a2, nworld1, nworld2, prefix + f1.name + ".")
      continue

    if isinstance(f1.type, wp.array) or isinstance(f2.type, wp.array):
      s1, s2 = a1.shape[0], a2.shape[0]
      if s1 == s2:
        continue
      test_obj.assertEqual(s2, nworld2, full_name + f" has leading dim {s2} with nworld={nworld2}. {msg}")
      test_obj.assertEqual(s1, nworld1, full_name + f" has leading dim {s1} with nworld={nworld1}. {msg}")


class IOTest(parameterized.TestCase):
  def test_put_model_nworld_array(self):
    """Tests that put_model arrays with nworld leading dim have `_is_batched`."""
    mjm, *_ = test_data.fixture("pendula.xml")
    m1 = mjw.put_model(mjm)

    self.assertTrue(hasattr(m1.geom_pos, "_is_batched"))
    self.assertEqual(m1.geom_pos.shape[0], 1)
    self.assertEqual(m1.geom_pos.strides[0], 0)
    self.assertLen(m1.geom_pos.strides, m1.geom_pos.ndim)
    self.assertTrue(hasattr(m1.opt.gravity, "_is_batched"))
    self.assertEqual(m1.opt.gravity.shape[0], 1)
    self.assertEqual(m1.opt.gravity.strides[0], 0)
    self.assertLen(m1.opt.gravity.strides, m1.opt.gravity.ndim)
    self.assertFalse(hasattr(m1.body_parentid, "_is_batched"))
    self.assertGreater(m1.body_parentid.shape[0], 0)
    self.assertGreater(m1.body_parentid.strides[0], 0)
    self.assertLen(m1.body_parentid.strides, m1.body_parentid.ndim)

  @parameterized.parameters(*_IO_TEST_MODELS)
  def test_put_data_nworld_array(self, xml):
    """Tests that put_data arrays that scale with nworld have leading dim nworld."""
    mjm, mjd, _, _ = test_data.fixture(xml)
    d1 = mjw.put_data(mjm, mjd, nworld=2, nconmax=1, njmax=40)
    dn = mjw.put_data(mjm, mjd, nworld=7, nconmax=1, njmax=40)
    _leading_dims_scale_w_nworld(self, d1, dn, 2, 7)

  def test_public_api_jax_compat(self):
    """Tests that annotations meet a set of criteria for JAX compat."""
    _check_annotation_compat(mjw.Model.__annotations__, "Model.")
    _check_annotation_compat(mjw.Data.__annotations__, "Data.")

  @parameterized.parameters(*_IO_TEST_MODELS)
  def test_types_match_annotations(self, xml):
    """Tests that the types of dataclass fields match the annotations."""
    mjm, _, m, d = test_data.fixture(xml)

    _check_type_matches_annotation(self, m, "Model.")
    _check_type_matches_annotation(self, d, "Data.")

    d = mjw.make_data(mjm, nworld=2)
    _check_type_matches_annotation(self, d, "Data.")

  @parameterized.parameters(*_IO_TEST_MODELS)
  def test_make_put_data_dims_match(self, xml):
    """Tests that make_data and put_data have matching dimensions."""
    mjm, mjd, _, _ = test_data.fixture(xml)
    dm2 = mjw.make_data(mjm, nworld=2, nconmax=30, njmax=40)
    dm3 = mjw.make_data(mjm, nworld=3, nconmax=30, njmax=40)

    dp2 = mjw.put_data(mjm, mjd, nworld=2, nconmax=30, njmax=40)
    dp3 = mjw.put_data(mjm, mjd, nworld=3, nconmax=30, njmax=40)

    _dims_match(self, dm2, dp2)
    _dims_match(self, dm3, dp3)


if __name__ == "__main__":
  wp.init()
  absltest.main()
