"""
Probability distribution.
"""

import numpy as np
cimport numpy as np


cdef class Distribution:
    """Probability distribution."""

    def sample(self, num_samples):
        """Draw num_samples from distribution."""
        raise NotImplementedError

    def log_likelihood(self, x):
        """Returns the log likelihood of 'x'.

        Args:
            x: Single or multiple points at which to evaluate
                the distribution.

        Returns:
            The log-likelihood of 'x' as a double.
        """
        raise NotImplementedError


cdef class RealDistribution(Distribution):
    """Probability distribution over R or R^n."""

    def log_likelihoods(self, x):
        """Returns log of evaluation at each point in 'x'.

        Args:
            x: Points at which to evaluate distribution.

        Returns:
            (N, ) array with log of value of distribution at every point.
        """
        raise NotImplementedError

    def log_likelihood(self, x, weights=None):
        """Returns the log-likelihood of x.

        Args:
            x: (N, ) or (N, D) array of data points.
            weights: (N, ) array of weights. If 'None', then
                all weights are assumed to be '1'.

        Returns:
            The log-likelihood of x as a double.
        """

        values = self.log_likelihoods(x) # (N, ) array
        if weights:
            return np.sum(values * weights)
        else:
            return np.sum(values)

    def __call__(self, x):
        """Returns values of distribution at given points.

        Args:
            x: (N, ) or (N, D) array of data points.

        Returns:
            (N, ) array with values of Gaussian at points in 'x'.
        """

        return np.exp(self.log_likelihoods(x))


class DiracDistribution(Distribution):
    """A Dirac delta function."""

    def __init__(self, center):
        """
        Args:
            center: The point at which the distribution is concentrated.
        """

        self.center = center

    def sample(self):
        """Returns the center."""

        return self.center

    def log_likelihood(self, x):
        """Returns the log-likelihood of 'x', constant in this case."""

        return 0.0


