import numpy as np
import pandas as pd
import os
import argparse

import tf_slim
import trfl
from utility import *
from SASRecModules import *

import tensorflow._api.v2.compat.v1 as tf

tf.disable_v2_behavior()


def parse_args():
    parser = argparse.ArgumentParser(description="Run supervised GRU.")

    parser.add_argument('--epoch', type=int, default=30,
                        help='Number of max epochs.')
    parser.add_argument('--data', nargs='?', default='../../data/RC15/',
                        help='data directory')
    parser.add_argument('--batch_size', type=int, default=256,
                        help='Batch size.')
    parser.add_argument('--hidden_factor', type=int, default=64,
                        help='Number of hidden factors, i.e., embedding size.')
    parser.add_argument('--r_click', type=float, default=0.2,
                        help='reward for the click behavior.')
    parser.add_argument('--r_buy', type=float, default=1.0,
                        help='reward for the purchase behavior.')
    parser.add_argument('--lr', type=float, default=0.01,
                        help='Learning rate.')
    parser.add_argument('--num_heads', default=1, type=int)
    parser.add_argument('--num_blocks', default=1, type=int)
    parser.add_argument('--discount', type=float, default=0.5,
                        help='Discount factor for RL.')
    parser.add_argument('--dropout_rate', default=0.1, type=float)

    parser.add_argument('--cuda', type=int, default=1,
                        help='cuda device.')

    parser.add_argument('--num_aug', type=int, default=2,
                        help='Number of state augmentations.')
    parser.add_argument('--item_num', type=int, default=0,
                        help='')
    parser.add_argument('--sigma', type=float, default=0.003,
                        help='Number of state augmentations.')
    parser.add_argument('--q_w', type=float, default=1.0,
                        help='weight of qloss.')

    parser.add_argument('--alpha', type=float, default=0.001,
                        help='')
    parser.add_argument('--beta', type=float, default=0.005,
                        help='')

    return parser.parse_args()


class SASRecnetwork:
    def __init__(self, hidden_size, learning_rate, item_num, state_size, args, name='DQNetwork'):
        self.state_size = state_size
        self.learning_rate = learning_rate
        self.hidden_size = hidden_size
        self.item_num = int(item_num)
        self.is_training = tf.placeholder(tf.bool, shape=())
        self.name = name
        with tf.variable_scope(self.name):
            self.all_embeddings = self.initialize_embeddings()

            self.inputs = tf.placeholder(tf.int32, [None, state_size], name='inputs')
            self.len_state = tf.placeholder(tf.int32, [None], name='len_state')
            self.target = tf.placeholder(tf.int32, [None], name='target')  # target item, to calculate ce loss

            self.noise = tf.placeholder(tf.float32, [args.num_aug, None, self.hidden_size], name='noise')

            self.input_emb = tf.nn.embedding_lookup(self.all_embeddings['state_embeddings'], self.inputs)
            # Positional Encoding
            pos_emb = tf.nn.embedding_lookup(self.all_embeddings['pos_embeddings'],
                                             tf.tile(tf.expand_dims(tf.range(tf.shape(self.inputs)[1]), 0),
                                                     [tf.shape(self.inputs)[0], 1]))
            self.seq = self.input_emb + pos_emb

            mask = tf.expand_dims(tf.to_float(tf.not_equal(self.inputs, item_num)), -1)
            # Dropout
            self.seq = tf.layers.dropout(self.seq,
                                         rate=args.dropout_rate,
                                         training=tf.convert_to_tensor(self.is_training))
            self.seq *= mask

            # Build blocks

            for i in range(args.num_blocks):
                with tf.variable_scope("num_blocks_%d" % i):
                    # Self-attention
                    self.seq = multihead_attention(queries=normalize(self.seq),
                                                   keys=self.seq,
                                                   num_units=self.hidden_size,
                                                   num_heads=args.num_heads,
                                                   dropout_rate=args.dropout_rate,
                                                   is_training=self.is_training,
                                                   causality=True,
                                                   scope="self_attention")

                    # Feed forward
                    self.seq = feedforward(normalize(self.seq), num_units=[self.hidden_size, self.hidden_size],
                                           dropout_rate=args.dropout_rate,
                                           is_training=self.is_training)
                    self.seq *= mask

            self.seq = normalize(self.seq)
            self.state_hidden = extract_axis_1(self.seq, self.len_state - 1)

            self.states_hidden_aug0 = tf.expand_dims(self.state_hidden, 0)
            self.states_hidden_aug1 = tf.concat([self.states_hidden_aug0,
                                                 tf.convert_to_tensor(self.noise) + self.states_hidden_aug0], axis=0)
            self.states_hidden_aug2 = tf.reshape(self.states_hidden_aug1, [-1, self.hidden_size])

            self.output1_ = tf_slim.fully_connected(self.states_hidden_aug2, self.item_num,
                                                    activation_fn=None, scope="q-value")  # all q-values
            self.output2 = tf_slim.fully_connected(self.state_hidden, self.item_num,
                                                   activation_fn=None, scope="ce-logits")  # all logits
            self.output1 = tf.reshape(self.output1_, [args.num_aug + 1, -1, self.item_num])


            # TRFL way
            self.actions = tf.placeholder(tf.int32, [None])
            self.targetQs_ = tf.placeholder(tf.float32, [args.num_aug + 1,None, item_num])
            self.targetQs_selector = tf.placeholder(tf.float32, [args.num_aug + 1, None,
                                                                 item_num])  # used for select best action for double q learning
            self.reward = tf.placeholder(tf.float32, [None])
            self.discount = tf.placeholder(tf.float32, [None])

            self.qloss = 0
            for i in range(args.num_aug + 1):
                qloss, q_learning = trfl.double_qlearning(self.output1[i], self.actions, self.reward, self.discount,
                                                          self.targetQs_[i], self.targetQs_selector[i])
                self.qloss += qloss

            self.celoss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=self.actions, logits=self.output2)
            self.loss = tf.reduce_mean(self.qloss + args.q_w * self.celoss)
            self.opt = tf.train.AdamOptimizer(learning_rate).minimize(self.loss)


    def initialize_embeddings(self):
        all_embeddings = dict()
        state_embeddings = tf.Variable(tf.random_normal([self.item_num + 1, self.hidden_size], 0.0, 0.01),
                                       name='state_embeddings')
        pos_embeddings = tf.Variable(tf.random_normal([self.state_size, self.hidden_size], 0.0, 0.01),
                                     name='pos_embeddings')
        all_embeddings['state_embeddings'] = state_embeddings
        all_embeddings['pos_embeddings'] = pos_embeddings
        return all_embeddings


def evaluate(sess):
    eval_sessions = pd.read_pickle(os.path.join(data_directory, 'val_sessions.df'))
    eval_ids = eval_sessions.session_id.unique()
    groups = eval_sessions.groupby('session_id')
    batch = 100
    evaluated = 0
    total_clicks = 0.0
    total_purchase = 0.0
    total_reward = [0, 0, 0, 0]
    hit_clicks = [0, 0, 0, 0]
    ndcg_clicks = [0, 0, 0, 0]
    hit_purchase = [0, 0, 0, 0]
    ndcg_purchase = [0, 0, 0, 0]
    while evaluated < len(eval_ids):
        states, len_states, actions, rewards = [], [], [], []
        for i in range(batch):
            id = eval_ids[evaluated]
            group = groups.get_group(id)
            history = []
            for index, row in group.iterrows():
                state = list(history)
                len_states.append(state_size if len(state) >= state_size else 1 if len(state) == 0 else len(state))
                state = pad_history(state, state_size, item_num)
                states.append(state)
                action = row['item_id']
                is_buy = row['is_buy']
                reward = reward_buy if is_buy == 1 else reward_click
                if is_buy == 1:
                    total_purchase += 1.0
                else:
                    total_clicks += 1.0
                actions.append(action)
                rewards.append(reward)
                history.append(row['item_id'])
            evaluated += 1
            if evaluated >= len(eval_ids):
                break
        noise = np.zeros([args.num_aug, len(actions), args.hidden_factor])
        prediction = sess.run(SASRec1.output2, feed_dict={SASRec1.inputs: states, SASRec1.len_state: len_states,
                                                          SASRec1.is_training: False, SASRec1.noise: noise})
        sorted_list = np.argsort(prediction)
        calculate_hit(sorted_list, topk, actions, rewards, reward_click, total_reward, hit_clicks, ndcg_clicks,
                      hit_purchase, ndcg_purchase)
    print('#############################################################')
    print('total clicks: %d, total purchase:%d' % (total_clicks, total_purchase))
    for i in range(len(topk)):
        hr_click = hit_clicks[i] / total_clicks
        hr_purchase = hit_purchase[i] / total_purchase
        ng_click = ndcg_clicks[i] / total_clicks
        ng_purchase = ndcg_purchase[i] / total_purchase
        print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
        print('cumulative reward @ %d: %f' % (topk[i], total_reward[i]))
        print('clicks hr ndcg @ %d : %f, %f' % (topk[i], hr_click, ng_click))
        print('purchase hr and ndcg @%d : %f, %f' % (topk[i], hr_purchase, ng_purchase))
    print('#############################################################')


if __name__ == '__main__':
    # Network parameters
    args = parse_args()
    os.environ["CUDA_VISIBLE_DEVICES"] = str(args.cuda)
    data_directory = args.data
    data_statis = pd.read_pickle(
        os.path.join(data_directory, 'data_statis.df'))  # read data statistics, includeing state_size and item_num
    state_size = data_statis['state_size'][0]  # the length of history to define the state
    item_num = data_statis['item_num'][0]  # total number of items
    reward_click = args.r_click
    reward_buy = args.r_buy
    topk = [5, 10, 15, 20]
    # save_file = 'pretrain-GRU/%d' % (hidden_size)

    tf.reset_default_graph()

    SASRec1 = SASRecnetwork(hidden_size=args.hidden_factor, learning_rate=args.lr, item_num=item_num,
                            state_size=state_size, args=args, name='SASRec1')
    SASRec2 = SASRecnetwork(hidden_size=args.hidden_factor, learning_rate=args.lr, item_num=item_num,
                            state_size=state_size, args=args, name='SASRec2')

    replay_buffer = pd.read_pickle(os.path.join(data_directory, 'replay_buffer.df'))
    saver = tf.train.Saver()

    total_step = 0
    gpu_options = tf.GPUOptions(allow_growth=True)
    train_noise_ = tf.random_uniform([args.num_aug, args.batch_size, args.hidden_factor],
                                     minval=args.alpha, maxval=args.beta, seed=2022)
    train_noise_1 = tf.random_uniform([args.num_aug, args.batch_size, args.hidden_factor],
                                      minval=args.alpha, maxval=args.beta, seed=10)
    with tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)) as sess:
        train_noise = train_noise_.eval()
        train_noise1 = train_noise_1.eval()
        # Initialize variables
        sess.run(tf.global_variables_initializer())
        # evaluate(sess)
        num_rows = replay_buffer.shape[0]
        num_batches = int(num_rows / args.batch_size)
        for i in range(args.epoch):
            for j in range(num_batches):
                batch = replay_buffer.sample(n=args.batch_size).to_dict()

                next_state = list(batch['next_state'].values())
                len_next_state = list(batch['len_next_states'].values())
                # double q learning, pointer is for selecting which network  is target and which is main
                pointer = np.random.randint(0, 2)
                if pointer == 0:
                    mainQN = SASRec1
                    target_QN = SASRec2
                else:
                    mainQN = SASRec2
                    target_QN = SASRec1
                target_Qs = sess.run(target_QN.output1,
                                     feed_dict={target_QN.inputs: next_state,
                                                target_QN.len_state: len_next_state,
                                                target_QN.is_training: True,
                                                target_QN.noise: train_noise})
                target_Qs_selector = sess.run(mainQN.output1,
                                              feed_dict={mainQN.inputs: next_state,
                                                         mainQN.len_state: len_next_state,
                                                         mainQN.is_training: True,
                                                         mainQN.noise: train_noise})

                # Set target_Qs to 0 for states where episode ends
                is_done = list(batch['is_done'].values())
                for index in range(target_Qs.shape[1]):
                    if is_done[index]:
                        target_Qs[:, index, :] = np.zeros([item_num])

                state = list(batch['state'].values())
                len_state = list(batch['len_state'].values())
                action = list(batch['action'].values())
                is_buy = list(batch['is_buy'].values())
                reward = []
                for k in range(len(is_buy)):
                    reward.append(reward_buy if is_buy[k] == 1 else reward_click)
                discount = [args.discount] * len(action)

                loss, _ = sess.run([mainQN.loss, mainQN.opt],
                                   feed_dict={mainQN.inputs: state,
                                              mainQN.len_state: len_state,
                                              mainQN.targetQs_: target_Qs,
                                              mainQN.reward: reward,
                                              mainQN.discount: discount,
                                              mainQN.actions: action,
                                              mainQN.targetQs_selector: target_Qs_selector,
                                              mainQN.is_training: True,
                                              mainQN.noise: train_noise1})
                total_step += 1
                if total_step % 200 == 0:
                    print("the loss in %dth batch is: %f" % (total_step, loss))
                if total_step % 2000 == 0:
                    evaluate(sess)
