from model import *
import argparse
import sys
import os
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=1,type=int)
    args = parser.parse_args()

    max_episode_length = 1000
    gamma = .995 # discount rate for advantage estimation and reward discounting
    s_size = 41
    a_size = 18 # Agent can move Left, Right, or Straight
    model_path = './models'
    load_model = args.load_model
    noisy=False
    num_workers = args.num_workers
    print(" num_workers = %d" % num_workers)
    print(" noisy_net_enabled = %s" % str(noisy))
    print(" load_model = %s" % str(args.load_model))

    tf.reset_default_graph()

    if not os.path.exists(model_path):
        os.makedirs(model_path)
        

    with tf.device("/cpu:0"): 
        global_episodes = tf.Variable(0,dtype=tf.int32,name='global_episodes',trainable=False)
        trainer = tf.train.AdamOptimizer(learning_rate=1e-4)
        master_network = AC_Network(s_size,a_size,'global',None,noisy) # Generate global network
        num_cpu = multiprocessing.cpu_count() # Set workers ot number of available CPU threads
        workers = []
            # Create worker classes
        for i in range(args.num_workers):
            worker = Worker(i,s_size,a_size,trainer,model_path,global_episodes,noisy,is_training= True)
            workers.append(worker)

        saver = tf.train.Saver()
        
    '''networks = ['global'] + ['worker_'+i for i in str(range(num_workers))]
    print(networks)'''
    #key = print_tensors_in_checkpoint_file('./tmp/checkpoints/mobilenet_v1_0.50_160.ckpt', tensor_name='',all_tensors=True)
    #print(key)
    
    with tf.Session() as sess:
        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:
            '''
            dict = {}
            value = slim.get_model_variables('global'+'/MobilenetV1')
            for variable in value:
                name = variable.name.replace('global'+'/','').split(':')[0]
                    #print(name)
                if name in key:
                    dict[name] = variable
                #print(dict)
                #print(dict)
            init_fn = slim.assign_from_checkpoint_fn(
                                os.path.join(checkpoints_dir, 'mobilenet_v1_0.50_160.ckpt'),
                                dict)
            init_fn(sess)'''
            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(max_episode_length,gamma,sess,coord,saver)
	    #worker.start(setting=0,vis=True)
            t = threading.Thread(target=(worker_work))
            t.daemon = True
            t.start()
            worker_threads.append(t)
        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