import linear_env_small as linear_env
import sim_env
from actor import Actor
from critic import Critic
from replay_buffer import ReplayBuffer
import numpy as np
import tensorflow as tf
import keras.backend as kbck
import json
import time
import argparse
import matplotlib.pylab as plt
import os.path
def ou(x, mu, theta, sigma):
    return theta * (mu - x) + sigma * np.random.randn(np.shape(x)[0])

def simulate(control, swmm ,flows):   
    best_reward = -1*np.inf
    BUFFER_SIZE = 100000
    BATCH_SIZE = 120
    GAMMA = 0.99
    TAU = 0.01     #Target Network HyperParameters
    LRA = 0.0001    #Learning rate for Actor
    LRC = 0.001     #Lerning rate for Critic

    action_dim = 4  
    state_dim = 5  
    max_steps = 6000
    np.random.seed(9501)

    EXPLORE = 100000.
    episode_count = 1000
    done = False
    step = 0
    epsilon = 1
    
    if swmm:
        print("No support")
    else:
        # Constants for the linear environment
        Hs = 2400
        A1 = 0.0020
        mu1 = 250
        sigma1 = 70
        A2 = 0.0048
        mu2 = 250
        sigma2 = 70
        dt = 1
        x = np.arange(Hs)
        d = np.zeros((2,Hs))
        
        if control:
            #Tensorflow GPU optimization
            config = tf.ConfigProto()
            config.gpu_options.allow_growth = True
            sess = tf.Session(config=config)
            kbck.set_session(sess)
            
            # Actor, critic and replay buffer creation
            actor = Actor(sess, state_dim, action_dim, BATCH_SIZE, TAU, LRA,flows)
            critic = Critic(sess, state_dim, action_dim, BATCH_SIZE, TAU, LRC)
            buff = ReplayBuffer(BUFFER_SIZE)


            # Get the linear environment

            reward_hist = []
            for i in range(episode_count):

                print("Episode : " + str(i) + " Replay Buffer " + str(buff.count()))
                A1 += 0.0004*np.random.rand()
                mu1 += 50*np.random.rand()
                sigma1 += 14*np.random.rand()
                A2 += 0.00096*np.random.rand()
                mu2 += 50*np.random.rand()
                sigma2 += 14*np.random.rand()
                d[0,:] = A1*np.exp((-1*(x-mu1)**2)/(2*sigma1**2))
                d[1,:] = A2*np.exp((-1*(x-mu2)**2)/(2*sigma2**2))
                vref = np.zeros((state_dim,))
                env  = linear_env.env(dt,d,vref)
                s_t = np.divide(env.reset(),env.vmax)
                total_reward = 0.
                
                for j in range(max_steps):
                    
                    ## Noise addition for exploration
                    ## Ornstein-Uhlenbeck process
                    loss = 0 
                    epsilon -= 1.0 / EXPLORE
                    a_t = np.zeros([1,action_dim])
                    noise_t = np.zeros([1,action_dim])
                    a_t_original = actor.munet.predict(s_t.reshape(1, s_t.shape[0]))  
                    noise_t[0,:] =  max(epsilon, 0) * ou(a_t_original[0,:], 0.5 , 1 , 1.5)
                    #noise_t[0,4:] =  max(epsilon, 0) * ou(a_t_original[0,4:],  0.5 , 1 , 1.5)
                    
                    a_t[0]  = a_t_original[0] + noise_t[0]
                    
                    #Act over the system and get info of the next states
                    s_t1 , r_t, done, _ = env.step(a_t[0],flows=flows)
                    s_t1 = np.divide(s_t1,env.vmax)

                    
                    #Add replay buffer
                    buff.add(s_t, a_t[0], r_t, s_t1, done)   
                    
                    #Do the batch update
                    batch = buff.getBatch(BATCH_SIZE)
                    states = np.asarray([e[0] for e in batch])
                    actions = np.asarray([e[1] for e in batch])
                    rewards = np.asarray([e[2] for e in batch])
                    next_states = np.asarray([e[3] for e in batch])
                    dones = np.asarray([e[4] for e in batch])
                    
                    # Get estimated q-values of the pair (next_state,mu(next_state))
                    actions_next = actor.target_munet.predict(next_states)
                    target_q_values = critic.target_qnet.predict([next_states, actions_next]) 
                    y_t = np.zeros(np.shape(actions))
                
                    for k in range(len(batch)):
                        if dones[k]:
                            y_t[k] = rewards[k]
                        else:
                            y_t[k] = rewards[k] + GAMMA*target_q_values[k]
            

                    loss += critic.qnet.train_on_batch([states,actions], y_t) 
                    a_for_grad = actor.munet.predict(states)
                    grads = critic.gradients(states, a_for_grad)
                    actor.train(states, grads)
                    actor.target_train()
                    critic.target_train()

                    total_reward = total_reward + GAMMA*r_t
                    s_t = s_t1
                    if j%100==0:
                        print("Episode", i, "Step", j, "Reward", r_t, "Loss", loss)

                    if done:
                        break

                reward_hist.append(total_reward)
                np.save("reward_small_history_flows_"+str(flows).lower()+".npy",np.array(reward_hist))

                if i%20 == 0:
                    print("Saving the networks...")
                    actor.munet.save_weights("./actors_small/anetwork_flows_"+str(flows).lower()+"_it_"+str(i)+".h5", overwrite=True)
                    critic.qnet.save_weights("./critics_small/cnetwork_flows_"+str(flows).lower()+"_it_"+str(i)+".h5", overwrite=True)

                if total_reward > best_reward:
                    print("Saving Best Actor...")
                    np.save("best_reward"+"_flows_"+str(flows)+".npy",np.array(total_reward))
                    actor.munet.save_weights("./actors_small/best_anetwork_flows_"+str(flows).lower()+".h5", overwrite=True)
                    critic.qnet.save_weights("./critics_small/best_cnetwork_flows_"+str(flows).lower()+".h5", overwrite=True)            
                    best_reward = total_reward
            
                print("TOTAL REWARD @ " + str(i) +"-th Episode  : Reward " + str(total_reward))
                print("Total Step: " + str(step))
                print("")