import json
import os
import re

import gradio as gr
import numpy as np
import pandas as pd
from apscheduler.schedulers.background import BackgroundScheduler
from huggingface_hub import HfApi

from src.backend import backend_routine

from src.logging import configure_root_logger, setup_logger


configure_root_logger()
logger = setup_logger(__name__)

API = HfApi(token=os.environ.get("TOKEN"))
RESULTS_REPO = f"open-rl-leaderboard/results"
ALL_ENV_IDS = {
    "Atari": [
        "Adventure",
        "AirRaid",
        "Alien",
        "Amidar",
        "Assault",
        "Asterix",
        "Asteroids",
        "Atlantis",
        "BankHeist",
        "BattleZone",
        "BeamRider",
        "Berzerk",
        "Bowling",
        "Boxing",
        "Breakout",
        "Carnival",
        "Centipede",
        "ChopperCommand",
        "CrazyClimber",
        "Defender",
        "DemonAttack",
        "DoubleDunk",
        "ElevatorAction",
        "Enduro",
        "FishingDerby",
        "Freeway",
        "Frostbite",
        "Gopher",
        "Gravitar",
        "Hero",
        "IceHockey",
        "Jamesbond",
        "JourneyEscape",
        "Kangaroo",
        "Krull",
        "KungFuMaster",
        "MontezumaRevenge",
        "MsPacman",
        "NameThisGame",
        "Phoenix",
        "Pitfall",
        "Pong",
        "Pooyan",
        "PrivateEye",
        "Qbert",
        "Riverraid",
        "RoadRunner",
        "Robotank",
        "Seaquest",
        "Skiing",
        "Solaris",
        "SpaceInvaders",
        "StarGunner",
        "Tennis",
        "TimePilot",
        "Tutankham",
        "UpNDown",
        "Venture",
        "VideoPinball",
        "WizardOfWor",
        "YarsRevenge",
        "Zaxxon",
    ],
    "Box2D": [
        "BipedalWalker-v3",
        "BipedalWalkerHardcore-v3",
        "CarRacing-v2",
        "LunarLander-v2",
        "LunarLanderContinuous-v2",
    ],
    "Toy text": [
        "Blackjack-v1",
        "CliffWalking-v0",
        "FrozenLake-v1",
        "FrozenLake8x8-v1",
    ],
    "Classic control": [
        "Acrobot-v1",
        "CartPole-v1",
        "MountainCar-v0",
        "MountainCarContinuous-v0",
        "Pendulum-v1",
    ],
    "MuJoCo": [
        "Ant-v4",
        "HalfCheetah-v4",
        "Hopper-v4",
        "Humanoid-v4",
        "HumanoidStandup-v4",
        "InvertedDoublePendulum-v4",
        "InvertedPendulum-v4",
        "Pusher-v4",
        "Reacher-v4",
        "Swimmer-v4",
        "Walker2d-v4",
    ],
}


def get_leaderboard_df():
    # List all results files in results repo
    pattern = re.compile(r"^[^/]*/[^/]*/[^/]*results_[a-f0-9]+\.json$")
    filenames = API.list_repo_files(RESULTS_REPO, repo_type="dataset")
    filenames = [filename for filename in filenames if pattern.match(filename)]

    data = []
    for filename in filenames:
        path = API.hf_hub_download(repo_id=RESULTS_REPO, filename=filename, repo_type="dataset")
        with open(path) as fp:
            report = json.load(fp)
        user_id, model_id = report["config"]["model_id"].split("/")
        row = {"user_id": user_id, "model_id": model_id}
        if report["status"] == "DONE" and len(report["results"]) > 0:
            env_ids = list(report["results"].keys())
            assert len(env_ids) == 1, "Only one environment supported for the moment"
            row["env_id"] = env_ids[0]
            row["mean_episodic_return"] = np.mean(report["results"][env_ids[0]]["episodic_returns"])
        data.append(row)

    df = pd.DataFrame(data)  # create DataFrame
    df = df.fillna("")  # replace NaN values with empty strings
    return df


def select_env(df: pd.DataFrame, env_id: str):
    df = df[df["env_id"] == env_id]
    df = df.sort_values("mean_episodic_return", ascending=False)
    df["ranking"] = np.arange(1, len(df) + 1)
    return df


def format_df(df: pd.DataFrame):
    # Add hyperlinks
    df = df.copy()
    for index, row in df.iterrows():
        user_id = row["user_id"]
        model_id = row["model_id"]
        df.loc[index, "user_id"] = f"[{user_id}](https://huggingface.co/{user_id})"
        df.loc[index, "model_id"] = f"[{model_id}](https://huggingface.co/{user_id}/{model_id})"

    # Keep only the relevant columns
    df = df[["ranking", "user_id", "model_id", "mean_episodic_return"]]
    return df.values.tolist()


TITLE = """
🚀 Open RL Leaderboard
"""

INTRODUCTION_TEXT = """
Welcome to the Open RL Leaderboard! This is a community-driven benchmark for reinforcement learning models.
"""

ABOUT_TEXT = r"""
The Open RL Leaderboard is a community-driven benchmark for reinforcement learning models.

## 🔌 How to have your agent evaluated?

The Open RL Leaderboard constantly scans the 🤗 Hub to detect new models to be evaluated. For your model to be evaluated, it must meet the following criteria.

1. The model must be public on the 🤗 Hub
2. The model must contain an `agent.pt` file.
3. The model must be [tagged](https://huggingface.co/docs/hub/model-cards#model-cards) `reinforcement-learning`
4. The model must be [tagged](https://huggingface.co/docs/hub/model-cards#model-cards) with the name of the environment you want to evaluate (for example `MountainCar-v0`)

Once your model meets these criteria, it will be automatically evaluated on the Open RL Leaderboard. That's it!

## 🏗️ How do I build the `agent.pt`?

The `agent.pt` file is a [TorchScript module](https://pytorch.org/docs/stable/jit.html#). It must be loadable using `torch.jit.load`.
The module must take batch of observations as input and return batch of actions. To check if your model is compatible with the Open RL Leaderboard, you can run the following code:

```python
import gymnasium as gym
import numpy as np
import torch

agent_path = "path/to/agent.pt"
env_id = ... # e.g. "MountainCar-v0"

agent = torch.jit.load(agent_path)
env = gym.make(env_id)
observations = np.array([env.observation_space.sample()])
observations = torch.from_numpy(observations)
actions = agent(observations)
actions = actions.numpy()[0]
assert env.action_space.contains(actions)
```

## 🕵 How are the models evaluated?

The evaluation is done by running the agent on the environment for 100 episodes. 

For further information, please refer to the [Open RL Leaderboard evaulation script](https://huggingface.co/spaces/open-rl-leaderboard/leaderboard/blob/main/src/evaluation.py).

### The particular case of Atari environments

Atari environments are evaluated on the `NoFrameskip-v4` version of the environment. For example, to evaluate an agent on the `Pong` environment, you must tag your model with `PongNoFrameskip-v4`. The environment is then wrapped to match the standard Atari preprocessing pipeline.

- No-op reset with a maximum of 30 no-ops
- Max and skip with a skip of 4
- Episodic life (although the reported score is for the full episode, not the life)
- Fire reset
- Clip reward (although the reported score is not clipped)
- Resize observation to 84x84
- Grayscale observation
- Frame stack of 4

## 🚑 Troubleshooting

If you encounter any issue, please open an issue on the [Open RL Leaderboard repository](https://huggingface.co/spaces/open-rl-leaderboard/leaderboard/discussions/new).

## 📜 Citation

```bibtex
@misc{open-rl-leaderboard,
  author = {Quentin Gallouédec and TODO},
  title = {Open RL Leaderboard},
  year = {2024},
  publisher = {Hugging Face},
  howpublished = "\url{https://huggingface.co/spaces/open-rl-leaderboard/leaderboard}",
}
```
"""


with gr.Blocks() as demo:
    gr.HTML(TITLE)
    gr.Markdown(INTRODUCTION_TEXT, elem_classes="markdown-text")
    with gr.Tabs(elem_classes="tab-buttons") as tabs:
        with gr.TabItem("🏅 Leaderboard"):
            df = get_leaderboard_df()
            for env_domain, env_ids in ALL_ENV_IDS.items():
                with gr.TabItem(env_domain):
                    for env_id in env_ids:
                        with gr.TabItem(env_id):
                            with gr.Row(equal_height=False):
                                if env_domain == "Atari":
                                    env_id = f"{env_id}NoFrameskip-v4"
                                env_df = select_env(df, env_id)
                                gr.components.Dataframe(
                                    value=format_df(env_df),
                                    headers=["🏆 Ranking", "🧑 User", "🤖 Model id", "📊 Mean episodic return"],
                                    datatype=["number", "markdown", "markdown", "number"],
                                    row_count=(10, "fixed"),
                                    scale=3,
                                )
                                # Get the best model and
                                if not env_df.empty:
                                    user_id = env_df.iloc[0]["user_id"]
                                    model_id = env_df.iloc[0]["model_id"]
                                    video_path = API.hf_hub_download(
                                        repo_id=f"{user_id}/{model_id}",
                                        filename="replay.mp4",
                                        revision="main",
                                        repo_type="model",
                                    )
                                    video = gr.PlayableVideo(
                                        video_path,
                                        autoplay=True,
                                        scale=1,
                                        min_width=50,
                                        show_download_button=False,
                                        label=model_id,
                                    )
                                    # Doesn't loop for the moment, see https://github.com/gradio-app/gradio/issues/7689

        with gr.TabItem("📝 About", elem_id="llm-benchmark-tab-table", id=2):
            gr.Markdown(ABOUT_TEXT)


scheduler = BackgroundScheduler()
scheduler.add_job(func=backend_routine, trigger="interval", seconds=10 * 60, max_instances=1)
scheduler.start()


if __name__ == "__main__":
    demo.queue().launch()