# Path: deeppy/base.py
# class Model(ParamMixin):
#     def setup(self, *array_shapes):
#         pass
# 
#     def update(self, *arrays):
#         raise NotImplementedError()
# 
# class CollectionMixin(ParamMixin, PhaseMixin):
#     collection = []
# 
#     @property
#     def params(self):
#         params = []
#         for obj in self.collection:
#             if isinstance(obj, ParamMixin):
#                 params.extend(obj.params)
#         return params
# 
#     @params.setter
#     def params(self, params):
#         idx = 0
#         for obj in self.collection:
#             if isinstance(obj, ParamMixin):
#                 n_params = len(obj.params)
#                 obj.params = params[idx:idx+n_params]
#                 idx += n_params
# 
#     @property
#     def phase(self):
#         return self._phase
# 
#     @phase.setter
#     def phase(self, phase):
#         if self._phase == phase:
#             return
#         self._phase = phase
#         for obj in self.collection:
#             if isinstance(obj, PhaseMixin):
#                 obj.phase = phase
# 
# Path: deeppy/filler.py
# class AutoFiller(Filler):
#     def __init__(self, gain=1.0):
#         self.gain = gain
# 
#     def array(self, shape):
#         ndim = len(shape)
#         if ndim == 2:
#             # Affine weights
#             scale = 1.0 / np.sqrt(shape[0])
#         elif ndim == 4:
#             # Convolution filter
#             scale = 1.0 / np.sqrt(np.prod(shape[1:]))
#         else:
#             raise ValueError('AutoFiller does not support ndim %i' % ndim)
#         scale = self.gain * scale / np.sqrt(3)
#         array = np.random.uniform(low=-scale, high=scale, size=shape)
#         return ca.array(array)
# 
# Path: deeppy/feed.py
# class Feed(object):
#     def __init__(self, x, batch_size=None, epoch_size=None):
#         self.x = x
#         self.n_samples = x.shape[0]
#         if batch_size is None:
#             batch_size = 128 if self.n_samples > 512 else self.n_samples
#         if epoch_size is None:
#             epoch_size = int(np.ceil(float(self.n_samples) / batch_size))
#         self.epoch_size = epoch_size
#         self.batch_size = batch_size
#         self.epoch_idx = 0
# 
#     @classmethod
#     def from_any(cls, arg):
#         if isinstance(arg, Feed):
#             arg.reset()
#             return arg
#         elif isinstance(arg, np.ndarray):
#             return cls(arg)
#         elif isinstance(arg, ca.ndarray):
#             return cls(np.array(arg))
#         elif isinstance(arg, tuple):
#             return SupervisedFeed(arg[0], arg[1])
#         raise ValueError('Invalid arguments.')
# 
#     def reset(self):
#         self.epoch_idx = 0
# 
#     def _batch_slices(self):
#         for b in range(self.epoch_size):
#             start = (self.epoch_idx * self.epoch_size + b) * self.batch_size
#             start = start % self.n_samples
#             stop = (start + self.batch_size) % self.n_samples
#             yield start, stop
#         self.epoch_idx += 1
# 
#     def batches(self):
#         x = ca.empty(self.x_shape, dtype=self.x.dtype)
#         for start, stop in self._batch_slices():
#             if stop > start:
#                 x_np = self.x[start:stop]
#             else:
#                 x_np = np.concatenate((self.x[start:], self.x[:stop]))
#             ca.copyto(x, x_np)
#             yield x,
# 
#     @property
#     def x_shape(self):
#         return (self.batch_size,) + self.x.shape[1:]
# 
#     @property
#     def shapes(self):
#         return self.x_shape,


import numpy as np
import cudarray as ca
import deeppy.expr as expr
from ..base import Model, CollectionMixin
from ..filler import AutoFiller
from ..feed import Feed

class KLDivergence(expr.Op):
    def __call__(self, mu, log_sigma):
        self.mu = mu
        self.log_sigma = log_sigma
        self.inputs = [mu, log_sigma]
        return self

    def setup(self):
        self.shape = (1,)
        self.array = ca.empty(self.shape)
        self.grad_array = ca.empty(self.shape)

    def fprop(self):
        tmp1 = self.mu.array**2
        ca.negative(tmp1, tmp1)
        tmp1 += self.log_sigma.array
        tmp1 += 1
        tmp1 -= ca.exp(self.log_sigma.array)
        self.array = ca.sum(tmp1)
        self.array *= -0.5

    def bprop(self):
        ca.multiply(self.mu.array, self.grad_array, self.mu.grad_array)
        ca.exp(self.log_sigma.array, out=self.log_sigma.grad_array)
        self.log_sigma.grad_array -= 1
        self.log_sigma.grad_array *= 0.5
        self.log_sigma.grad_array *= self.grad_array



# Path: deeppy/base.py
# class Model(ParamMixin):
#     def setup(self, *array_shapes):
#         pass
# 
#     def update(self, *arrays):
#         raise NotImplementedError()
# 
# class CollectionMixin(ParamMixin, PhaseMixin):
#     collection = []
# 
#     @property
#     def params(self):
#         params = []
#         for obj in self.collection:
#             if isinstance(obj, ParamMixin):
#                 params.extend(obj.params)
#         return params
# 
#     @params.setter
#     def params(self, params):
#         idx = 0
#         for obj in self.collection:
#             if isinstance(obj, ParamMixin):
#                 n_params = len(obj.params)
#                 obj.params = params[idx:idx+n_params]
#                 idx += n_params
# 
#     @property
#     def phase(self):
#         return self._phase
# 
#     @phase.setter
#     def phase(self, phase):
#         if self._phase == phase:
#             return
#         self._phase = phase
#         for obj in self.collection:
#             if isinstance(obj, PhaseMixin):
#                 obj.phase = phase
# 
# Path: deeppy/filler.py
# class AutoFiller(Filler):
#     def __init__(self, gain=1.0):
#         self.gain = gain
# 
#     def array(self, shape):
#         ndim = len(shape)
#         if ndim == 2:
#             # Affine weights
#             scale = 1.0 / np.sqrt(shape[0])
#         elif ndim == 4:
#             # Convolution filter
#             scale = 1.0 / np.sqrt(np.prod(shape[1:]))
#         else:
#             raise ValueError('AutoFiller does not support ndim %i' % ndim)
#         scale = self.gain * scale / np.sqrt(3)
#         array = np.random.uniform(low=-scale, high=scale, size=shape)
#         return ca.array(array)
# 
# Path: deeppy/feed.py
# class Feed(object):
#     def __init__(self, x, batch_size=None, epoch_size=None):
#         self.x = x
#         self.n_samples = x.shape[0]
#         if batch_size is None:
#             batch_size = 128 if self.n_samples > 512 else self.n_samples
#         if epoch_size is None:
#             epoch_size = int(np.ceil(float(self.n_samples) / batch_size))
#         self.epoch_size = epoch_size
#         self.batch_size = batch_size
#         self.epoch_idx = 0
# 
#     @classmethod
#     def from_any(cls, arg):
#         if isinstance(arg, Feed):
#             arg.reset()
#             return arg
#         elif isinstance(arg, np.ndarray):
#             return cls(arg)
#         elif isinstance(arg, ca.ndarray):
#             return cls(np.array(arg))
#         elif isinstance(arg, tuple):
#             return SupervisedFeed(arg[0], arg[1])
#         raise ValueError('Invalid arguments.')
# 
#     def reset(self):
#         self.epoch_idx = 0
# 
#     def _batch_slices(self):
#         for b in range(self.epoch_size):
#             start = (self.epoch_idx * self.epoch_size + b) * self.batch_size
#             start = start % self.n_samples
#             stop = (start + self.batch_size) % self.n_samples
#             yield start, stop
#         self.epoch_idx += 1
# 
#     def batches(self):
#         x = ca.empty(self.x_shape, dtype=self.x.dtype)
#         for start, stop in self._batch_slices():
#             if stop > start:
#                 x_np = self.x[start:stop]
#             else:
#                 x_np = np.concatenate((self.x[start:], self.x[:stop]))
#             ca.copyto(x, x_np)
#             yield x,
# 
#     @property
#     def x_shape(self):
#         return (self.batch_size,) + self.x.shape[1:]
# 
#     @property
#     def shapes(self):
#         return self.x_shape,


# Path: deeppy/model/variational_autoencoder.py
import numpy as np
import cudarray as ca
import deeppy.expr as expr
from ..base import Model, CollectionMixin
from ..filler import AutoFiller
from ..feed import Feed


class KLDivergence(expr.Op):
    def __call__(self, mu, log_sigma):
        self.mu = mu
        self.log_sigma = log_sigma
        self.inputs = [mu, log_sigma]
        return self

    def setup(self):
        self.shape = (1,)
        self.array = ca.empty(self.shape)
        self.grad_array = ca.empty(self.shape)

    def fprop(self):
        tmp1 = self.mu.array**2
        ca.negative(tmp1, tmp1)
        tmp1 += self.log_sigma.array
        tmp1 += 1
        tmp1 -= ca.exp(self.log_sigma.array)
        self.array = ca.sum(tmp1)
        self.array *= -0.5

    def bprop(self):
        ca.multiply(self.mu.array, self.grad_array, self.mu.grad_array)
        ca.exp(self.log_sigma.array, out=self.log_sigma.grad_array)
        self.log_sigma.grad_array -= 1
        self.log_sigma.grad_array *= 0.5
        self.log_sigma.grad_array *= self.grad_array



class VariationalAutoencoder(Model, CollectionMixin):