# -*- coding: utf-8 -*-
"""
mcmc_difussion.py: Provides facilities for Markov-Chain Monte Carlo
integration to convolute an arbitrary function with an isotropic gaussian
kernel of varying width (σ), to implement a technique called coarse to
fine smoothing[1].


[1]: J. Mobahi and J. W. Fished: On the Link Between Gaussian Homotopy
     Continuation and Convex Envelopes, 2015.

@author: vharisop
"""
from __future__ import division
import numpy
import theano.tensor as T
import theano

def iso_gauss_kern(disps, sigma):
    """
    An implementation of an isotropic gaussian kernel in Theano.
    Accepts a [M x N] tensor of matrix type and results in a 1xM
    vector, which is the sampled isotropic gaussian.

    Arguments
    ---------
    disps : theano.tensor.var.TensorVariable
        The input tensor, which is expected to hold a matrix with
        the sampled points as rows
    sigma : float
        The standard deviation of the kernel

    Returns
    -------
    theano.tensor.var.TensorVariable
        A tensor representing the gaussian kernel, which is supposed
        to produce a 1xM vector of values
    """
    numer = T.exp(-T.sum(disps ** 2, axis=1) / (2 * (sigma ** 2)))
    norm_fact = 1 / (numpy.sqrt(2 * numpy.pi * sigma) ** disps.shape[1])
    return norm_fact * numer


def iso_gauss_single(disp, sigma):
    """
    An implementation of an isotropic gaussian kernel in Theano.
    Accepts a [1 x N] tensor of matrix type and results in a 1x1
    vector, which is the sampled isotropic gaussian.

    Arguments
    ---------
    disp : theano.tensor.var.TensorVariable
        The input tensor, which is expected to hold a vector
        with a sample point
    sigma : float
        The standard deviation of the kernel

    Returns
    -------
    theano.tensor.var.TensorVariable
        A tensor representing the gaussian kernel, which is supposed
        to produce a single value
    """
    numer = T.exp(-T.sum(disp ** 2) / (2 * (sigma ** 2)))
    norm_fact = 1 / (numpy.sqrt(2 * numpy.pi * sigma) ** disp.shape)
    return norm_fact * numer


def iso_gauss_kern_grad(disps, sigma):
    """
    An implementation of an isotropic gaussian kernel in Theano.
    Accepts a [M x N] tensor of matrix type and results in a NxM
    vector, which is the gradient of the sampled isotropic gaussian
    kernel.

    Arguments
    ---------
    disps : theano.tensor.var.TensorVariable
        The input tensor, which is expected to hold a matrix with
        the sampled points as rows
    sigma : float
        The standard deviation of the kernel

    Returns
    -------
    theano.tensor.var.TensorVariable
        A tensor representing the gaussian kernel, which is supposed
        to produce a 1xM vector of values
    """
    numer = T.exp(-T.sum(disps ** 2, axis=1) / (2 * (sigma ** 2)))
    norm_fact = 1 / (numpy.sqrt(2 * numpy.pi * sigma) ** disps.shape[1])
    g_kern = norm_fact * numer
    f_mul = lambda vec, val: vec * val
    grad_vec, _ = theano.scan(fn=f_mul, sequences=[disps, g_kern])
    return grad_vec


def mcmc_single_convolution(f, wt, disps, sigma):
    """
    Creates a tensor holding the markov-chain monte carlo convolution of
    a given function and the isotropic gaussian kernel. The output is to
    be a 1xM vector.

    Arguments
    ---------
    f : theano.tensor.var.TensorVariable
        The function to be convoluted
    wt : theano.tensor.var.TensorVariable
        The tensor of inputs (params) assumed to be providing f. In the
        context of a neural network, these inputs are the parameters of
        the network, rather than the actual input patterns
    disps : theano.tensor.var.TensorVariable
        The tensor of sampled displacements
    sigma : float
        The deviation of the gaussian kernel

    Returns
    -------
    theano.tensor.var.TensorVariable
        A tensor holding the numerical convolution of `f` with the
        gaussian kernel
    """
    # NOTE: here f must be a tensor that evaluates to a scalar
    convolution, _ = theano.scan(
        fn=(lambda d: f(wt - d) * iso_gauss_single(d, sigma)),
        sequences=disps)
    return T.mean(convolution)


def mcmc_single_gradient(f, pt, disps, sigma):
    """
    Creates a tensor holding the markov-chain monte carlo convolution of
    a given function with the derivative of the isotropic gaussian kernel.
    The output is to be a 1xN vector.

    Arguments
    ---------
    f : theano.tensor.var.TensorVariable
        The function to be convolved
    pt : theano.tensor.var.TensorVariable
        The tensor of inputs, assumed to be the source of f
    disps : theano.tensor.var.TensorVariable
        The tensor of sampled displacements
    sigma : float
        The gaussian kernel's deviation

    Returns
    -------
    theano.tensor.var.TensorVariable
        A tensor holding the numerical convolution of `f` with the
        derivative of the gaussian kernel
    """
    # gradient of gaussian: (-v / sigma^2) * G(v)
    # f(x - v) * (-v/ σ^2)
    fn = lambda d: \
        f(pt - d) * iso_gauss_single(d, sigma) * (-d / (sigma ** 2))
    convolution, _ = theano.scan(fn=fn, sequences=disps)
    # convolution is now a N x M vector, where N is the number of points
    # and M is the number of samples. We multiply elemwise with the dis-
    # -placement vector.
    return T.mean(convolution, axis=0)


def mcmc_convolution(f, pts, disps, sigma):
    """
    Creates a tensor holding the markov-chain monte carlo convolution of
    a given function and the isotropic gaussian kernel. The output is to
    be a 1xM vector.

    Arguments
    ---------
    f : theano.tensor.var.TensorVariable
        The function to be convoluted
    pts : theano.tensor.var.TensorVariable
        The tensor of inputs, assumed to be providing f. In the context
        of a neural network, these inputs are the network's parameters,
        rather than the actual input patterns
    disps : theano.tensor.var.TensorVariable
        The tensor of sampled displacements
    sigma : float
        The deviation of the gaussian kernel

    Returns
    -------
    theano.tensor.var.TensorVariable
        A tensor holding the numerical convolution of `f` with the
        gaussian kernel
    """
    # NOTE: here f must be a tensor that evaluates to a scalar
    fn = lambda x: T.dot(f(x - disps), iso_gauss_kern(disps, sigma))
    convolution, _ = theano.scan(fn=fn, sequences=pts)
    return T.mean(convolution, axis=1)


def mcmc_gradient(f, pts, disps, sigma):
    """
    Creates a tensor holding the markov-chain monte carlo convolution of
    a given function with the derivative of the isotropic gaussian kernel.
    The output is to be a NxM vector.

    Arguments
    ---------
    f : theano.tensor.var.TensorVariable
        The function to be convolved
    pts : theano.tensor.var.TensorVariable
        The tensor of inputs, assumed to be the source of f
    disps : theano.tensor.var.TensorVariable
        The tensor of sampled displacements
    sigma : float
        The gaussian kernel's deviation

    Returns
    -------
    theano.tensor.var.TensorVariable
        A tensor holding the numerical convolution of `f` with the
        derivative of the gaussian kernel
    """
    # gradient of gaussian: (-v / sigma^2) * G(v)
    # f(x - v) * (-v/ σ^2)
    fn = lambda x: T.dot(
        f(x - disps),
        iso_gauss_kern_grad(disps, sigma))
    convolution, _ = theano.scan(fn=fn, sequences=pts)
    # convolution is now a N x M vector, where N is the number of points
    # and M is the number of samples. We multiply elemwise with the dis-
    # -placement vector.
    return T.mean(convolution, axis=1)

# TODO: Check if we actually implemented this correctly:
# return (1 / num_pts) * numpy.sum(f(x - pt) * kern(pt) for pt in sampled)

if __name__ == '__main__':
    print("Example: X is a series of points")
    print("We will find the gaussian kernel on it")
    X = T.matrix('X')  # points
    D = T.matrix('D')  # displacements
    f = lambda x: x.norm(L=2)
    conv = theano.function(
        [X, D],
        [mcmc_convolution(f, X, D, 1.0),
         mcmc_gradient(f, X, D, 1.0)])

    Xs = numpy.random.randn(10, 2)
    Ds = numpy.random.multivariate_normal(
        [0.0, 0.0], 3.0 * numpy.eye(2), 100)
    # conv(Xs, Ds) should return a 1 x 10 vector, containing the mcmc
    # convolution at each of the 10 data points, each of them having been
    # convoluted with a 100 point kernel, as well as a 10x2 vector,
    # containing the gradient of the above convolution
    Vs, GVs = conv(Xs, Ds)
