# Path: sparse/util/util.py
# def prod(tuple):
#     result = 1
# 
#     for v in tuple:
#         result *= v
# 
#     return result


import torch
import torch.nn.functional as F
import util, sys
from torch import FloatTensor, LongTensor
from torch.autograd import Variable
from sparse.util import prod
from util import d
    :param indices: Long tensor
    :param in_rank:
    :return: (1) A matrix of size N by 2, (2) the dimensions of M
    """

    batchsize, n, rank = indices.size()

    inrank = len(in_shape)
    outrank = len(out_shape)

    result = torch.cuda.LongTensor(batchsize, n, 2) if indices.is_cuda else LongTensor(batchsize, n, 2)

    left = fi_matrix(indices[:, :, 0:outrank], out_shape)   # i index of the weight matrix
    right = fi_matrix(indices[:, :, outrank:rank], in_shape) # j index

    result = torch.cat([left.unsqueeze(2), right.unsqueeze(2)], dim=2)

    return result, LongTensor((prod(out_shape), prod(in_shape)))

def fi_matrix(indices, shape):
    batchsize, rows, rank = indices.size()

    prod = torch.LongTensor(rank).fill_(1)

    if indices.is_cuda:
        prod = prod.cuda()

    for i in range(rank):
        prod[i] = 1
        for j in range(i + 1, len(shape)):

# Path: sparse/util/util.py
# def prod(tuple):
#     result = 1
# 
#     for v in tuple:
#         result *= v
# 
#     return result


# Path: sparse/tensors.py
import torch
import torch.nn.functional as F
import util, sys
from torch import FloatTensor, LongTensor
from torch.autograd import Variable
from sparse.util import prod
from util import d

    :param indices: Long tensor
    :param in_rank:
    :return: (1) A matrix of size N by 2, (2) the dimensions of M
    """

    batchsize, n, rank = indices.size()

    inrank = len(in_shape)
    outrank = len(out_shape)

    result = torch.cuda.LongTensor(batchsize, n, 2) if indices.is_cuda else LongTensor(batchsize, n, 2)

    left = fi_matrix(indices[:, :, 0:outrank], out_shape)   # i index of the weight matrix
    right = fi_matrix(indices[:, :, outrank:rank], in_shape) # j index

    result = torch.cat([left.unsqueeze(2), right.unsqueeze(2)], dim=2)

    return result, LongTensor((prod(out_shape), prod(in_shape)))

def fi_matrix(indices, shape):
    batchsize, rows, rank = indices.size()

    prod = torch.LongTensor(rank).fill_(1)

    if indices.is_cuda:
        prod = prod.cuda()

    for i in range(rank):
        prod[i] = 1
        for j in range(i + 1, len(shape)):

            prod[i] *= shape[j]