# Copyright 2023 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
#
#     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.
# ==============================================================================
"""Utilities for testing."""

import sys
from typing import Dict, List, Tuple
from xml.etree import ElementTree as ET

from etils import epath
import mujoco
import numpy as np

TEST_FILES: List[str] = [
    'constraints.xml',
    'convex.xml',
    'pendula.xml',
]

_ACTUATOR_TYPES = ['motor', 'velocity', 'position', 'general', 'intvelocity']
_JOINT_TYPES = ['free', 'hinge', 'slide', 'ball']
_JOINT_AXES = ['1 0 0', '0 1 0', '0 0 1']
_FRICTIONS = ['1.2 0.003 0.0002', '0.2 0.0001 0.0005']
_KP_POS = ['1', '2']
_KP_INTVEL = ['10000', '2000']
_KV_VEL = ['123', '1']
_PAIR_FRICTIONS = ['1.2 0.9 0.003 0.0002 0.0001']
_SOLREFS = ['0.04 1.01', '0.05 1.02', '0.03 1.1', '0.015 1.0']
_SOLIMPS = [
    '0.75 0.94 0.002 0.2 2',
    '0.8 0.99 0.001 0.3 6',
    '0.6 0.9 0.003 0.1 1',
]
_DIMS = ['3']
_MARGINS = ['0.0', '0.01', '0.02']
_GAPS = ['0.0', '0.005']
_GEARS = ['20', '50', '100']


def p(pct: int) -> bool:
  assert 0 <= pct <= 100
  return np.random.uniform(low=0, high=100) < pct


def _make_joint(joint_type: str, name: str) -> Dict[str, str]:
  """Returns attributes for a joint."""
  joint_attr = {'type': joint_type, 'name': name}

  if joint_type not in ('free', 'ball'):
    joint_attr['axis'] = np.random.choice(_JOINT_AXES)
    lb, ub = -np.random.uniform() * 90, np.random.uniform() * 90
    joint_attr['range'] = f'{lb:.2f} {ub:.2f}'
  elif joint_type == 'ball':
    joint_attr['axis'] = '1 0 0'
    ub = np.random.uniform() * 90
    joint_attr['range'] = f'0.0 {ub:.2f}'

  if p(50) and joint_type != 'free':
    lb, ub = -np.random.uniform(), np.random.uniform()
    joint_attr['actuatorfrcrange'] = f'{lb:.2f} {ub:.2f}'

  if joint_type not in ('free',):
    joint_attr['damping'] = '{:.2f}'.format(np.random.uniform() * 20)
    joint_attr['stiffness'] = '{:.2f}'.format(np.random.uniform() * 20)

  return joint_attr


def _geom_solparams(
    pair: bool = False, enable_contact: bool = True
) -> Dict[str, str]:
  """Returns geom solver parameters."""
  params = {
      'contype': np.random.choice(['0', '1']) if enable_contact else '0',
      'conaffinity': np.random.choice(['0', '1']) if enable_contact else '0',
      'priority': np.random.choice(['-1', '2']),
      'solmix': np.random.choice(['0.0', '1.6']),
      'friction': np.random.choice(_FRICTIONS),
      'condim': np.random.choice(_DIMS),
  }
  pair_params = {
      'solreffriction': np.random.choice(_SOLREFS),
      'friction': np.random.choice(_PAIR_FRICTIONS),
      'condim': np.random.choice(_DIMS),
  }
  params = pair_params if pair else params
  params.update({
      'solimp': np.random.choice(_SOLIMPS),
      'solref': np.random.choice(_SOLREFS),
      'margin': np.random.choice(_MARGINS),
      'gap': np.random.choice(_GAPS),
  })

  return params


def _make_geom(
    pos: str, size: float, name: str, enable_contact: bool = True
) -> Dict[str, str]:
  """Returns attributes for a sphere geom."""
  attr = {
      'pos': pos,
      'type': 'sphere',
      'name': name,
      'size': f'{size:.2f}',
      'mass': '1',
  }
  attr.update(_geom_solparams(pair=False, enable_contact=enable_contact))

  return attr


def _make_actuator(actuator_type: str, joint: str) -> Dict[str, str]:
  """Returns attributes for an actuator."""
  attr = {'joint': joint}
  if actuator_type == 'motor':
    attr['gear'] = np.random.choice(_GEARS)
  elif actuator_type == 'position':
    attr['kp'] = np.random.choice(_KP_POS)
  elif actuator_type == 'general':
    attr['biastype'] = 'affine'
    attr['gainprm'] = '35 0 0'
    attr['biasprm'] = '0 -35 -0.65'
  elif actuator_type == 'intvelocity':
    attr['kp'] = np.random.choice(_KP_INTVEL)
    lb, ub = -np.random.uniform(), np.random.uniform()
    attr['actrange'] = f'{lb:.2f} {ub:.2f}'
  elif actuator_type == 'velocity':
    attr['kv'] = np.random.choice(_KV_VEL)

  if p(50) and actuator_type != 'intvelocity':
    lb, ub = -np.random.uniform(), np.random.uniform()
    attr['ctrlrange'] = f'{lb:.2f} {ub:.2f}'

  if p(50):
    lb, ub = -np.random.uniform(), np.random.uniform()
    attr['forcerange'] = f'{lb*10:.2f} {ub*10:.2f}'

  return attr


def create_mjcf(
    seed: int,
    min_trees: int = 1,
    max_trees: int = 1,
    max_tree_depth: int = 5,
    body_pos: Tuple[float, float, float] = (0.0, 0.0, -0.5),
    geom_pos: Tuple[float, float, float] = (0.0, 0.0, 0.0),
    max_stacked_joints=4,
    max_geoms_per_body=2,
    max_contact_excludes=1,
    max_contact_pairs=4,
    disable_actuation_pct: int = 0,
    add_actuators: bool = False,
    root_always_free: bool = False,
    enable_contact: bool = True,
) -> str:
  """Creates a random MJCF for testing.

  Args:
    seed: seed for rng
    min_trees: minimum number of kinematic trees to generate
    max_trees: maximum number of kinematic trees to generate
    max_tree_depth: the maximum tree depth
    body_pos: the default body position relative to the parent
    geom_pos: the default geom position in the body frame
    max_stacked_joints: maximum number of joints to stack for each body
    max_geoms_per_body: maximum number of geoms per body
    max_contact_excludes: maximum number of bodies to exlude from contact
    max_contact_pairs: maximum number of explicit geom contact pairs in the xml
    disable_actuation_pct: the percentage of time to disable actuation via the
      disable flag
    add_actuators: whether to add actuators
    root_always_free: if True, the root body of each kinematic tree has a free
      joint with the world
    enable_contact: if False, disables all contacts via contype/conaffinity

  Returns:
    an XML string for the MuJoCo config
  Raises:
    AssertionError when args are not in the correct ranges
  """
  np.random.seed(seed)

  assert min_trees <= max_trees
  assert max_tree_depth >= 1
  assert 0 <= disable_actuation_pct <= 100
  assert max_stacked_joints >= 1
  assert max_geoms_per_body >= 1
  assert max_contact_excludes >= 1
  assert max_contact_pairs >= 1

  mjcf = ET.Element('mujoco')
  opt = ET.SubElement(mjcf, 'option', {'timestep': '0.005', 'solver': 'CG'})
  world = ET.SubElement(mjcf, 'worldbody')
  ET.SubElement(mjcf, 'compiler', {'autolimits': 'true'})

  # disable flags
  if p(disable_actuation_pct):
    ET.SubElement(opt, 'flag', {'actuation': 'disable'})

  ET.SubElement(
      world,
      'geom',
      {
          'name': 'plane',
          'type': 'plane',
          'contype': '1' if enable_contact else '0',
          'conaffinity': '1' if enable_contact else '0',
          'size': '40 40 40',
      },
  )

  # kinematic trees
  tree_depth = np.random.randint(1, max_tree_depth + 1)

  def make_tree(body: ET.Element, depth: int) -> None:
    if depth >= tree_depth:
      return

    z_pos = np.random.uniform(low=-1, high=1) * 0.01  # small jitter
    pos = f'{body_pos[0]:.3f} {body_pos[1]:.3f} {body_pos[2] + z_pos:.3f}'
    n_bodies = len(list(mjcf.iter('body')))
    child = ET.SubElement(body, 'body', {'pos': pos, 'name': f'body{n_bodies}'})

    n_joints = len(list(mjcf.iter('joint')))
    for nj in range(np.random.randint(1, max_stacked_joints + 1)):
      joint_type = np.random.choice(_JOINT_TYPES)
      if nj == 0 and depth == 0 and root_always_free:
        joint_type = 'free'

      # free joint only allowed at top level
      while joint_type == 'free' and (depth > 0 or nj > 0):
        joint_type = np.random.choice(_JOINT_TYPES)

      joint_attr = _make_joint(joint_type, name=f'joint{n_joints + nj}')
      ET.SubElement(child, 'joint', joint_attr)

      prev_joints = child.findall('joint')
      had_ball_or_free = any(
          [j.get('type') in ('ball', 'free') for j in prev_joints]
      )
      if had_ball_or_free:
        break  # do not stack more joints

    n_geoms = len(list(mjcf.iter('geom')))
    for _ in range(np.random.randint(1, max_geoms_per_body + 1)):
      pos = ('{:.2f} ' * 3).format(*geom_pos).strip()
      size = 0.2 + np.random.uniform(low=-1, high=1) * 0.02
      geom_attr = _make_geom(
          pos, size, name=f'geom{n_geoms}', enable_contact=enable_contact
      )
      ET.SubElement(child, 'geom', geom_attr)
      n_geoms += 1

    make_tree(child, depth + 1)

  num_trees = np.random.randint(min_trees, max_trees + 1)
  for _ in range(num_trees):
    make_tree(world, 0)

  # actuators
  if add_actuators:
    actuator = ET.SubElement(mjcf, 'actuator')
    n_joints = len(list(mjcf.iter('joint')))
    nu = np.random.randint(1, n_joints + 1)
    actuators = []
    for i in range(nu):
      actuator_type = np.random.choice(_ACTUATOR_TYPES)
      attr = _make_actuator(actuator_type, joint=f'joint{i}')
      actuators.append((actuator_type, attr))

    np.random.shuffle(actuators)
    for typ, attr in actuators:
      ET.SubElement(actuator, typ, attr)

  # contact pairs
  contact = ET.SubElement(mjcf, 'contact')
  geoms = list(mjcf.iter('geom'))
  geom_names = [geom.get('name') for geom in geoms]
  n_geoms = len(geoms)
  pairs = set()
  for _ in range(min(max_contact_pairs, n_geoms * (n_geoms - 1) // 2)):
    if p(80):
      continue

    geom1, geom2 = np.random.choice(geom_names, replace=False, size=2)
    if geom1 > geom2:
      geom1, geom2 = geom2, geom1

    if (geom1, geom2) in pairs:
      continue

    pairs.add((geom1, geom2))
    attr = {'geom1': geom1, 'geom2': geom2}
    attr.update(_geom_solparams(pair=True))
    ET.SubElement(contact, 'pair', attr)

  # exclude contacts
  bodies = list(mjcf.iter('body'))
  body_names = [b.get('name') for b in bodies]
  n_bodies = len(bodies)
  for _ in range(min(max_contact_excludes, (n_bodies * (n_bodies - 1) // 2))):
    if p(50):
      continue

    body1, body2 = np.random.choice(body_names, replace=False, size=2)
    ET.SubElement(contact, 'exclude', {'body1': body1, 'body2': body2})

  # ElementTree.indent is not available before Python 3.9
  if sys.version_info.minor >= 9:
    ET.indent(mjcf)

  return ET.tostring(mjcf).decode('utf-8')


def load_test_file(name: str) -> mujoco.MjModel:
  """Loads a mujoco.MjModel based on the file name."""
  path = epath.resource_path('mujoco.mjx') / 'test_data' / name
  m = mujoco.MjModel.from_xml_path(path.as_posix())
  return m
