import gradio as gr
import os
from moviepy.editor import *

def replay(option):
	path = ""
	# Get the correct model
	if (option == "LunarLander-v2 🚀👩‍🚀"):
  		path = "./LunarLander-v2.mp4"
	elif(option == "CartPole-v1 🕹️"):
  		path = "./CartPole-v1.mp4"
	elif(option == "Atari Space Invaders 👾"):
  		path = "./SpaceInvadersNoFrameskip-v4.mp4"
	
	
  	# The only turnaround I found (Base64 video pb)
	videoclip = VideoFileClip(path)
	videoclip.write_videofile("new_filename.mp4")
	return 'new_filename.mp4'

"""
TODO: Next version with live video generation
def replay_classical(hf_model_filename, hf_model_id):
	import gym
	from stable_baselines3 import PPO
	from stable_baselines3.common.evaluation import evaluate_policy

	model = PPO.load_from_huggingface(hf_model_id,hf_model_filename)

	eval_env = gym.make(option)


def replay_atari(hf_model_filename, hf_model_id):
"""

iface = gr.Interface(
    replay,
    [
        gr.inputs.Dropdown(["Atari Space Invaders 👾", "CartPole-v1 🕹️", "LunarLander-v2 🚀👩‍🚀"]),
    ],
    "video",
     title = 'Stable Baselines 3 with 🤗',
            description = '',
             article = 
                        '''<div>
                        	<p style="text-align: center">This version of the RL library allows you to load models directly from the Hugging Face Hub</p>
                            <p style="text-align: center"> Select the trained agent you want to watch perform.
                            These models are from <a href="https://github.com/araffin/rl-baselines-zoo">Stable Baseline Zoo</a></p>
                            <p>
                            There are currently 3 models:
                            <ul>
                            	<li><a href="https://huggingface.co/ThomasSimonini/stable-baselines3-ppo-SpaceInvadersNoFrameskip-v4">PPO SpaceInvadersNoFrameskip-v4</a></li>
                            	<li><a href="https://huggingface.co/ThomasSimonini/stable-baselines3-ppo-LunarLander-v2">PPO LunarLander-v2</a></li>
                            	<li><a href="https://huggingface.co/ThomasSimonini/stable-baselines3-ppo-CartPole-v1">PPO CartPole-v1</a></li>
                            </ul>
                        </div>'''
            )
   

iface.launch()