# Copyright (c) Anand Patil, 2007

__docformat__ = 'reStructuredText'

__all__ = ['GaussianProcess', 'GPEvaluation', 'GPSubmodel']

import pymc as pm
import copy
import numpy as np
from .Realization import Realization


__all__ = ['GaussianProcess', 'GPEvaluation', 'GPSubmodel']


def gp_logp(x, M, C, mesh, f_eval, M_obs, C_obs):
    raise TypeError('GP objects have no logp function')


def gp_rand(M, C, mesh, f_eval, M_obs, C_obs, size=None):
    # M and C are input pre-observed, so no need to
    out = pm.gp.Realization(M_obs, C_obs)
    out.x_sofar = mesh
    out.f_sofar = f_eval
    out.M = M
    out.C = C
    return out


class GaussianProcess(pm.Stochastic):

    """
    G=GaussianProcess(name, submodel, **kwds)

    A stochastic variable valued as a Gaussian process realization.

    :Arguments:

        -   `name`: The name of the variable.

        -   `submodel`: The Gaussian process submodel to which the
            variable belongs.

        -   `relative_precision`: See documentation.

    :SeeAlso: GPSubmodel, Realization, GPEvaluation
    """

    def __init__(self, name, submodel, trace=True, value=None, rseed=False,
                 observed=False, cache_depth=2, plot=None, verbose=-1):

        self.submodel = submodel

        pm.Stochastic.__init__(self, gp_logp, GaussianProcess.__doc__, name,
                               {'M': submodel.M,
                                'C': submodel.C,
                                'mesh': submodel.mesh,
                                'f_eval': submodel.f_eval,
                                'M_obs': submodel.M_obs,
                                'C_obs': submodel.C_obs},
                               gp_rand, trace, value, np.dtype('object'), rseed, observed, cache_depth, plot, verbose, False)

        self.rand()

    def gen_lazy_function(self):
        pass

    def get_logp(self):
        raise TypeError('Gaussian process %s has no logp.' % self.__name__)

    def set_logp(self, new_logp):
        raise TypeError('Gaussian process %s has no logp.' % self.__name__)

    def set_value(self, new_value):
        # If this is a new value pulled off the trace, save it the trouble of
        # observing
        if np.all(new_value.x_sofar == pm.utils.value(self.parents['mesh'])):
            if new_value.need_init_obs:
                new_value.M_internal = copy.copy(
                    pm.utils.value(
                        self.parents[
                            'M_obs']))
                new_value.C_internal = copy.copy(
                    pm.utils.value(
                        self.parents[
                            'C_obs']))
                new_value.need_init_obs = False
        pm.Stochastic.set_value(self, new_value)

    logp = property(fget=get_logp, fset=set_logp)

# for attr in ['neg','abs','invert'] + ['add', 'mul', 'sub'] + ['iadd','isub','imul','idiv','itruediv','ifloordiv','imod','ipow','ilshift','irshift','iand','ixor','ior','unicode']:
#     GaussianProcess.__delattr__('__%s__'%attr)


class GPEvaluation(pm.MvNormalChol):
    pass


class GPSubmodel(pm.ObjectContainer):

    """
    G=GaussianProcess(name, M, C, mesh[, init_vals, obs_on_mesh, tally_all])

    A Gaussian process submodel, which is a container for GaussianProcess and
    GPEvaluation objects.

    :Arguments:

        -   `name`: A prefix for the names of all the variables.

        -   `M`, `C`: The mean and covariance serving as parents of
            the submodel.

        -   `mesh`: See documentation.

        -   `init_vals`: The initial value of self.f_eval

        -   `obs_on_mesh`: Whether self.f_eval's value is observed.

        -   `tally_all`: Whether all variables should be tallied, or just
                f and f_eval.

    :Attributes:

        -   `f`: A GaussianProcess object, which is valued as a GP realization
            with prior mean self.M and prior covariance self.C.

        -   `f_eval`: The evaluation of self.f on self.mesh. This is a
            multivariate normal variable with a logp attribute.

        -   `mesh`: The mesh that was input, wrapped in a deterministic.

        -   `C_eval`: The evaluation of self.C on self.mesh.

        -   `S_eval`: The lower-triangular Cholesky factor of C_eval.

        -   `M_obs`, `C_obs`: Versions of self.M and self.C, observed on
                self.mesh with value self.f_eval.value.

        -   `fr_check`: A potential that enforces the constraint that
                self.C_eval must be full-rank.


    :SeeAlso: GaussianProcess, Realization, GPEvaluation
    """

    def __init__(self, name, M, C, mesh, init_vals=None,
                 obs_on_mesh=False, tally_all=False, **kwds):

        if isinstance(mesh, pm.Variable):
            mesh = pm.Lambda(
                '%s_mesh' %
                name,
                lambda mesh=mesh: pm.gp.regularize_array(
                    mesh),
                trace=False)
        else:
            mesh = pm.gp.regularize_array(mesh)
        self.name = name

        self.mesh = mesh
        self.M = M
        self.C = C

        if len(mesh) == 0:
            self.M_obs = M
            self.C_obs = C
            self.f_eval = None
        else:

            @pm.deterministic(
                trace=tally_all, name='%s_covariance_bits' % name)
            def covariance_bits(C=C, mesh=mesh):
                """
                Both the realization 'f' and the on-mesh evaluation 'f_eval' need the
                Cholesky factor of the covariance evaluation. The Gibbs step method
                also needs the full covariance evaluation. The mean needs a certain other
                function of the full covariance evaluation.

                All these things can be got as byproducts of Covariance.observe. Keeping the
                observed covariance and using it as the parent of f means the computations only
                get done once.
                """
                C_obs = copy.copy(C)
                try:
                    U, C_eval, Uo_Cxo = C_obs.observe(
                        mesh, np.zeros(mesh.shape[0]), output_type='s')
                    return U.T.copy('F'), C_eval, C_obs, Uo_Cxo
                except np.linalg.LinAlgError:
                    return None

            S_eval = pm.Lambda(
                '%s_S_eval' %
                name,
                lambda cb=covariance_bits: cb[
                    0] if cb else None,
                doc="The lower triangular Cholesky factor of %s.C_eval" %
                name,
                trace=tally_all or kwds.get(
                    'tally_S_eval',
                    False))
            C_eval = pm.Lambda(
                '%s_C_eval' % name,
                lambda cb=covariance_bits: cb[
                    1] if cb else None,
                doc="The evaluation %s.C(%s.mesh, %s.mesh)" % (
                    name,
                    name,
                    name),
                trace=tally_all or kwds.get(
                    'tally_C_eval',
                    False))
            C_obs = pm.Lambda(
                '%s_C_obs' % name,
                lambda cb=covariance_bits: cb[
                    2] if cb else None,
                doc="%s.C, observed on %s.mesh" % (
                    name,
                    name),
                trace=tally_all or kwds.get(
                    'tally_C_obs',
                    False))
            Uo_Cxo = pm.Lambda(
                '%s_Uo_Cxo' % name,
                lambda cb=covariance_bits: cb[
                    3] if cb else None,
                doc="A byproduct of observation of %s.C that can be used by %s.M" % (
                    name,
                    name),
                trace=tally_all or kwds.get(
                    'tally_Uo_Cxo',
                    False))
            M_eval = pm.Lambda(
                '%s_M_eval' % name,
                lambda M=M,
                mesh=mesh,
                Uo_Cxo=Uo_Cxo: M(
                    mesh,
                    Uo_Cxo=Uo_Cxo),
                trace=tally_all or kwds.get(
                    'tally_M_eval',
                    False),
                doc="The evaluation %s.M(%s.mesh)" % (
                    name,
                    name))

            @pm.potential(name='%s_fr_check' % name)
            def fr_check(S_eval=S_eval):
                """
                Forbids non-positive-definite C_evals.
                """
                if S_eval is None:
                    return -np.inf
                else:
                    return 0
            fr_check = fr_check

            f_eval = GPEvaluation(
                '%s_f_eval' % name, mu=M_eval, sig=S_eval, value=init_vals, trace=kwds.get('tally_f_eval', True), observed=obs_on_mesh,
                doc="The evaluation %s.f(%s.mesh).\nThis is a multivariate normal variable with mean %s.M_eval and covariance %s.C_eval." % (name, name, name, name))

            @pm.deterministic(
                trace=tally_all or kwds.get('tally_M_obs', False), name='%s_M_obs' % name)
            def M_obs(M=M, f_eval=f_eval, C_obs=C_obs, mesh=mesh):
                """
                Creates an observed mean object to match %sC_obs.
                """ % name
                M_obs = copy.copy(M)
                M_obs.observe(C_obs, mesh, f_eval)
                return M_obs

            self.M_obs = M_obs
            self.C_obs = C_obs
            self.f_eval = f_eval

        f = GaussianProcess(
            '%s_f' %
            name,
            self,
            trace=tally_all or kwds.get(
                'tally_f',
                True))

        f.rand()
        l = locals()
        lk = filter(lambda k: isinstance(l[k], pm.Node), l.keys())
        l = dict([(k, l[k]) for k in lk])
        pm.ObjectContainer.__init__(self, l)

    def getobjects(self):
        names = ['M_eval', 'C_eval', 'S_eval', 'f_eval', 'f', 'fr_check']
        return dict(zip(['%s_%s' % (self.name, name) for name in names], [
                    getattr(self, name) for name in names]))
