import gym
import torch
from easydict import EasyDict

from ding.config import compile_config
from ding.envs import DingEnvWrapper
from ding.model import DQN
from ding.policy import DQNPolicy, single_env_forward_wrapper
from dizoo.cliffwalking.config.cliffwalking_dqn_config import create_config, main_config
from dizoo.cliffwalking.envs.cliffwalking_env import CliffWalkingEnv


def main(main_config: EasyDict, create_config: EasyDict, ckpt_path: str):
    main_config.exp_name = f'cliffwalking_dqn_seed0_deploy'
    cfg = compile_config(main_config, create_cfg=create_config, auto=True)
    env = CliffWalkingEnv(cfg.env)
    env.enable_save_replay(replay_path=f'./{main_config.exp_name}/video')
    model = DQN(**cfg.policy.model)
    state_dict = torch.load(ckpt_path, map_location='cpu')
    model.load_state_dict(state_dict['model'])
    policy = DQNPolicy(cfg.policy, model=model).eval_mode
    forward_fn = single_env_forward_wrapper(policy.forward)
    obs = env.reset()
    returns = 0.
    while True:
        action = forward_fn(obs)
        obs, rew, done, info = env.step(action)
        returns += rew
        if done:
            break
    print(f'Deploy is finished, final epsiode return is: {returns}')


if __name__ == "__main__":
    main(
        main_config=main_config,
        create_config=create_config,
        ckpt_path=f'./cliffwalking_dqn_seed0/ckpt/ckpt_best.pth.tar'
    )