import os
import json

import click

import gym

import pandas as pd
from pyvirtualdisplay import Display
from gym.wrappers.monitoring.video_recorder import VideoRecorder

from rampwf.utils.importing import import_module_from_source


def make_video(root_dir, env, submission, agent, seed, epoch):
    """Make a video from a trace.

    The trace is read from
    <env>/submissions/<submission>/mbrl_outputs/<agent>/seed_<seed>/
    epoch_<epoch>/trace.csv

    and the generated video saved in
    <env>/submissions/<submission>/mbrl_outputs/<agent>/seed_<seed>/
    epoch_<epoch>/video/

    Parameters
    ----------
    submission : string
        Submission name.
    agent : string
        Name of the agent.
    seed : int
        The seed used to generate the trace.
    epoch : int
        The epoch of the trace used to generate the video.
    """

    display = Display(visible=False, size=(1400, 900))
    _ = display.start()
    env_path = os.path.join(env, 'env.py')
    env_module = import_module_from_source(env_path, 'env')
    env = env_module.Env()

    epoch_path = os.path.join(root_dir,
        'submissions', submission, 'mbrl_outputs',
        agent, f'seed_{seed}', f'epoch_{epoch}')
    data_path = os.path.join(epoch_path, 'trace.csv')

    video_path = os.path.join(epoch_path, 'video.mp4')


    env.reset()

    data = pd.read_csv(data_path)

    metadata_path = os.path.join("data", "metadata.json")
    with open(metadata_path) as f:
        metadata = json.load(f)

    action_names = metadata["action"]
    X = data.dropna(axis=0)

    states_mask = X.columns.str.startswith("state_")

    video = VideoRecorder(env, video_path)
    # returns an initial observation
    env.reset()

    for i in range(len(X) - 1):
        env.render()
        video.capture_frame()
        current_state = X.loc[:, states_mask].to_numpy()[i]
        env.set_numpy_state(current_state)
        current_action = X[action_names].to_numpy()[i]
        env.step(current_action)

    video.close()
    env.close()


@click.command()
@click.option("--root_dir", default=".", show_default=True,
              type=click.STRING,
              help="Root directory from where to search for submission and agent data trace")
@click.option("--env", required=True,
              type=click.STRING,
              help="Env path to create the video")
@click.option("--submission", default="real_system", show_default=True,
              type=click.STRING,
              help="Model submission. Choose 'real_system' if you want to "
                   "use the real environment.")
@click.option('--agent', default='random_shooting', show_default=True,
              type=click.STRING, help="Agent.")
@click.option("--seed", default=0, show_default=True,
              help="The seed used to generate the trace.")
@click.option("--epoch", default=0, show_default=True,
              help="The epoch of the trace used to generate the video.")

def make_video_command(root_dir, env, submission, agent, seed, epoch):
    return make_video(root_dir, env, submission, agent, seed, epoch)


if __name__ == '__main__':


    make_video_command()
