# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/transforms/transforms.bipartite.ipynb (unless otherwise specified).

__all__ = ['BipartiteDataset']

# Cell
import numpy as np
import time
from tqdm.notebook import tqdm

import torch
from torch.utils.data import Dataset

# Cell
class BipartiteDataset(Dataset):
    def __init__(self, args, train, neg_dist, offset, num_u, num_v,K):
        self.args = args
        self.edge_1 = torch.tensor(train[self.args.user_col].values-1)
        self.edge_2 = torch.tensor(train[self.args.item_col].values-1) +num_u
        self.edge_3 = torch.tensor(train[self.args.feedback_col].values) - offset
        self.neg_dist = neg_dist
        self.K = K;
        self.num_u = num_u
        self.num_v = num_v
        self.tot = np.arange(num_v)
        self.train = train

    def negs_gen_(self):
        print('negative sampling...'); st=time.time()
        self.edge_4 = torch.empty((len(self.edge_1),self.K),dtype=torch.long)
        prog = tqdm(desc='negative sampling for each epoch...',total=len(set(self.train[self.args.user_col].values)),position=0)
        for j in set(self.train[self.args.user_col].values):
            pos=self.train[self.train[self.args.user_col]==j][self.args.item_col].values-1
            neg = np.setdiff1d(self.tot,pos)
            temp = (torch.tensor(np.random.choice(neg,len(pos)*self.K,replace=True,p=self.neg_dist[neg]/self.neg_dist[neg].sum()))+self.num_u).long()
            self.edge_4[self.edge_1==j-1]=temp.view(int(len(temp)/self.K),self.K)
            prog.update(1)
        prog.close()
        self.edge_4 = torch.tensor(self.edge_4).long()
        print('complete ! %s'%(time.time()-st))

    def negs_gen_EP(self,epoch):
        print('negative sampling for next epochs...'); st=time.time()
        self.edge_4_tot = torch.empty((len(self.edge_1),self.K,epoch),dtype=torch.long)
        prog = tqdm(desc='negative sampling for next epochs...',total=len(set(self.train[self.args.user_col].values)),position=0)
        for j in set(self.train[self.args.user_col].values):
            pos=self.train[self.train[self.args.user_col]==j][self.args.item_col].values-1
            neg = np.setdiff1d(self.tot,pos)
            temp = (torch.tensor(np.random.choice(neg,len(pos)*self.K*epoch,replace=True,p=self.neg_dist[neg]/self.neg_dist[neg].sum()))+self.num_u).long()
            self.edge_4_tot[self.edge_1==j-1]=temp.view(int(len(temp)/self.K/epoch),self.K,epoch)
            prog.update(1)
        prog.close()
        self.edge_4_tot = torch.tensor(self.edge_4_tot).long()
        print('complete ! %s'%(time.time()-st))

    def __len__(self):
        return len(self.edge_1)

    def __getitem__(self,idx):
        u = self.edge_1[idx]
        v = self.edge_2[idx]
        w = self.edge_3[idx]
        negs = self.edge_4[idx]
        return u,v,w,negs