import datetime
import math
import numpy as np
import torch
from torch import nn
from tqdm import tqdm
from aggregator import LocalAggregator, GlobalAggregator
from np_model import DeterministicNet, LatentNet, Decoder
from torch.nn import Module, Parameter
import torch.nn.functional as F


class CombineGraph(Module):
    def __init__(self, opt, num_node, adj_all, num):
        super(CombineGraph, self).__init__()
        self.opt = opt

        self.batch_size = opt.batch_size
        self.num_node = num_node
        self.dim = opt.hiddenSize
        self.dropout_local = opt.dropout_local
        self.dropout_global = opt.dropout_global
        self.hop = opt.n_iter
        self.sample_num = opt.n_sample
        self.adj_all = trans_to_cuda(torch.Tensor(adj_all)).long()
        self.num = trans_to_cuda(torch.Tensor(num)).float()

        # Aggregator
        self.local_agg = LocalAggregator(self.dim, self.opt.alpha, dropout=0.0)
        self.global_agg = []
        for i in range(self.hop):
            if opt.activate == 'relu':
                agg = GlobalAggregator(self.dim, opt.dropout_gcn, act=torch.relu)
            else:
                agg = GlobalAggregator(self.dim, opt.dropout_gcn, act=torch.tanh)
            self.add_module('agg_gcn_{}'.format(i), agg)
            self.global_agg.append(agg)

        # Item representation & Position representation
        self.embedding = nn.Embedding(num_node, self.dim, padding_idx=0)
        self.pos_embedding = nn.Embedding(200, self.dim)

        # Parameters
        self.w_1 = nn.Parameter(torch.Tensor(2 * self.dim, self.dim))
        self.w_2 = nn.Parameter(torch.Tensor(self.dim, 1))
        self.glu1 = nn.Linear(self.dim, self.dim)
        self.glu2 = nn.Linear(self.dim, self.dim, bias=False)
        self.linear_transform = nn.Linear(self.dim, self.dim, bias=False)

        self.leakyrelu = nn.LeakyReLU(opt.alpha)
        self.loss_function = nn.CrossEntropyLoss()
        self.optimizer = torch.optim.Adam(self.parameters(), lr=opt.lr, weight_decay=opt.l2)
        self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, step_size=opt.lr_dc_step, gamma=opt.lr_dc)

        self.reset_parameters()

    def reset_parameters(self):
        stdv = 1.0 / math.sqrt(self.dim)
        for weight in self.parameters():
            weight.data.uniform_(-stdv, stdv)

    def sample(self, target, n_sample):
        return self.adj_all[target.view(-1)], self.num[target.view(-1)]

    def compute_scores(self, hidden, mask):
        mask = mask.float().unsqueeze(-1)

        batch_size = hidden.shape[0]
        len = hidden.shape[1]
        pos_emb = self.pos_embedding.weight[:len]
        pos_emb = pos_emb.unsqueeze(0).repeat(batch_size, 1, 1)

        hs = torch.sum(hidden * mask, -2) / torch.sum(mask, 1)
        hs = hs.unsqueeze(-2).repeat(1, len, 1)
        nh = torch.matmul(torch.cat([pos_emb, hidden], -1), self.w_1)
        nh = torch.tanh(nh)
        nh = torch.sigmoid(self.glu1(nh) + self.glu2(hs))
        beta = torch.matmul(nh, self.w_2)
        beta = beta * mask
        select = torch.sum(beta * hidden, 1)
        #select = F.normalize(select, p=2.0, dim=-1)

        # [b, d]
        return select

    def forward(self, inputs, adj, mask_item, item):
        batch_size = inputs.shape[0]
        seqs_len = inputs.shape[1]
        h = self.embedding(inputs)

        # local
        h_local = self.local_agg(h, adj, mask_item)

        # global
        item_neighbors = [inputs]
        weight_neighbors = []
        support_size = seqs_len

        for i in range(1, self.hop + 1):
            item_sample_i, weight_sample_i = self.sample(item_neighbors[-1], self.sample_num)
            support_size *= self.sample_num
            item_neighbors.append(item_sample_i.view(batch_size, support_size))
            weight_neighbors.append(weight_sample_i.view(batch_size, support_size))

        entity_vectors = [self.embedding(i) for i in item_neighbors]
        weight_vectors = weight_neighbors

        session_info = []
        item_emb = self.embedding(item) * mask_item.float().unsqueeze(-1)
        
        # mean 
        sum_item_emb = torch.sum(item_emb, 1) / torch.sum(mask_item.float(), -1).unsqueeze(-1)

        
        sum_item_emb = sum_item_emb.unsqueeze(-2)
        for i in range(self.hop):
            session_info.append(sum_item_emb.repeat(1, entity_vectors[i].shape[1], 1))

        for n_hop in range(self.hop):
            entity_vectors_next_iter = []
            shape = [batch_size, -1, self.sample_num, self.dim]
            for hop in range(self.hop - n_hop):
                aggregator = self.global_agg[n_hop]
                vector = aggregator(self_vectors=entity_vectors[hop],
                                    neighbor_vector=entity_vectors[hop+1].view(shape),
                                    masks=None,
                                    batch_size=batch_size,
                                    neighbor_weight=weight_vectors[hop].view(batch_size, -1, self.sample_num),
                                    extra_vector=session_info[hop])
                entity_vectors_next_iter.append(vector)
            entity_vectors = entity_vectors_next_iter

        h_global = entity_vectors[0].view(batch_size, seqs_len, self.dim)

        # combine
        h_local = F.dropout(h_local, self.dropout_local, training=self.training)
        h_global = F.dropout(h_global, self.dropout_global, training=self.training)
        output = h_local + h_global

        return output


def trans_to_cuda(variable):
    if torch.cuda.is_available():
        return variable.cuda()
    else:
        return variable


def trans_to_cpu(variable):
    if torch.cuda.is_available():
        return variable.cpu()
    else:
        return variable







class NP_MISR(nn.Module):
    def __init__(self, opt,num_node, gnn, gnn_weights_path="NP-MISR-main\pretrain\retailrocketweight\gnn_weights19.pth"):
        super(NP_MISR, self).__init__()
        self.num_class = num_node
        self.num_latent = opt.num_latent
        self.batch_size = opt.batch_size
        self.gnn = gnn
        self.gnn.load_state_dict(torch.load(gnn_weights_path))
        self.gnn.eval()
        for param in self.gnn.parameters():
            param.requires_grad = False

        d_model = opt.hiddenSize
        self.deterministicNet = DeterministicNet(d_model, d_model ,d_model)
        self.latent_encoder = LatentNet(d_model, d_model, self.num_latent)
        self.decoder = Decoder(d_model * 3, d_model)
        self.optimizer = torch.optim.Adam(
            filter(lambda p: p.requires_grad, self.parameters()), 
            lr=0.001, 
            weight_decay=1e-5
        )
        self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, step_size=3, gamma=0.1)
        self.batch_kl_losses = []

    def train_model(self, data):
        targets, sessionEmb = forward(self.gnn, data)
        Loss = 0
        KL = 0
        batch_size = len(targets)

        for i in range(batch_size):
            label = trans_to_cuda(targets[i].unsqueeze(0).long())
            current_session = sessionEmb[i].squeeze(0)
            d = current_session.size(-1)
            remaining_sessions = torch.cat([sessionEmb[:i], sessionEmb[i + 1:]],
                                           dim=0).squeeze(1)
            remaining_labels = torch.cat([targets[:i], targets[i+1:]], dim=0)
            current_label = targets[i]

            similarities = F.cosine_similarity(current_session.unsqueeze(0), remaining_sessions, dim=1)
            num_similar = max(1, int(0.05 * (batch_size - 1)))
            top_indices = torch.topk(similarities, num_similar).indices
            context = remaining_sessions[top_indices]

            remaining_labels = trans_to_cuda(remaining_labels)
            current_label = trans_to_cuda(current_label)
            contextlabel = self.gnn.embedding(remaining_labels[top_indices])
            current_label = self.gnn.embedding(current_label).unsqueeze(0)
            current_session = current_session.unsqueeze(0)

            determ_repr = self.deterministicNet(context, contextlabel)
            # [b, t, d]
            prior_mu, prior_log_var, prior = self.latent_encoder(context, contextlabel)
            posterior_mu, posterior_log_var, posterior = self.latent_encoder(current_session, current_label)
            context_mean = prior_mu.squeeze(0)
            context_var = prior_log_var.squeeze(0)

            target_emb = self.gnn.embedding(label)
            target_mean = torch.mean(target_emb, dim=0)
            target_var = torch.var(target_emb, dim=0)

            mean_diff = torch.mean(torch.abs(context_mean - target_mean))
            var_diff = torch.mean(torch.abs(context_var - target_var))

            _, t, _ = posterior.size()
            determ_repr = determ_repr.unsqueeze(1).expand(-1, t, -1)
            x = current_session.unsqueeze(1).expand(-1, t, -1)

            decoder_input = torch.cat([determ_repr, posterior, x], dim=-1)
            decoder_output = self.decoder(decoder_input)
            weight_emb = self.gnn.embedding.weight
            weight_emb = F.normalize(weight_emb, p=2.0, dim=-1)
            logit = torch.matmul(decoder_output, weight_emb.transpose(1, 0))
            logit = logit.view(t, self.num_class)
            label = label.unsqueeze(-1).expand(-1, t).contiguous().view(-1)
            loss = F.cross_entropy(logit, label, ignore_index=0)
            kl = self.kl_divergence(prior_mu, prior_log_var, posterior_mu, posterior_log_var)
            Loss = Loss + loss
            KL = KL + kl
        total_kl = KL.item()
        self.batch_kl_losses.append(total_kl)
        return (Loss + KL) / batch_size

    def inference(self, data):
        _, _, _, _, _, inputs = data
        inputs = inputs[:,0]
        
        targets, sessionEmb = forward(self.gnn, data)
        batch_size = len(targets)
        probs = []
        for i in range(batch_size):
            current_session = sessionEmb[i].squeeze(0)
            d = current_session.size(-1)
            remaining_sessions = torch.cat([sessionEmb[:i], sessionEmb[i + 1:]],
                                           dim=0).squeeze(1)
            remaining_labels = torch.cat([inputs[:i], inputs[i + 1:]], dim=0)
            current_label = targets[i]

            similarities = F.cosine_similarity(current_session.unsqueeze(0), remaining_sessions, dim=1)
            num_similar = max(1, int(0.05 * (batch_size - 1)))
            top_indices = torch.topk(similarities, num_similar).indices
            context = remaining_sessions[top_indices]

            remaining_labels = trans_to_cuda(remaining_labels)
            current_label = trans_to_cuda(current_label)
            contextlabel = self.gnn.embedding(remaining_labels[top_indices])
            current_label = current_label.long().unsqueeze(0)
            current_session = current_session.unsqueeze(0)

            determ_repr = self.deterministicNet(context, contextlabel)
            # [b, t, d]
            prior_mu, prior_log_var, prior = self.latent_encoder(context, contextlabel)
            
            _, t, _ = prior.size()
            determ_repr = determ_repr.unsqueeze(1).expand(-1, t, -1)
            x = current_session.unsqueeze(1).expand(-1, t, -1)

            decoder_input = torch.cat([determ_repr, prior, x], dim=-1)
            decoder_output = self.decoder(decoder_input)
            weight_emb = self.gnn.embedding.weight
            weight_emb = F.normalize(weight_emb, p=2.0, dim=-1)
            logit = torch.matmul(decoder_output, weight_emb.transpose(1, 0))
            logit = logit.mean(dim=0, keepdim=True)
            prob = F.softmax(logit, dim=-1)
            probs.append(prob)
        probs = torch.cat(probs, dim=0)
        return targets, probs

    def save(self, path):
        torch.save(self.state_dict(), path)

    def load(self, path):
        self.load_state_dict(torch.load(path))

    def kl_divergence(self, prior_mu, prior_var, posterior_mu, posterior_var):
        kl_div = (torch.exp(posterior_var) + (posterior_mu - prior_mu) ** 2) / torch.exp(prior_var) - 1. + (
                prior_var - posterior_var)
        kl_div = 0.5 * kl_div.sum()
        return kl_div


def forward(model, data ):
    alias_inputs, adj, items, mask, targets, inputs = data
    alias_inputs = trans_to_cuda(alias_inputs).long()
    items = trans_to_cuda(items).long()
    adj = trans_to_cuda(adj).float()
    mask = trans_to_cuda(mask).long()
    inputs = trans_to_cuda(inputs).long()

    hidden = model(items, adj, mask, inputs)
    get = lambda index: hidden[index][alias_inputs[index]]
    seq_hidden = torch.stack([get(i) for i in torch.arange(len(alias_inputs)).long()])
    return targets, model.compute_scores(seq_hidden, mask)


def train_test(model, train_data, test_data):
   
    print('start training: ', datetime.datetime.now())
    model.train()
    total_loss = 0.0
    train_loader = torch.utils.data.DataLoader(train_data, num_workers=4, batch_size=model.batch_size,
                                               shuffle=True, pin_memory=True)
    for data in train_loader:
        model.optimizer.zero_grad()
        loss = model.train_model(data)
        loss.backward()
        model.optimizer.step()
        total_loss += loss.item()
    print('\tLoss:\t%.3f' % total_loss)
    model.scheduler.step()
    

    k_values = [3, 5, 10, 20]
    print('start predicting: ', datetime.datetime.now())
    model.eval()
    test_loader = torch.utils.data.DataLoader(test_data, num_workers=4, batch_size=model.batch_size,
                                              shuffle=True, pin_memory=True)
    results = {k: [] for k in k_values}
    hits = {k: [] for k in k_values}
    mrrs = {k: [] for k in k_values}
    covs = {k: set() for k in k_values}
    
    for data in test_loader:
        targets, (scores) = model.inference(data)
        targets = targets.numpy()
        for k in k_values:
            sub_scores = scores.topk(k)[1]
            sub_scores = trans_to_cpu(sub_scores).detach().numpy()
            for score, target in zip(sub_scores, targets):
                hits[k].append(np.isin(target, score))
                for item in score:
                    if isinstance(item, np.ndarray):
                        item = tuple(item)
                    covs[k].update(item)
                if len(np.where(score == target)[0]) == 0:
                    mrrs[k].append(0)
                else:
                    mrrs[k].append(1 / (np.where(score == target)[0][0] + 1))
    
    for k in k_values:
        results[k].append(np.mean(hits[k]) * 100)
        results[k].append(np.mean(mrrs[k]) * 100)
        results[k].append(len(covs[k]))
        print(f"k = {k}: Recall@{k} = {results[k][0]:.4f}, MRR@{k} = {results[k][1]:.4f}, Cov@{k} = {results[k][2]}")

    return results

