import numpy as np

from PuzzleLib import Config

from PuzzleLib.Backend import gpuarray, Blas
from PuzzleLib.Backend.Dnn import batchNormNd, batchNormNdBackward

from PuzzleLib.Variable import Variable
from PuzzleLib.Modules.Module import ModuleError, Module


class BatchNormND(Module):
	def __init__(self, nd, maps, epsilon=1e-5, initFactor=1.0, minFactor=0.1, sscale=0.01, affine=True, name=None,
				 empty=False, inplace=False):
		super().__init__(name)
		self.inplace = inplace

		if inplace and Config.showWarnings:
			Config.getLogger().info("Warning: %s is using inplace flag", self)

		self.maps = maps
		self.epsilon = epsilon
		self.initFactor = initFactor
		self.minFactor = minFactor
		self.numOfProps = 0

		self.affine = affine

		self.scale, self.bias, self.mean, self.var = None, None, None, None
		self.savemean, self.saveinvvar, self.scalegrad, self.biasgrad = None, None, None, None

		if empty:
			return

		shape = (1, maps) + self.repeat(1, nd)

		scale = np.random.normal(1.0, sscale if affine else 0.0, shape).astype(self.calctype)
		var = np.ones(shape, dtype=self.calctype)

		self.setVar("scale", Variable(gpuarray.to_gpu(scale)))
		self.setVar("bias", Variable(gpuarray.zeros(shape, dtype=self.calctype)))

		self.setAttr("mean", gpuarray.zeros(shape, dtype=self.calctype))
		self.setAttr("var", gpuarray.to_gpu(var))


	def updateData(self, data):
		if self.train:
			if self.inplace:
				raise ModuleError("%s: using inplace flag in train mode is prohibited" % self)

			self.numOfProps += 1
			factor = max(self.initFactor / self.numOfProps, self.minFactor)

			self.data, self.savemean, self.saveinvvar = batchNormNd(
				data, self.scale, self.bias, self.mean, self.var, self.epsilon, factor, False
			)
		else:
			self.data = batchNormNd(
				data, self.scale, self.bias, self.mean, self.var, self.epsilon, 0, True,
				out=data if self.inplace else None
			)


	def updateGrad(self, grad):
		tup = batchNormNdBackward(self.inData, grad, self.scale, self.savemean, self.saveinvvar, self.epsilon)

		if self.affine:
			self.grad, self.scalegrad, self.biasgrad = tup
		else:
			self.grad, _, _ = tup


	def accGradParams(self, grad, scale=1.0, momentum=0.0):
		if self.affine:
			Blas.addVectorToVector(
				self.scalegrad.ravel(), self.vars["scale"].grad.ravel(), out=self.vars["scale"].grad.ravel(),
				alpha=scale, beta=momentum
			)
			Blas.addVectorToVector(
				self.biasgrad.ravel(), self.vars["bias"].grad.ravel(), out=self.vars["bias"].grad.ravel(),
				alpha=scale, beta=momentum
			)


	def dataShapeFrom(self, shape):
		return shape


	def gradShapeFrom(self, shape):
		return shape


	def reset(self):
		super().reset()
		self.savemean, self.saveinvvar = None, None

		if self.affine:
			self.scalegrad, self.biasgrad = None, None


	def calcMode(self, T):
		if Config.backend == Config.Backend.cuda:
			if T not in {np.float16, np.float32}:
				raise ModuleError("Unsupported dtype %s" % T)

		elif T != np.float32:
			raise ModuleError("Unsupported dtype %s" % T)

		self.calctype = T

[((38, 8, 38, 43), 'numpy.ones', 'np.ones', (), '', True, 'import numpy as np\n'), ((66, 8, 66, 104), 'PuzzleLib.Backend.Dnn.batchNormNdBackward', 'batchNormNdBackward', ({(66, 28, 66, 39): 'self.inData', (66, 41, 66, 45): 'grad', (66, 47, 66, 57): 'self.scale', (66, 59, 66, 72): 'self.savemean', (66, 74, 66, 89): 'self.saveinvvar', (66, 91, 66, 103): 'self.epsilon'}, {}), '(self.inData, grad, self.scale, self.savemean, self.\n    saveinvvar, self.epsilon)', False, 'from PuzzleLib.Backend.Dnn import batchNormNd, batchNormNdBackward\n'), ((43, 23, 43, 65), 'PuzzleLib.Backend.gpuarray.zeros', 'gpuarray.zeros', (), '', False, 'from PuzzleLib.Backend import gpuarray, Blas\n'), ((44, 22, 44, 42), 'PuzzleLib.Backend.gpuarray.to_gpu', 'gpuarray.to_gpu', ({(44, 38, 44, 41): 'var'}, {}), '(var)', False, 'from PuzzleLib.Backend import gpuarray, Blas\n'), ((55, 47, 57, 4), 'PuzzleLib.Backend.Dnn.batchNormNd', 'batchNormNd', ({(56, 4, 56, 8): 'data', (56, 10, 56, 20): 'self.scale', (56, 22, 56, 31): 'self.bias', (56, 33, 56, 42): 'self.mean', (56, 44, 56, 52): 'self.var', (56, 54, 56, 66): 'self.epsilon', (56, 68, 56, 74): 'factor', (56, 76, 56, 81): 'False'}, {}), '(data, self.scale, self.bias, self.mean, self.var, self.epsilon,\n    factor, False)', False, 'from PuzzleLib.Backend.Dnn import batchNormNd, batchNormNdBackward\n'), ((59, 15, 62, 4), 'PuzzleLib.Backend.Dnn.batchNormNd', 'batchNormNd', (), '', False, 'from PuzzleLib.Backend.Dnn import batchNormNd, batchNormNdBackward\n'), ((37, 10, 37, 65), 'numpy.random.normal', 'np.random.normal', ({(37, 27, 37, 30): '1.0', (37, 32, 37, 57): 'sscale if affine else 0.0', (37, 59, 37, 64): 'shape'}, {}), '(1.0, sscale if affine else 0.0, shape)', True, 'import numpy as np\n'), ((40, 32, 40, 54), 'PuzzleLib.Backend.gpuarray.to_gpu', 'gpuarray.to_gpu', ({(40, 48, 40, 53): 'scale'}, {}), '(scale)', False, 'from PuzzleLib.Backend import gpuarray, Blas\n'), ((41, 31, 41, 73), 'PuzzleLib.Backend.gpuarray.zeros', 'gpuarray.zeros', (), '', False, 'from PuzzleLib.Backend import gpuarray, Blas\n'), ((50, 10, 50, 82), 'PuzzleLib.Modules.Module.ModuleError', 'ModuleError', ({(50, 22, 50, 81): "('%s: using inplace flag in train mode is prohibited' % self)"}, {}), "('%s: using inplace flag in train mode is prohibited' % self)", False, 'from PuzzleLib.Modules.Module import ModuleError, Module\n'), ((105, 10, 105, 49), 'PuzzleLib.Modules.Module.ModuleError', 'ModuleError', ({(105, 22, 105, 48): "('Unsupported dtype %s' % T)"}, {}), "('Unsupported dtype %s' % T)", False, 'from PuzzleLib.Modules.Module import ModuleError, Module\n'), ((108, 9, 108, 48), 'PuzzleLib.Modules.Module.ModuleError', 'ModuleError', ({(108, 21, 108, 47): "('Unsupported dtype %s' % T)"}, {}), "('Unsupported dtype %s' % T)", False, 'from PuzzleLib.Modules.Module import ModuleError, Module\n'), ((19, 3, 19, 21), 'PuzzleLib.Config.getLogger', 'Config.getLogger', ({}, {}), '()', False, 'from PuzzleLib import Config\n')]