#!/usr/bin/env python
import os, logging, gym
from baselines import logger
from baselines.common import set_global_seeds
from baselines.common.misc_util import boolean_flag
from baselines import bench
from baselines.a2c.a2c import learn
from baselines.common.vec_env.subproc_vec_env import SubprocVecEnv
from baselines.common.vec_env.vec_frame_stack import VecFrameStack
import time

import gym
from gym import spaces
import pybullet as p
from itertools import cycle
import numpy as np

camTargetPos = [0, 0, 0]
cameraUp = [0, 0, 1]
cameraPos = [1, 1, 1]
pitch = -10.0
roll = 0
upAxisIndex = 2
camDistance = 4
pixelWidth = 320
pixelHeight = 200
nearPlane = 0.01
farPlane = 100
fov = 60


class TestEnv(gym.Env):

  def __init__(
      self,
      renderer='tiny',  # ('tiny', 'egl', 'debug')
  ):
    self.action_space = spaces.Discrete(2)
    self.iter = cycle(range(0, 360, 10))

    # how we want to show
    assert renderer in ('tiny', 'egl', 'debug', 'plugin')
    self._renderer = renderer
    self._render_width = 84
    self._render_height = 84
    # connecting
    if self._renderer == "tiny" or self._renderer == "plugin":
      optionstring = '--width={} --height={}'.format(self._render_width, self._render_height)
      p.connect(p.DIRECT, options=optionstring)

      if self._renderer == "plugin":
        plugin_fn = os.path.join(
            p.__file__.split("bullet3")[0],
            "bullet3/build/lib.linux-x86_64-3.5/eglRenderer.cpython-35m-x86_64-linux-gnu.so")
        plugin = p.loadPlugin(plugin_fn, "_tinyRendererPlugin")
        if plugin < 0:
          print("\nPlugin Failed to load! Try installing via `pip install -e .`\n")
          sys.exit()
        print("plugin =", plugin)

    elif self._renderer == "egl":
      optionstring = '--width={} --height={}'.format(self._render_width, self._render_height)
      optionstring += ' --window_backend=2 --render_device=0'
      p.connect(p.GUI, options=optionstring)

    elif self._renderer == "debug":
      #print("Connection: SHARED_MEMORY")
      #cid = p.connect(p.SHARED_MEMORY)
      #if (cid<0):
      cid = p.connect(p.GUI)
      p.resetDebugVisualizerCamera(1.3, 180, -41, [0.52, -0.2, -0.33])

    p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0)
    p.configureDebugVisualizer(p.COV_ENABLE_SEGMENTATION_MARK_PREVIEW, 0)
    p.configureDebugVisualizer(p.COV_ENABLE_DEPTH_BUFFER_PREVIEW, 0)
    p.configureDebugVisualizer(p.COV_ENABLE_RGB_BUFFER_PREVIEW, 0)

  def __del__(self):
    p.disconnect()

  def reset(self):
    pass

  def step(self, action):
    p.stepSimulation()
    start = time.time()
    yaw = next(self.iter)
    viewMatrix = p.computeViewMatrixFromYawPitchRoll(camTargetPos, camDistance, yaw, pitch, roll,
                                                     upAxisIndex)
    aspect = pixelWidth / pixelHeight
    projectionMatrix = p.computeProjectionMatrixFOV(fov, aspect, nearPlane, farPlane)
    img_arr = p.getCameraImage(pixelWidth,
                               pixelHeight,
                               viewMatrix,
                               projectionMatrix,
                               shadow=1,
                               lightDirection=[1, 1, 1],
                               renderer=p.ER_BULLET_HARDWARE_OPENGL)
    #renderer=pybullet.ER_TINY_RENDERER)
    self._observation = img_arr[2]
    return np.array(self._observation), 0, 0, {}

  def seed(self, seed=None):
    pass


def train(env_id, num_timesteps=300, seed=0, num_env=2, renderer='tiny'):

  def make_env(rank):

    def _thunk():
      if env_id == "TestEnv":
        env = TestEnv(renderer=renderer)  #gym.make(env_id)
      else:
        env = gym.make(env_id)
      env.seed(seed + rank)
      env = bench.Monitor(env, logger.get_dir() and os.path.join(logger.get_dir(), str(rank)))
      gym.logger.setLevel(logging.WARN)
      # only clip rewards when not evaluating
      return env

    return _thunk

  set_global_seeds(seed)
  env = SubprocVecEnv([make_env(i) for i in range(num_env)])

  env.reset()
  start = time.time()
  for i in range(num_timesteps):
    action = [env.action_space.sample() for _ in range(num_env)]
    env.step(action)
  stop = time.time()
  duration = (stop - start)
  if (duration):
    fps = num_timesteps / duration
  else:
    fps = 0
  env.close()
  return num_env, fps


if __name__ == "__main__":
  env_id = "TestEnv"
  res = []

  for renderer in ('tiny', 'plugin', 'egl'):
    for i in (1, 8):
      tmp = train(env_id, num_env=i, renderer=renderer)
      print(renderer, tmp)
      res.append((renderer, tmp))
  print()
  print("rendertest_sync.py")
  print("back nenv fps fps_tot")
  for renderer, i in res:
    print(renderer, '\t', i[0], round(i[1]), '\t', round(i[0] * i[1]))