//  [python]
def PREMA(X, O1, O2, mask, R, iteration, device):
    """
    PREMA:
        notation mapping (to the paper):
            A1 -> U, A2 -> V, A3 -> W, tmp1 -> Q1, tmp2 -> Q2
            O1 -> C1, O2 -> C2
    """

    # preparation
    new_d1, d2, d3 = O1.shape
    d1, new_d2, d3 = O2.shape
    A1, A2, A3, tmp1, tmp2 = PREMA_init(X, O1, O2, R, device)
    print ('finish initialization')
    print ()
    loss_list, rec_list = [], []
    
    # intermediate constant data (for fast gradient computation)
    X_1 = X.reshape(d1, -1)
    X_2 = X.permute(1,0,2).reshape(d2, -1)
    X_3 = X.permute(2,0,1).reshape(d3, -1)
    M1 = mask.reshape(d1, -1)
    M2 = mask.permute(1,0,2).reshape(d2, -1)
    M3 = mask.permute(2,0,1).reshape(d3, -1)

    def obj_func(X, coord, grad):
        """
        X = (A1, A2, A3, tmp1, tmp2)
        coord decides which one get the gradient
        """
        X[coord] -= grad
        A1, A2, A3, tmp1, tmp2 = X

        rec = torch.einsum('ir,jr,kr->ijk', A1, A2, A3)
        LOSS1 = torch.norm(O1 - torch.einsum('ir,jr,kr->ijk',tmp1,A2,A3)) / torch.norm(O1)
        LOSS2 = torch.norm(O2 - torch.einsum('ir,jr,kr->ijk',A1,tmp2,A3)) / torch.norm(O2)
        LOSS3 = torch.norm(mask * (rec - X)) / torch.norm(mask * X)
        LOSS = Lambda * LOSS1 ** 2 + Lambda * LOSS2 ** 2 + LOSS3 ** 2
        
        return LOSS

    # Kharti-rao product
    def kr(A, B):
        return torch.cat([torch.outer(A[:,i], B[:,i]).view(-1, 1) for i in range(A.shape[1])], dim=1)

    def grad_tmp1(Lambda):
        return - torch.einsum('ijk,jr,kr->ir',O1,A2,A3) + torch.einsum('ar,br,bj,cr,cj->aj',tmp1,A2,A2,A3,A3)
    
    def grad_tmp2(Lambda):
        return - torch.einsum('ijk,ir,kr->jr',O2,A1,A3) + torch.einsum('ar,br,bj,cr,cj->aj',tmp2,A1,A1,A3,A3)

    def grad_A1(Lambda):
        return - Lambda * torch.einsum('ijk,jr,kr->ir',O2,tmp2,A3) + Lambda * torch.einsum('ar,br,bj,cr,cj->aj',A1,tmp2,tmp2,A3,A3)\
                            + (M1 * (A1 @ kr(A2, A3).T - X_1)) @ kr(A2, A3) + beta * torch.ones((d1, new_d1)).to(device) @ tmp1
    def grad_A2(Lambda):
        return - Lambda * torch.einsum('ijk,ir,kr->jr',O1,tmp1,A3) + Lambda * torch.einsum('ar,br,bj,cr,cj->aj',A2,tmp1,tmp1,A3,A3)\
                            + (M2 * (A2 @ kr(A1, A3).T - X_2)) @ kr(A1, A3) + beta * torch.ones((d2, new_d2)).to(device) @ tmp2
    def grad_A3(Lambda):
        return - Lambda * torch.einsum('ijk,ir,jr->kr',O1,tmp1,A2) + Lambda * torch.einsum('ar,br,bj,cr,cj->aj',A3,tmp1,tmp1,A2,A2)\
                - Lambda * torch.einsum('ijk,ir,jr->kr',O2,A1,tmp2) + Lambda * torch.einsum('ar,br,bj,cr,cj->aj',A3,A1,A1,tmp2,tmp2)\
                           + (M3 * (A3 @ kr(A1, A2).T - X_3)) @ kr(A1, A2)

    def line_search(obj_func, X, coord, grad, alpha=[1e-9, 3e-7], visit=[-1, -1], depth=3):

        if depth == 0:
            return alpha[0]

        result = []
        for i, j in zip(alpha, visit):
            result.append(obj_func(X, coord, i * grad)) if j < 0 else result.append(j)
        mid = sum(alpha) / 2

        if result[0] >= result[1]:
            return line_search(obj_func, X, coord, grad, alpha=[mid, alpha[1]], visit=[-1, result[1]], depth=depth-1)
        else:
            return line_search(obj_func, X, coord, grad, alpha=[alpha[0], mid], visit=[result[0], -1], depth=depth-1)


    alpha = 2e-8
    beta = 1e-3
    
    # start gradient descent
    tic = time.time()
    for i in range(iteration):
        Lambda = max(1. * np.exp(-i / 50), 1e-3)

        GRAD = grad_tmp1(Lambda)
        # alpha = line_search(obj_func, [A1, A2, A3, tmp1, tmp2], 3, GRAD)
        tmp1 -= alpha * GRAD

        GRAD = grad_tmp2(Lambda)
        # alpha = line_search(obj_func, [A1, A2, A3, tmp1, tmp2], 4, GRAD)
        tmp2 -= alpha * GRAD

        GRAD = grad_A1(Lambda)
        # alpha = line_search(obj_func, [A1, A2, A3, tmp1, tmp2], 0, GRAD)
        A1 -= alpha * GRAD

        GRAD = grad_A2(Lambda)
        # alpha = line_search(obj_func, [A1, A2, A3, tmp1, tmp2], 1, GRAD)
        A2 -= alpha * GRAD
        
        GRAD = grad_A3(Lambda)
        # alpha = line_search(obj_func, [A1, A2, A3, tmp1, tmp2], 2, GRAD)
        A3 -= alpha * GRAD
        
        # loss
        rec = torch.einsum('ir,jr,kr->ijk', A1, A2, A3)
        LOSS1 = torch.norm(O1 - torch.einsum('ir,jr,kr->ijk',tmp1,A2,A3)) / torch.norm(O1)
        LOSS2 = torch.norm(O2 - torch.einsum('ir,jr,kr->ijk',A1,tmp2,A3)) / torch.norm(O2)
        LOSS3 = torch.norm(mask * (rec - X)) / torch.norm(mask * X)
        LOSS = Lambda * LOSS1 ** 2 + Lambda * LOSS2 ** 2 + LOSS3 ** 2
        RecLOSS = 1 - torch.norm(rec - X) / torch.norm(X)
        if i % 20 == 0:
            print ('{}/{}'.format(i, iteration), 'loss:', LOSS.item(), 'rec:', RecLOSS.item(), 'time span: ', time.time() - tic)

        # collect loss
        tic = time.time()

        loss_list.append(LOSS); 
        rec_list.append(RecLOSS)

        # rescaling
        norm1 = (A1**2).sum(dim=0) ** .5
        norm2 = (A2**2).sum(dim=0) ** .5
        norm3 = (A3**2).sum(dim=0) ** .5
        norm = (norm1 * norm2 * norm3) ** (1/3.)

        A1 *= norm / norm1
        A2 *= norm / norm2
        A3 *= norm / norm3
        tmp1 *= norm / norm1 
        tmp2 *= norm / norm2
        
    return A1, A2, A3, tmp1, tmp2, loss_list, rec_list