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

[((1067, 1102), 'numpy.ones', 'np.ones', (['shape'], {'dtype': 'self.calctype'}), '(shape, dtype=self.calctype)\n', (1074, 1102), True, 'import numpy as np\n'), ((1936, 2037), 'PuzzleLib.Backend.Dnn.batchNormNdBackward', 'batchNormNdBackward', (['self.inData', 'grad', 'self.scale', 'self.savemean', 'self.saveinvvar', 'self.epsilon'], {}), '(self.inData, grad, self.scale, self.savemean, self.\n    saveinvvar, self.epsilon)\n', (1955, 2037), False, 'from PuzzleLib.Backend.Dnn import batchNormNd, batchNormNdBackward\n'), ((1261, 1303), 'PuzzleLib.Backend.gpuarray.zeros', 'gpuarray.zeros', (['shape'], {'dtype': 'self.calctype'}), '(shape, dtype=self.calctype)\n', (1275, 1303), False, 'from PuzzleLib.Backend import gpuarray, Blas\n'), ((1327, 1347), 'PuzzleLib.Backend.gpuarray.to_gpu', 'gpuarray.to_gpu', (['var'], {}), '(var)\n', (1342, 1347), False, 'from PuzzleLib.Backend import gpuarray, Blas\n'), ((1640, 1734), 'PuzzleLib.Backend.Dnn.batchNormNd', 'batchNormNd', (['data', 'self.scale', 'self.bias', 'self.mean', 'self.var', 'self.epsilon', 'factor', '(False)'], {}), '(data, self.scale, self.bias, self.mean, self.var, self.epsilon,\n    factor, False)\n', (1651, 1734), False, 'from PuzzleLib.Backend.Dnn import batchNormNd, batchNormNdBackward\n'), ((1763, 1887), 'PuzzleLib.Backend.Dnn.batchNormNd', 'batchNormNd', (['data', 'self.scale', 'self.bias', 'self.mean', 'self.var', 'self.epsilon', '(0)', '(True)'], {'out': '(data if self.inplace else None)'}), '(data, self.scale, self.bias, self.mean, self.var, self.epsilon,\n    0, True, out=data if self.inplace else None)\n', (1774, 1887), False, 'from PuzzleLib.Backend.Dnn import batchNormNd, batchNormNdBackward\n'), ((981, 1036), 'numpy.random.normal', 'np.random.normal', (['(1.0)', '(sscale if affine else 0.0)', 'shape'], {}), '(1.0, sscale if affine else 0.0, shape)\n', (997, 1036), True, 'import numpy as np\n'), ((1136, 1158), 'PuzzleLib.Backend.gpuarray.to_gpu', 'gpuarray.to_gpu', (['scale'], {}), '(scale)\n', (1151, 1158), False, 'from PuzzleLib.Backend import gpuarray, Blas\n'), ((1192, 1234), 'PuzzleLib.Backend.gpuarray.zeros', 'gpuarray.zeros', (['shape'], {'dtype': 'self.calctype'}), '(shape, dtype=self.calctype)\n', (1206, 1234), False, 'from PuzzleLib.Backend import gpuarray, Blas\n'), ((1427, 1499), 'PuzzleLib.Modules.Module.ModuleError', 'ModuleError', (["('%s: using inplace flag in train mode is prohibited' % self)"], {}), "('%s: using inplace flag in train mode is prohibited' % self)\n", (1438, 1499), False, 'from PuzzleLib.Modules.Module import ModuleError, Module\n'), ((2901, 2940), 'PuzzleLib.Modules.Module.ModuleError', 'ModuleError', (["('Unsupported dtype %s' % T)"], {}), "('Unsupported dtype %s' % T)\n", (2912, 2940), False, 'from PuzzleLib.Modules.Module import ModuleError, Module\n'), ((2975, 3014), 'PuzzleLib.Modules.Module.ModuleError', 'ModuleError', (["('Unsupported dtype %s' % T)"], {}), "('Unsupported dtype %s' % T)\n", (2986, 3014), False, 'from PuzzleLib.Modules.Module import ModuleError, Module\n'), ((527, 545), 'PuzzleLib.Config.getLogger', 'Config.getLogger', ([], {}), '()\n', (543, 545), False, 'from PuzzleLib import Config\n')]