import numpy as np
import tensorflow as tf
import random


class PolicyGradientAgent(tf.keras.Model):
    def __init__(self, datum_size, num_stocks, past_num):
        """
        this class inherits from tf.keras.Model
        a general class of policy gradient RL agents using actor-critic
        assume this agent can manage num_stocks stocks at the same time, and the agent's actions have no impact on
        the stock environment. When performing each action, assume the agent can only sell/buy/hold one of the
        num_stocks stocks.

        (outdated, don't read, kept for documentation purpose)
        # the agent's action space is a tensor of size [3 * num_stocks], where indexes [3i, 3i+1, 3i+2] refers to the
        # i-th stock (zero-indexing). Each element of this tensor is in [0, 1] and all elements should sum up to 1 (this
        # tensor is a probability of taking each action). For the i-th stock, action[3i] is the probability of holding,
        # action[3i+1] is the probability of buying, action[3i+2] is the probability of selling.
        # To choose an action, we sample from this probability tensor and get an index k. if (k mod 3) == 0, then the
        # agent will do nothing. if (k mod 3) == 1, then the agent will buy (action[k] * self.max_buy) amount of the
        # corresponding stock. If (k mod 3) == 2, then the agent will sell (action[k] * self.max_sell) amount of the
        # corresponding stock.

        :param datum_size: the size of the state space, passed in from pre-processing
        :param num_stocks: the number of stocks the agent manages at once
        :param past_num: the past number of days to consider when making the current decision
        """
        super(PolicyGradientAgent, self).__init__()

        # stock market decision hyper params
        self.max_sell = 100  # the dollar amount the agent is willing to sell at most in one action
        self.max_buy = 100  # the dollar amount the agent is willing to buy at most in one action
        self.capital = 1000  # the initial amount of capital the agent holds
        self.past_num = past_num  # the number of days of historical financial data the agent will consider when
                                    # making the current decision.

        # RL agent params
        self.datum_size = datum_size
        self.num_stocks = num_stocks
        self.batch_size = 60
        self.buffer = []  # initialize the memory replay buffer
        self.buffer_size = 100  # maximum episodes the buffer can hold
        self.buffer_num_elt = 0  # the number of current elements in the buffer
        self.buffer_episode_lens = []  # a list of `episode length` of experience that were stored in the buffer
                                        # used to facilitate self.forget()
        self.num_actions = 3 * self.num_stocks
        self.actor_H1 = 64  # hidden layer output sizes
        self.actor_H2 = 64
        self.actor_H3 = 64
        self.actor_H4 = 64
        self.critic_H1 = 64
        self.critic_H2 = 32
        # self.lr_schedule = tf.keras.optimizers.schedules.PiecewiseConstantDecay(boundaries=[30, 110, 220, 300],
        #                                                                         values=[0.01, 0.005, 0.003, 0.002,
        #                                                                                 0.001])
        # self.lr_schedule = 0.003
        self.lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(initial_learning_rate=0.01,
                                                                          decay_rate=0.98, decay_steps=100000)
        self.optimizer = tf.keras.optimizers.Adam(learning_rate=self.lr_schedule)

        # model layers
        initializer = tf.keras.initializers.GlorotUniform()
        self.lift = tf.keras.layers.Dense(datum_size * num_stocks * 3)
        self.actor_gru_1 = tf.keras.layers.GRU(self.actor_H1, return_sequences=True, return_state=True, kernel_initializer=initializer)
        self.actor_dropout_1 = tf.keras.layers.Dropout(rate=0.1)
        self.actor_gru_2 = tf.keras.layers.GRU(self.actor_H2, kernel_initializer=initializer)
        self.actor_dropout_2 = tf.keras.layers.Dropout(rate=0.1)
        self.actor_dense_1 = tf.keras.layers.Dense(self.actor_H3, activation='relu', kernel_initializer=initializer)
        # self.actor_dense_2 = tf.keras.layers.Dense(self.actor_H4, activation='relu', kernel_initializer=initializer)
        self.actor_dense_3 = tf.keras.layers.Dense(self.num_actions, kernel_initializer=initializer)
        self.normalizor = tf.keras.layers.LayerNormalization()
        self.critic_dense_1 = tf.keras.layers.Dense(self.critic_H1, activation='relu', kernel_initializer=initializer)
        self.critic_dense_2 = tf.keras.layers.Dense(self.critic_H2, activation='relu', kernel_initializer=initializer)
        self.critic_dense_3 = tf.keras.layers.Dense(1, kernel_initializer=initializer)

    def call(self, states):
        """
        performs the forward pass on a bunch of states, generate probabilities for taking each action
        uses 2 GRU to summarize sequential stock data, and a dense layer to classify to actions

        :param states: a batch of state tuples (self.batch_sz, ), where each tuple is (<price_history>, <portfolio>)
                        where price_history is of shape (model.num_stocks, model.past_num, datum_size),
                        and portfolio is of shape (model.num_stocks+1,)
        :return a tensor of size (batch_sz, self.num_stock, 3), representing the action probability of taking each
                action. The tensor always sums to 1 on the last dimension (the prob of buying/selling/holding a
                particular stock on a particular day sums to one)
                a tensor that summarize the states tensor of shape [batch_sz * self.actor_H2]
        """
        # extract info from states
        price_history = [state[0] for state in states]  # (batch_sz, num_stocks, past_num, datum_size)
        price_history = tf.transpose(price_history, perm=[0, 2, 1, 3])  # (batch_sz, past_num, num_stocks, datum_size)
        price_history = tf.reshape(price_history, (-1, self.past_num, self.num_stocks * self.datum_size))  # (batch_sz, past_num, num_stocks * datum_size)
        portfolio = [state[1] for state in states]  # (batch_sz, num_stock + 1)

        # pass through layers
        price_history = self.lift(price_history)
        gru_1_out_whole_seq, _ = self.actor_gru_1(price_history)  # (batch_sz, past_num, actor_H1)
        gru_1_out_whole_seq = self.actor_dropout_1(gru_1_out_whole_seq)
        gru_2_out = self.actor_gru_2(gru_1_out_whole_seq)  # (batch_sz, actor_H2)
        gru_2_out = self.actor_dropout_2(gru_2_out)
        past_and_current_info = tf.concat([gru_2_out, portfolio], axis=1)  # (batch_sz, actor_H2 + num_stock + 1)

        actor_out = self.actor_dense_1(past_and_current_info)
        # actor_out = self.actor_dense_2(actor_out)
        actor_out = self.actor_dense_3(actor_out)  # (batch_sz * self.num_actions)

        # reshape the output and softmax
        actor_out = tf.reshape(actor_out, (-1, self.num_stocks, 3))
        actor_out = self.normalizor(actor_out)
        action_probs = tf.nn.softmax(actor_out)  # softmaxes across last dimension only
        return action_probs, past_and_current_info

    def value(self, states_summary):
        """
        this is the critic, it estimates the value of being in each state using the value network
        pass the summary of states info through 3 dense layers

        :param states_summary: a batch of summary of input states (self.batch_sz, self.num_actions), this is the second
                                output of the call function
        :return a tensor of size (batch_sz, 1) representing the value of each inputted states
        """
        hidden_1 = self.critic_dense_1(states_summary)  # (batch_sz * self.critic_H1)
        hidden_2 = self.critic_dense_2(hidden_1)  # (batch_sz * self.critic_H2)
        value_out = self.critic_dense_3(hidden_2)  # (batch_sz * 1)
        return value_out

    def loss(self, states, actions_taken, discounted_reward):
        """
        computes the loss of a single forward pass for the agent

        :param states: a batch of state tuples (self.batch_sz, ), where each tuple is (<price_history>, <portfolio>)
                        where price_history is of shape (model.num_stocks, model.past_num, datum_size),
                        and portfolio is of shape (model.num_stocks+1,)
        :param actions_taken: a batch of the sequence of actions that the agent actually took in the episode.(batch_sz,)
                                Each action is of shape (model.num_stocks,) containing values of
                                0 (hold), 1 (buy), or 2 (sell), corresponding to the action taken for a stock.
                              if batch_sz==2 and num_stocks==3, example: [[2, 2, 1], [1, 0, 2]]
        :param discounted_reward: discounted rewards through the batch. (batch_sz, )
        :return a scalar loss of the whole batch
        """
        action_probs, states_summary = self.call(states)  # action_probs: (batch_sz, self.num_stock, 3)
        values = self.value(states_summary)  # (batch_sz, 1)
        values = tf.reshape(values, (-1))  # (batch_sz,)
        hot_tens = tf.one_hot(actions_taken, 3)  # depth of 3, since we have 3 possible actions
        # only desired indices will be multiplied by 1, others are multiplied by 0
        filtered_tens = tf.multiply(action_probs, hot_tens) 
        # add the last dimension to get rid of 0's from filtered_tens. result: (batch_sz, num_stocks)
        # probs_action_taken_each_stock = tf.keras.backend.sum(filtered_tens, axis=-1)  # option1
        probs_action_taken_each_stock = tf.reduce_sum(filtered_tens, axis=-1, keepdims=True)  # option2
        # probs_of_action_taken = tf.reduce_sum(probs_action_taken_each_stock, axis=1)
        # assume each stock is independent
        probs_of_action_taken = tf.reduce_prod(probs_action_taken_each_stock, axis=1)
        probs_of_action_taken = tf.reshape(probs_of_action_taken, (probs_of_action_taken.shape[0],)) #option2
        
        advantage = discounted_reward - values  # (batch_sz,)
        actor_loss = - tf.reduce_sum(tf.math.multiply(tf.math.log(probs_of_action_taken), tf.stop_gradient(advantage)))
        critic_loss = tf.reduce_sum(tf.math.square(advantage))
        loss = actor_loss + critic_loss  # scalar loss of the batch
        return loss

    def remember(self, states, actions, discounted_rewards):
        """
        put an entire episode of states, actions, rewards into the memory-replay buffer
        call this function right after an entire episode of (s,a,r) is generated to remember the episode

        :param states: an episode_len of state tuples (self.batch_sz, ), where each tuple is
                        (<price_history>, <portfolio>), where price_history is of shape
                        (model.num_stocks, model.past_num, datum_size), and portfolio is of shape (model.num_stocks+1,)
        :param actions: an episode_len of the sequence of actions that the agent actually took in the episode.
                        (episode_len, self.num_stocks)
                        Each action is of shape (model.num_stocks,) containing values of 0 (hold), 1 (buy), or 2 (sell),
                        corresponding to the action taken for a stock.
        :param discounted_rewards: the sequence of rewards received in an episode. (episode_len, )
        :return Nothing
        """
        if self.buffer_num_elt >= self.buffer_size:
            self.forget()
        episode_len = len(states)
        list_of_pairs = list(zip(states, actions, discounted_rewards))
        self.buffer += list_of_pairs
        self.buffer_episode_lens.append(episode_len)

    def forget(self):
        """
        get rid of the earliest episode in the memory buffer, because the experience is no longer useful.
        """
        earliest_episode_len = self.buffer_episode_lens.pop(0)  # number of elements to remove
        del self.buffer[:earliest_episode_len]  # remove the first n elements

    def experience_replay(self):
        """
        sample from the current memory buffer to get (state, action, discounted_reward) pairs and train on these
        experience replay will make the stock data more iid and stablize training
        This function is used to sample from self.buffer

        :return (states, actions_taken, discounted_reward), each of which is a list of length [batch_sz]
                if the buffer is large enough, else the length is [self.buffer_size]
        """
        if len(self.buffer) < self.batch_size:
            # not enough experience to sample from, just use all exp
            states, actions_taken, discounted_reward = list(zip(*self.buffer))  # unzip the buffer
            return states, actions_taken, discounted_reward
        else:
            # enough experience to sample from
            list_of_pairs = random.sample(self.buffer, self.batch_size)  # sample batch_size of samples from buffer
            states, actions_taken, discounted_reward = list(zip(*list_of_pairs))
            return states, actions_taken, discounted_reward


def save_model(model, file_path):
    """
    save the weights of a trained model to the current directory

    :param model: a trained model
    :param file_path: a string, specifying the file path to save the model
    """
    model.save(file_path)


def load_model(file_path):
    """
    load the model with a previously trained model

    :param file_path: a string, specifying the file path to the saved model
    :return model: a trained model
    """
    model = tf.keras.models.load_model(file_path)
    return model
