"""Modules for maintaining statistics based on exponential decay"""
__docformat__ = "restructuredtext en"

import copy
import numpy
import theano
import theano.tensor

class ExponentialMean(theano.Module):
    """Maintain an exponentially-decaying estimate of the mean

    This module computes the exact mean of the first `max_denom` values of `x`.
    After the first `max_denom` values, it tracks the mean using the formula:

    :math:`self.curval = (1.0 - (1.0/max_denom)) * self.old_curval + (1.0/max_denom) * x`


    The symbolic buffer containing the running mean is called `old_curval`.  (This has a value
    in the ModuleInstance).

    The symbolic variable for the updated running mean is called `curval`.

    """

    max_denom = None
    """The average will be updated as if the current estimated average was estimated from at
    most `max_denom-1` values."""

    ival = None
    """The initial value for the estimated average."""

    def __init__(self, x, max_denom, ival):
        """
        :param x: track the mean of this Variable

        :param max_denom: see `self.max_denom`
    
        :param ival: This should be a tensor of zeros with a shape that matches `x`'s runtime
        value.

        """
        super(ExponentialMean, self).__init__()

        self.max_denom = max_denom
        self.ival = ival

        if len(ival.shape) == 0:
            x_type = theano.tensor.dscalar
        elif len(ival.shape) == 1:
            x_type = theano.tensor.dvector
        elif len(ival.shape) == 2:
            x_type = theano.tensor.dmatrix
        else:
            #TODO: x_type = theano.tensor.TensorType(...)
            raise NotImplementedError()

        self.old_curval = (x_type())
        # TODO: making this an lscalar caused different optimizations, followed by integer
        # division somewhere were i wanted float division.... and the wrong answer.
        self.denom = (theano.tensor.dscalar())

        alpha = 1.0 / self.denom
        self.curval = (1.0 - alpha) * self.old_curval + alpha * x

    def updates(self):
        """
        :returns: the symbolic updates necessary to refresh the mean estimate
        :rtype: dict Variable -> Variable
        """
        return {
                self.old_curval: self.curval,
                self.denom: theano.tensor.smallest(self.denom + 1, self.max_denom)
                }

    def _instance_initialize(self, obj):
        obj.denom = 1
        obj.old_curval = copy.copy(self.ival)

def exp_mean(x, x_shape, max_denom=100):
    """Return an `ExponentialMean` to track a Variable `x` with given shape
    
    :type x: Variable
    :type x_shape: tuple
    :type max_denom: int

    :rtype: ExponentialMean instance
    """
    return ExponentialMean(x, 
            max_denom=max_denom,
            ival=numpy.zeros(x_shape))

def exp_mean_sqr(x, x_shape, max_denom=100):
    """Return an `ExponentialMean` to track a Variable `x`'s square with given shape
    
    :type x: Variable
    :type x_shape: tuple
    :type max_denom: int

    :rtype: ExponentialMean instance
    """
    return ExponentialMean(x**2, 
            max_denom=max_denom,
            ival=numpy.zeros(x_shape))

class exp_var(theano.Module):
    """Module with interface similar to `ExponentialMean` for tracking elementwise variance

    """

    mean = None
    """`ExponentialMean`: current estimate of mean of `x` """

    mean_sqr = None
    """`ExponentialMean`: current estimate of mean sqr of `x` """

    curval = None
    """Variable: Current estimate of the variance of `x` """

    def __init__(self, x, x_shape, max_denom=100):
        """
        :param x: track the variance of this Variable

        :param max_denom: see `self.max_denom`
    
        :param ival: This should be a tensor of zeros with a shape that matches `x`'s runtime
        value.

        """
        super(exp_var,self).__init__()

        self.mean = exp_mean(x, x_shape, max_denom)
        self.mean_sqr = exp_mean_sqr(x, x_shape, max_denom)

        self.curval = self.mean_sqr.curval - self.mean.curval**2

    def updates(self):
        """
        :returns: the symbolic updates necessary to refresh the variance estimate
        :rtype: dict Variable -> Variable
        """
        rval = {}
        rval.update(self.mean.updates())
        rval.update(self.mean_sqr.updates())
        return rval

    def _instance_initialize(self, obj):
        obj.mean.initialize()
        obj.mean_sqr.initialize()

class DynamicNormalizer(theano.Module):
    """
    Normalizes `input` using geometric-decaying estimates of the mean and variance.  The
    `output` should mean near zero, and variance near 1.
    """
    def __init__(self, input, input_shape, max_denom=100, eps=1.0e-8):
        super(DynamicNormalizer, self).__init__()
        self.input = input
        self.d_mean = exp_mean(input, input_shape, max_denom=max_denom)
        self.d_var = exp_var(input, input_shape, max_denom=max_denom)
        self.output = (input - self.d_mean.curval) / theano.tensor.sqrt(self.d_var.curval+eps)

    def updates(self):
        rval = {}
        rval.update(self.d_mean.updates())
        rval.update(self.d_var.updates())
        return rval

    def _instance_initialize(self, obj):
        obj.d_mean.initialize()
        obj.d_var.initialize()