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

import copy
import numpy
import theano
import theano.tensor
from theano.compile.sandbox import shared

class ExponentialMean(object):
    """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.running <- (1.0 - (1.0/max_denom)) * self.running + (1.0/max_denom) * x`

    """

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

    running = None
    """Shared: The running mean statistic from which the output is computed."""

    denom = None
    """Shared: The number of examples we've updated from so far
    """

    def __init__(self, input, max_denom, ival):
        """
        :param input: 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 `input`'s runtime
        value.

        """
        dtype=ival.dtype #dtype is an actual numpy dtype object, not a string
        self.max_denom = max_denom

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

        self.running = shared(numpy.array(ival, copy=True))
        # TODO: making this an lscalar caused different optimizations, followed by integer
        # division somewhere were I wanted float division.... and the wrong answer.
        self.denom = shared(numpy.asarray(1, dtype=dtype))

        alpha = 1.0 / self.denom
        self.output = (1.0 - alpha) * self.running + theano.tensor.cast(alpha * input, str(dtype))

        self.updates = [
                (self.running, self.output),
                (self.denom, theano.tensor.smallest(self.denom + 1, self.max_denom)),
                ]

        assert self.output.type.dtype == dtype

    @classmethod
    def new(cls, x, x_shape, max_denom, dtype='float64'):
        """Return an `ExponentialMean` to track a Variable `x` with given shape
        
        :type x: Variable
        :type x_shape: tuple
        :type max_denom: int
        :type dtype: string
        :param dtype: the running average will be computed at this precision

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