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

import tf_slim
import trfl
from utility import *
from Modules import *

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

tf.disable_v2_behavior()


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

    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.001,
                        help='Learning rate.')
    parser.add_argument('--num_heads', default=1, type=int)
    parser.add_argument('--num_blocks', default=1, type=int)
    parser.add_argument('--hidden_dropout_prob', default=0.1, type=float)
    parser.add_argument('--cuda', type=int, default=1,
                        help='cuda device.')
    parser.add_argument('--q_w', type=float, default=1.0,
                        help='weight of qloss.')

    parser.add_argument("--no_filters", action="store_true",
                        help="if no filters, filter layers transform to self-attention")
    return parser.parse_args()


class FMLPRecModel:
    def __init__(self, hidden_size, learning_rate, item_num, state_size):
        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=())

        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

        # Embedding Layer
        self.input_emb = tf.nn.embedding_lookup(all_embeddings['state_embeddings'], self.inputs)
        # Positional Encoding
        self.pos_emb = tf.nn.embedding_lookup(all_embeddings['pos_embeddings'],
                                              tf.tile(tf.expand_dims(tf.range(tf.shape(self.inputs)[1]), 0),
                                                      [tf.shape(self.inputs)[0], 1]))
        self.sequence_emb = self.input_emb + self.pos_emb
        # Normalize & DropOut
        self.sequence_emb = normalize(self.sequence_emb)
        self.sequence_emb = tf.layers.dropout(self.sequence_emb, rate=args.hidden_dropout_prob)

        attention_mask = tf.cast((self.inputs < item_num), dtype=tf.float32)
        extended_attention_mask = tf.expand_dims(tf.expand_dims(attention_mask, dim=1), dim=2)
        max_len = attention_mask.shape[-1]  # 64
        attn_shape = (1, max_len, max_len)  # (1,64,64)
        subsequent_mask = tf.expand_dims(tf.linalg.band_part(tf.ones(attn_shape), -1, 0), dim=1)
        extended_attention_mask = extended_attention_mask * subsequent_mask
        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0

        self.complex_weight = tf.Variable(tf.random_normal([1, self.state_size // 2 + 1, args.hidden_factor, 2],
                                                           0.0, 1.0), name='complex_weight')

        item_encoded_layers = item_encoder(args, self.sequence_emb, extended_attention_mask, self.complex_weight,
                                           output_all_encoded_layers=True)
        sequence_output = item_encoded_layers[-1]

        self.state_hidden = extract_axis_1(sequence_output, self.len_state - 1)

        self.output = tf_slim.fully_connected(self.state_hidden, self.item_num, activation_fn=None, scope='fc')

        self.loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=self.target, logits=self.output)
        self.loss = tf.reduce_mean(self.loss)
        self.opt = tf.train.AdamOptimizer(self.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
        prediction = sess.run(FMLPRec.output, feed_dict={FMLPRec.inputs: states, FMLPRec.len_state: len_states,
                                                        FMLPRec.is_training: False})
        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, including 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()

    FMLPRec = FMLPRecModel(hidden_size=args.hidden_factor, learning_rate=args.lr, item_num=item_num,
                           state_size=state_size)

    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)
    with tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)) as sess:
        # 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()
                state = list(batch['state'].values())
                len_state = list(batch['len_state'].values())
                target = list(batch['action'].values())
               
                loss, _ = sess.run([FMLPRec.loss, FMLPRec.opt],
                                   feed_dict={FMLPRec.inputs: state,
                                              FMLPRec.len_state: len_state,
                                              FMLPRec.target: target,
                                              FMLPRec.is_training: True})
                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)
        # saver.save(sess, save_file)
