# Copyright 2021 The Brax Authors.
#
# 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.

"""Trains a ur5e robot arm to move its end effector to a sequence of targets.

The 6 joints have been placed faithfully to the actual joint locations of the
ur5e robot arm model.  Because brax does not yet support meshes, the capsule
locations for collider boundaries are only approximate.

See https://www.universal-robots.com/products/ur5-robot/ for more details."""

from typing import Tuple

import dataclasses
import jax
import jax.numpy as jnp
import brax
from brax.envs import env
from brax.physics import math

from google.protobuf import text_format


class Ur5e(env.Env):
  """Trains a UR5E robotic arm to touch a sequence of random targets."""

  def __init__(self, **kwargs):
    config = text_format.Parse(_SYSTEM_CONFIG, brax.Config())
    super().__init__(config, **kwargs)
    self.target_idx = self.sys.body_idx['Target']
    self.torso_idx = self.sys.body_idx['wrist_3_link']
    self.target_radius = .02
    self.target_distance = .5

  def reset(self, rng: jnp.ndarray) -> env.State:
    qp = self.sys.default_qp()
    rng, target = self._random_target(rng)
    pos = jax.ops.index_update(qp.pos, jax.ops.index[self.target_idx], target)
    qp = dataclasses.replace(qp, pos=pos)
    info = self.sys.info(qp)
    obs = self._get_obs(qp, info)
    reward, done, steps, zero = jnp.zeros(4)
    metrics = {
        'hits': zero,
        'weightedHits': zero,
        'movingToTarget': zero,
    }
    return env.State(rng, qp, info, obs, reward, done, steps, metrics)

  def step(self, state: env.State, action: jnp.ndarray) -> env.State:
    rng = state.rng
    qp, info = self.sys.step(state.qp, action)
    obs = self._get_obs(qp, info)

    # small reward for end effector moving towards target
    torso_delta = qp.pos[self.torso_idx] - state.qp.pos[self.torso_idx]
    target_rel = qp.pos[self.target_idx] - qp.pos[self.torso_idx]
    target_dist = jnp.linalg.norm(target_rel)
    target_dir = target_rel / (1e-6 + target_dist)
    moving_to_target = .1 * jnp.dot(torso_delta, target_dir)

    # big reward for reaching target
    target_hit = jnp.where(target_dist < self.target_radius, 1.0, 0.0)
    weighted_hit = target_hit

    reward = moving_to_target + weighted_hit

    steps = state.steps + self.action_repeat
    done = jnp.where(steps >= self.episode_length, 1.0, 0.0)
    metrics = {
        'hits': target_hit,
        'weightedHits': weighted_hit,
        'movingToTarget': moving_to_target,
    }

    # teleport any hit targets
    rng, target = self._random_target(rng)
    target = jnp.where(target_hit, target, qp.pos[self.target_idx])
    pos = jax.ops.index_update(qp.pos, jax.ops.index[self.target_idx], target)
    qp = dataclasses.replace(qp, pos=pos)
    return env.State(rng, qp, info, obs, reward, done, steps, metrics)

  def _get_obs(self, qp: brax.QP, info: brax.Info) -> jnp.ndarray:
    """Egocentric observation of target and arm body."""
    torso_fwd = math.rotate(jnp.array([1., 0., 0.]), qp.rot[self.torso_idx])
    torso_up = math.rotate(jnp.array([0., 0., 1.]), qp.rot[self.torso_idx])

    v_inv_rotate = jax.vmap(math.inv_rotate, in_axes=(0, None))

    pos_local = qp.pos - qp.pos[self.torso_idx]
    pos_local = v_inv_rotate(pos_local, qp.rot[self.torso_idx])
    vel_local = v_inv_rotate(qp.vel, qp.rot[self.torso_idx])

    target_local = pos_local[self.target_idx]
    target_local_mag = jnp.reshape(jnp.linalg.norm(target_local), -1)
    target_local_dir = target_local / (1e-6 + target_local_mag)

    pos_local = jnp.reshape(pos_local, -1)
    vel_local = jnp.reshape(vel_local, -1)

    contact_mag = jnp.sum(jnp.square(info.contact.vel), axis=-1)
    contacts = jnp.where(contact_mag > 0.00001, 1, 0)

    return jnp.concatenate([
        torso_fwd, torso_up, target_local_mag, target_local_dir, pos_local,
        vel_local, contacts
    ])

  def _random_target(self, rng: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray]:
    """Returns a target location in a random circle slightly above xy plane."""
    rng, rng1, rng2 = jax.random.split(rng, 3)
    dist = self.target_radius + self.target_distance * jax.random.uniform(rng1)
    ang = jnp.pi * 2. * jax.random.uniform(rng2)
    target_x = dist * jnp.cos(ang)
    target_y = dist * jnp.sin(ang)
    target_z = .5
    target = jnp.array([target_x, target_y, target_z]).transpose()
    return rng, target

_SYSTEM_CONFIG = """
bodies {
  name: "floor"
  colliders {
    plane {
    }
  }
  inertia {
    x: 1.0
    y: 1.0
    z: 1.0
  }
  mass: 1.0
  frozen {
    position {
      x: 1.0
      y: 1.0
      z: 1.0
    }
    rotation {
      x: 1.0
      y: 1.0
      z: 1.0
    }
    all: true
  }
}
bodies {
  name: "shoulder_link"
  colliders {
    position {
      y: 0.06682991981506348
    }
    rotation {
      x: 90.0
    }
    capsule {
      radius: 0.05945208668708801
      length: 0.13365983963012695
    }
  }
  inertia {
    x: 1.0
    y: 1.0
    z: 1.0
  }
  mass: 1.0
}
bodies {
  name: "upper_arm_link"
  colliders {
    position {
      z: 0.21287038922309875
    }
    rotation {
    }
    capsule {
      radius: 0.05968618765473366
      length: 0.5446449518203735
    }
  }
  inertia {
    x: 1.0
    y: 1.0
    z: 1.0
  }
  mass: 1.0
}
bodies {
  name: "forearm_link"
  colliders {
    position {
      z: 0.1851803958415985
    }
    rotation {
    }
    capsule {
      radius: 0.05584339052438736
      length: 0.48926496505737305
    }
  }
  inertia {
    x: 1.0
    y: 1.0
    z: 1.0
  }
  mass: 1.0
}
bodies {
  name: "wrist_1_link"
  colliders {
    position {
      y: 0.10467606782913208
    }
    rotation {
      x: 90.0
    }
    capsule {
      radius: 0.038744933903217316
      length: 0.10467606782913208
    }
  }
  inertia {
    x: 1.0
    y: 1.0
    z: 1.0
  }
  mass: 1.0
}
bodies {
  name: "wrist_2_link"
  colliders {
    position {
      z: 0.052344050258398056
    }
    rotation {
    }
    capsule {
      radius: 0.03879201412200928
      length: 0.10468810051679611
    }
  }
  inertia {
    x: 1.0
    y: 1.0
    z: 1.0
  }
  mass: 1.0
}
bodies {
  name: "wrist_3_link"
  colliders {
    position {
      y: -0.04025782644748688
    }
    rotation {
      x: 90.0
    }
    capsule {
      radius: 0.01725015603005886
      length: 0.08051565289497375
    }
  }
  inertia {
    x: 1.0
    y: 1.0
    z: 1.0
  }
  mass: 1.0
}
bodies {
  name: "Target"
  colliders {
    sphere {
      radius: 0.1
    }
  }
  inertia {
    x: 1.0
    y: 1.0
    z: 1.0
  }
  mass: 1.0
  frozen {
    position {
      x: 1.0
      y: 1.0
      z: 1.0
    }
    rotation {
      x: 1.0
      y: 1.0
      z: 1.0
    }
    all: true
  }
}
joints {
  name: "shoulder_pan_joint"
  stiffness: 40000.0
  parent: "floor"
  child: "shoulder_link"
  parent_offset {
    z: 0.163
  }
  child_offset {
  }
  rotation {
    y: -90.0
  }
  angular_damping: 50.0
  angle_limit {
    min: -360.0
    max: 360.0
  }
}
joints {
  name: "shoulder_lift_joint"
  stiffness: 40000.0
  parent: "shoulder_link"
  child: "upper_arm_link"
  parent_offset {
    y: 0.138
  }
  child_offset {
  }
  rotation {
    z: 90.0
  }
  angular_damping: 50.0
  angle_limit {
    min: -360.0
    max: 360.0
  }
}
joints {
  name: "elbow_joint"
  stiffness: 40000.0
  parent: "upper_arm_link"
  child: "forearm_link"
  parent_offset {
    y: -0.13
    z: 0.425
  }
  child_offset {
  }
  rotation {
    z: 90.0
  }
  angular_damping: 50.0
  angle_limit {
    min: -360.0
    max: 360.0
  }
}
joints {
  name: "wrist_1_joint"
  stiffness: 40000.0
  parent: "forearm_link"
  child: "wrist_1_link"
  parent_offset {
    z: 0.3919999897480011
  }
  child_offset {
  }
  rotation {
    z: 90.0
  }
  angular_damping: 50.0
  angle_limit {
    min: -360.0
    max: 360.0
  }
}
joints {
  name: "wrist_2_joint"
  stiffness: 40000.0
  parent: "wrist_1_link"
  child: "wrist_2_link"
  parent_offset {
    y: 0.12700000405311584
  }
  child_offset {
  }
  rotation {
    y: -90.0
  }
  angular_damping: 50.0
  angle_limit {
    min: -360.0
    max: 360.0
  }
}
joints {
  name: "wrist_3_joint"
  stiffness: 40000.0
  parent: "wrist_2_link"
  child: "wrist_3_link"
  parent_offset {
    z: 0.1
  }
  child_offset {
  }
  rotation {
    z: 90.0
  }
  angular_damping: 50.0
  angle_limit {
    min: -360.0
    max: 360.0
  }
}
actuators {
  name: "shoulder_pan_joint"
  joint: "shoulder_pan_joint"
  strength: 100.0
  torque {
  }
}
actuators {
  name: "shoulder_lift_joint"
  joint: "shoulder_lift_joint"
  strength: 100.0
  torque {
  }
}
actuators {
  name: "elbow_joint"
  joint: "elbow_joint"
  strength: 100.0
  torque {
  }
}
actuators {
  name: "wrist_1_joint"
  joint: "wrist_1_joint"
  strength: 100.0
  torque {
  }
}
actuators {
  name: "wrist_2_joint"
  joint: "wrist_2_joint"
  strength: 100.0
  torque {
  }
}
actuators {
  name: "wrist_3_joint"
  joint: "wrist_3_joint"
  strength: 100.0
  torque {
  }
}
friction: 0.6
gravity {
  z: -9.81
}
angular_damping: -0.05
baumgarte_erp: 0.1
dt: 0.02
substeps: 8
"""
