from model import *

import sys

import os

import multiprocessing

import threading

import argparse

worker_threads = []

def main():

    parser = argparse.ArgumentParser(description='Train or test neural net motor controller')

    parser.add_argument('--load_model', dest='load_model', action='store_true', default=False)

    parser.add_argument('--num_workers', dest='num_workers',action='store',default=3,type=int)

    parser.add_argument('--visualize', dest='vis', action='store_true', default=False)

    args = parser.parse_args()

    load_model = args.load_model

    num_workers = args.num_workers

    vis = args.vis

    training = True#not load_model

    model_path = './models'

    if not os.path.exists(model_path):

        os.makedirs(model_path)

    # hyperparameters

    explore = 1000

    batch_size = 32

    gamma = 0.995

    n_step = 3

    tf.reset_default_graph()

    with tf.Session() as sess:

        with tf.device("/cpu:0"): 

            global_episodes = tf.Variable(0,dtype=tf.int32,name='global_episodes',trainable=False)

            global_actor_network = ActorNetwork(sess,41+14+3,18,'global'+'/actor')

            num_cpu = multiprocessing.cpu_count() # Set workers ot number of available CPU threads

            workers = []

            # Create worker classes

            for i in range(num_workers):

                worker = Worker(sess,i,model_path,global_episodes,explore,training,vis,batch_size,gamma,n_step,global_actor_network.net)

            workers.append(worker)

            saver = tf.train.Saver()

            coord = tf.train.Coordinator()

            if load_model == True:

                print ('Loading Model...')

                ckpt = tf.train.get_checkpoint_state(model_path)

                saver.restore(sess,ckpt.model_checkpoint_path)

                print ('Loading Model succeeded...')

            else:

                sess.run(tf.global_variables_initializer())

            # This is where the asynchronous magic happens.

            # Start the "work" process for each worker in a separate thread.

            for worker in workers:

                worker_work = lambda: worker.work(coord,saver)

                t = threading.Thread(target=(worker_work))

                t.daemon = True

                t.start()

                worker_threads.append(t)

                sleep(0.05)

            coord.join(worker_threads)

if __name__ == "__main__":

    try:

        main()

    except KeyboardInterrupt:

        print("Ctrl-c received! Sending kill to threads...")

        for t in worker_threads:

            t.kill_received = True