# -*- coding: utf-8 -*-
import numpy as np
import torch
torch.manual_seed(2020)
from torch import nn
import torch.nn.functional as F
from math import sqrt
import pdb
from sklearn.metrics import roc_auc_score

mse_func = lambda x,y: np.mean((x-y)**2)
acc_func = lambda x,y: np.sum(x == y) / len(x)

def generate_total_sample(num_user, num_item):
    sample = []
    for i in range(num_user):
        sample.extend([[i,j] for j in range(num_item)])
    return np.array(sample)

def sigmoid(x):
    return 1.0 / (1 + np.exp(-x))


class MF(nn.Module):
    def __init__(self, num_users, num_items, batch_size, embedding_k=4, *args, **kwargs):
        super(MF, self).__init__()
        self.num_users = num_users
        self.num_items = num_items
        self.embedding_k = embedding_k
        self.W = torch.nn.Embedding(self.num_users, self.embedding_k)
        self.H = torch.nn.Embedding(self.num_items, self.embedding_k)
        self.batch_size = batch_size
        self.sigmoid = torch.nn.Sigmoid()
        self.xent_func = torch.nn.BCELoss()

    def forward(self, x, is_training=False):
        user_idx = torch.LongTensor(x[:,0])
        item_idx = torch.LongTensor(x[:,1])
        U_emb = self.W(user_idx)
        V_emb = self.H(item_idx)

        out = torch.sum(U_emb.mul(V_emb), 1)

        if is_training:
            return out, U_emb, V_emb
        else:
            return out
        
    def complete(self):
        for i in range(self.embedding_k):
            complete = torch.zeros([self.num_users,self.num_items])
            for m in range(self.num_users):                   
                 complete[m, :] = self.W(torch.LongTensor([m]))[0][i]
            if i == 0:
                complete1 = complete.reshape(self.num_users*self.num_items, 1)
            else:
                complete1 = torch.cat((complete1 ,complete.reshape(self.num_users*self.num_items, 1)), dim = 1)                
                
        for i in range(self.embedding_k):
            complete = torch.zeros([self.num_users,self.num_items])
            for k in range(self.num_items):
                complete[:,k] = self.H(torch.LongTensor([k]))[0][i]
            if i == 0:
                complete2 = complete.reshape(self.num_users*self.num_items, 1)
            else:
                complete2 = torch.cat((complete2 ,complete.reshape(self.num_users*self.num_items, 1)), dim = 1)
        
        complete = torch.cat((complete1, complete2),dim = 1)

        return complete
    
    def fit(self, x, y, stop = 5,
        num_epoch=1000, lr=0.05, lamb=0, 
        tol=1e-4, verbose=False):

        optimizer = torch.optim.Adam(self.parameters(), lr=lr, weight_decay=lamb)
        last_loss = 1e9

        num_sample = len(x)
        total_batch = num_sample // self.batch_size

        early_stop = 0
        for epoch in range(num_epoch):
            all_idx = np.arange(num_sample)
            np.random.shuffle(all_idx)
            epoch_loss = 0

            for idx in range(total_batch):
                # mini-batch training
                selected_idx = all_idx[self.batch_size*idx:(idx+1)*self.batch_size]
                sub_x = x[selected_idx]
                sub_y = y[selected_idx]
                sub_y = torch.Tensor(sub_y)

                pred, u_emb, v_emb = self.forward(sub_x, True)
                pred = self.sigmoid(pred)

                xent_loss = self.xent_func(pred,sub_y)

                optimizer.zero_grad()
                xent_loss.backward()
                optimizer.step()
                
                epoch_loss += xent_loss.detach().cpu().numpy()

            relative_loss_div = (last_loss-epoch_loss)/(last_loss+1e-10)
            if  relative_loss_div < tol:
                if early_stop > stop:
                    print("[MF] epoch:{}, xent:{}".format(epoch, epoch_loss))
                    break
                early_stop += 1
                
            last_loss = epoch_loss

            if epoch % 10 == 0 and verbose:
                print("[MF] epoch:{}, xent:{}".format(epoch, epoch_loss))

            if epoch == num_epoch - 1:
                print("[MF] Reach preset epochs, it seems does not converge.")

    def predict(self, x):
        pred = self.forward(x)
        pred = self.sigmoid(pred)
        return pred.detach().cpu().numpy()

    
class MF_BaseModel(nn.Module):
    def __init__(self, num_users, num_items, batch_size, embedding_k=4, *args, **kwargs):
        super(MF_BaseModel, self).__init__()
        self.num_users = num_users
        self.num_items = num_items
        self.embedding_k = embedding_k
        self.W = torch.nn.Embedding(self.num_users, self.embedding_k)
        self.H = torch.nn.Embedding(self.num_items, self.embedding_k)
        self.batch_size = batch_size
        self.sigmoid = torch.nn.Sigmoid()
        self.xent_func = torch.nn.BCELoss()

    def forward(self, x, is_training=False):
        user_idx = torch.LongTensor(x[:, 0])
        item_idx = torch.LongTensor(x[:, 1])
        U_emb = self.W(user_idx)
        V_emb = self.H(item_idx)

        out = torch.sum(U_emb.mul(V_emb), 1)

        if is_training:
            return out, U_emb, V_emb
        else:
            return out

    def predict(self, x):
        pred = self.forward(x)
        return pred.detach().cpu()

    
class MF_IPS(nn.Module):
    def __init__(self, num_users, num_items, batch_size, embedding_k=4, *args, **kwargs):
        super().__init__()
        self.num_users = num_users
        self.num_items = num_items
        self.embedding_k = embedding_k
        self.batch_size = batch_size
        self.prediction_model = MF_BaseModel(
            num_users=self.num_users, num_items=self.num_items, batch_size = self.batch_size, embedding_k=self.embedding_k)
        self.original_model = MF(
            num_users=self.num_users, num_items=self.num_items, batch_size = self.batch_size, embedding_k=self.embedding_k)
        
        self.linear_1 = torch.nn.Linear(self.embedding_k*2, 1)
        
        self.sigmoid = torch.nn.Sigmoid()
        self.xent_func = torch.nn.BCELoss()


    def _compute_IPS(self, x, y, gamma):

        x_train = torch.zeros([self.num_users,self.num_items])

        for i in range(len(x)):
            x_train[x[i][0],x[i][1]] = 1 # o
        prediction = (x_train.reshape(self.num_users*self.num_items,1)).type(torch.FloatTensor)      
        
        x_train = self.original_model.complete().type(torch.FloatTensor).detach()     
        
        optimizer = torch.optim.SGD([self.linear_1.weight, self.linear_1.bias], lr=1e-3, momentum=0.9)

        last_loss = 1e9
        early_stop = 0
        for epoch in range(1000):
            all_idx = np.arange(self.num_users*self.num_items)
            np.random.shuffle(all_idx)
            total_batch = (self.num_users*self.num_items)// self.batch_size
            for idx in range(total_batch):
                selected_idx = all_idx[self.batch_size*idx:(idx+1)*self.batch_size]
                sub_x_train = x_train[selected_idx].detach()
                sub_prediction = prediction[selected_idx]
             
                out = self.linear_1(sub_x_train)
                out = self.sigmoid(out)
                loss = self.xent_func(out, sub_prediction)
                
                xent_loss = loss
                optimizer.zero_grad()
                xent_loss.backward()
                optimizer.step()
                
            if (epoch + 1) % 15 == 0:
                print('*'*10)
                print('epoch {}'.format(epoch+1))
            
            if loss > last_loss:
                early_stop += 1 
            else:
                last_loss = loss
            
            if early_stop >= 5:
                break
                
        x_train = x_train.detach()
        propensity = self.sigmoid(self.linear_1(x_train.detach()))
        propensity[np.where(propensity.cpu() <= gamma)] = gamma
        
        one_over_zl = 1 / propensity        
        return prediction, propensity 

    
    def fit(self, x, y, gamma, stop = 5,
        num_epoch=1000, lr=0.05, lamb=0, 
        tol=1e-4, verbose=False):

        optimizer_prediction = torch.optim.Adam(self.prediction_model.parameters(), lr=lr, weight_decay=lamb)
        
        last_loss = 1e9

        num_sample = len(x)
        total_batch = num_sample // self.batch_size

        early_stop = 0
        observation, one_over_zl = self._compute_IPS(x, y, gamma)        

        one_over_zl_obs = one_over_zl[np.where(observation.cpu() == 1)]

        for epoch in range(num_epoch):
            all_idx = np.arange(num_sample)
            np.random.shuffle(all_idx)
            epoch_loss = 0

            for idx in range(total_batch):
                # mini-batch training
                selected_idx = all_idx[self.batch_size*idx:(idx+1)*self.batch_size]
                sub_x = x[selected_idx]
                sub_y = y[selected_idx]

                # propensity score
                inv_prop = one_over_zl_obs[selected_idx].detach()

                sub_y = torch.Tensor(sub_y)

                pred, u_emb, v_emb = self.prediction_model.forward(sub_x, True)
                pred = self.sigmoid(pred)

                xent_loss = F.binary_cross_entropy(pred, sub_y,
                    weight=inv_prop)

                loss = xent_loss

                optimizer_prediction.zero_grad()
                loss.backward()
                optimizer_prediction.step()
                
                epoch_loss += xent_loss.detach().cpu().numpy()

            relative_loss_div = (last_loss-epoch_loss)/(last_loss+1e-10)
            if  relative_loss_div < tol:
                if early_stop > stop:
                    print("[MF-IPS] epoch:{}, xent:{}".format(epoch, epoch_loss))
                    break
                early_stop += 1
                
            last_loss = epoch_loss

            if epoch % 10 == 0 and verbose:
                print("[MF-IPS] epoch:{}, xent:{}".format(epoch, epoch_loss))

            if epoch == num_epoch - 1:
                print("[MF-IPS] Reach preset epochs, it seems does not converge.")

    def predict(self, x):
        pred = self.prediction_model.forward(x)
        pred = self.sigmoid(pred)
        return pred.detach().cpu().numpy()        
    
class MF_IPS_AT(nn.Module):
    def __init__(self, num_users, num_items, batch_size, embedding_k=4, *args, **kwargs):
        super().__init__()
        self.num_users = num_users
        self.num_items = num_items
        self.embedding_k = embedding_k
        self.batch_size = batch_size
        self.prediction1_model = MF_BaseModel(
            num_users=self.num_users, num_items=self.num_items, batch_size = self.batch_size, embedding_k=self.embedding_k)
        self.prediction2_model = MF_BaseModel(
            num_users=self.num_users, num_items=self.num_items, batch_size = self.batch_size, embedding_k=self.embedding_k)
        self.prediction_model = MF_BaseModel(
            num_users=self.num_users, num_items=self.num_items, batch_size = self.batch_size, embedding_k=self.embedding_k)
        self.original_model = MF(
            num_users=self.num_users, num_items=self.num_items, batch_size = self.batch_size, embedding_k=self.embedding_k)
        
        self.linear_1 = torch.nn.Linear(self.embedding_k*2, 1)
        
        self.sigmoid = torch.nn.Sigmoid()
        self.xent_func = torch.nn.BCELoss()


    def _compute_IPS(self, x, y, gamma):
        
#         self.original_model.fit(x = x, y = y, num_epoch=1000, lr=0.05, lamb=0, 
#             tol=1e-4, verbose=False)

        x_train = torch.zeros([self.num_users,self.num_items])

        for i in range(len(x)):
            x_train[x[i][0],x[i][1]] = 1 # o
        prediction = (x_train.reshape(self.num_users*self.num_items,1)).type(torch.FloatTensor)     
        
        x_train = self.original_model.complete().type(torch.FloatTensor).detach()     
        
        optimizer = torch.optim.SGD([self.linear_1.weight, self.linear_1.bias], lr=1e-3, momentum=0.9)
        #for name, param in self.named_parameters():
           # param = param.cuda()
        last_loss = 1e9    #print(name, param)
        early_stop = 0
        for epoch in range(1000):
            all_idx = np.arange(self.num_users*self.num_items)
            np.random.shuffle(all_idx)
            total_batch = (self.num_users*self.num_items)// self.batch_size
            #print(total_batch)
            #print(type(total_batch))
            #total_batch = int(total_batch)
            #print(early_stop)
            for idx in range(total_batch):
                selected_idx = all_idx[self.batch_size*idx:(idx+1)*self.batch_size]
                sub_x_train = x_train[selected_idx].detach()
                sub_prediction = prediction[selected_idx]
             
                out = self.linear_1(sub_x_train)
                #print(out)
                out = self.sigmoid(out)
                loss = self.xent_func(out, sub_prediction)
                
                xent_loss = loss
                #print()
                optimizer.zero_grad()
                #print(self.linear_1.weight.grad)
                xent_loss.backward()
                #print(self.linear_1.weight.grad)
                optimizer.step()
                #print("1")
            if (epoch + 1) % 15 == 0:
                print('*'*10)
                print('epoch {}'.format(epoch+1)) # 训练轮数

            
            if loss > last_loss:
                early_stop += 1 
            else:
                last_loss = loss
            
            if early_stop >= 5:
                break
                
        x_train = x_train.detach()
        propensity = self.sigmoid(self.linear_1(x_train.detach()))
        propensity[np.where(propensity.cpu() <= gamma)] = gamma
        
        one_over_zl = 1 / propensity
        #one_over_zl = torch.Tensor(one_over_zl)
        
        return prediction, propensity 

    
    def fit(self, x, y, gamma, tao, stop = 5, G = 4,
        num_epoch=1000, lr=0.05, lamb=0, 
        tol=1e-4, verbose=False):

        optimizer_prediction1 = torch.optim.Adam(
            self.prediction1_model.parameters(), lr=lr, weight_decay=lamb)
        optimizer_prediction2 = torch.optim.Adam(
            self.prediction2_model.parameters(), lr=lr, weight_decay=lamb)
        optimizer_prediction = torch.optim.Adam(
            self.prediction_model.parameters(), lr=lr, weight_decay=lamb)
        
        last_loss = 1e9
        x_all = generate_total_sample(self.num_users, self.num_items)
        num_sample = len(x)
        total_batch = num_sample // self.batch_size

        early_stop = 0
        
        observation, one_over_zl = self._compute_IPS(x, y, gamma)        

        one_over_zl_obs = one_over_zl[np.where(observation.cpu() == 1)]
        for epoch in range(num_epoch):                   
            all_idx = np.arange(num_sample)
            np.random.shuffle(all_idx)
            epoch_loss = 0

            for idx in range(total_batch):
                # mini-batch training
                selected_idx = all_idx[self.batch_size*idx:(idx+1)*self.batch_size]
                sub_x = x[selected_idx]
                sub_y = y[selected_idx]

                # propensity score
                inv_prop = one_over_zl_obs[selected_idx]

                sub_y = torch.Tensor(sub_y)

                pred, u_emb, v_emb = self.prediction1_model.forward(sub_x, True)
                pred = self.sigmoid(pred)
                
                xent_loss = F.binary_cross_entropy(pred, sub_y,
                    weight=inv_prop.detach())

                loss = xent_loss

                optimizer_prediction1.zero_grad()
                loss.backward()
                optimizer_prediction1.step()
                
                epoch_loss += xent_loss.detach().cpu().numpy()

            relative_loss_div = (last_loss-epoch_loss)/(last_loss+1e-10)
            if  relative_loss_div < tol:
                if early_stop > 5:
                    print("[MF-IPS-Pred1] epoch:{}, xent:{}".format(epoch, epoch_loss))
                    break
                early_stop += 1
                
            last_loss = epoch_loss

            if epoch % 10 == 0 and verbose:
                print("[MF-IPS-Pred1] epoch:{}, xent:{}".format(epoch, epoch_loss))

            if epoch == num_epoch - 1:
                print("[MF-IPS-Pred1] Reach preset epochs, it seems does not converge.")

        early_stop = 0
        for epoch in range(num_epoch):                   
            all_idx = np.arange(num_sample)
            np.random.shuffle(all_idx)
            epoch_loss = 0

            for idx in range(total_batch):
                # mini-batch training
                selected_idx = all_idx[self.batch_size*idx:(idx+1)*self.batch_size]
                sub_x = x[selected_idx]
                sub_y = y[selected_idx]

                # propensity score
                inv_prop = one_over_zl_obs[selected_idx]

                sub_y = torch.Tensor(sub_y)

                pred, u_emb, v_emb = self.prediction2_model.forward(sub_x, True)
                pred = self.sigmoid(pred)
                
                xent_loss = F.binary_cross_entropy(pred, sub_y,
                    weight=inv_prop.detach())

                loss = xent_loss

                optimizer_prediction2.zero_grad()
                loss.backward()
                optimizer_prediction2.step()
                
                epoch_loss += xent_loss.detach().cpu().numpy()

            relative_loss_div = (last_loss-epoch_loss)/(last_loss+1e-10)
            if  relative_loss_div < tol:
                if early_stop > 5:
                    print("[MF-IPS-Pred2] epoch:{}, xent:{}".format(epoch, epoch_loss))
                    break
                early_stop += 1
                
            last_loss = epoch_loss

            if epoch % 10 == 0 and verbose:
                print("[MF-IPS-Pred2] epoch:{}, xent:{}".format(epoch, epoch_loss))

            if epoch == num_epoch - 1:
                print("[MF-IPS-Pred2] Reach preset epochs, it seems does not converge.")
        
        early_stop = 0
        for epoch in range(num_epoch):
            all_idx = np.arange(num_sample) # observation
            np.random.shuffle(all_idx)

            # sampling counterfactuals
            ul_idxs = np.arange(x_all.shape[0]) # all
            np.random.shuffle(ul_idxs)

            epoch_loss = 0

            for idx in range(total_batch):                
                x_sampled = x_all[ul_idxs[G*idx* self.batch_size : G*(idx+1)*self.batch_size]]
                pred_u1 = self.prediction1_model.forward(x_sampled)
                pred_u2 = self.prediction2_model.forward(x_sampled)
                pred_u1 = self.sigmoid(pred_u1)
                pred_u2 = self.sigmoid(pred_u2)
                x_sampled_common = x_sampled[(pred_u1.detach().cpu().numpy() - pred_u2.detach().cpu().numpy()) < tao]

                pred_u3 = self.prediction_model.forward(x_sampled_common)
                pred_u3 = self.sigmoid(pred_u3)

                sub_y = self.prediction1_model.forward(x_sampled_common)
                sub_y = self.sigmoid(sub_y)
                #print(sub_y)
                #sub_y = torch.Tensor(sub_y).cuda()
                
                xent_loss = F.binary_cross_entropy(pred_u3, sub_y.detach())

                loss = xent_loss

                optimizer_prediction.zero_grad()
                loss.backward()
                optimizer_prediction.step()
                
                epoch_loss += xent_loss.detach().cpu().numpy()

            relative_loss_div = (last_loss-epoch_loss)/(last_loss+1e-10)
            if  relative_loss_div < tol:
                if early_stop > stop:
                    print("[MF-IPS-AT] epoch:{}, xent:{}".format(epoch, epoch_loss))
                    break
                early_stop += 1
                
            last_loss = epoch_loss

            if epoch % 10 == 0 and verbose:
                print("[MF-IPS-AT] epoch:{}, xent:{}".format(epoch, epoch_loss))

            if epoch == num_epoch - 1:
                print("[MF-IPS-AT] Reach preset epochs, it seems does not converge.")

    def predict(self, x):
        pred = self.prediction_model.forward(x)
        pred = self.sigmoid(pred)
        return pred.detach().cpu().numpy()    
    
class MF_SNIPS(nn.Module):
    def __init__(self, num_users, num_items, batch_size, embedding_k=4, *args, **kwargs):
        super().__init__()
        self.num_users = num_users
        self.num_items = num_items
        self.embedding_k = embedding_k
        self.batch_size = batch_size
        self.prediction_model = MF_BaseModel(
            num_users = self.num_users, num_items = self.num_items, batch_size = self.batch_size, embedding_k=self.embedding_k, *args, **kwargs)
        self.original_model = MF(
            num_users = self.num_users, num_items = self.num_items, batch_size = self.batch_size, embedding_k=self.embedding_k, *args, **kwargs)

        self.linear_1 = torch.nn.Linear(self.embedding_k*2, 1)

        self.sigmoid = torch.nn.Sigmoid()
        self.xent_func = torch.nn.BCELoss()

    def _compute_IPS(self, x, y, gamma):
        
        x_train = torch.zeros([self.num_users,self.num_items])

        for i in range(len(x)):
            x_train[x[i][0],x[i][1]] = 1 # o
        prediction = (x_train.reshape(self.num_users*self.num_items,1)).type(torch.FloatTensor)

        x_train = self.original_model.complete().type(torch.FloatTensor).detach()
               
        optimizer = torch.optim.SGD([self.linear_1.weight, self.linear_1.bias], lr=1e-3, momentum=0.9)

        last_loss = 1e9
        early_stop = 0
        for epoch in range(1000):
            all_idx = np.arange(self.num_users*self.num_items)
            np.random.shuffle(all_idx)
            total_batch = (self.num_users*self.num_items)// self.batch_size
            for idx in range(total_batch):
                selected_idx = all_idx[self.batch_size*idx:(idx+1)*self.batch_size]
                sub_x_train = x_train[selected_idx].detach()
                sub_prediction = prediction[selected_idx]
             
                out = self.linear_1(sub_x_train)
                out = self.sigmoid(out)
                loss = self.xent_func(out, sub_prediction)
                
                xent_loss = loss

                optimizer.zero_grad()
                xent_loss.backward()
                optimizer.step()
                
            if (epoch + 1) % 15 == 0:
                print('*'*10)
                print('epoch {}'.format(epoch+1))  
            
            if loss > last_loss:
                early_stop += 1 
            else:
                last_loss = loss
            
            if early_stop >= 5:
                break
                
        x_train = x_train.detach()
        propensity = self.sigmoid(self.linear_1(x_train.detach())) 
        propensity[np.where(propensity.cpu() <= gamma)] = gamma

        one_over_zl = 1 / propensity
        return prediction, one_over_zl  
    
    def fit(self, x, y, gamma, stop = 5,
        num_epoch=1000, lr=0.05, lamb=0, 
        tol=1e-4, verbose=True):

        optimizer_prediction = torch.optim.Adam(self.prediction_model.parameters(), lr=lr, weight_decay=lamb)
        
        last_loss = 1e9

        num_sample = len(x)
        total_batch = num_sample // self.batch_size

        early_stop = 0
        observation, one_over_zl = self._compute_IPS(x, y, gamma)
        one_over_zl_obs = one_over_zl[np.where(observation.cpu() == 1)]
        
        for epoch in range(num_epoch):
            all_idx = np.arange(num_sample)
            np.random.shuffle(all_idx)
            epoch_loss = 0

            for idx in range(total_batch):
                # mini-batch training
                selected_idx = all_idx[self.batch_size*idx:(idx+1)*self.batch_size]
                sub_x = x[selected_idx]
                sub_y = y[selected_idx]

                # propensity score
                inv_prop = one_over_zl_obs[selected_idx].detach()

                sub_y = torch.Tensor(sub_y)

                pred, u_emb, v_emb = self.prediction_model.forward(sub_x, True)
                pred = self.sigmoid(pred)

                xent_loss = F.binary_cross_entropy(pred, sub_y,
                    weight=inv_prop, reduction = "sum")
                
                xent_loss = xent_loss / (inv_prop.detach().sum())

                loss = xent_loss

                optimizer_prediction.zero_grad()
                loss.backward()
                optimizer_prediction.step()
                
                epoch_loss += xent_loss.detach().cpu().numpy()

            relative_loss_div = (last_loss-epoch_loss)/(last_loss+1e-10)
            if  relative_loss_div < tol:
                if early_stop > stop:
                    print("[MF-SNIPS] epoch:{}, xent:{}".format(epoch, epoch_loss))
                    break
                early_stop += 1
                
            last_loss = epoch_loss

            if epoch % 10 == 0 and verbose:
                print("[MF-SNIPS] epoch:{}, xent:{}".format(epoch, epoch_loss))

            if epoch == num_epoch - 1:
                print("[MF-SNIPS] Reach preset epochs, it seems does not converge.")

    def predict(self, x):
        pred = self.prediction_model.forward(x)
        pred = self.sigmoid(pred)
        return pred.detach().cpu().numpy()        
    
    
class MF_CVIB(nn.Module):
    def __init__(self, num_users, num_items, embedding_k=4, *args, **kwargs):
        super(MF_CVIB, self).__init__()
        self.num_users = num_users
        self.num_items = num_items
        self.embedding_k = embedding_k
        self.W = torch.nn.Embedding(self.num_users, self.embedding_k)
        self.H = torch.nn.Embedding(self.num_items, self.embedding_k)

        self.sigmoid = torch.nn.Sigmoid()
        self.xent_func = torch.nn.BCELoss()

    def forward(self, x, is_training=False):
        user_idx = torch.LongTensor(x[:,0])
        item_idx = torch.LongTensor(x[:,1])
        U_emb = self.W(user_idx)
        V_emb = self.H(item_idx)

        out = torch.sum(U_emb.mul(V_emb), 1)

        if is_training:
            return out, U_emb, V_emb
        else:
            return out

    def fit(self, x, y, 
        num_epoch=1000, batch_size=128, lr=0.05, lamb=0, 
        alpha=0.1, gamma=0.01, stop = 1,
        tol=1e-4, verbose=True):

        self.alpha = alpha
        self.gamma = gamma

        optimizer = torch.optim.Adam(self.parameters(), lr=lr, weight_decay=lamb)
        last_loss = 1e9

        # generate all counterfactuals and factuals for info reg
        x_all = generate_total_sample(self.num_users, self.num_items)

        num_sample = len(x)
        total_batch = num_sample // batch_size
        early_stop = 0

        for epoch in range(num_epoch):
            all_idx = np.arange(num_sample)
            np.random.shuffle(all_idx)

            # sampling counterfactuals
            ul_idxs = np.arange(x_all.shape[0])
            np.random.shuffle(ul_idxs)

            epoch_loss = 0
            for idx in range(total_batch):
                # mini-batch training
                selected_idx = all_idx[batch_size*idx:(idx+1)*batch_size]
                sub_x = x[selected_idx]
                sub_y = y[selected_idx]
                sub_y = torch.Tensor(sub_y)

                pred, u_emb, v_emb = self.forward(sub_x, True)
                pred = self.sigmoid(pred)
                xent_loss = self.xent_func(pred,sub_y)

                # pair wise loss
                x_sampled = x_all[ul_idxs[idx* batch_size:(idx+1)*batch_size]]

                pred_ul,_,_ = self.forward(x_sampled, True)
                pred_ul = self.sigmoid(pred_ul)

                logp_hat = pred.log()

                pred_avg = pred.mean()
                pred_ul_avg = pred_ul.mean()

                info_loss = self.alpha * (- pred_avg * pred_ul_avg.log() - (1-pred_avg) * (1-pred_ul_avg).log()) + self.gamma* torch.mean(pred * logp_hat)

                loss = xent_loss + info_loss

                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
                
                epoch_loss += xent_loss.detach().numpy()

            relative_loss_div = (last_loss-epoch_loss)/(last_loss+1e-10)
            if  relative_loss_div < tol:
                if early_stop > stop:
                    print("[MF-CVIB] epoch:{}, xent:{}".format(epoch, epoch_loss))
                    break
                early_stop += 1
                
            last_loss = epoch_loss

            if epoch % 10 == 0 and verbose:
                print("[MF-CVIB] epoch:{}, xent:{}".format(epoch, epoch_loss))

            if epoch == num_epoch - 1:
                print("[MF-CVIB] Reach preset epochs, it seems does not converge.")

    def predict(self, x):
        pred = self.forward(x)
        pred = self.sigmoid(pred)
        return pred.detach().numpy()
    
    
class MF_DR(nn.Module):
    def __init__(self, num_users, num_items, batch_size, embedding_k=4, *args, **kwargs):
        super().__init__()
        self.num_users = num_users
        self.num_items = num_items
        self.embedding_k = embedding_k
        self.batch_size = batch_size
        self.prediction_model = MF_BaseModel(
            num_users = self.num_users, num_items = self.num_items, batch_size = self.batch_size, embedding_k=self.embedding_k, *args, **kwargs)
        self.original_model = MF(
            num_users = self.num_users, num_items = self.num_items, batch_size = self.batch_size, embedding_k=self.embedding_k, *args, **kwargs)

        self.linear_1 = torch.nn.Linear(self.embedding_k*2, 1)

        self.sigmoid = torch.nn.Sigmoid()
        self.xent_func = torch.nn.BCELoss()

    def _compute_IPS(self, x, y, gamma):
        
        x_train = torch.zeros([self.num_users,self.num_items])

        for i in range(len(x)):
            x_train[x[i][0],x[i][1]] = 1 # o
        prediction = (x_train.reshape(self.num_users*self.num_items,1)).type(torch.FloatTensor)

        x_train = self.original_model.complete().type(torch.FloatTensor).detach()
        
        
        optimizer = torch.optim.SGD([self.linear_1.weight, self.linear_1.bias], lr=1e-3, momentum=0.9)

        last_loss = 1e9
        early_stop = 0
        for epoch in range(1000):
            all_idx = np.arange(self.num_users*self.num_items)
            np.random.shuffle(all_idx)
            total_batch = (self.num_users*self.num_items)// self.batch_size
            for idx in range(total_batch):
                selected_idx = all_idx[self.batch_size*idx:(idx+1)*self.batch_size]
                sub_x_train = x_train[selected_idx].detach()
                sub_prediction = prediction[selected_idx]
             
                out = self.linear_1(sub_x_train)
                out = self.sigmoid(out)
                loss = self.xent_func(out, sub_prediction)
                
                xent_loss = loss

                optimizer.zero_grad()
                xent_loss.backward()
                optimizer.step()

            if (epoch + 1) % 15 == 0:
                print('*'*10)
                print('epoch {}'.format(epoch+1))
            
            if loss > last_loss:
                early_stop += 1 
            else:
                last_loss = loss
            
            if early_stop >= 5:
                break
        
        x_train = x_train.detach()
        propensity = self.sigmoid(self.linear_1(x_train.detach())) 
        propensity[np.where(propensity.cpu() <= gamma)] = gamma
        
        one_over_zl = 1 / propensity
        return prediction, one_over_zl  

    def fit(self, x, y, prior_y, gamma, stop = 5,
        num_epoch=1000, lr=0.05, lamb=0, 
        tol=1e-4, G = 1, verbose=True): 

        optimizer_prediction = torch.optim.Adam(self.prediction_model.parameters(), lr=lr, weight_decay=lamb)

        last_loss = 1e9

        # generate all counterfactuals and factuals
        x_all = generate_total_sample(self.num_users, self.num_items) 

        num_sample = len(x)
        total_batch = num_sample // self.batch_size

        observation, one_over_zl = self._compute_IPS(x, y, gamma)
        one_over_zl_obs = one_over_zl[np.where(observation.cpu() == 1)]
        
        prior_y = prior_y.mean()
        early_stop = 0
        for epoch in range(num_epoch):
            all_idx = np.arange(num_sample) 
            np.random.shuffle(all_idx)

            # sampling counterfactuals
            ul_idxs = np.arange(x_all.shape[0])
            np.random.shuffle(ul_idxs)

            epoch_loss = 0

            for idx in range(total_batch):
                # mini-batch training
                
                selected_idx = all_idx[self.batch_size*idx:(idx+1)*self.batch_size]
                sub_x = x[selected_idx]
                sub_y = y[selected_idx]

                # propensity score
                inv_prop = one_over_zl_obs[selected_idx].detach()

                sub_y = torch.Tensor(sub_y)

                pred, u_emb, v_emb = self.prediction_model.forward(sub_x, True)  
                pred = self.sigmoid(pred)

                x_sampled = x_all[ul_idxs[G * idx* self.batch_size: G * (idx+1)*self.batch_size]]

                pred_ul,_,_ = self.prediction_model.forward(x_sampled, True)
                pred_ul = self.sigmoid(pred_ul)

                xent_loss = F.binary_cross_entropy(pred, sub_y, weight=inv_prop, reduction="sum")
                
                imputation_y = torch.Tensor([prior_y]* G *selected_idx.shape[0])
                imputation_loss = F.binary_cross_entropy(pred, imputation_y[0:self.batch_size], reduction="sum")

                ips_loss = (xent_loss - imputation_loss)

                # direct loss
                direct_loss = F.binary_cross_entropy(pred_ul, imputation_y,reduction="sum")

                loss = (ips_loss + direct_loss)/x_sampled.shape[0]

                optimizer_prediction.zero_grad()
                loss.backward()
                optimizer_prediction.step()
                
                epoch_loss += xent_loss.detach().cpu().numpy()

            relative_loss_div = (last_loss-epoch_loss)/(last_loss+1e-10)
            if  relative_loss_div < tol:
                if early_stop > stop:
                    print("[MF-DR] epoch:{}, xent:{}".format(epoch, epoch_loss))
                    break
                early_stop += 1
                
            last_loss = epoch_loss

            if epoch % 10 == 0 and verbose:
                print("[MF-DR] epoch:{}, xent:{}".format(epoch, epoch_loss))

            if epoch == num_epoch - 1:
                print("[MF-DR] Reach preset epochs, it seems does not converge.")

    def predict(self, x):
        pred = self.prediction_model.forward(x)
        pred = self.sigmoid(pred)
        return pred.detach().cpu().numpy()


class MF_TDR(nn.Module):
    def __init__(self, num_users, num_items, batch_size, embedding_k=4, *args, **kwargs):
        super().__init__()
        self.num_users = num_users
        self.num_items = num_items
        self.embedding_k = embedding_k
        self.batch_size = batch_size
        self.prediction_model = MF_BaseModel(
            num_users = self.num_users, num_items = self.num_items, batch_size = self.batch_size, embedding_k=self.embedding_k, *args, **kwargs)
        self.original_model = MF(
            num_users = self.num_users, num_items = self.num_items, batch_size = self.batch_size, embedding_k=self.embedding_k, *args, **kwargs)

        self.linear_1 = torch.nn.Linear(self.embedding_k*2, 1)

        self.sigmoid = torch.nn.Sigmoid()
        self.xent_func = torch.nn.BCELoss()

    def _compute_IPS(self, x, y, gamma):
        
        x_train = torch.zeros([self.num_users,self.num_items])

        for i in range(len(x)):
            x_train[x[i][0],x[i][1]] = 1
        prediction = (x_train.reshape(self.num_users*self.num_items,1)).type(torch.FloatTensor)

        x_train = self.original_model.complete().type(torch.FloatTensor).detach()
        
        optimizer = torch.optim.SGD([self.linear_1.weight, self.linear_1.bias], lr=1e-3, momentum=0.9)

        last_loss = 1e9
        early_stop = 0
        for epoch in range(1000):
            all_idx = np.arange(self.num_users*self.num_items)
            np.random.shuffle(all_idx)
            total_batch = (self.num_users*self.num_items) // self.batch_size
            for idx in range(total_batch):
                selected_idx = all_idx[self.batch_size*idx:(idx+1)*self.batch_size]
                sub_x_train = x_train[selected_idx].detach()
                sub_prediction = prediction[selected_idx]
             
                out = self.linear_1(sub_x_train)
                out = self.sigmoid(out)
                loss = self.xent_func(out, sub_prediction)
                
                xent_loss = loss
                
                optimizer.zero_grad()
                xent_loss.backward()
                optimizer.step()

            if (epoch + 1) % 15 == 0:
                print('*'*10)
                print('epoch {}'.format(epoch+1))
              
            if loss > last_loss:
                early_stop += 1 
            else:
                last_loss = loss
            
            if early_stop >= 5:
                break
        
        x_train = x_train.detach()
        propensity = self.sigmoid(self.linear_1(x_train.detach())) 
        propensity[np.where(propensity.cpu() <= gamma)] = gamma
        one_over_zl = 1 / propensity
        return prediction, one_over_zl  

    def fit(self, x, y, prior_y, gamma, stop = 5,
        num_epoch=1000, lr=0.05, lamb=0, 
        tol=1e-4, G = 1, verbose=True): 

        optimizer_prediction = torch.optim.Adam(self.prediction_model.parameters(), lr=lr, weight_decay=lamb) 

        last_loss = 1e9

        # generate all counterfactuals and factuals
        x_all = generate_total_sample(self.num_users, self.num_items)

        num_sample = len(x) #6960 
        total_batch = num_sample // self.batch_size

        observation, one_over_zl = self._compute_IPS(x, y, gamma)
        one_over_zl_obs = one_over_zl[np.where(observation.cpu() == 1)]

        prior_y = prior_y.mean()
        early_stop = 0
        for epoch in range(num_epoch):
            all_idx = np.arange(num_sample)
            np.random.shuffle(all_idx)

            # sampling counterfactuals
            ul_idxs = np.arange(x_all.shape[0])
            np.random.shuffle(ul_idxs)

            epoch_loss = 0

            for idx in range(total_batch):
                
                # mini-batch training                
                selected_idx = all_idx[self.batch_size*idx:(idx+1)*self.batch_size]
                sub_x = x[selected_idx]
                sub_y = y[selected_idx]

                # propensity score
                inv_prop = one_over_zl_obs[selected_idx].detach()

                sub_y = torch.Tensor(sub_y)

                pred, u_emb, v_emb = self.prediction_model.forward(sub_x, True)  
                pred = self.sigmoid(pred)

                x_sampled = x_all[ul_idxs[G * idx* self.batch_size: G * (idx+1)*self.batch_size]]

                pred_ul,_,_ = self.prediction_model.forward(x_sampled, True)
                pred_ul = self.sigmoid(pred_ul)

                xent_loss = F.binary_cross_entropy(pred, sub_y, weight=inv_prop, reduction="sum")               
                
                imputation_y = torch.Tensor([prior_y]* G *selected_idx.shape[0])
                imputation_loss = F.binary_cross_entropy(pred, imputation_y[0:self.batch_size], reduction="sum")

                ips_loss = (xent_loss - imputation_loss)

                # direct loss
                direct_loss = F.binary_cross_entropy(pred_ul, imputation_y,reduction="sum")

                loss = (ips_loss + direct_loss)/x_sampled.shape[0]

                optimizer_prediction.zero_grad()
                loss.backward()
                optimizer_prediction.step()
                
                epoch_loss += xent_loss.detach().cpu().numpy()

            relative_loss_div = (last_loss-epoch_loss)/(last_loss+1e-10)
            if  relative_loss_div < tol:
                if early_stop > stop:
                    
                    print("[MF-TDR] epoch:{}, xent:{}".format(epoch, epoch_loss))

                    e_loss = F.binary_cross_entropy(pred, sub_y, reduction = "none")                    
                    e_hat_loss = F.binary_cross_entropy(pred, imputation_y[0:self.batch_size], reduction = "none")
                    
                    TMLE_beta = inv_prop-1
                    TMLE_alpha = e_loss - e_hat_loss
                    TMLE_epsilon = ((TMLE_alpha * TMLE_beta).sum()/(TMLE_beta * TMLE_beta).sum())
                    e_hat_TMLE = TMLE_epsilon.item() * (one_over_zl.float()- torch.tensor([1.])) 
                    e_hat_TMLE_obs = e_hat_TMLE[np.where(observation.cpu() == 1)]
                    
                    np.random.shuffle(all_idx)
                    np.random.shuffle(x_all)
                    
                    selected_idx = all_idx[0:self.batch_size]
                    sub_x = x[selected_idx]
                    sub_y = y[selected_idx]

                    # propensity score
                    inv_prop = one_over_zl_obs[selected_idx].detach()

                    sub_y = torch.Tensor(sub_y)

                    pred, u_emb, v_emb = self.prediction_model.forward(sub_x, True)  
                    pred = self.sigmoid(pred)

                    x_sampled = x_all[ul_idxs[0: G * self.batch_size]]

                    pred_ul,_,_ = self.prediction_model.forward(x_sampled, True)
                    pred_ul = self.sigmoid(pred_ul)

                    xent_loss = ((F.binary_cross_entropy(pred, sub_y, reduction="none") ** 2) * inv_prop).sum() # o*eui/pui
                        
                    imputation_loss = ((F.binary_cross_entropy(pred, imputation_y[0:self.batch_size], reduction="none") + e_hat_TMLE_obs[selected_idx].squeeze().detach()) ** 2).sum()
 
                    ips_loss = (xent_loss - imputation_loss)

                    sub_x_sampled_number = []
                    for i in x_sampled:
                        sub_x_sampled_number.append((self.num_items * i[0] + i[1]))
                    sub_x_sampled_number = np.array(sub_x_sampled_number)
                    
                    direct_loss = ((F.binary_cross_entropy(pred_ul, imputation_y, reduction="none") + e_hat_TMLE[sub_x_sampled_number].squeeze().detach()) ** 2).sum()                    
                    
                    loss = (ips_loss + direct_loss)/sub_x_sampled_number.shape[0]

                    optimizer_prediction.zero_grad()
                    loss.backward()
                    optimizer_prediction.step()
                    
                    break
                else:
                    early_stop += 1
                
            last_loss = epoch_loss

            if epoch % 10 == 0 and verbose:
                print("[MF-TDR] epoch:{}, xent:{}".format(epoch, epoch_loss))

            if epoch == num_epoch - 1:
                print("[MF-TDR] Reach preset epochs, it seems does not converge.")
                        
    def predict(self, x):
        pred = self.prediction_model.forward(x)
        pred = self.sigmoid(pred)
        return pred.detach().cpu().numpy()


class MF_DR_JL(nn.Module):
    def __init__(self, num_users, num_items, batch_size, embedding_k=4, *args, **kwargs):
        super().__init__()
        self.num_users = num_users
        self.num_items = num_items
        self.embedding_k = embedding_k
        self.batch_size = batch_size
        self.prediction_model = MF_BaseModel(
            num_users = self.num_users, num_items = self.num_items, batch_size = self.batch_size, embedding_k=self.embedding_k, *args, **kwargs)
        self.imputation = MF_BaseModel(
            num_users=self.num_users, num_items=self.num_items, batch_size = self.batch_size, embedding_k=self.embedding_k)
        self.original_model = MF(
            num_users = self.num_users, num_items = self.num_items, batch_size = self.batch_size, embedding_k=self.embedding_k, *args, **kwargs)

        self.linear_1 = torch.nn.Linear(self.embedding_k*2, 1)

        self.sigmoid = torch.nn.Sigmoid()
        self.xent_func = torch.nn.BCELoss()

    def _compute_IPS(self, x, y, gamma):
        
        x_train = torch.zeros([self.num_users ,self.num_items])

        for i in range(len(x)):
            x_train[x[i][0],x[i][1]] = 1 # o
        prediction = (x_train.reshape(self.num_users *self.num_items,1)).type(torch.FloatTensor)

        x_train = self.original_model.complete().type(torch.FloatTensor)
        
        optimizer = torch.optim.SGD([self.linear_1.weight, self.linear_1.bias], lr=1e-3, momentum=0.9)
        
        last_loss = 1e9
        early_stop = 0
        for epoch in range(1000):
            all_idx = np.arange(self.num_users *self.num_items)
            np.random.shuffle(all_idx)
            total_batch = (self.num_users *self.num_items)// self.batch_size
            for idx in range(total_batch):
                selected_idx = all_idx[self.batch_size*idx:(idx+1)*self.batch_size]
                sub_x_train = x_train[selected_idx].detach()
                sub_prediction = prediction[selected_idx]
             
                out = self.linear_1(sub_x_train)
                out = self.sigmoid(out)
                loss = self.xent_func(out, sub_prediction)
                
                xent_loss = loss

                optimizer.zero_grad()
                xent_loss.backward()
                optimizer.step()

            if (epoch + 1) % 15 == 0:
                print('*'*10)
                print('epoch {}'.format(epoch+1))
            
            if loss > last_loss:
                early_stop += 1 
            else:
                last_loss = loss
            
            if early_stop >= 5:
                break
        
        x_train = x_train.detach()
        propensity = self.sigmoid(self.linear_1(x_train.detach())) 
        propensity[np.where(propensity.cpu() <= gamma)] = gamma
        one_over_zl = 1 / propensity

        return prediction, one_over_zl  
        
    def fit(self, x, y, x_test, y_test, stop = 1,
        num_epoch=1000, lr=0.05, lamb=0, gamma = 0.1,
        tol=1e-4, G=1, verbose=True): 

        optimizer_prediction = torch.optim.Adam(
            self.prediction_model.parameters(), lr=lr, weight_decay=lamb)
        optimizer_imputation = torch.optim.Adam(
            self.imputation.parameters(), lr=lr, weight_decay=lamb)
        
        last_loss = 1e9
           
        # generate all counterfactuals and factuals
        x_all = generate_total_sample(self.num_users, self.num_items)

        num_sample = len(x) #6960 
        total_batch = num_sample // self.batch_size
        observation, one_over_zl = self._compute_IPS(x, y, gamma)

        early_stop = 0

        one_over_zl_obs = one_over_zl[np.where(observation.cpu() == 1)]
        
        for epoch in range(num_epoch): 
            all_idx = np.arange(num_sample)
            np.random.shuffle(all_idx)

            # sampling counterfactuals
            ul_idxs = np.arange(x_all.shape[0]) 
            np.random.shuffle(ul_idxs)

            epoch_loss = 0

            for idx in range(total_batch):

                # mini-batch training
                selected_idx = all_idx[self.batch_size*idx:(idx+1)*self.batch_size]
                sub_x = x[selected_idx] 
                sub_y = y[selected_idx]

                # propensity score
                inv_prop = one_over_zl_obs[selected_idx].detach()                
                
                sub_y = torch.Tensor(sub_y)
                        
                pred = self.prediction_model.forward(sub_x)
                imputation_y = self.imputation.predict(sub_x)
                pred = self.sigmoid(pred)
                imputation_y = self.sigmoid(imputation_y)
                                
                x_sampled = x_all[ul_idxs[G*idx* self.batch_size : G*(idx+1)*self.batch_size]]
                                       
                pred_u = self.prediction_model.forward(x_sampled) 
                imputation_y1 = self.imputation.predict(x_sampled)
                pred_u = self.sigmoid(pred_u)     
                imputation_y1 = self.sigmoid(imputation_y1)              
                               
                xent_loss = F.binary_cross_entropy(pred, sub_y, weight=inv_prop, reduction="sum") # o*eui/pui
                imputation_loss = F.binary_cross_entropy(pred, imputation_y, reduction="sum")

                ips_loss = (xent_loss - imputation_loss) # batch size

                direct_loss = F.binary_cross_entropy(pred_u, imputation_y1, reduction="sum")

                loss = (ips_loss + direct_loss)/x_sampled.shape[0]

                optimizer_prediction.zero_grad()
                loss.backward()
                optimizer_prediction.step()
                               
                epoch_loss += xent_loss.detach().cpu().numpy()                

                pred = self.prediction_model.predict(sub_x)
                pred = self.sigmoid(pred)
                imputation_y = self.imputation.forward(sub_x)
                imputation_y = self.sigmoid(imputation_y)
                
                e_loss = F.binary_cross_entropy(pred, sub_y, reduction="none")
                e_hat_loss = F.binary_cross_entropy(imputation_y, pred, reduction="none")
                imp_loss = (((e_loss.detach() - e_hat_loss) ** 2) * inv_prop).sum()

                optimizer_imputation.zero_grad()
                imp_loss.backward()
                optimizer_imputation.step()                
                
            relative_loss_div = (last_loss-epoch_loss)/(last_loss+1e-10)
            if  relative_loss_div < tol:
                if early_stop > stop:
                    print("[MF-DR-JL] epoch:{}, xent:{}".format(epoch, epoch_loss))
                    break
                else:
                    early_stop += 1
                
            last_loss = epoch_loss

            if epoch % 10 == 0 and verbose:
                print("[MF-DR-JL] epoch:{}, xent:{}".format(epoch, epoch_loss))

            if epoch == num_epoch - 1:
                print("[MF-DR-JL] Reach preset epochs, it seems does not converge.")
    
    def predict(self, x):
        pred = self.prediction_model.predict(x)
        pred = self.sigmoid(pred)
        return pred.detach().cpu().numpy()            

    
class MF_TDR_JL(nn.Module):
    def __init__(self, num_users, num_items, batch_size, embedding_k=4, *args, **kwargs):
        super().__init__()
        self.num_users = num_users
        self.num_items = num_items
        self.embedding_k = embedding_k
        self.batch_size = batch_size
        self.prediction_model = MF_BaseModel(
            num_users = self.num_users, num_items = self.num_items, batch_size = self.batch_size, embedding_k=self.embedding_k, *args, **kwargs)
        self.imputation = MF_BaseModel(
            num_users=self.num_users, num_items=self.num_items, batch_size = self.batch_size, embedding_k=self.embedding_k)
        self.original_model = MF(
            num_users = self.num_users, num_items = self.num_items, batch_size = self.batch_size, embedding_k=self.embedding_k, *args, **kwargs)

        self.linear_1 = torch.nn.Linear(self.embedding_k*2, 1)

        self.sigmoid = torch.nn.Sigmoid()
        self.xent_func = torch.nn.BCELoss()

    def _compute_IPS(self, x, y, gamma):
        
        x_train = torch.zeros([self.num_users,self.num_items])

        for i in range(len(x)):
            x_train[x[i][0],x[i][1]] = 1 # o
        prediction = (x_train.reshape(self.num_users*self.num_items,1)).type(torch.FloatTensor)

        x_train = self.original_model.complete().type(torch.FloatTensor).detach()
        
        optimizer = torch.optim.SGD([self.linear_1.weight, self.linear_1.bias], lr=1e-3, momentum=0.9)

        last_loss = 1e9 
        early_stop = 0
        for epoch in range(1000):
            all_idx = np.arange(self.num_users*self.num_items)
            np.random.shuffle(all_idx)
            total_batch = (self.num_users*self.num_items)// self.batch_size
            #print(early_stop)
            for idx in range(total_batch):
                selected_idx = all_idx[self.batch_size*idx:(idx+1)*self.batch_size]
                sub_x_train = x_train[selected_idx].detach()
                sub_prediction = prediction[selected_idx]
             
                out = self.linear_1(sub_x_train)
                out = self.sigmoid(out)
                loss = self.xent_func(out, sub_prediction)
                
                xent_loss = loss
                optimizer.zero_grad()
                xent_loss.backward()

                optimizer.step()

            if (epoch + 1) % 15 == 0:
                print('*'*10)
                print('epoch {}'.format(epoch+1))
            
            if loss > last_loss:
                early_stop += 1 
            else:
                last_loss = loss
            
            if early_stop >= 5:
                break
        
        x_train = x_train.detach()
        propensity = self.sigmoid(self.linear_1(x_train.detach())) 
        propensity[np.where(propensity.cpu() <= gamma)] = gamma
        one_over_zl = 1 / propensity
                      
        return prediction, one_over_zl  
        
    def fit(self, x, y, stop = 1,
        num_epoch=1000,lr=0.05, lamb=0, gamma = 0.1,
        tol=1e-4, G=1, verbose=True): 

        optimizer_prediction = torch.optim.Adam(
            self.prediction_model.parameters(), lr=lr, weight_decay=lamb)
        optimizer_imputation = torch.optim.Adam(
            self.imputation.parameters(), lr=lr, weight_decay=lamb)
     
        last_loss = 1e9
   
        # generate all counterfactuals and factuals
        x_all = generate_total_sample(self.num_users, self.num_items)

        num_sample = len(x) 
        total_batch = num_sample // self.batch_size
        observation, one_over_zl = self._compute_IPS(x, y, gamma)

        early_stop = 0
        one_over_zl_obs = one_over_zl[np.where(observation.cpu() == 1)]

        for epoch in range(num_epoch):            
            all_idx = np.arange(num_sample)
            np.random.shuffle(all_idx)

            # sampling counterfactuals
            ul_idxs = np.arange(x_all.shape[0])
            np.random.shuffle(ul_idxs)

            epoch_loss = 0

            for idx in range(total_batch):

                # mini-batch training
                selected_idx = all_idx[self.batch_size*idx:(idx+1)*self.batch_size]
                sub_x = x[selected_idx] 
                sub_y = y[selected_idx]

                # propensity score

                inv_prop = one_over_zl_obs[selected_idx].detach()                
                
                sub_y = torch.Tensor(sub_y)
                        
                pred = self.prediction_model.forward(sub_x)
                imputation_y = self.imputation.predict(sub_x)
                pred = self.sigmoid(pred)
                imputation_y = self.sigmoid(imputation_y)
                                
                x_sampled = x_all[ul_idxs[G*idx* self.batch_size : G*(idx+1)*self.batch_size]]
                                       
                pred_u = self.prediction_model.forward(x_sampled)
                imputation_y1 = self.imputation.predict(x_sampled)
                pred_u = self.sigmoid(pred_u)     
                imputation_y1 = self.sigmoid(imputation_y1)          
                                
                xent_loss = F.binary_cross_entropy(pred, sub_y, weight=inv_prop, reduction="sum") # o*eui/pui
                imputation_loss = F.binary_cross_entropy(pred, imputation_y, reduction="sum")
                        
                ips_loss = xent_loss - imputation_loss 
                                
                # direct loss                                
                direct_loss = F.binary_cross_entropy(pred_u, imputation_y1, reduction="sum")             

                loss = (ips_loss + direct_loss)/x_sampled.shape[0]

                optimizer_prediction.zero_grad()
                loss.backward()
                optimizer_prediction.step()
                                     
                epoch_loss += xent_loss.detach().cpu().numpy()                

                pred = self.prediction_model.predict(sub_x)
                pred = self.sigmoid(pred)
                imputation_y = self.imputation.forward(sub_x)
                imputation_y = self.sigmoid(imputation_y)
                    
                e_loss = F.binary_cross_entropy(pred, sub_y, reduction="none")
                e_hat_loss = F.binary_cross_entropy(imputation_y, pred, reduction="none")
                imp_loss = (((e_loss.detach() - e_hat_loss) ** 2) * inv_prop).sum()

                optimizer_imputation.zero_grad()
                imp_loss.backward()
                optimizer_imputation.step()                
                
            relative_loss_div = (last_loss-epoch_loss)/(last_loss+1e-10)
            if  relative_loss_div < tol:
                if early_stop > stop:
                    print("[MF-TDR-JL] epoch:{}, xent:{}".format(epoch, epoch_loss))
                    
                    e_loss = F.binary_cross_entropy(pred, sub_y, reduction="none")
                    e_hat_loss = F.binary_cross_entropy(imputation_y, pred, reduction="none")   
                    
                    TMLE_beta = inv_prop-1
                    TMLE_alpha = e_loss - e_hat_loss
                    TMLE_epsilon = ((TMLE_alpha * TMLE_beta).sum()/(TMLE_beta * TMLE_beta).sum())
                    e_hat_TMLE = TMLE_epsilon.item() * (one_over_zl.float()- torch.tensor([1.]))
                    e_hat_TMLE_obs = e_hat_TMLE[np.where(observation.cpu() == 1)]

                    np.random.shuffle(x_all)
                    np.random.shuffle(all_idx)
                    
                    selected_idx = all_idx[0:self.batch_size]
                    sub_x = x[selected_idx] 
                    sub_y = y[selected_idx]

                    inv_prop = one_over_zl_obs[selected_idx].detach()                
                
                    sub_y = torch.Tensor(sub_y)
                       
                    pred = self.prediction_model.forward(sub_x)
                    imputation_y = self.imputation.predict(sub_x)
                    pred = self.sigmoid(pred)
                    imputation_y = self.sigmoid(imputation_y)
                               
                    x_sampled = x_all[ul_idxs[0 : G*self.batch_size]]
                                       
                    pred_u = self.prediction_model.forward(x_sampled) 
                    imputation_y1 = self.imputation.predict(x_sampled)
                    pred_u = self.sigmoid(pred_u)     
                    imputation_y1 = self.sigmoid(imputation_y1)                             
                
                    xent_loss = ((F.binary_cross_entropy(pred, sub_y, reduction ="none") ** 2) * inv_prop).sum()
                    imputation_loss = ((F.binary_cross_entropy(pred, imputation_y, reduction="none")
                                        + e_hat_TMLE_obs[selected_idx].squeeze().detach()) ** 2).sum()
                        
                    ips_loss = xent_loss - imputation_loss
                    
                    # direct loss
                    sub_x_sampled_number = []
                    for i in x_sampled:
                        sub_x_sampled_number.append((self.num_items * i[0] + i[1]))
                    sub_x_sampled_number = np.array(sub_x_sampled_number)                 
                
                    direct_loss = ((F.binary_cross_entropy(pred_u, imputation_y1, reduction="none") + e_hat_TMLE[sub_x_sampled_number].squeeze().detach()) ** 2).sum()
                    
                    loss = (ips_loss + direct_loss)/sub_x_sampled_number.shape[0]
                    
                    optimizer_prediction.zero_grad()
                    loss.backward()
                    optimizer_prediction.step()
                    break
                else:
                    early_stop += 1
                
            last_loss = epoch_loss

            if epoch % 10 == 0 and verbose:
                print("[MF-TDR-JL] epoch:{}, xent:{}".format(epoch, epoch_loss))

            if epoch == num_epoch - 1:
                print("[MF-TDR-JL] Reach preset epochs, it seems does not converge.")

    def predict(self, x):
        pred = self.prediction_model.predict(x)
        pred = self.sigmoid(pred)
        return pred.detach().cpu().numpy()        

    
class MF_DR_CL(nn.Module):
    def __init__(self, num_users, num_items, batch_size, embedding_k=4, *args, **kwargs):
        super().__init__()
        self.num_users = num_users
        self.num_items = num_items
        self.embedding_k = embedding_k
        self.batch_size = batch_size
        self.prediction_model = MF_BaseModel(
            num_users=self.num_users, num_items=self.num_items, batch_size = self.batch_size, embedding_k=self.embedding_k)
        self.imputation = MF_BaseModel(
            num_users=self.num_users, num_items=self.num_items, batch_size = self.batch_size, embedding_k=self.embedding_k)
        self.original_model = MF(
            num_users=self.num_users, num_items=self.num_items, batch_size = self.batch_size, embedding_k=self.embedding_k)
        
        self.linear_1 = torch.nn.Linear(self.embedding_k*2, 1)
        
        self.sigmoid = torch.nn.Sigmoid()
        self.xent_func = torch.nn.BCELoss()
        
    def _compute_IPS(self, x, y):
        
        x_train = torch.zeros([self.num_users,self.num_items])

        for i in range(len(x)):
            x_train[x[i][0],x[i][1]] = 1 # o
        prediction = (x_train.reshape(self.num_users*self.num_items,1)).type(torch.FloatTensor)

        x_train = self.original_model.complete().type(torch.FloatTensor)
                
        optimizer = torch.optim.SGD([self.linear_1.weight, self.linear_1.bias], lr=1e-3, momentum=0.9)

        last_loss = 1e9
        early_stop = 0
        for epoch in range(1000):
            all_idx = np.arange(self.num_users*self.num_items)
            np.random.shuffle(all_idx)
            total_batch = (self.num_users*self.num_items)// self.batch_size                      
            for idx in range(total_batch):
                selected_idx = all_idx[self.batch_size*idx:(idx+1)*self.batch_size]
                sub_x_train = x_train[selected_idx].detach()
                sub_prediction = prediction[selected_idx]
             
                out = self.linear_1(sub_x_train)
                out = self.sigmoid(out)
                loss = self.xent_func(out, sub_prediction)
                
                xent_loss = loss

                optimizer.zero_grad()
                xent_loss.backward()
                optimizer.step()
                      
            if (epoch + 1) % 15 == 0:
                print('*'*10)
                print('epoch {}'.format(epoch+1))
            
            if loss > last_loss:
                early_stop += 1 
            else:
                last_loss = loss
            
            if early_stop >= 5:
                break
        
        x_train = x_train.detach()
        return prediction, x_train          
        
    def fit(self, x, y, x_test, y_test, stop = 1,
        num_epoch=1000, lr=0.05, lamb=0, gamma = 0.1,
        tol=1e-4, G=1, verbose=True): 

        optimizer_prediction = torch.optim.Adam(
            self.prediction_model.parameters(), lr=lr, weight_decay=lamb)
        optimizer_imputation = torch.optim.Adam(
            self.imputation.parameters(), lr=lr, weight_decay=lamb)
        optimizer_propensity = torch.optim.Adam(
            [self.linear_1.weight, self.linear_1.bias], lr=lr, weight_decay=lamb)
     
        last_loss = 1e9

        # generate all counterfactuals and factuals
        x_all = generate_total_sample(self.num_users, self.num_items)

        num_sample = len(x) #6960 
        total_batch = num_sample // self.batch_size
        observation, x_train = self._compute_IPS(x, y)

        early_stop = 0
        
        for epoch in range(num_epoch):            
            all_idx = np.arange(num_sample)
            np.random.shuffle(all_idx)

            # sampling counterfactuals
            ul_idxs = np.arange(x_all.shape[0])
            np.random.shuffle(ul_idxs)

            epoch_loss = 0
            a = self.sigmoid(self.linear_1(x_train.detach()))
            b = a.clone()
            b[np.where(b.cpu() <= gamma)] = gamma
            propensity = b

            one_over_zl = 1 / propensity
            one_over_zl_obs = one_over_zl[np.where(observation.cpu() == 1)]
            one_over_zl = torch.squeeze(one_over_zl)

            for idx in range(total_batch):
                # mini-batch training
                selected_idx = all_idx[self.batch_size*idx:(idx+1)*self.batch_size]
                inv_prop = one_over_zl_obs[selected_idx]                
                inv_prop = torch.squeeze(inv_prop)
                sub_x = x[selected_idx] 
                sub_y = y[selected_idx]

                # propensity score                
                sub_y = torch.Tensor(sub_y)
                        
                pred = self.prediction_model.forward(sub_x)
                imputation_y = self.imputation.predict(sub_x)
                pred = self.sigmoid(pred)
                imputation_y = self.sigmoid(imputation_y)

                x_sampled = x_all[ul_idxs[G*idx* self.batch_size : G*(idx+1)*self.batch_size]]

                pred_u = self.prediction_model.forward(x_sampled)
                imputation_y1 = self.imputation.predict(x_sampled)
                pred_u = self.sigmoid(pred_u)
                imputation_y1 = self.sigmoid(imputation_y1)
                
                xent_loss = F.binary_cross_entropy(pred, sub_y, weight=inv_prop.detach(), reduction="sum")
                imputation_loss = F.binary_cross_entropy(pred, imputation_y, reduction="sum")
                        
                ips_loss = (xent_loss - imputation_loss)                
                
                # direct loss
                direct_loss = F.binary_cross_entropy(pred_u, imputation_y1, reduction="sum")
                 
                # propensity loss                
                pred_prop = 1/(one_over_zl)
                observation = torch.squeeze(observation)
                prop_loss = F.binary_cross_entropy(pred_prop, observation)               
                
                loss = (ips_loss + direct_loss)/x_sampled.shape[0] + prop_loss
                      
                optimizer_prediction.zero_grad()
                optimizer_propensity.zero_grad()
                loss.backward()    
                optimizer_prediction.step()
                optimizer_propensity.step()

                epoch_loss += xent_loss.detach().cpu().numpy()
                
                a = self.sigmoid(self.linear_1(x_train.detach()))
                b = a.clone()
                b[np.where(b.cpu() <= gamma)] = gamma
                propensity = b

                one_over_zl = 1 / propensity
                one_over_zl_obs = one_over_zl[np.where(observation.cpu() == 1)]
                one_over_zl = torch.squeeze(one_over_zl)
                
                inv_prop = one_over_zl_obs[selected_idx]                
                inv_prop = torch.squeeze(inv_prop)
                
                pred = self.prediction_model.predict(sub_x)
                pred = self.sigmoid(pred)
                imputation_y = self.imputation.forward(sub_x)
                imputation_y = self.sigmoid(imputation_y)

                e_loss = F.binary_cross_entropy(pred, sub_y, reduction="none")
                e_hat_loss = F.binary_cross_entropy(imputation_y, pred, reduction="none")
                imp_loss = (((e_loss.detach() - e_hat_loss) ** 2) * inv_prop.detach()).sum()

                optimizer_imputation.zero_grad()
                imp_loss.backward()
                optimizer_imputation.step()                

            relative_loss_div = (last_loss-epoch_loss)/(last_loss+1e-10)
            if  relative_loss_div < tol:
                if early_stop > stop:
                    print("[MF-DR-CL] epoch:{}, xent:{}".format(epoch, epoch_loss))
                    break
                early_stop += 1
                
            last_loss = epoch_loss

            if epoch % 5 == 0 and verbose:
                print("[MF-DR-CL] epoch:{}, xent:{}".format(epoch, epoch_loss))

            if epoch == num_epoch - 1:
                print("[MF-DR-CL] Reach preset epochs, it seems does not converge.")      
    
    def predict(self, x):
        pred = self.prediction_model.predict(x)
        pred = self.sigmoid(pred)
        return pred.detach().cpu().numpy()    

    
class MF_TDR_CL(nn.Module):
    def __init__(self, num_users, num_items, batch_size, embedding_k=4, *args, **kwargs):
        super().__init__()
        self.num_users = num_users
        self.num_items = num_items
        self.embedding_k = embedding_k
        self.batch_size = batch_size
        self.prediction_model = MF_BaseModel(
            num_users=self.num_users, num_items=self.num_items, batch_size = self.batch_size, embedding_k=self.embedding_k)
        self.imputation = MF_BaseModel(
            num_users=self.num_users, num_items=self.num_items, batch_size = self.batch_size, embedding_k=self.embedding_k)
        self.original_model = MF(
            num_users=self.num_users, num_items=self.num_items, batch_size = self.batch_size, embedding_k=self.embedding_k)
        
        self.linear_1 = torch.nn.Linear(self.embedding_k*2, 1)
        
        self.sigmoid = torch.nn.Sigmoid()
        self.xent_func = torch.nn.BCELoss()

        
    def _compute_IPS(self, x, y):
                      
        x_train = torch.zeros([self.num_users,self.num_items])

        for i in range(len(x)):
            x_train[x[i][0],x[i][1]] = 1 # o
        prediction = (x_train.reshape(self.num_users*self.num_items,1)).type(torch.FloatTensor)

        x_train = self.original_model.complete().type(torch.FloatTensor)

        optimizer = torch.optim.SGD([self.linear_1.weight, self.linear_1.bias], lr=1e-3, momentum=0.9)
                      
        last_loss = 1e9
        early_stop = 0
        for epoch in range(1000):
            all_idx = np.arange(self.num_users*self.num_items)
            np.random.shuffle(all_idx)
            total_batch = (self.num_users*self.num_items)// self.batch_size
            for idx in range(total_batch):
                selected_idx = all_idx[self.batch_size*idx:(idx+1)*self.batch_size]
                sub_x_train = x_train[selected_idx].detach()
                sub_prediction = prediction[selected_idx]
             
                out = self.linear_1(sub_x_train)
                out = self.sigmoid(out)
                loss = self.xent_func(out, sub_prediction)
                
                xent_loss = loss

                optimizer.zero_grad()
                xent_loss.backward()
                optimizer.step()

            if (epoch + 1) % 15 == 0:
                print('*'*10)
                print('epoch {}'.format(epoch+1))
            
            if loss > last_loss:
                early_stop += 1 
            else:
                last_loss = loss
            
            if early_stop >= 5:
                break
                
        x_train = x_train.detach()

        return prediction, x_train          
        
    def fit(self, x, y, x_test, y_test, stop = 1, 
        num_epoch=1000, lr=0.05, lamb=0, gamma = 0.1,
        tol=1e-4, G=1, verbose=True): 

        optimizer_prediction = torch.optim.Adam(
            self.prediction_model.parameters(), lr=lr, weight_decay=lamb)
        optimizer_imputation = torch.optim.Adam(
            self.imputation.parameters(), lr=lr, weight_decay=lamb)
        optimizer_propensity = torch.optim.Adam(
            [self.linear_1.weight, self.linear_1.bias], lr=lr, weight_decay=lamb)
     
        last_loss = 1e9

        x_all = generate_total_sample(self.num_users, self.num_items)

        num_sample = len(x)
        total_batch = num_sample // self.batch_size
        observation, x_train = self._compute_IPS(x, y)

        early_stop = 0
        
        e_hat_TMLE = torch.Tensor([0]* self.num_users*self.num_items)
        TMLE_epsilon = 0
        for epoch in range(num_epoch):
            
            all_idx = np.arange(num_sample)
            np.random.shuffle(all_idx)

            ul_idxs = np.arange(x_all.shape[0])
            np.random.shuffle(ul_idxs)

            epoch_loss = 0
            
            a = self.sigmoid(self.linear_1(x_train.detach()))
            b = a.clone()
            b[np.where(b.cpu() <= gamma)] = gamma
            propensity = b

            one_over_zl = 1 / propensity
            one_over_zl_obs = one_over_zl[np.where(observation.cpu() == 1)]
            one_over_zl = torch.squeeze(one_over_zl)
            
            for idx in range(total_batch):

                # mini-batch training
                selected_idx = all_idx[self.batch_size*idx:(idx+1)*self.batch_size]
                sub_x = x[selected_idx] 
                sub_y = y[selected_idx]

                # propensity score
                inv_prop = one_over_zl_obs[selected_idx]                
                inv_prop = torch.squeeze(inv_prop)
                
                sub_y = torch.Tensor(sub_y)
                        
                pred = self.prediction_model.forward(sub_x)
                imputation_y = self.imputation.predict(sub_x)
                pred = self.sigmoid(pred)
                imputation_y = self.sigmoid(imputation_y)
                                
                x_sampled = x_all[ul_idxs[G*idx* self.batch_size : G*(idx+1)*self.batch_size]]
                                       
                pred_u = self.prediction_model.forward(x_sampled) 
                imputation_y1 = self.imputation.predict(x_sampled)
                pred_u = self.sigmoid(pred_u)     
                imputation_y1 = self.sigmoid(imputation_y1)
                
                sub_x_number = []
                for i in sub_x:
                    sub_x_number.append((self.num_items * i[0] + i[1]))
                sub_x_number = np.array(sub_x_number)
                
                xent_loss = F.binary_cross_entropy(pred, sub_y, weight=inv_prop.detach(), reduction="sum")             
                
                imputation_loss = (F.binary_cross_entropy(pred, imputation_y, reduction="none") +
                                    e_hat_TMLE[sub_x_number].squeeze().detach()).sum()
                
                ips_loss = (xent_loss - imputation_loss)                
            
                sub_x_sampled_number = []
                for i in x_sampled:
                    sub_x_sampled_number.append((self.num_items * i[0] + i[1]))
                sub_x_sampled_number = np.array(sub_x_sampled_number)
                
                direct_loss = (F.binary_cross_entropy(pred_u, imputation_y1, reduction="none") + 
                               e_hat_TMLE[sub_x_sampled_number].squeeze().detach()).sum()
                
                # propensity loss                
                pred_prop = 1/(one_over_zl)

                observation = torch.squeeze(observation)
                prop_loss = F.binary_cross_entropy(pred_prop, observation, reduction = "mean")
                
                loss = (ips_loss + direct_loss)/x_sampled.shape[0] + prop_loss
                
                optimizer_prediction.zero_grad()
                optimizer_propensity.zero_grad()                
                loss.backward()                                           
                optimizer_prediction.step()               
                optimizer_propensity.step()
                                
                epoch_loss += xent_loss.detach().cpu().numpy()    
                
                a = self.sigmoid(self.linear_1(x_train.detach()))
                b = a.clone()
                b[np.where(b.cpu() <= gamma)] = gamma
                propensity = b

                one_over_zl = 1 / propensity
                one_over_zl_obs = one_over_zl[np.where(observation.cpu() == 1)]
                one_over_zl = torch.squeeze(one_over_zl)
                
                inv_prop = one_over_zl_obs[selected_idx]
                inv_prop = torch.squeeze(inv_prop)
                
                pred = self.prediction_model.predict(sub_x)
                pred = self.sigmoid(pred)
                imputation_y = self.imputation.forward(sub_x)
                imputation_y = self.sigmoid(imputation_y)
                    
                e_loss = F.binary_cross_entropy(pred, sub_y, reduction="none")
                e_hat_loss = F.binary_cross_entropy(imputation_y, pred, reduction="none")
                imp_loss = (((e_loss.detach() - e_hat_loss - e_hat_TMLE[sub_x_number].squeeze().detach()) ** 2) * inv_prop.detach()).sum()

                optimizer_imputation.zero_grad()
                imp_loss.backward()
                optimizer_imputation.step()                
                with torch.no_grad():
                    e_loss = F.binary_cross_entropy(pred, sub_y, reduction="none")
                    e_hat_loss = F.binary_cross_entropy(imputation_y, pred, reduction="none")                

                    TMLE_beta = inv_prop-1
                    TMLE_alpha = e_loss - e_hat_loss
                    TMLE_epsilon = ((TMLE_alpha * TMLE_beta).sum()/(TMLE_beta * TMLE_beta).sum())

                    one_over_zl -= 1

                    e_hat_TMLE += TMLE_epsilon.item() * (one_over_zl)
                
                #print("[MF-DR-TMLE-TL] batch:{}, mse:{}, auc:{}".format(total_batch*epoch + idx, mse_mfips, auc_mfips))        
            relative_loss_div = (last_loss-epoch_loss)/(last_loss+1e-10)
            if  relative_loss_div < tol:
                if early_stop > stop:
                    print("[MF-TDR-CL] epoch:{}, xent:{}".format(epoch, epoch_loss))
                    break
                early_stop += 1
                
            last_loss = epoch_loss

            if epoch % 10 == 0 and verbose:
                print("[MF-TDR-CL] epoch:{}, xent:{}".format(epoch, epoch_loss))

            if epoch == num_epoch - 1:
                print("[MF-TDR-CL] Reach preset epochs, it seems does not converge.")
        
    def predict(self, x):
        pred = self.prediction_model.predict(x)
        pred = self.sigmoid(pred)
        return pred.detach().cpu().numpy()        
