import gym
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
from q_learning_bins import plot_running_avg

# so you can test different architectures
class HiddenLayer:
    def __init__(self, M1, M2, f=tf.nn.tanh, use_bias=True):
        self.W = tf.Variable(tf.random.normal(shape=(M1, M2)))
        self.use_bias = use_bias
        if use_bias:
            self.b = tf.Variable(np.zeros(M2).astype(np.float32))
        self.f = f

    def forward(self, X):
        if self.use_bias:
            a = tf.matmul(X, self.W) + self.b
        else:
            a = tf.matmul(X, self.W)
        return self.f(a)


# approximates pi(a | s)
class PolicyModel:
    # D = input layer size
    # K = output layer size (number of actions)
    def __init__(self, D, K, hidden_layer_sizes):
        # create the graph
        self.layers = []
        M1 = D
        for M2 in hidden_layer_sizes:
            layer = HiddenLayer(M1, M2)
            self.layers.append(layer)
            M1 = M2

        # final layer, usees softmax
        layer = HiddenLayer(M1, K, tf.nn.softmax, use_bias=False)
        self.layers.append(layer)

        # inputs and targets
        self.X = tf.placeholder(tf.float32, shape=(None, D), name='X')
        self.actions = tf.placeholder(tf.int32, shape = (None,), name = 'actions')
        self.advantages = tf.placeholder(tf.float32, shape = (None,), name = 'advantages')

        # calculate output and cost
        Z = self.X
        for layer in self.layers:
            Z = layer.forward(Z)
        p_a_given_s = Z

        self.predict_op = p_a_given_s

        # get the probabilities of the actions selected
        # e.g.: p_a_given_s = [0.25, 0.4, 0.35], actions = [1, 1, 0, 0, 2, 1]
        # p_a_given_s * tf.one_hot(actions, 3) =
        # [[0, 0.40, 0],
        #  [0, 0.40, 0],
        #  [0.25, 0, 0],
        #  [0.25, 0, 0],
        #  [0, 0, 0.35],
        #  [0, 0.40, 0]]
        # using tf.reduce_sum gives [0.4, 0.4, 0.25, 0.25, 0.35, 0.4]
        selected_probs = tf.log(
            tf.reduce_sum(
                p_a_given_s * tf.one_hot(self.actions, K),
                reduction_indices=[1]
            )
        )

        cost = -tf.reduce_sum(self.advantages * selected_probs)

        self.train_op = tf.train.AdagradOptimizer(1e-1).minimize(cost)

    def set_session(self, session):
        self.session = session

    def partial_fit(self, X, actions, advantages):
        X = np.atleast_2d(X)
        actions = np.atleast_1d(actions)
        advantages = np.atleast_1d(advantages)
        self.session.run(
            self.train_op,
            feed_dict={
                self.X:X,
                self.actions:actions,
                self.advantages:advantages
            }
        )

    def predict(self, X):
        X = np.atleast_2d(X)
        return self.session.run(self.predict_op, feed_dict = {self.X:X})

    def sample_action(self, X):
        p = self.predict(X)[0]
        return np.random.choice(len(p), p=p)

# approximates V(s)
class ValueModel:
    def __init__(self, D, hidden_layer_sizes):
        # create the graph
        self.layers = []
        M1 = D
        for M2 in hidden_layer_sizes:
            layer = HiddenLayer(M1, M2)
            self.layers.append(layer)
            M1 = M2

        # final layer
        layer = HiddenLayer(M1, 1, lambda x:x)
        self.layers.append(layer)

        # inputs and targets:
        self.X = tf.placeholder(tf.float32, shape=(None,D), name='X')
        self.Y = tf.placeholder(tf.float32, shape=(None,), name='Y')

        # calculate output and cost
        Z = self.X
        for layer in self.layers:
            Z = layer.forward(Z)

        Y_hat = tf.reshape(Z, [-1]) # the output
        self.predict_op = Y_hat

        cost = tf.reduce_sum(tf.square(self.Y - Y_hat))

        self.train_op = tf.train.GradientDescentOptimizer(1e-4).minimize(cost)

    def set_session(self, session):
        self.session = session

    def partial_fit(self, X, Y):
        X = np.atleast_2d(X)
        Y = np.atleast_1d(Y)
        self.session.run(
            self.train_op,
            feed_dict={self.X:X, self.Y:Y}
        )

    def predict(self, X):
        X = np.atleast_2d(X)
        return self.session.run(self.predict_op, feed_dict={self.X:X})


def play_one_td(env, p_model, v_model, gamma):
    observation = env.reset()
    done = False
    total_reward = 0
    iters = 0

    while not done:
        action = p_model.sample_action(observation)
        prev_observation = observation
        observation, reward, done, info = env.step(action)

        # update the models
        V_next = v_model.predict(observation)
        G = reward + gamma * np.max(V_next) # Q learning
        advantage = G - v_model.predict(prev_observation)
        p_model.partial_fit(prev_observation, action, advantage)
        v_model.partial_fit(prev_observation, G)

        if reward == 1:
            total_reward += reward
        iters += 1

    return total_reward

def play_one_mc(env, p_model, v_model, gamma):
    observation = env.reset()
    done = False
    total_reward = 0
    iters = 0

    states = []
    actions = []
    rewards = []

    reward = 0
    while not done:
        action = p_model.sample_action(observation)

        states.append(observation)
        actions.append(action)
        rewards.append(reward)

        observation, reward, done, info = env.step(action)

        if done:
            reward = -200

        if reward == 1:
            total_reward += reward
        iters += 1

    # save the final (s, a, r) tuple
    action = p_model.sample_action(observation)
    states.append(observation)
    actions.append(action)
    rewards.append(reward)

    # Calculate returns and advantages
    returns = []
    advantages = []
    G = 0
    for s, r in zip(reversed(states), reversed(rewards)):
        returns.append(G)
        advantages.append(G - v_model.predict(s)[0])
        G = r + gamma * G
    returns.reverse()
    advantages.reverse()

    # update the models
    p_model.partial_fit(states, actions, advantages)
    v_model.partial_fit(states, returns)

    return total_reward

def main():
    env = gym.make('CartPole-v0')
    D = env.observation_space.shape[0]
    K = env.action_space.n
    p_model = PolicyModel(D, K, [10])
    v_model = ValueModel(D, [5, 5])
    init = tf.global_variables_initializer()
    session = tf.InteractiveSession()
    session.run(init)
    p_model.set_session(session)
    v_model.set_session(session)
    gamma = 0.99

    N = 3000
    total_rewards = np.empty(N)
    for n in range(N):
        total_reward = play_one_mc(env, p_model, v_model, gamma)
        total_rewards[n] = total_reward
        if n % 100 == 0:
            print('episode:', n, 'total reward:', total_reward, 'avg reward (last 100):', total_rewards[max(0, n - 100):n+1].mean())

    print('avg reward for last 100 episodes:', total_rewards[-100:].mean())
    print('total steps:', total_rewards.sum())

    plt.plot(total_rewards)
    plt.title("Rewards")
    plt.show()

    plot_running_avg(total_rewards)

main()