# %%
import pathlib
import pickle
import tempfile

import stable_baselines3 as sb3
from stable_baselines3.common.env_util import make_vec_env

from imitation.algorithms import adversarial, bc
from imitation.data import rollout
from imitation.util import logger

from intersim.envs.intersimple import IntersimpleReward, speed_reward

from gail.discriminator import MlpDiscriminator
import numpy as np
import functools
from stable_baselines3.common.evaluation import evaluate_policy
from ray import tune
import os
import torch

model_name = 'gail_flat'

# %%
# Load pickled test demonstrations.
#with open("data/NormalizedIntersimpleExpert_IntersimpleRewardAgent51.pkl", "rb") as f:
with open("data/NormalizedIntersimpleExpert_IntersimpleRewardAgent51Mu.001.pkl", "rb") as f:
    # This is a list of `imitation.data.types.Trajectory`, where
    # every instance contains observations and actions for a single expert
    # demonstration.
    trajectories = pickle.load(f)

# %%
# Convert List[types.Trajectory] to an instance of `imitation.data.types.Transitions`.
# This is a more general dataclass containing unordered
# (observation, actions, next_observation) transitions.
transitions = rollout.flatten_trajectories(trajectories)

venv = make_vec_env(IntersimpleReward, n_envs=2, env_kwargs={'agent': 51})

tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
tempdir_path = pathlib.Path(tempdir.name)
print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")

def training_function(config, checkpoint_dir=None):
    logger.configure(tempdir_path / "GAIL/")

    discriminator = MlpDiscriminator()
    if checkpoint_dir:
        discriminator.load_state_dict(torch.load(os.path.join(checkpoint_dir, 'disc_checkpoint')))
        generator = sb3.PPO.load(os.path.join(checkpoint_dir, 'gen_checkpoint'))
    else:
        generator = sb3.PPO("MlpPolicy", venv, verbose=1, n_steps=config['n_steps'])

    gail_trainer = adversarial.GAIL(
        venv,
        expert_data=transitions,
        expert_batch_size=config['expert_batch_size'],
        n_disc_updates_per_round=config['n_disc_updates_per_round'],
        discrim_kwargs={'discrim_net': MlpDiscriminator()},
        gen_algo=generator,
        allow_variable_horizon=True,
    )

    def callback(epoch):
        print("callback")
        eval_env = IntersimpleReward(agent=51, reward=functools.partial(speed_reward, collision_penalty=0.))
        #sync_envs_normalization(self.training_env, self.eval_env)
        episode_rewards, episode_lengths = evaluate_policy(generator, eval_env, return_episode_rewards=True)
        tune.report(
            reward=np.mean(episode_rewards),
            length=np.mean(episode_lengths),
            training_iteration=epoch,
        )

        with tune.checkpoint_dir(step=epoch) as checkpoint_dir:
            gail_trainer.gen_algo.save(os.path.join(checkpoint_dir, 'gen_checkpoint'))
            torch.save(discriminator.state_dict(), os.path.join(checkpoint_dir, 'disc_checkpoint'))

    gail_trainer.train(total_timesteps=40000, callback=callback)

analysis = tune.run(
    training_function,
    config = {
        'expert_batch_size': tune.randint(1, 22), #220,
        'n_disc_updates_per_round': tune.randint(2, 100), #16,
        'n_steps': tune.randint(1, 10000), #4096,
    },
    resources_per_trial={
        'cpu': 1,
        # 'gpu': 1,
    },
    local_dir='ray',
    num_samples=10,
)

print('Best config', analysis.get_best_config(metric='progress', mode='max'))

#del gail_trainer

# %%
model = sb3.PPO.load(model_name)

env = IntersimpleReward(agent=51)

obs = env.reset()
while True:
    action, _states = model.predict(obs)
    obs, rewards, done, info = env.step(action)
    env.render(mode='post')
    if done:
        break

env.close(filestr='render/'+model_name)