# Author: Pearu Peterson
# Created: July 2020

import itertools
import torch

def generate_indices(shape):
    for indices in itertools.product(*[list(range(d)) for d in shape]):
        yield indices

def generate_coo_tensor(shape):
    indices_lst = []
    values_lst = []
    for indices in generate_indices(shape):
        indices_lst.append(indices)
        values_lst.append(torch.randn(()))
        yield torch.sparse_coo_tensor(torch.tensor(indices_lst).T, torch.tensor(values_lst), shape).coalesce()


def get_storage_size(a):
    if a.is_sparse:
        return get_storage_size(a._indices()) + get_storage_size(a._values())
    return a.numel() * a.element_size()

def find_sparse_dense_equivalence(N, last_sparsity=None):

    shape = (N,) * 2

    # sparse tensors using about the same amount of memory as dense
    # tensors in the context of matrix multiplication
    same_storage_coalesce = None
    same_storage_uncoalesce = None

    for A in generate_coo_tensor(shape):
        assert A.is_coalesced()
        density = A._values().numel() / A.numel()

        B = A.to_dense()
        AA = torch.sparse.mm(A, A)
        BB = torch.mm(B, B)

        data = dict(
            sparse_mm_storage_uncoalesce = 2*get_storage_size(A) + get_storage_size(AA),
            sparse_mm_storage_coalesce = 2*get_storage_size(A) + get_storage_size(AA.coalesce()),
            dense_mm_storage = 2*get_storage_size(A) + 2*get_storage_size(B) + get_storage_size(BB),
            sparse_vs_dense = get_storage_size(A) / get_storage_size(B),
            sparsity = 1 - density
        )
        if same_storage_coalesce is None and data['dense_mm_storage'] < data['sparse_mm_storage_coalesce']:
            data['kind'] = 'coalesce'
            same_storage_coalesce = data
            break
        if same_storage_uncoalesce is None and data['dense_mm_storage'] < data['sparse_mm_storage_uncoalesce']:
            data['kind'] = 'uncoalesce'
            same_storage_uncoalesce = data
    return same_storage_coalesce, same_storage_uncoalesce

for N in [10, 20, 30, 40, 50, 100]:
    coalesce, uncoalesce = find_sparse_dense_equivalence(N)
    print(f'N={N}, coalesce sparsity={coalesce["sparsity"]:.4f}, uncoalesce sparsity={uncoalesce["sparsity"]:.4f}')
