#@title Install Rendering Dependencies { display-mode: "form" }
#@markdown (You only need to run this code when using Colab's hosted runtime)

import os
from IPython.display import HTML, display

def progress(value, max=100):
    return HTML("""
        <progress
            value='{value}'
            max='{max}',
            style='width: 100%'
        >
            {value}
        </progress>
    """.format(value=value, max=max))

pro_bar = display(progress(0, 100), display_id=True)

try:
  import google.colab
  INSTALL_XVFB = True
except ImportError:
  INSTALL_XVFB = 'COLAB_ALWAYS_INSTALL_XVFB' in os.environ

if INSTALL_XVFB:
  !sudo apt-get update -qq
  pro_bar.update(progress(50, 100))
  !sudo DEBIAN_FRONTEND=noninteractive apt-get install -y -qq xvfb
  pro_bar.update(progress(90, 100))
  import subprocess
  subprocess.Popen(['Xvfb', ':1', '-screen', '0', '1024x768x24', '-ac', '+extension', 'GLX', '+render', '-noreset'], 
                   stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
  os.environ["DISPLAY"] = ":1"
pro_bar.update(progress(100, 100))

try:
  import mlagents
  print("ml-agents already installed")
except ImportError:
  !python -m pip install -q mlagents==1.1.0
  print("Installed ml-agents")

# -----------------
# This code is used to close an env that might not have been closed before
try:
  env.close()
except:
  pass
# -----------------

from mlagents_envs.registry import default_registry
env_id = "GridWorld"

# Import the EngineConfigurationChannel class
from mlagents_envs.side_channel.engine_configuration_channel import EngineConfigurationChannel

# Create the side channel
engine_config_channel = EngineConfigurationChannel()

# Pass the side channel to the make method
# Note, the make method takes a LIST of SideChannel as input
env = default_registry[env_id].make(side_channels = [engine_config_channel])

# Configure the Unity Engine
engine_config_channel.set_configuration_parameters(target_frame_rate = 30)

env.reset()

# ...
# Perform experiment on environment
# ...

env.close()

import matplotlib.pyplot as plt
import numpy as np
%matplotlib inline

# -----------------
# This code is used to close an env that might not have been closed before
try:
  env.close()
except:
  pass
# -----------------

from mlagents_envs.registry import default_registry
env_id = "GridWorld"

# Import the EngineConfigurationChannel class
from mlagents_envs.side_channel.environment_parameters_channel import EnvironmentParametersChannel

# Create the side channel
env_parameters = EnvironmentParametersChannel()

# Pass the side channel to the make method
# Note, the make method takes a LIST of SideChannel as input
env = default_registry[env_id].make(side_channels = [env_parameters])

env.reset()
behavior_name = list(env.behavior_specs)[0]

print("Observation without changing the environment parameters")
decision_steps, terminal_steps = env.get_steps(behavior_name)
plt.imshow(np.moveaxis(decision_steps.obs[0][0,:,:,:], 0, -1))
plt.show()

print("Increasing the dimensions of the grid from 5 to 7")
env_parameters.set_float_parameter("gridSize", 7)
print("Increasing the number of X from 1 to 5")
env_parameters.set_float_parameter("numObstacles", 5)

# Any change to a SideChannel will only be effective after a step or reset
# In the GridWorld Environment, the grid's dimensions can only change at reset
env.reset()


decision_steps, terminal_steps = env.get_steps(behavior_name)
plt.imshow(np.moveaxis(decision_steps.obs[0][0,:,:,:], 0, -1))
plt.show()



env.close()