import random


def generateIndices(n_blocks, N, D):
	"""
	generates indices for block matrix computation.
	Checked.
	Input:
	n_blocks: number of blocks to use.
	N: number of samples.
	D: number of genes.
	Output:
	y_indices_to_use[i][j] is the indices of block j in sample i.
	"""
	y_indices_to_use = []
	idxs = list(range(D))
	n_in_block = int(1. * D / n_blocks)

	for i in range(N):
		partition = []
		random.shuffle(idxs)
		n_added = 0

		for block in range(n_blocks):
			start = n_in_block * block
			end = start + n_in_block

			if block < n_blocks - 1:
				idxs_in_block = idxs[start:end]
			else:
				idxs_in_block = idxs[start:]

			partition.append(sorted(idxs_in_block))
			n_added += len(idxs_in_block)

		y_indices_to_use.append(partition)

		if i == 0:
			print('Block sizes', [len(a) for a in partition])

		assert(n_added == D)

	return y_indices_to_use