import gradio as gr
import requests.exceptions
from huggingface_hub import HfApi, hf_hub_download
from huggingface_hub.repocard import metadata_load

def load_agent(model_id_1, model_id_2):
    """
    This function load the agent's video and results
    :return: video_path
    """
    # Load the metrics
    metadata_1 = get_metadata(model_id_1)
    
    # Get the accuracy
    results_1 = parse_metrics_accuracy(metadata_1)
    
    # Load the video
    video_path_1 = hf_hub_download(model_id_1, filename="replay.mp4")
    
    # Load the metrics
    metadata_2 = get_metadata(model_id_2)
    
    # Get the accuracy
    results_2 = parse_metrics_accuracy(metadata_2)
    
    # Load the video
    video_path_2 = hf_hub_download(model_id_2, filename="replay.mp4")
    
    return model_id_1, video_path_1, results_1, model_id_2, video_path_2, results_2

def parse_metrics_accuracy(meta):
    if "model-index" not in meta:
        return None
    result = meta["model-index"][0]["results"]
    metrics = result[0]["metrics"]
    accuracy = metrics[0]["value"]
    return accuracy

def get_metadata(model_id):
    """
    Get the metadata of the model repo
    :param model_id:
    :return: metadata
    """
    try:
        readme_path = hf_hub_download(model_id, filename="README.md")
        metadata = metadata_load(readme_path)
        print(metadata)
        return metadata
    except requests.exceptions.HTTPError:
        return None


gr.Interface(load_agent, 
[
        gr.Textbox(
            label="Model 1",
        ),
        gr.Textbox(
            label="Model 2",
        ),
    ],
    [ "text", "video", gr.Textbox(
            label="Mean Reward +/- Std Reward",
        ), "text", "video",  gr.Textbox(
            label="Mean Reward +/- Std Reward",
        )],
        examples=[
        ["sb3/a2c-AntBulletEnv-v0","sb3/ppo-AntBulletEnv-v0"],
        ["ThomasSimonini/a2c-AntBulletEnv-v0", "sb3/a2c-AntBulletEnv-v0"],
        ["sb3/dqn-SpaceInvadersNoFrameskip-v4", "sb3/a2c-SpaceInvadersNoFrameskip-v4"], 
        ["ThomasSimonini/ppo-QbertNoFrameskip-v4","sb3/ppo-QbertNoFrameskip-v4"],
    ],
    title="Compare Deep Reinforcement Learning agents",
    description="Type two models id you want to compare"
    ).launch()