import roboschool
import gym
import numpy as np
import cv2
import torch
import os
import subprocess
import threading
from socket import *
import matplotlib.pyplot as plt
from torchup.logging.factory import LogFactory
from torchup.agents.DoubleDQN import DoubleDQNAgent
from torchup.models import DuelingDQNModel
from torchup.wrappers import (FrameHistoryWrapper,
                              FrameSkipWrapper,
                              TensorboardWrapper,
                              VideoRecorderWrapper,
                              DiscretizeWrapper,
                              DQNImgWrapper)


port = 6000

# meta data
meta_data = {
    'env': 'RoboschoolInvertedPendulumSwingup-v1',
    'experiment': 'discrete to continuous',
    'frame_skip': 4,
    'history': 4,
    'action_space': None,
    'port': None,
    'training_frames': int(10e7)
}

if __name__ == '__main__':
    s = socket(AF_INET, SOCK_STREAM, 0)
    logger = LogFactory()
    env = gym.make('RoboschoolInvertedPendulumSwingup-v1')
    env = DiscretizeWrapper(env, n_actions=3)
    env = TensorboardWrapper(
            FrameHistoryWrapper(
                FrameSkipWrapper(
                    DQNImgWrapper(
                        VideoRecorderWrapper(env, directory=logger.exp_dir_path)))),
            logger=logger.writer)

    meta_data['action_space'] = int(env.action_space.n)

    # Configure agent and logging
    agent = DoubleDQNAgent(model=DuelingDQNModel, env=env)
    logger.agent_cfg(agent)

    agent.logger = logger

    # Tensorboard
    meta_data['port'] = port
    logger.add_meta_data(meta_data)

    args = ['tensorboard',
            '--logdir={}'.format(logger.exp_dir_path),
            '--port={}'.format(port)]
    p = subprocess.Popen(args)

    agent.train()
