import numpy as np
import torch
from nltk.translate.bleu_score import sentence_bleu
import os
from metric import get_bleu
import torch.nn.functional as F
import torch.nn as nn
import datetime

"""
use learned prior as the initial states
"""


class _EVAL(object):
    def __init__(self, vocab_obj, args, device):
        super().__init__()

        self.m_sos_idx = vocab_obj.sos_idx
        self.m_eos_idx = vocab_obj.eos_idx
        self.m_pad_idx = vocab_obj.pad_idx
        self.m_i2w = vocab_obj.m_i2w

        self.m_epoch = args.epochs
        self.m_batch_size = args.batch_size 
        self.m_mean_loss = 0

        self.m_x0 = args.x0
        self.m_k = args.k

        self.m_latent_size = args.latent_size
        self.m_anneal_func = args.anneal_func
        self.m_device = device
        self.m_model_path = args.model_path

    def f_init_eval(self, network, model_file=None, reload_model=False):
        if reload_model:
            print("reload model")
            if not model_file:
                model_file = "model_best.pt"
            model_name = os.path.join(self.m_model_path, model_file)
            print("model name", model_name)
            check_point = torch.load(model_name)
            network.load_state_dict(check_point['model'])

        self.m_network = network

    def f_eval(self, train_data, eval_data):
        self.f_eval_new(train_data, eval_data)
        self.f_eval_rec(train_data, eval_data)

    def f_init_user_item(self, eval_data):
        user2uid = eval_data.m_user2uid
        item2iid = eval_data.m_item2iid

        user_num = len(user2uid)
        item_num = len(item2iid)
        latent_size = self.m_latent_size

        self.m_user_embedding = torch.zeros(user_num, latent_size)
        self.m_item_embedding = torch.zeros(item_num, latent_size)

        self.m_user_num = torch.zeros((user_num, 1))
        self.m_item_num = torch.zeros((item_num, 1))

        self.m_local_embedding = torch.zeros(1, latent_size)
        self.m_local_num = 0

    def f_get_user_item(self, train_data, eval_data):
        s_time = datetime.datetime.now()

        print("loading user & item & local from prior")

        # eval_user2uid = eval_data.m_user2uid
        # eval_item2iid = eval_data.m_item2iid

        # self.m_network.eval()
        # with torch.no_grad():
        #     for input_batch, input_length_batch, user_batch, item_batch, target_batch, target_length_batch, random_flag in train_data:

        #         input_batch_gpu = input_batch.to(self.m_device)
        #         input_length_batch_gpu = input_length_batch.to(self.m_device)

        #         user_batch_gpu = user_batch.to(self.m_device)
        #         item_batch_gpu = item_batch.to(self.m_device)

        #         target_batch_gpu = target_batch.to(self.m_device)
        #         target_length_batch_gpu = target_length_batch.to(self.m_device)

        #         input_de_batch_gpu = target_batch_gpu[:, :-1]
        #         input_de_length_batch_gpu = target_length_batch_gpu-1

        #         logits_gpu, z_prior_gpu, z_mean_gpu, z_logv_gpu, z_gpu, s_prior_gpu, s_mean_gpu, s_logv_gpu, s_gpu, l_mean_gpu, l_logv_gpu, l_gpu, variational_hidden_gpu = self.m_network(input_batch_gpu, input_length_batch_gpu, input_de_batch_gpu, input_de_length_batch_gpu, user_batch_gpu, item_batch_gpu, random_flag)

        #         z_mean = z_mean_gpu.cpu()
        #         s_mean = s_mean_gpu.cpu()
        #         l_mean = l_mean_gpu.cpu()

        #         for i, user_idx in enumerate(user_batch):
        #             user_idx = user_idx.item()
        #             if user_idx in eval_user2uid:
        #                 uid = eval_user2uid[user_idx]
        #                 z_mean_i = z_mean[i]

        #                 self.m_user_embedding[uid] += z_mean_i.detach()
        #                 self.m_user_num[uid] += 1.0

        #             item_idx = item_batch[i].item()
        #             if item_idx in eval_item2iid:
        #                 iid = eval_item2iid[item_idx]
        #                 s_mean_i = s_mean[i]
        #                 self.m_item_embedding[iid] += s_mean_i.detach()
        #                 self.m_item_num[iid] += 1.0

        #         # self.m_local_embedding += torch.mean(l_mean.detach(), dim=0).unsqueeze(0)
        #         self.m_local_embedding += torch.mean(l_mean.detach(), dim=0)

        #         self.m_local_num += 1
        # for 

        # self.m_user_embedding = self.m_user_embedding/self.m_user_num
        # self.m_item_embedding = self.m_item_embedding/self.m_item_num
        
        # if torch.isnan(self.m_user_num).any():
        #     print("self.m_user_num", self.m_user_num)

        # if torch.isnan(self.m_item_num).any():
        #     print("self.m_item_num", self.m_item_num)

        # print("self.m_local_num", self.m_local_num)
        # # print("local size", self.m_local_embedding.size(), self.m_local_embedding)
        # self.m_local_embedding = self.m_local_embedding/self.m_local_num

        self.m_user_embedding = self.m_network.m_user_embedding
        self.m_item_embedding = self.m_network.m_item_embedding

        e_time = datetime.datetime.now()
        print("load user item duration", e_time-s_time)

    def f_eval_new(self, train_data, eval_data):
        self.f_init_user_item(eval_data)

        self.f_get_user_item(train_data, eval_data)

        eval_user2uid = eval_data.m_user2uid
        eval_item2iid = eval_data.m_item2iid

        batch_index = 0
        bleu_score_list = []

        self.m_network.eval()
        with torch.no_grad():
            for input_batch, input_length_batch, user_batch, item_batch, target_batch, target_length_batch, random_flag in eval_data:
                # print("batch_index", batch_index)

                batch_size = input_batch.size(0)

                input_batch_gpu = input_batch.to(self.m_device)
                input_length_batch_gpu = input_length_batch.to(self.m_device)

                user_batch_gpu = user_batch.to(self.m_device)
                item_batch_gpu = item_batch.to(self.m_device)

                target_batch_gpu = target_batch.to(self.m_device)
                target_length_batch_gpu = target_length_batch.to(self.m_device)

                input_de_batch_gpu = target_batch_gpu[:, :-1]
                input_de_length_batch_gpu = target_length_batch_gpu-1

                tmp = torch.zeros((batch_size, self.m_latent_size))
                z_mean_gpu = torch.randn_like(tmp).to(self.m_device)
                s_mean_gpu = torch.randn_like(tmp).to(self.m_device)
                # z_mean_gpu = self.m_user_embedding(user_batch_gpu)
                # s_mean_gpu = self.m_item_embedding(item_batch_gpu)
                # l_mean_gpu = self.m_local_embedding.to(self.m_device)
                # l_mean_gpu = torch.cat(batch_size*[self.m_local_embedding])
                l_mean_gpu = torch.randn_like(s_mean_gpu)

                mean = z_mean_gpu+s_mean_gpu+l_mean_gpu
                max_seq_len = max(target_length_batch-1)
                
                # print("z_mean_gpu", z_mean_gpu)
                # if torch.isnan(z_mean_gpu).any():
                # print("z size", z_mean_gpu.size())
                # print("s size", s_mean_gpu.size())
                # print("l size", l_mean_gpu.size())
                
                samples, z = self.f_decode_text(z_mean_gpu, s_mean_gpu, l_mean_gpu, max_seq_len)

                lens = target_length_batch-1
                lens = lens.tolist()
                preds = samples.cpu().tolist()
                target_batch = target_batch[:, 1:].tolist()

                preds = [pred_i[:lens[index]] for index, pred_i in enumerate(preds)]
                targets = [target_i[:lens[index]] for index, target_i in enumerate(target_batch)]

                bleu_score_batch = get_bleu(preds, targets)

                bleu_score_list.append(bleu_score_batch)

                batch_index += 1
        
        mean_bleu_score = np.mean(bleu_score_list)
        print("bleu score", mean_bleu_score)

    def f_eval_rec(self, train_data, eval_data):
        self.m_mean_loss = 0
        # for epoch_i in range(self.m_epoch):
        # batch_size = args.batch_size

        infer_loss_list = []
        
        batch_index = 0

        bleu_score_list = []
        with torch.no_grad():
            for input_batch, input_length_batch, user_batch, item_batch, target_batch, target_length_batch, random_flag in eval_data:
                
                # if batch_index > 0:
                #     break

                # batch_index += 1
                input_batch_gpu = input_batch.to(self.m_device)
                input_length_batch_gpu = input_length_batch.to(self.m_device)

                user_batch_gpu = user_batch.to(self.m_device)
                item_batch_gpu = item_batch.to(self.m_device)

                target_batch_gpu = target_batch.to(self.m_device)
                target_length_batch_gpu = target_length_batch.to(self.m_device)

                input_de_batch_gpu = target_batch_gpu[:, :-1]
                input_de_length_batch_gpu = target_length_batch_gpu-1

                logits, z_prior, z_mean, z_logv, z, s_prior, s_mean, s_logv, s, l_mean, l_logv, l, variational_hidden = self.m_network(input_batch_gpu, input_length_batch_gpu, input_de_batch_gpu, input_de_length_batch_gpu, user_batch_gpu, item_batch_gpu, random_flag)       

                # mean = torch.cat([z_mean, s_mean], dim=1)
                # mean = z_mean+s_mean+l_mean
                # mean = z_mean+s_mean
                # mean = s_mean+l_mean
                if random_flag == 0:
                    mean = z_mean+s_mean+l_mean
                elif random_flag == 1:
                    mean = z_mean+s_mean
                elif random_flag == 2:
                    mean = s_mean+l_mean
                elif random_flag == 3:
                    mean = z_mean+l_mean

                max_seq_len = max(target_length_batch-1)
                samples, z = self.f_decode_text(z_mean, s_mean, l_mean, max_seq_len)

                lens = target_length_batch-1
                lens = lens.tolist()
                preds = samples.cpu().tolist()
                target_batch = target_batch[:, 1:].tolist()

                preds = [pred_i[:lens[index]]for index, pred_i in enumerate(preds)]
                targets = [target_i[:lens[index]]for index, target_i in enumerate(target_batch)]

                bleu_score_batch = get_bleu(preds, targets)

                bleu_score_list.append(bleu_score_batch)

        mean_bleu_score = np.mean(bleu_score_list)
        print("bleu score", mean_bleu_score)

    def f_decode_text(self, z, s, l, max_seq_len, n=4):
        if z is None:
            assert "z is none"

        batch_size = z.size(0)

        seq_idx = torch.arange(0, batch_size).long().to(self.m_device)

        seq_running = torch.arange(0, batch_size).long().to(self.m_device)

        seq_mask = torch.ones(batch_size).bool().to(self.m_device)

        running_seqs = torch.arange(0, batch_size).long().to(self.m_device)

        generations = torch.zeros(batch_size, max_seq_len).fill_(self.m_pad_idx).to(self.m_device).long()

        t = 0

        var_de = self.m_network.m_latent2hidden(z+s+l)

        # print("hidden size", hidden.size())
        hidden = None

        while(t < max_seq_len and len(running_seqs)>0):
            # print("t", t)
            if t == 0:
                input_seq = torch.zeros(batch_size).fill_(self.m_sos_idx).long().to(self.m_device)
            
            # input_seq = input_seq.unsqueeze(1)
            # print("input seq size", input_seq.size())
            input_embedding = self.m_network.m_embedding(input_seq)
            
            input_embedding = input_embedding+var_de

            input_embedding = input_embedding.unsqueeze(1)

            if torch.isnan(input_embedding).any():
                print("input_embedding", input_embedding)

            # print("input_embedding", input_embedding.size())
            output, hidden = self.m_network.m_decoder_rnn(input_embedding, hidden)

            logits = self.m_network.m_output2vocab(output)

            input_seq = self._sample(logits)

            if len(input_seq.size()) < 1:
                input_seq = input_seq.unsqueeze(0)

            generations = self._save_sample(generations, input_seq, seq_running, t)

            seq_mask[seq_running] = (input_seq != self.m_eos_idx).bool()
            seq_running = seq_idx.masked_select(seq_mask)

            running_mask = (input_seq != self.m_eos_idx).bool()
            running_seqs = running_seqs.masked_select(running_mask)

            if len(running_seqs) > 0:
                input_seq = input_seq[running_seqs]

                hidden = hidden[:, running_seqs]
                # var_de = var_de[running_seqs]
                # repeat_hidden_0 = repeat_hidden_0[running_seqs]
                output = output[running_seqs].squeeze(1)

                running_seqs = torch.arange(0, len(running_seqs)).long().to(self.m_device)

                z = z[running_seqs]
                s = s[running_seqs]
                l = l[running_seqs]
                
                z_attn = torch.cat([output, z], dim=-1)
                s_attn = torch.cat([output, s], dim=-1)
                l_attn = torch.cat([output, l], dim=-1)

                z_attn_score = self.m_network.m_attn(z_attn)
                s_attn_score = self.m_network.m_attn(s_attn)
                l_attn_score = self.m_network.m_attn(l_attn)

                attn_score = F.softmax(torch.cat([z_attn_score, s_attn_score, l_attn_score], dim=-1), dim=-1)
                # print("attn_score", attn_score)
                
                var_de = attn_score[:, 0].unsqueeze(1)*z
                var_de = var_de + attn_score[:, 1].unsqueeze(1)*s
                var_de = var_de + attn_score[:, 2].unsqueeze(1)*l

                var_de = self.m_network.m_latent2hidden(var_de)

                # var_de_flag = self.m_network.m_decoder_gate(output.squeeze(1))
                # var_de = self.m_network.m_latent2hidden((1-var_de_flag)*z+var_de_flag*s+l)

            t += 1

        return generations, z

    def _sample(self, dist, mode="greedy"):
        if mode == 'greedy':
            _, sample = torch.topk(dist, 1, dim=-1)
        sample = sample.squeeze()

        return sample

    def _save_sample(self, save_to, sample, running_seqs, t):
        # select only still running
        running_latest = save_to[running_seqs]
        # update token at position t
        running_latest[:,t] = sample.data
        # save back
        save_to[running_seqs] = running_latest

        return save_to

def idx2word(idx, i2w, pad_idx):

    sent_str = [str()]*len(idx)

    for i, sent in enumerate(idx):
        # print(" "*10, "*"*10)
        for word_id in sent:

            if word_id == pad_idx:
                break
            # print('word_id', word_id.item())
            sent_str[i] += i2w[str(word_id.item())] + " "

        sent_str[i] = sent_str[i].strip()

    return sent_str