from imageio import imread
import numpy as np


class Dataset(object):
	def __init__(self, uuids, annotations):
		super(Dataset, self).__init__()
		self.uuids = uuids
		self._annot = annotations

	def __len__(self):
		return len(self.uuids)

	def _get(self, method, i):
		return getattr(self._annot, method)(self.uuids[i])



	def get_example(self, i, mode="RGB"):
		methods = ["image", "parts", "label"]
		im_path, parts, label = [self._get(m, i) for m in methods]
		return imread(im_path, pilmode=mode), parts, label

	__getitem__  = get_example


# some convention functions

DEFAULT_RATIO = np.sqrt(49 / 400)

def __expand_parts(p):
	return p[:, 0], p[:, 1:3], p[:, 3].astype(bool)



def uniform_part_locs(im, ratio=DEFAULT_RATIO, round_op=np.floor):
	h, w, c = im.shape

	part_w = round_op(w * ratio).astype(np.int32)
	part_h = round_op(h * ratio).astype(np.int32)

	n, m = w // part_w, h // part_h

	idxs = np.arange(n*m)
	locs = np.zeros((2, n*m), dtype=np.int32)


	for x in range(n):
		for y in range(m):
			i = y * n + x
			x0, y0 = x * part_w, y * part_h
			locs[:, i] = [x0 + part_w // 2, y0 + part_h // 2]

	return idxs, locs

def visible_part_locs(p):
	idxs, locs, vis = __expand_parts(p)
	return idxs[vis], locs[vis].T


def crops(im, xy, ratio=DEFAULT_RATIO, padding_mode="edge"):
	assert im.ndim == 3, "Only RGB images are currently supported!"

	h, w, c = im.shape
	crop_h, crop_w = int(h * ratio), int(w * ratio)
	crops = np.zeros((xy.shape[1], crop_h, crop_w, c), dtype=im.dtype)

	pad_h, pad_w = crop_h // 2, crop_w // 2

	padded_im = np.pad(im, [(pad_h, pad_h), (pad_w, pad_w), [0,0]], mode=padding_mode)

	for i, (x, y) in enumerate(xy.T):
		x0, y0 = x - crop_w // 2 + pad_w, y - crop_h // 2 + pad_h
		crops[i] = padded_im[y0:y0+crop_h, x0:x0+crop_w]

	return crops

def visible_crops(im, p, *args, **kw):
	idxs, locs, vis = __expand_parts(p)
	parts = crops(im, locs[vis].T, *args, **kw)
	res = np.zeros((len(idxs),) + parts.shape[1:], dtype=parts.dtype)
	res[vis] = parts
	return res

def reveal_parts(im, xy, ratio=DEFAULT_RATIO):
	h, w, c = im.shape
	crop_h, crop_w = int(h * ratio), int(w * ratio)

	res = np.zeros_like(im)
	for x, y in xy.T:
		x0, y0 = max(x - crop_w // 2, 0), max(y - crop_h // 2, 0)
		res[y0:y0+crop_h, x0:x0+crop_w] = im[y0:y0+crop_h, x0:x0+crop_w]

	return res