from functools import wraps
import os
import pickle

import numpy as np
from scipy.sparse import coo_matrix

def pickle_cache(func):
    @wraps(func)
    def function_wrapper(*args, cache, reprocess=False, **kwargs):
        if os.path.exists(cache) and not reprocess:
            with open(cache, 'rb') as f:
                output = pickle.load(f)
            return output
        else:
            output = func(*args, **kwargs)
            with open(cache, 'wb') as f:
                pickle.dump(output, f)
            return output

    return function_wrapper

@pickle_cache
def load_mtx_format(filename):
    with open(filename) as f:
        # Skip 2
        f.readline()
        f.readline()
        n_col, n_row, n_count = map(int, f.readline().split())
        data = []
        for line in f.readlines():
            col, row, count = map(int, line.split())
            data.append((count, row - 1, col - 1))
        data, i, j = [np.array(x) for x in zip(*data)]
        expression = coo_matrix((data, (i, j)), shape=(n_row, n_col))

    return expression
