"""This file implements the functionalities of an example simulated system using pybullet.

"""
import copy
import math
import numpy as np
import os
import pybullet_data
import random


class BoxStackPyBulletSim(object):
  """The ExamplePyBulletSim class that adds some objects to the scene, steps the sim and return a reward.

  """

  def __init__(self, pybullet_client, urdf_root=pybullet_data.getDataPath(), time_step=0.01):
    """Constructs an example simulation and reset it to the initial states.

    Args:
      pybullet_client: The instance of BulletClient to manage different
        simulations.
      urdf_root: The path to the urdf folder.
      time_step: The time step of the simulation.
    """
    self._pybullet_client = pybullet_client
    self._urdf_root = urdf_root
    self.m_actions_taken_since_reset = 0
    self.time_step = time_step
    self.stateId = -1
    self.Reset(reload_urdf=True)

  def Reset(self, reload_urdf=False):
    """Reset the minitaur to its initial states.

    Args:
      reload_urdf: Whether to reload the urdf file. If not, Reset() just place
        the minitaur back to its starting position.
    """
    self.m_actions_taken_since_reset = 0
    xPosRange = 0.025
    yPosRange = 0.025
    boxHalfExtents = 0.025

    if reload_urdf:
      camInfo = self._pybullet_client.getDebugVisualizerCamera()
      cameraDistance = camInfo[10]
      print("cameraYaw=", camInfo[8])
      print("cameraPitch=", camInfo[9])
      print("camtarget=", camInfo[11])
      print("projectionMatrix=", camInfo[3])
      self._pybullet_client.resetDebugVisualizerCamera(cameraDistance=0.3,
                                                       cameraYaw=camInfo[8],
                                                       cameraPitch=camInfo[9],
                                                       cameraTargetPosition=camInfo[11])

      plane = self._pybullet_client.loadURDF("plane.urdf")
      texUid = self._pybullet_client.loadTexture("checker_blue.png")
      self._pybullet_client.changeVisualShape(plane, -1, textureUniqueId=texUid)

      self._numObjects = 4  #random number?

      self._cubes = []

      red = [0.97, 0.25, 0.25, 1]
      green = [0.41, 0.68, 0.31, 1]
      yellow = [0.92, 0.73, 0, 1]
      blue = [0, 0.55, 0.81, 1]
      colors = [red, green, yellow, blue]

      for i in range(self._numObjects):
        pos = [0, 0, boxHalfExtents + i * 2 * boxHalfExtents]
        orn = self._pybullet_client.getQuaternionFromEuler([0, 0, 0])
        orn = [0, 0, 0, 1]
        cube = self._pybullet_client.loadURDF("cube_small.urdf", pos, orn)
        self._pybullet_client.changeVisualShape(cube, -1, rgbaColor=colors[i])
        self._cubes.append(cube)

      self._pybullet_client.setGravity(0, 0, -10)
      self.stateId = self._pybullet_client.saveState()
    else:
      if (self.stateId >= 0):
        self._pybullet_client.restoreState(self.stateId)
    index = 0
    for i in self._cubes:
      posX = random.uniform(-xPosRange, xPosRange)
      posY = random.uniform(-yPosRange, yPosRange)
      yaw = random.uniform(-math.pi, math.pi)
      pos = [posX, posY, boxHalfExtents + index * 2 * boxHalfExtents]
      index += 1
      orn = self._pybullet_client.getQuaternionFromEuler([0, 0, yaw])
      self._pybullet_client.resetBasePositionAndOrientation(i, pos, orn)

  def GetActionDimension(self):
    """Get the length of the action list.

    Returns:
      The length of the action list.
    """
    return 4  #self.num_motors

  def GetObservationUpperBound(self):
    """Get the upper bound of the observation.

    Returns:
      The upper bound of an observation. See GetObservation() for the details
        of each element of an observation.
    """
    upper_bound = np.array([0.0] * self.GetObservationDimension())
    upper_bound[0:] = 1.0  # Quaternion of base orientation.
    return upper_bound

  def GetObservationLowerBound(self):
    """Get the lower bound of the observation."""
    return -self.GetObservationUpperBound()

  def GetObservationDimension(self):
    """Get the length of the observation list.

    Returns:
      The length of the observation list.
    """
    sz = len(self.GetObservation())
    print("sz=", sz)
    return sz

  def GetObservation(self):
    """Get the observations of minitaur.

    Returns:
      The observation list. observation[0:4] is the base orientation in quaternion form.
    """
    observation = []
    for i in self._cubes:
      pos, orn = self._pybullet_client.getBasePositionAndOrientation(i)
      observation.extend(list(pos))
      observation.extend(list(orn))
    return observation

  def ApplyAction(self, action):
    """Set the desired action.
    """
    self.m_actions_taken_since_reset += 1

  def Termination(self):
    return self.m_actions_taken_since_reset >= 200


def CreateSim(pybullet_client, urdf_root, time_step):
  sim = BoxStackPyBulletSim(pybullet_client=pybullet_client,
                            urdf_root=urdf_root,
                            time_step=time_step)
  return sim