import glob
import json
import logging
import os

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

from src.backend import backend_routine
from src.logging import configure_root_logger, setup_logger

configure_root_logger()
logger = setup_logger(__name__)
logging.getLogger("absl").setLevel(logging.WARNING)

API = HfApi(token=os.environ.get("TOKEN"))
RESULTS_REPO = "open-rl-leaderboard/results"
REFRESH_RATE = 5 * 60  # 5 minutes
ALL_ENV_IDS = {
    "Atari": [
        "AdventureNoFrameskip-v4",
        "AirRaidNoFrameskip-v4",
        "AlienNoFrameskip-v4",
        "AmidarNoFrameskip-v4",
        "AssaultNoFrameskip-v4",
        "AsterixNoFrameskip-v4",
        "AsteroidsNoFrameskip-v4",
        "AtlantisNoFrameskip-v4",
        "BankHeistNoFrameskip-v4",
        "BattleZoneNoFrameskip-v4",
        "BeamRiderNoFrameskip-v4",
        "BerzerkNoFrameskip-v4",
        "BowlingNoFrameskip-v4",
        "BoxingNoFrameskip-v4",
        "BreakoutNoFrameskip-v4",
        "CarnivalNoFrameskip-v4",
        "CentipedeNoFrameskip-v4",
        "ChopperCommandNoFrameskip-v4",
        "CrazyClimberNoFrameskip-v4",
        "DefenderNoFrameskip-v4",
        "DemonAttackNoFrameskip-v4",
        "DoubleDunkNoFrameskip-v4",
        "ElevatorActionNoFrameskip-v4",
        "EnduroNoFrameskip-v4",
        "FishingDerbyNoFrameskip-v4",
        "FreewayNoFrameskip-v4",
        "FrostbiteNoFrameskip-v4",
        "GopherNoFrameskip-v4",
        "GravitarNoFrameskip-v4",
        "HeroNoFrameskip-v4",
        "IceHockeyNoFrameskip-v4",
        "JamesbondNoFrameskip-v4",
        "JourneyEscapeNoFrameskip-v4",
        "KangarooNoFrameskip-v4",
        "KrullNoFrameskip-v4",
        "KungFuMasterNoFrameskip-v4",
        "MontezumaRevengeNoFrameskip-v4",
        "MsPacmanNoFrameskip-v4",
        "NameThisGameNoFrameskip-v4",
        "PhoenixNoFrameskip-v4",
        "PitfallNoFrameskip-v4",
        "PongNoFrameskip-v4",
        "PooyanNoFrameskip-v4",
        "PrivateEyeNoFrameskip-v4",
        "QbertNoFrameskip-v4",
        "RiverraidNoFrameskip-v4",
        "RoadRunnerNoFrameskip-v4",
        "RobotankNoFrameskip-v4",
        "SeaquestNoFrameskip-v4",
        "SkiingNoFrameskip-v4",
        "SolarisNoFrameskip-v4",
        "SpaceInvadersNoFrameskip-v4",
        "StarGunnerNoFrameskip-v4",
        "TennisNoFrameskip-v4",
        "TimePilotNoFrameskip-v4",
        "TutankhamNoFrameskip-v4",
        "UpNDownNoFrameskip-v4",
        "VentureNoFrameskip-v4",
        "VideoPinballNoFrameskip-v4",
        "WizardOfWorNoFrameskip-v4",
        "YarsRevengeNoFrameskip-v4",
        "ZaxxonNoFrameskip-v4",
    ],
    "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 iqm(x):
    score_dict = {"a": np.expand_dims(np.array(x), 1)}

    def aggregate_func(x):
        return np.array([metrics.aggregate_iqm(x)])

    aggregate_scores, aggregate_score_cis = rly.get_interval_estimates(score_dict, aggregate_func, reps=1000)
    iqm, _, _ = aggregate_scores["a"][0], aggregate_score_cis["a"][0][0], aggregate_score_cis["a"][1][0]
    return iqm


def get_leaderboard_df():
    dir_path = API.snapshot_download(repo_id=RESULTS_REPO, repo_type="dataset")
    pattern = os.path.join(dir_path, "**", "results_*.json")
    filenames = glob.glob(pattern, recursive=True)

    data = []
    for filename in filenames:
        with open(filename) as fp:
            report = json.load(fp)
        user_id, model_id = report["config"]["model_id"].split("/")
        row = {"user_id": user_id, "model_id": model_id, "model_sha": report["config"]["model_sha"]}
        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["iqm_episodic_return"] = iqm(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("iqm_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", "iqm_episodic_return"]]
    return df.values.tolist()


def refresh_dataframes():
    df = get_leaderboard_df()
    all_dfs = [format_df(select_env(df, env_id)) for env_id in all_env_ids]
    return all_dfs


def refresh_videos():
    df = get_leaderboard_df()
    outputs = []
    for env_id in all_env_ids:
        env_df = select_env(df, env_id)
        if not env_df.empty:
            user_id = env_df.iloc[0]["user_id"]
            model_id = env_df.iloc[0]["model_id"]
            model_sha = env_df.iloc[0]["model_sha"]
            repo_id = f"{user_id}/{model_id}"
            video_path = API.hf_hub_download(repo_id=repo_id, filename="replay.mp4", revision=model_sha, repo_type="model")
            outputs.append(video_path)
        else:
            outputs.append(None)
    return outputs


def refresh_winners():
    df = get_leaderboard_df()
    outputs = []
    for env_id in all_env_ids:
        env_df = select_env(df, env_id)
        if not env_df.empty:
            winner = f'{env_df.iloc[0]["user_id"]}/{env_df.iloc[0]["model_id"]}'
            outputs.append(
                f"""## {env_id}

### 🏆 [{winner}](https://huggingface.co/{winner}) 🏆"""
            )
            # # Or in HTML:
            # outputs.append(f'<h3>🏆 <a href="https://huggingface.co/{model}">{model}</a> 🏆</h3>')
        else:
            outputs.append(
                f"""## {env_id}

### 🤷‍♂️ No winner yet"""
            )
    return outputs


HEADING = """
# 🥇 Open RL Leaderboard 🥇

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. It usually takes a few minutes for the evaluation to be completed.
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](https://huggingface.co/spaces/open-rl-leaderboard/leaderboard/discussions/new) on the Open RL Leaderboard repository.

## 🏃 Next steps

We are working on adding more environments and metrics to the Open RL Leaderboard.
If you have any suggestions, please [open an discussion](https://huggingface.co/spaces/open-rl-leaderboard/leaderboard/discussions/new) on the Open RL Leaderboard repository.

## 📜 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}",
}
```
"""

css = """
.generating {
    border: none;
}
h2 {
    text-align: center;
}
h3 {
    text-align: center;
}

"""

with gr.Blocks(css=css) as demo:
    gr.Markdown(HEADING)
    with gr.Tabs(elem_classes="tab-buttons") as tabs:
        with gr.TabItem("🏅 Leaderboard"):
            df = get_leaderboard_df()
            all_env_ids = []
            all_gr_dfs = []
            all_gr_videos = []
            all_gr_winners = []
            for env_domain, env_ids in ALL_ENV_IDS.items():
                with gr.TabItem(env_domain):
                    for env_id in env_ids:
                        # If the env_id envs with "NoFrameskip-v4", we remove it
                        tab_env_id = env_id[: -len("NoFrameskip-v4")] if env_id.endswith("NoFrameskip-v4") else env_id
                        with gr.TabItem(tab_env_id):
                            logger.info(f"Creating tab for {env_id}")
                            with gr.Row(equal_height=False):
                                with gr.Column(scale=3):
                                    # Display the leaderboard
                                    gr_df = gr.components.Dataframe(
                                        headers=["🏆", "🧑 User", "🤖 Model id", "📊 IQM episodic return"],
                                        datatype=["number", "markdown", "markdown", "number"],
                                        row_count=(20, "fixed"),
                                    )
                                with gr.Column(scale=1):
                                    with gr.Row():  # Display the env_id and the winner
                                        gr_winner = gr.Markdown()
                                    with gr.Row():  # Play the video of the best model
                                        gr_video = gr.PlayableVideo(  # Doesn't loop for the moment, see https://github.com/gradio-app/gradio/issues/7689
                                            min_width=50,
                                            autoplay=True,
                                            show_download_button=False,
                                            show_share_button=False,
                                            show_label=False,
                                        )

                                all_env_ids.append(env_id)
                                all_gr_dfs.append(gr_df)
                                all_gr_winners.append(gr_winner)
                                all_gr_videos.append(gr_video)

        with gr.TabItem("📝 About"):
            gr.Markdown(ABOUT_TEXT)

    demo.load(refresh_dataframes, outputs=all_gr_dfs, every=REFRESH_RATE)
    demo.load(refresh_videos, outputs=all_gr_videos, every=REFRESH_RATE)
    demo.load(refresh_winners, outputs=all_gr_winners, every=REFRESH_RATE)


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


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