from rbm.matrix import Mat, np

def normalize(image: Mat):
	mean = np.mean(image, axis=1, keepdims=True)
	std  = np.std(image,  axis=1, keepdims=True)
	std  = np.where(std < 1e-8, np.ones_like(std), std)
	return (image - mean) / std

class SubImage:
	def __init__(self, sx: int, sy: int, px: int, py: int):
		self.sx = sx
		self.sy = sy
		self.px = px
		self.py = py

	def check(self, image: Mat):
		nx = (image.shape[0] - self.sx)
		ny = (image.shape[1] - self.sy)
		if nx % self.px:
			raise ValueError
		if ny % self.py:
			raise ValueError

		nx = int(nx / self.px) + 1
		ny = int(ny / self.py) + 1
		return nx, ny

	def __call__(self, images: Mat, n_ch: int = 3):
		if len(images.shape) == 3:
			return self.process_image(np.reshape(images, (images.shape[0], n_ch) + images.shape[1:3]))

		if len(images.shape) == 2:
			return self.process_image(np.reshape(images, (1,n_ch,) + images.shape)), images.shape

		return self.process_image(images)

	def process_image(self, image: Mat):
		nx, ny = self.check(image[0, 0, :, :])
		result = np.zeros((image.shape[0]*nx*ny, image.shape[1], self.sx, self.sy))
		k = 0
		for n in range(image.shape[0]):
			for x in range(0, image.shape[2] - self.sx + 1, self.px):
				for y in range(0, image.shape[3] - self.sy + 1, self.py):
					for c in range(image.shape[1]):
						si = image[n, c, x:x+self.sx, y:y+self.sy]
						result[k, c] = si
					k += 1

		return result

def conv2d(image: Mat, kernel: Mat, stride: tuple[int,int]):
	pad_x = int((kernel.shape[0] - 1)/2)
	pad_y = int((kernel.shape[1] - 1)/2)
	image_padded = np.pad(image, pad_width=(pad_x, pad_y))

	nx = (image_padded.shape[0] - kernel.shape[0])
	ny = (image_padded.shape[1] - kernel.shape[1])
	if nx % stride[0]:
		raise ValueError
	if ny % stride[1]:
		raise ValueError
	pass

	result = np.zeros((image_padded.shape[0] + kernel.shape[0] - 1, image_padded.shape[1] + kernel.shape[1] - 1))
	k = 0
	for r in range(pad_y, image_padded.shape[0]-pad_y):
		result[r,:] = np.convolve(image_padded[r,:], kernel[k,:])
		k += 1

	k = 0
	for c in range(pad_x, image_padded.shape[1]-pad_x):
		result[:, c] += np.convolve(image_padded[:,c], kernel[:,k])
		k += 1

	return result
