import requests
import pandas as pd
from tqdm.auto import tqdm

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


# Based on Omar Sanseviero work
# Make model clickable link
def make_clickable_model(model_name):
    link = "https://huggingface.co/" + model_name
    return f'<a target="_blank" href="{link}">{model_name}</a>'

# Make user clickable link
def make_clickable_user(user_id):
    link = "https://huggingface.co/" + user_id
    return f'<a  target="_blank" href="{link}">{user_id}</a>'
    
def get_model_ids(rl_env):
    api = HfApi()
    models = api.list_models(filter=rl_env)
    model_ids = [x.modelId for x in models]
    return model_ids
    
def get_metadata(model_id):
    try:
        readme_path = hf_hub_download(model_id, filename="README.md")
        return metadata_load(readme_path)
    except requests.exceptions.HTTPError:
        # 404 README.md not found
        return None
        
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"]
    #print("ACCURACY", accuracy)
    return accuracy

# We keep the worst case episode
def parse_rewards(accuracy):
    if accuracy !=  None:
        parsed = accuracy.split(' +/- ')
        mean_reward = float(parsed[0])
        std_reward =  float(parsed[1])
    else:
        mean_reward = -1000
        std_reward = -1000
    return mean_reward, std_reward

def get_data(rl_env):
    data = []
    model_ids = get_model_ids(rl_env)
    for model_id in tqdm(model_ids):
        meta = get_metadata(model_id)
        if meta is None:
            continue
        user_id = model_id.split('/')[0]
        row = {}
        row["User"] = user_id
        row["Model"] = model_id
        accuracy = parse_metrics_accuracy(meta)
        #print("RETURNED ACCURACY", accuracy)
        mean_reward, std_reward = parse_rewards(accuracy)
        #print("MEAN REWARD", mean_reward)
        row["Results"] = mean_reward - std_reward
        row["Mean Reward"] = mean_reward
        row["Std Reward"] = std_reward
        data.append(row)
    return pd.DataFrame.from_records(data)


def get_data_per_env(rl_env):
    dataframe = get_data(rl_env)
    dataframe = dataframe.fillna("")

    if not dataframe.empty:
        # turn the model ids into clickable links
        dataframe["User"] = dataframe["User"].apply(make_clickable_user)
        dataframe["Model"] = dataframe["Model"].apply(make_clickable_model)
        dataframe = dataframe.sort_values(by=['Results'], ascending=False)
        table_html = dataframe.to_html(escape=False, index=False,justify = 'left')
     
        #table_html = table_html.replace("<thead>", '<thead align="left">')  # left-align the headers
        return table_html,dataframe,dataframe.empty
    else: 
        html = """<div style="color: green">
                <p> ⌛ Please wait. Results will be out soon... </p>
                </div>
               """
        return html,dataframe,dataframe.empty   



RL_ENVS = ['LunarLander-v2','CarRacing-v0','MountainCar-v0']
           
RL_DETAILS ={'CarRacing-v0':{'title':" The Car Racing 🏎️ Leaderboard 🚀",'data':get_data_per_env('CarRacing-v0')},
            'MountainCar-v0':{'title':"The Mountain Car ⛰️ 🚗 Leaderboard 🚀",'data':get_data_per_env('MountainCar-v0')},
            'LunarLander-v2':{'title':" The Lunar Lander 🌕 Leaderboard 🚀",'data':get_data_per_env('LunarLander-v2')}
            }


def get_info_display(len_dataframe,env_name,name_leaderboard,is_empty):
    if not is_empty:
        markdown = """
        <div class='infoPoint'>
        <h1> {name_leaderboard} </h1>
        <br>
        <p> This is a leaderboard of <b>{len_dataframe}</b> agents playing {env_name} 👩‍🚀. </p>
        <br>
        <p> We use lower bound result to sort the models: mean_reward - std_reward. </p>
        <br>    
        <p> You can click on the model's name to be redirected to its model card which includes documentation. </p>
        <br>
        <p> You want to try your model? Read this <a href="https://github.com/huggingface/deep-rl-class/blob/Unit1/unit1/README.md" target="_blank">Unit 1</a> of Deep Reinforcement Learning Class.
        </p>
        </div>
        """.format(len_dataframe = len_dataframe,env_name = env_name,name_leaderboard = name_leaderboard)

    else:
        markdown = """
        <div class='infoPoint'>
        <h1> {name_leaderboard} </h1>
        <br>
        </div>                  
        """.format(name_leaderboard =  name_leaderboard)
    return markdown    

def reload_leaderboard(rl_env):
    global RL_DETAILS
    RL_DETAILS[rl_env]['data'] = get_data_per_env(rl_env)
  
    data_html,data_dataframe,is_empty = RL_DETAILS[rl_env]['data'] 
    data_html = data_html.replace('thead align="left"', '<thead align="left" style="color:red">')  # left-align the headers

    markdown = get_info_display(len(data_dataframe),rl_env,RL_DETAILS[rl_env]['title'],is_empty)            
    
    return markdown,data_html     
            
with open('app.css','r') as f:
    BLOCK_CSS = f.read() 

block = gr.Blocks(css=BLOCK_CSS)
with block:
    
    with gr.Tabs():
        for rl_env in RL_ENVS:
            with gr.TabItem(rl_env):
                data_html,data_dataframe,is_empty = RL_DETAILS[rl_env]['data'] 
                markdown = get_info_display(len(data_dataframe),rl_env,RL_DETAILS[rl_env]['title'],is_empty)            
                reload = gr.Button('Reload Leaderboard')
                env_state =gr.Variable(default_value=rl_env)  
                output_markdown = gr.HTML(markdown)
                output_html = gr.HTML(data_html)
                reload.click(reload_leaderboard,inputs=[env_state],outputs=[output_markdown,output_html])
            

block.launch()