] if full_cov and full_output_cov",
)
def mix_latent_gp(
    W: tf.Tensor, g_mean: tf.Tensor, g_var: tf.Tensor, full_cov: bool, full_output_cov: bool
) -> MeanAndVariance:
    r"""Takes the mean and variance of an uncorrelated L-dimensional latent GP
    and returns the mean and the variance of the mixed GP, `f = W g`,
    where both f and g are GPs.

    :return: f_mean and f_var
    """
    f_mean = tf.tensordot(g_mean, W, [[-1], [-1]])  # [..., N, P]

    if full_cov and full_output_cov:  # g_var is [L, ..., N, N]
        # this branch is practically never taken
        g_var = rollaxis_left(g_var, 1)  # [..., N, N, L]

        g_var = tf.expand_dims(g_var, axis=-2)  # [..., N, N, 1, L]
        g_var_W = g_var * W  # [..., N, P, L]
        f_var = tf.tensordot(g_var_W, W, [[-1], [-1]])  # [..., N, N, P, P]
        f_var = leading_transpose(f_var, [..., -4, -2, -3, -1])  # [..., N, P, N, P]

    elif full_cov and not full_output_cov:  # g_var is [L, ..., N, N]
        # this branch is practically never taken
        f_var = tf.tensordot(g_var, W ** 2, [[0], [-1]])  # [..., N, N, P]
        f_var = leading_transpose(f_var, [..., -1, -3, -2])  # [..., P, N, N]

    elif not full_cov and full_output_cov:  # g_var is [..., N, L]
        g_var = tf.expand_dims(g_var, axis=-2)  # [..., N, 1, L]
        g_var_W = g_var * W  # [..., N, P, L]
        f_var = tf.tensordot(g_var_W, W, [[-1], [-1]])  # [..., N, P, P]

    elif not full_cov and not full_output_cov:  # g_var is [..., N, L]
        W_squared = W ** 2  # [P, L]
        f_var = tf.tensordot(g_var, W_squared, [[-1], [-1]])  # [..., N, P]

    return f_mean, f_var


@check_shapes(
    "Kmns: [P, M, batch..., N]",
    "Kmms: [P, M, M]",
    "Knns: [P, batch..., N, N] if full_cov",
    "Knns: [P, batch..., N] if not full_cov",
    "f: [M, P]",
    "q_sqrt: [M_R_or_R_M_M...]",
    "return[0]: [batch..., N, R]",
    "return[1]: [batch..., R, N, N] if full_cov",
    "return[1]: [batch..., N, R] if not full_cov",
)
def separate_independent_conditional_implementation(
    Kmns: tf.Tensor,
    Kmms: tf.Tensor,
    Knns: tf.Tensor,
    f: tf.Tensor,
    *,
    full_cov: bool = False,
    q_sqrt: Optional[tf.Tensor] = None,
    white: bool = False,
) -> MeanAndVariance:
    """
    Multi-output GP with independent GP priors.

    Number of latent processes equals the number of outputs (L = P).

    Further reference:

    - See `gpflow.conditionals._conditional` for a detailed explanation of
      conditional in the single-output case.
    - See the multioutput notebook for more information about the multioutput framework.
    - See above for the parameters and the return value.
    """
    fs = tf.transpose(f)[:, :, None]  # [P, M, 1]
    # [P, 1, M, M]  or  [P, M, 1]

    if q_sqrt is not None:
        q_sqrts = (
            tf.transpose(q_sqrt)[:, :, None] if q_sqrt.shape.ndims == 2 else q_sqrt[:, None, :, :]
        )
        base_conditional_args_to_map = (
            Kmms,
            Kmns,
            Knns,
            fs,
            q_sqrts,
        )  # type: Tuple[tf.Tensor, ...]

        def single_gp_conditional(
            t: Tuple[tf.Tensor, ...]
        ) -> MeanAndVariance:  # pragma: no cover - tf.map_fn is invisible to codecov
            Kmm, Kmn, Knn, f, q_sqrt = t
            return base_conditional(Kmn, Kmm, Knn, f, full_cov=full_cov, q_sqrt=q_sqrt, white=white)

    else:
        base_conditional_args_to_map = (Kmms, Kmns, Knns, fs)

        def single_gp_conditional(
            t: Tuple[tf.Tensor, ...]
        ) -> MeanAndVariance:  # pragma: no cover - tf.map_fn is invisible to codecov
            Kmm, Kmn, Knn, f = t
            return base_conditional(Kmn, Kmm, Knn, f, full_cov=full_cov, q_sqrt=q_sqrt, white=white)

    rmu, rvar = tf.map_fn(
        single_gp_conditional, base_conditional_args_to_map, (default_float(), default_float())
    )  # [P, N, 1], [P, 1, N, N] or [P, N, 1]

    fmu = rollaxis_left(tf.squeeze(rmu, axis=-1), 1)  # [N, P]

    if full_cov:
        fvar = tf.squeeze(rvar, axis=-3)  # [..., 0, :, :]  # [P, N, N]
    else:
        fvar = rollaxis_left(tf.squeeze(rvar, axis=-1), 1)  # [N, P]

    return fmu, fvar


# gpflow/config/__config__.py

# Copyright 2019-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
This is a private module that manages GPflow configuration.

The module provides functions to modify default settings of GPflow, such as:
- the standard float precision and integer type
- the type of positive transformation
- a value for a minimum shift from zero for the positive transformation
- an output format for `gpflow.utilities.print_summary`

The module holds global configuration :class:`Config` variable that stores all
setting values.

Environment variables are an alternative way for changing the default GPflow
configuration.

.. warning::
    The user has to set environment variables before running python
    interpreter to modify the configuration.

Full set of environment variables and available options:

* ``GPFLOW_INT``: "int16", "int32", or "int64"
* ``GPFLOW_FLOAT``: "float16", "float32", or "float64"
* ``GPFLOW_POSITIVE_BIJECTOR``: "exp" or "softplus"
* ``GPFLOW_POSITIVE_MINIMUM``: Any positive float number
* ``GPFLOW_SUMMARY_FMT``: "notebook" or any other format that :mod:`tabulate` can handle.
* ``GPFLOW_JITTER``: Any positive float number

The user can also change the GPflow configuration temporarily with a context
manager :func:`as_context`:

>>> config = Config(jitter=1e-5)
>>> with as_context(config):
>>>     # ...code here sees new config
"""

import contextlib
import enum
import os
from dataclasses import dataclass, field, replace
from typing import Any, Dict, Generator, List, Mapping, Optional, Union

import numpy as np
import tabulate
import tensorflow as tf
import tensorflow_probability as tfp

__all__ = [
    "Config",
    "as_context",
    "config",
    "default_float",
    "default_int",
    "default_jitter",
    "default_positive_bijector",
    "default_positive_minimum",
    "default_summary_fmt",
    "positive_bijector_type_map",
    "set_config",
    "set_default_float",
    "set_default_int",
    "set_default_jitter",
    "set_default_positive_bijector",
    "set_default_positive_minimum",
    "set_default_summary_fmt",
]

__config: Optional["Config"] = None


class _Values(enum.Enum):
    """Setting's names collection with default values. The `name` method returns name
    of the environment variable. E.g. for `SUMMARY_FMT` field the environment variable
    will be `GPFLOW_SUMMARY_FMT`."""

    INT = np.int32
    FLOAT = np.float64
    POSITIVE_BIJECTOR = "softplus"
    POSITIVE_MINIMUM = 0.0
    SUMMARY_FMT = "fancy_grid"
    JITTER = 1e-6

    @property
    def name(self) -> str:  # type: ignore[override]  # name is generated and has weird typing.
        return f"GPFLOW_{super().name}"


def _default(value: _Values) -> Any:
    """Checks if value is set in the environment."""
    return os.getenv(value.name, default=value.value)


def _default_numeric_type_factory(
    valid_types: Mapping[str, type], enum_key: _Values, type_name: str
) -> type:
    value: Union[str, type] = _default(enum_key)
    if isinstance(value, type) and (value in valid_types.values()):
        return value
    assert isinstance(value, str)  # Hint for mypy
    if value not in valid_types:
        raise TypeError(f"Config cannot recognize {type_name} type.")
    return valid_types[value]


def _default_int_factory() -> type:
    valid_types = dict(int16=np.int16, int32=np.int32, int64=np.int64)
    return _default_numeric_type_factory(valid_types, _Values.INT, "int")


def _default_float_factory() -> type:
    valid_types = dict(float16=np.float16, float32=np.float32, float64=np.float64)
    return _default_numeric_type_factory(valid_types, _Values.FLOAT, "float")


def _default_jitter_factory() -> float:
    value = _default(_Values.JITTER)
    try:
        return float(value)
    except ValueError:
        raise TypeError("Config cannot set the jitter value with non float type.")


def _default_positive_bijector_factory() -> str:
    bijector_type: str = _default(_Values.POSITIVE_BIJECTOR)
    if bijector_type not in positive_bijector_type_map().keys():
        raise TypeError(
            "Config cannot set the passed value as a default positive bijector."
            f"Available options: {set(positive_bijector_type_map().keys())}"
        )
    return bijector_type


def _default_positive_minimum_factory() -> float:
    value = _default(_Values.POSITIVE_MINIMUM)
    try:
        return float(value)
    except ValueError:
        raise TypeError("Config cannot set the positive_minimum value with non float type.")


def _default_summary_fmt_factory() -> Optional[str]:
    result: Optional[str] = _default(_Values.SUMMARY_FMT)
    return result


# The following type alias is for the Config class, to help a static analyser distinguish
# between the built-in 'float' type and the 'float' type defined in the that class.
Float = Union[float]


@dataclass(frozen=True)
class Config:
    """
    Immutable object for storing global GPflow settings
    """

    int: type = field(default_factory=_default_int_factory)
    """Integer data type, int32 or int64."""

    float: type = field(default_factory=_default_float_factory)
    """Float data type, float32 or float64"""

    jitter: Float = field(default_factory=_default_jitter_factory)
    """
    Jitter value. Mainly used for for making badly conditioned matrices more stable.

    Default value is `1e-6`.
    """

    positive_bijector: str = field(default_factory=_default_positive_bijector_factory)
    """
    Method for positive bijector, either "softplus" or "exp".

    Default is "softplus".
    """

    positive_minimum: Float = field(default_factory=_default_positive_minimum_factory)
    """Lower bound for the positive transformation."""

    summary_fmt: Optional[str] = field(default_factory=_default_summary_fmt_factory)
    """Summary format for module printing."""


def config() -> Config:
    """Returns current active config."""
    assert __config is not None, "__config is None. This should never happen."
    return __config


def default_int() -> type:
    """Returns default integer type"""
    return config().int


def default_float() -> type:
    """Returns default float type"""
    return config().float


def default_jitter() -> float:
    """
    The jitter is a constant that GPflow adds to the diagonal of matrices
    to achieve numerical stability of the system when the condition number
    of the associated matrices is large, and therefore the matrices nearly singular.
    """
    return config().jitter


def default_positive_bijector() -> str:
    """Type of bijector used for positive constraints: exp or softplus."""
    return config().positive_bijector


def default_positive_minimum() -> float:
    """Shift constant that GPflow adds to all positive constraints."""
    return config().positive_minimum


def default_summary_fmt() -> Optional[str]:
    """Summary printing format as understood by :mod:`tabulate` or a special case "notebook"."""
    return config().summary_fmt


def set_config(new_config: Config) -> None:
    """Update GPflow config with new settings from `new_config`."""
    global __config
    __config = new_config


def set_default_int(value_type: type) -> None:
    """
    Sets default integer type. Available options are ``np.int16``, ``np.int32``,
    or ``np.int64``.
    """
    try:
        tf_dtype = tf.as_dtype(value_type)  # Test that it's a tensorflow-valid dtype
    except TypeError:
        raise TypeError(f"{value_type} is not a valid tf or np dtype")

    if not tf_dtype.is_integer:
        raise TypeError(f"{value_type} is not an integer dtype")

    set_config(replace(config(), int=tf_dtype.as_numpy_dtype))


def set_default_float(value_type: type) -> None:
    """
    Sets default float type. Available options are `np.float16`, `np.float32`,
    or `np.float64`.
    """
    try:
        tf_dtype = tf.as_dtype(value_type)  # Test that it's a tensorflow-valid dtype
    except TypeError:
        raise TypeError(f"{value_type} is not a valid tf or np dtype")

    if not tf_dtype.is_floating:
        raise TypeError(f"{value_type} is not a float dtype")

    set_config(replace(config(), float=tf_dtype.as_numpy_dtype))


def set_default_jitter(value: float) -> None:
    """
    Sets constant jitter value.
    The jitter is a constant that GPflow adds to the diagonal of matrices
    to achieve numerical stability of the system when the condition number
    of the associated matrices is large, and therefore the matrices nearly singular.
    """
    if not (
        isinstance(value, (tf.Tensor, np.ndarray)) and len(value.shape) == 0
    ) and not isinstance(value, float):
        raise TypeError("Expected float32 or float64 scalar value")

    if value < 0:
        raise ValueError("Jitter must be non-negative")

    set_config(replace(config(), jitter=value))


def set_default_positive_bijector(value: str) -> None:
    """
    Sets positive bijector type.
    There are currently two options implemented: "exp" and "softplus".
    """
    type_map = positive_bijector_type_map()
    if isinstance(value, str):
        value = value.lower()
    if value not in type_map:
        raise ValueError(f"`{value}` not in set of valid bijectors: {sorted(type_map)}")

    set_config(replace(config(), positive_bijector=value))


def set_default_positive_minimum(value: float) -> None:
    """Sets shift constant for positive transformation."""
    if not (
        isinstance(value, (tf.Tensor, np.ndarray)) and len(value.shape) == 0
    ) and not isinstance(value, float):
        raise TypeError("Expected float32 or float64 scalar value")

    if value < 0:
        raise ValueError("Positive minimum must be non-negative")

    set_config(replace(config(), positive_minimum=value))


def set_default_summary_fmt(value: Optional[str]) -> None:
    formats: List[Optional[str]] = list(tabulate.tabulate_formats)
    formats.extend(["notebook", None])
    if value not in formats:
        raise ValueError(f"Summary does not support '{value}' format")

    set_config(replace(config(), summary_fmt=value))


def positive_bijector_type_map() -> Dict[str, type]:
    return {
        "exp": tfp.bijectors.Exp,
        "softplus": tfp.bijectors.Softplus,
    }


@contextlib.contextmanager
def as_context(temporary_config: Optional[Config] = None) -> Generator[None, None, None]:
    """Ensure that global configs defaults, with a context manager. Useful for testing."""
    current_config = config()
    temporary_config = replace(current_config) if temporary_config is None else temporary_config
    try:
        set_config(temporary_config)
        yield
    finally:
        set_config(current_config)


# Set global config.
set_config(Config())


# gpflow/config/__init__.py

from .__config__ import *
from .__config__ import __doc__  # needs explicit import

__all__ = [
    "Config",
    "as_context",
    "config",
    "default_float",
    "default_int",
    "default_jitter",
    "default_positive_bijector",
    "default_positive_minimum",
    "default_summary_fmt",
    "positive_bijector_type_map",
    "set_config",
    "set_default_float",
    "set_default_int",
    "set_default_jitter",
    "set_default_positive_bijector",
    "set_default_positive_minimum",
    "set_default_summary_fmt",
]


# gpflow/covariances/__init__.py

from . import kufs, kuus, multioutput
from .dispatch import Kuf, Kuu

__all__ = [
    "Kuf",
    "Kuu",
    "dispatch",
    "kufs",
    "kuus",
    "multioutput",
]


# gpflow/covariances/dispatch.py

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from ..utilities import Dispatcher

Kuu = Dispatcher("Kuu")
Kuf = Dispatcher("Kuf")


# gpflow/covariances/kufs.py

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np  # pylint: disable=unused-import  # Used by Sphinx to generate documentation.
import tensorflow as tf

from ..base import TensorLike, TensorType
from ..experimental.check_shapes import check_shapes
from ..inducing_variables import InducingPatches, InducingPoints, Multiscale
from ..kernels import Convolutional, Kernel, SquaredExponential
from .dispatch import Kuf


@Kuf.register(InducingPoints, Kernel, TensorLike)
@check_shapes(
    "inducing_variable: [M, D, 1]",
    "Xnew: [batch..., N, D]",
    "return: [M, batch..., N]",
)
def Kuf_kernel_inducingpoints(
    inducing_variable: InducingPoints, kernel: Kernel, Xnew: TensorType
) -> tf.Tensor:
    return kernel(inducing_variable.Z, Xnew)


@Kuf.register(Multiscale, SquaredExponential, TensorLike)
@check_shapes(
    "inducing_variable: [M, D, 1]",
    "Xnew: [batch..., N, D]",
    "return: [M, batch..., N]",
)
def Kuf_sqexp_multiscale(
    inducing_variable: Multiscale, kernel: SquaredExponential, Xnew: TensorType
) -> tf.Tensor:
    Xnew, _ = kernel.slice(Xnew, None)
    Zmu, Zlen = kernel.slice(inducing_variable.Z, inducing_variable.scales)
    idlengthscales = kernel.lengthscales + Zlen
    d = inducing_variable._cust_square_dist(Xnew, Zmu, idlengthscales[None, :, :])
    lengthscales = tf.reduce_prod(kernel.lengthscales / idlengthscales, 1)
    lengthscales = tf.reshape(lengthscales, (1, -1))
    return tf.transpose(kernel.variance * tf.exp(-0.5 * d) * lengthscales)


@Kuf.register(InducingPatches, Convolutional, object)
@check_shapes(
    "inducing_variable: [M, D, 1]",
    "Xnew: [batch..., N, D2]",
    "return: [M, batch..., N]",
)
def Kuf_conv_patch(
    inducing_variable: InducingPatches, kernel: Convolutional, Xnew: TensorType
) -> tf.Tensor:
    Xp = kernel.get_patches(Xnew)  # [N, num_patches, patch_len]
    bigKzx = kernel.base_kernel.K(
        inducing_variable.Z, Xp
    )  # [M, N, P] -- thanks to broadcasting of kernels
    Kzx = tf.reduce_sum(bigKzx * kernel.weights if hasattr(kernel, "weights") else bigKzx, [2])
    return Kzx / kernel.num_patches


# gpflow/covariances/kuus.py

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import tensorflow as tf

from ..config import default_float
from ..experimental.check_shapes import check_shapes
from ..inducing_variables import InducingPatches, InducingPoints, Multiscale
from ..kernels import Convolutional, Kernel, SquaredExponential
from .dispatch import Kuu


@Kuu.register(InducingPoints, Kernel)
@check_shapes(
    "inducing_variable: [M, D, 1]",
    "return: [M, M]",
)
def Kuu_kernel_inducingpoints(
    inducing_variable: InducingPoints, kernel: Kernel, *, jitter: float = 0.0
) -> tf.Tensor:
    Kzz = kernel(inducing_variable.Z)
    Kzz += jitter * tf.eye(inducing_variable.num_inducing, dtype=Kzz.dtype)
    return Kzz


@Kuu.register(Multiscale, SquaredExponential)
@check_shapes(
    "inducing_variable: [M, D, 1]",
    "return: [M, M]",
)
def Kuu_sqexp_multiscale(
    inducing_variable: Multiscale, kernel: SquaredExponential, *, jitter: float = 0.0
) -> tf.Tensor:
    Zmu, Zlen = kernel.slice(inducing_variable.Z, inducing_variable.scales)
    idlengthscales2 = tf.square(kernel.lengthscales + Zlen)
    sc = tf.sqrt(
        idlengthscales2[None, ...] + idlengthscales2[:, None, ...] - kernel.lengthscales ** 2
    )
    d = inducing_variable._cust_square_dist(Zmu, Zmu, sc)
    Kzz = kernel.variance * tf.exp(-d / 2) * tf.reduce_prod(kernel.lengthscales / sc, 2)
    Kzz += jitter * tf.eye(inducing_variable.num_inducing, dtype=Kzz.dtype)
    return Kzz


@Kuu.register(InducingPatches, Convolutional)
@check_shapes(
    "inducing_variable: [M, D, 1]",
    "return: [M, M]",
)
def Kuu_conv_patch(
    inducing_variable: InducingPatches, kernel: Convolutional, jitter: float = 0.0
) -> tf.Tensor:
    return kernel.base_kernel.K(inducing_variable.Z) + jitter * tf.eye(
        inducing_variable.num_inducing, dtype=default_float()
    )


# gpflow/expectations/__init__.py

from . import (
    cross_kernels,
    linears,
    mean_functions,
    misc,
    products,
    quadratures,
    squared_exponentials,
    sums,
)
from .expectations import expectation, quadrature_expectation

__all__ = [
    "cross_kernels",
    "dispatch",
    "expectation",
    "expectations",
    "linears",
    "mean_functions",
    "misc",
    "products",
    "quadrature_expectation",
    "quadratures",
    "squared_exponentials",
    "sums",
]


# gpflow/expectations/cross_kernels.py

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Union, cast

import tensorflow as tf

from .. import kernels
from ..base import TensorType
from ..experimental.check_shapes import check_shapes
from ..inducing_variables import InducingPoints
from ..probability_distributions import DiagonalGaussian, Gaussian
from . import dispatch
from .expectations import expectation


@dispatch.expectation.register(
    (Gaussian, DiagonalGaussian),
    kernels.SquaredExponential,
    InducingPoints,
    kernels.Linear,
    InducingPoints,
)
@check_shapes(
    "p: [N, D]",
    "feat1: [M1, D, P]",
    "feat2: [M2, D, P]",
    "return: [N, M1, M2]",
)
def _expectation_gaussian_sqe_inducingpoints__linear_inducingpoints(
    p: Union[Gaussian, DiagonalGaussian],
    sqexp_kern: kernels.SquaredExponential,
    feat1: InducingPoints,
    lin_kern: kernels.Linear,
    feat2: InducingPoints,
    nghp: None = None,
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <Ka_{Z1, x_n} Kb_{x_n, Z2}>_p(x_n)
        - K_lin_{.,.} :: SqExp kernel
        - K_sqexp_{.,.} :: Linear kernel
    Different Z1 and Z2 are handled if p is diagonal and K_lin and K_sqexp have disjoint
    active_dims, in which case the joint expectations simplify into a product of expectations

    :return: NxM1xM2
    """
    if sqexp_kern.on_separate_dims(lin_kern) and isinstance(
        p, DiagonalGaussian
    ):  # no joint expectations required
        eKxz1 = expectation(p, (sqexp_kern, feat1))
        eKxz2 = expectation(p, (lin_kern, feat2))
        return eKxz1[:, :, None] * eKxz2[:, None, :]

    if feat1 != feat2:
        raise NotImplementedError("inducing_variables have to be the same for both kernels.")

    if sqexp_kern.active_dims != lin_kern.active_dims:
        raise NotImplementedError("active_dims have to be the same for both kernels.")

    # use only active dimensions
    Xcov = sqexp_kern.slice_cov(tf.linalg.diag(p.cov) if isinstance(p, DiagonalGaussian) else p.cov)
    Z, Xmu = sqexp_kern.slice(feat1.Z, p.mu)

    N = tf.shape(Xmu)[0]
    D = tf.shape(Xmu)[1]

    def take_with_ard(value: TensorType) -> TensorType:
        if not sqexp_kern.ard:
            return tf.zeros((D,), dtype=value.dtype) + value
        return value

    lin_kern_variances = take_with_ard(lin_kern.variance)
    sqexp_kern_lengthscales = take_with_ard(sqexp_kern.lengthscales)

    chol_L_plus_Xcov = tf.linalg.cholesky(
        tf.linalg.diag(sqexp_kern_lengthscales ** 2) + Xcov
    )  # NxDxD

    Z_transpose = tf.transpose(Z)
    all_diffs = Z_transpose - tf.expand_dims(Xmu, 2)  # NxDxM
    exponent_mahalanobis = tf.linalg.triangular_solve(
        chol_L_plus_Xcov, all_diffs, lower=True
    )  # NxDxM
    exponent_mahalanobis = tf.reduce_sum(tf.square(exponent_mahalanobis), 1)  # NxM
    exponent_mahalanobis = tf.exp(-0.5 * exponent_mahalanobis)  # NxM

    sqrt_det_L = tf.reduce_prod(sqexp_kern_lengthscales)
    sqrt_det_L_plus_Xcov = tf.exp(
        tf.reduce_sum(tf.math.log(tf.linalg.diag_part(chol_L_plus_Xcov)), axis=1)
    )
    determinants = sqrt_det_L / sqrt_det_L_plus_Xcov  # N
    eKxz_sqexp = sqexp_kern.variance * (
        determinants[:, None] * exponent_mahalanobis
    )  ## NxM <- End RBF eKxz code

    tiled_Z = tf.tile(tf.expand_dims(Z_transpose, 0), (N, 1, 1))  # NxDxM
    z_L_inv_Xcov = tf.linalg.matmul(
        tiled_Z, Xcov / sqexp_kern_lengthscales[:, None] ** 2.0, transpose_a=True
    )  # NxMxD

    cross_eKzxKxz = tf.linalg.cholesky_solve(
        chol_L_plus_Xcov,
        cast(tf.Tensor, (lin_kern_variances * sqexp_kern_lengthscales ** 2.0))[..., None] * tiled_Z,
    )  # NxDxM

    cross_eKzxKxz = tf.linalg.matmul(
        (z_L_inv_Xcov + Xmu[:, None, :]) * eKxz_sqexp[..., None], cross_eKzxKxz
    )  # NxMxM
    return cross_eKzxKxz


@dispatch.expectation.register(
    (Gaussian, DiagonalGaussian),
    kernels.Linear,
    InducingPoints,
    kernels.SquaredExponential,
    InducingPoints,
)
@check_shapes(
    "p: [N, D]",
    "feat1: [M1, D, P]",
    "feat2: [M2, D, P]",
    "return: [N, M1, M2]",
)
def _expectation_gaussian_linear_inducingpoints__sqe_inducingpoints(
    p: Union[Gaussian, DiagonalGaussian],
    lin_kern: kernels.Linear,
    feat1: InducingPoints,
    sqexp_kern: kernels.SquaredExponential,
    feat2: InducingPoints,
    nghp: None = None,
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <Ka_{Z1, x_n} Kb_{x_n, Z2}>_p(x_n)
        - K_lin_{.,.} :: Linear kernel
        - K_sqexp_{.,.} :: sqexp kernel
    Different Z1 and Z2 are handled if p is diagonal and K_lin and K_sqexp have disjoint
    active_dims, in which case the joint expectations simplify into a product of expectations

    :return: NxM1xM2
    """
    return tf.linalg.adjoint(expectation(p, (sqexp_kern, feat2), (lin_kern, feat1)))


# gpflow/expectations/dispatch.py

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from ..utilities import Dispatcher

expectation = Dispatcher("expectation")
quadrature_expectation = Dispatcher("quadrature_expectation")
variational_expectation = Dispatcher("variational_expectation")


# gpflow/expectations/expectations.py

# Copyright 2018 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Optional, Tuple, Union, cast

import tensorflow as tf

from ..base import TensorType
from ..experimental.check_shapes import check_shapes
from ..inducing_variables import InducingVariables
from ..kernels import Kernel
from ..mean_functions import MeanFunction
from ..probability_distributions import (
    DiagonalGaussian,
    Gaussian,
    MarkovGaussian,
    ProbabilityDistribution,
)
from . import dispatch

ProbabilityDistributionLike = Union[ProbabilityDistribution, Tuple[TensorType, TensorType]]
"""
Either a prabability distribution, or a tuple of mean, covariance that is turned into an
appropriate Gaussian distribution, depending on the shape of the covariance.
"""

ExpectationObject = Union[Kernel, MeanFunction, None]
PackedExpectationObject = Union[ExpectationObject, Tuple[Kernel, InducingVariables]]


def expectation(
    p: ProbabilityDistributionLike,
    obj1: PackedExpectationObject,
    obj2: PackedExpectationObject = None,
    nghp: Optional[int] = None,
) -> tf.Tensor:
    """
    Compute the expectation <obj1(x) obj2(x)>_p(x)
    Uses multiple-dispatch to select an analytical implementation,
    if one is available. If not, it falls back to quadrature.

    :type p: (mu, cov) tuple or a `ProbabilityDistribution` object
    :type obj1: kernel, mean function, (kernel, inducing_variable), or None
    :type obj2: kernel, mean function, (kernel, inducing_variable), or None
    :param int nghp: passed to `_quadrature_expectation` to set the number
                     of Gauss-Hermite points used: `num_gauss_hermite_points`
    :return: a 1-D, 2-D, or 3-D tensor containing the expectation

    Allowed combinations

    - Psi statistics:
        >>> eKdiag = expectation(p, kernel)  (N)  # Psi0
        >>> eKxz = expectation(p, (kernel, inducing_variable))  (NxM)  # Psi1
        >>> exKxz = expectation(p, identity_mean, (kernel, inducing_variable))  (NxDxM)
        >>> eKzxKxz = expectation(p, (kernel, inducing_variable), (kernel, inducing_variable))  (NxMxM)  # Psi2

    - kernels and mean functions:
        >>> eKzxMx = expectation(p, (kernel, inducing_variable), mean)  (NxMxQ)
        >>> eMxKxz = expectation(p, mean, (kernel, inducing_variable))  (NxQxM)

    - only mean functions:
        >>> eMx = expectation(p, mean)  (NxQ)
        >>> eM1x_M2x = expectation(p, mean1, mean2)  (NxQ1xQ2)
        .. note:: mean(x) is 1xQ (row vector)

    - different kernels. This occurs, for instance, when we are calculating Psi2 for Sum kernels:
        >>> eK1zxK2xz = expectation(p, (kern1, inducing_variable), (kern2, inducing_variable))  (NxMxM)
    """
    p, obj1, feat1, obj2, feat2 = _init_expectation(p, obj1, obj2)
    try:
        return dispatch.expectation(p, obj1, feat1, obj2, feat2, nghp=nghp)
    except NotImplementedError:
        return dispatch.quadrature_expectation(p, obj1, feat1, obj2, feat2, nghp=nghp)


def quadrature_expectation(
    p: ProbabilityDistributionLike,
    obj1: PackedExpectationObject,
    obj2: PackedExpectationObject = None,
    nghp: Optional[int] = None,
) -> tf.Tensor:
    """
    Compute the expectation <obj1(x) obj2(x)>_p(x)
    Uses Gauss-Hermite quadrature for approximate integration.

    :type p: (mu, cov) tuple or a `ProbabilityDistribution` object
    :type obj1: kernel, mean function, (kernel, inducing_variable), or None
    :type obj2: kernel, mean function, (kernel, inducing_variable), or None
    :param int num_gauss_hermite_points: passed to `_quadrature_expectation` to set
                                         the number of Gauss-Hermite points used
    :return: a 1-D, 2-D, or 3-D tensor containing the expectation
    """
    print(f"2. p={p}, obj1={obj1}, obj2={obj2}")
    p, obj1, feat1, obj2, feat2 = _init_expectation(p, obj1, obj2)
    return dispatch.quadrature_expectation(p, obj1, feat1, obj2, feat2, nghp=nghp)


@check_shapes(
    "return[0]: [N, D]",
    "return[2]: [M1, D, P]",
    "return[4]: [M2, D, P]",
)
def _init_expectation(
    p: ProbabilityDistributionLike, obj1: PackedExpectationObject, obj2: PackedExpectationObject
) -> Tuple[
    ProbabilityDistribution,
    ExpectationObject,
    Optional[InducingVariables],
    ExpectationObject,
    Optional[InducingVariables],
]:
    if isinstance(p, tuple):
        mu, cov = p
        classes = [DiagonalGaussian, Gaussian, MarkovGaussian]
        p = classes[cov.ndim - 2](*p)  # type: ignore[abstract]

    obj1, feat1 = obj1 if isinstance(obj1, tuple) else (obj1, None)
    obj2, feat2 = obj2 if isinstance(obj2, tuple) else (obj2, None)
    return (
        # type-ignore instead of cast, because it dependes on versions whether a cast is necessary.
        p,  # type: ignore[return-value]
        cast(ExpectationObject, obj1),
        cast(Optional[InducingVariables], feat1),
        cast(ExpectationObject, obj2),
        cast(Optional[InducingVariables], feat2),
    )


# gpflow/expectations/linears.py

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Type, Union

import tensorflow as tf

from .. import kernels
from .. import mean_functions as mfn
from ..experimental.check_shapes import check_shapes
from ..inducing_variables import InducingPoints
from ..probability_distributions import DiagonalGaussian, Gaussian, MarkovGaussian
from . import dispatch
from .expectations import expectation

NoneType: Type[None] = type(None)


@dispatch.expectation.register(Gaussian, kernels.Linear, NoneType, NoneType, NoneType)
@check_shapes(
    "p: [N, D]",
    "return: [N]",
)
def _expectation_gaussian_linear(
    p: Gaussian, kernel: kernels.Linear, _: None, __: None, ___: None, nghp: None = None
) -> tf.Tensor:
    """
    Compute the expectation:
    <diag(K_{X, X})>_p(X)
        - K_{.,.} :: Linear kernel

    :return: N
    """
    # use only active dimensions
    Xmu, _ = kernel.slice(p.mu, None)
    Xcov = kernel.slice_cov(p.cov)

    return tf.reduce_sum(kernel.variance * (tf.linalg.diag_part(Xcov) + Xmu ** 2), 1)


@dispatch.expectation.register(Gaussian, kernels.Linear, InducingPoints, NoneType, NoneType)
@check_shapes(
    "p: [N, D]",
    "inducing_variable: [M, D, P]",
    "return: [N, M]",
)
def _expectation_gaussian_linear_inducingpoints(
    p: Gaussian,
    kernel: kernels.Linear,
    inducing_variable: InducingPoints,
    _: None,
    __: None,
    nghp: None = None,
) -> tf.Tensor:
    """
    Compute the expectation:
    <K_{X, Z}>_p(X)
        - K_{.,.} :: Linear kernel

    :return: NxM
    """
    # use only active dimensions
    Z, Xmu = kernel.slice(inducing_variable.Z, p.mu)

    return tf.linalg.matmul(Xmu, Z * kernel.variance, transpose_b=True)


@dispatch.expectation.register(Gaussian, kernels.Linear, InducingPoints, mfn.Identity, NoneType)
@check_shapes(
    "p: [N, D]",
    "inducing_variable: [M, D, P]",
    "return: [N, M, D]",
)
def _expectation_gaussian_linear_inducingpoints__identity(
    p: Gaussian,
    kernel: kernels.Linear,
    inducing_variable: InducingPoints,
    mean: mfn.Identity,
    _: None,
    nghp: None = None,
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <K_{Z, x_n} x_n^T>_p(x_n)
        - K_{.,.} :: Linear kernel

    :return: NxMxD
    """
    Xmu, Xcov = p.mu, p.cov

    N = tf.shape(Xmu)[0]
    var_Z = kernel.variance * inducing_variable.Z  # MxD
    tiled_Z = tf.tile(tf.expand_dims(var_Z, 0), (N, 1, 1))  # NxMxD
    return tf.linalg.matmul(tiled_Z, Xcov + (Xmu[..., None] * Xmu[:, None, :]))


@dispatch.expectation.register(
    MarkovGaussian, kernels.Linear, InducingPoints, mfn.Identity, NoneType
)
@check_shapes(
    "p: [N, D]",
    "inducing_variable: [M, D, P]",
    "return: [N, M, D]",
)
def _expectation_markov_linear_inducingpoints__identity(
    p: MarkovGaussian,
    kernel: kernels.Linear,
    inducing_variable: InducingPoints,
    mean: mfn.Identity,
    _: None,
    nghp: None = None,
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <K_{Z, x_n} x_{n+1}^T>_p(x_{n:n+1})
        - K_{.,.} :: Linear kernel
        - p       :: MarkovGaussian distribution (p.cov 2x(N+1)xDxD)

    :return: NxMxD
    """
    Xmu, Xcov = p.mu, p.cov

    N = tf.shape(Xmu)[0] - 1
    var_Z = kernel.variance * inducing_variable.Z  # MxD
    tiled_Z = tf.tile(tf.expand_dims(var_Z, 0), (N, 1, 1))  # NxMxD
    eXX = Xcov[1, :-1] + (Xmu[:-1][..., None] * Xmu[1:][:, None, :])  # NxDxD
    return tf.linalg.matmul(tiled_Z, eXX)


@dispatch.expectation.register(
    (Gaussian, DiagonalGaussian), kernels.Linear, InducingPoints, kernels.Linear, InducingPoints
)
@check_shapes(
    "p: [N, D]",
    "feat1: [M, D, P]",
    "feat2: [M, D, P]",
    "return: [N, M, M]",
)
def _expectation_gaussian_linear_inducingpoints__linear_inducingpoints(
    p: Union[Gaussian, DiagonalGaussian],
    kern1: kernels.Linear,
    feat1: InducingPoints,
    kern2: kernels.Linear,
    feat2: InducingPoints,
    nghp: None = None,
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <Ka_{Z1, x_n} Kb_{x_n, Z2}>_p(x_n)
        - Ka_{.,.}, Kb_{.,.} :: Linear kernels
    Ka and Kb as well as Z1 and Z2 can differ from each other, but this is supported
    only if the Gaussian p is Diagonal (p.cov NxD) and Ka, Kb have disjoint active_dims
    in which case the joint expectations simplify into a product of expectations

    :return: NxMxM
    """
    if kern1.on_separate_dims(kern2) and isinstance(
        p, DiagonalGaussian
    ):  # no joint expectations required
        eKxz1 = expectation(p, (kern1, feat1))
        eKxz2 = expectation(p, (kern2, feat2))
        return eKxz1[:, :, None] * eKxz2[:, None, :]

    if kern1 != kern2 or feat1 != feat2:
        raise NotImplementedError(
            "The expectation over two kernels has only an "
            "analytical implementation if both kernels are equal."
        )

    kernel = kern1
    inducing_variable = feat1

    # use only active dimensions
    Xcov = kernel.slice_cov(tf.linalg.diag(p.cov) if isinstance(p, DiagonalGaussian) else p.cov)
    Z, Xmu = kernel.slice(inducing_variable.Z, p.mu)

    N = tf.shape(Xmu)[0]
    var_Z = kernel.variance * Z
    tiled_Z = tf.tile(tf.expand_dims(var_Z, 0), (N, 1, 1))  # NxMxD
    XX = Xcov + tf.expand_dims(Xmu, 1) * tf.expand_dims(Xmu, 2)  # NxDxD
    return tf.linalg.matmul(tf.linalg.matmul(tiled_Z, XX), tiled_Z, transpose_b=True)


# gpflow/expectations/mean_functions.py

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# noqa: F811

from typing import Type, Union

import tensorflow as tf

from .. import mean_functions as mfn
from ..experimental.check_shapes import check_shapes
from ..probability_distributions import Gaussian
from . import dispatch
from .expectations import expectation

NoneType: Type[None] = type(None)


@dispatch.expectation.register(Gaussian, (mfn.Linear, mfn.Constant), NoneType, NoneType, NoneType)
@check_shapes(
    "p: [N, D]",
    "return: [N, Q]",
)
def _expectation_gaussian_linear(
    p: Gaussian,
    mean: Union[mfn.Linear, mfn.Constant],
    _: None,
    __: None,
    ___: None,
    nghp: None = None,
) -> tf.Tensor:
    """
    Compute the expectation:
    <m(X)>_p(X)
        - m(x) :: Linear, Identity or Constant mean function

    :return: NxQ
    """
    return mean(p.mu)


@dispatch.expectation.register(Gaussian, mfn.Constant, NoneType, mfn.Constant, NoneType)
@check_shapes(
    "p: [N, D]",
    "return: [N, Q1, Q2]",
)
def _expectation_gaussian_constant__constant(
    p: Gaussian, mean1: mfn.Constant, _: None, mean2: mfn.Constant, __: None, nghp: None = None
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <m1(x_n)^T m2(x_n)>_p(x_n)
        - m1(.), m2(.) :: Constant mean functions

    :return: NxQ1xQ2
    """
    return mean1(p.mu)[:, :, None] * mean2(p.mu)[:, None, :]


@dispatch.expectation.register(Gaussian, mfn.Constant, NoneType, mfn.MeanFunction, NoneType)
@check_shapes(
    "p: [N, D]",
    "return: [N, Q1, Q2]",
)
def _expectation_gaussian_constant__meanfunction(
    p: Gaussian, mean1: mfn.Constant, _: None, mean2: mfn.Constant, __: None, nghp: None = None
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <m1(x_n)^T m2(x_n)>_p(x_n)
        - m1(.) :: Constant mean function
        - m2(.) :: General mean function

    :return: NxQ1xQ2
    """
    e_mean2 = expectation(p, mean2)
    return mean1(p.mu)[:, :, None] * e_mean2[:, None, :]


@dispatch.expectation.register(Gaussian, mfn.MeanFunction, NoneType, mfn.Constant, NoneType)
@check_shapes(
    "p: [N, D]",
    "return: [N, Q1, Q2]",
)
def _expectation_gaussian_meanfunction__constant(
    p: Gaussian,
    mean1: mfn.MeanFunction,
    _: None,
    mean2: mfn.MeanFunction,
    __: None,
    nghp: None = None,
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <m1(x_n)^T m2(x_n)>_p(x_n)
        - m1(.) :: General mean function
        - m2(.) :: Constant mean function

    :return: NxQ1xQ2
    """
    e_mean1 = expectation(p, mean1)
    return e_mean1[:, :, None] * mean2(p.mu)[:, None, :]


@dispatch.expectation.register(Gaussian, mfn.Identity, NoneType, mfn.Identity, NoneType)
@check_shapes(
    "p: [N, D]",
    "return: [N, D, D]",
)
def _expectation_gaussian_identity__identity(
    p: Gaussian, mean1: mfn.Identity, _: None, mean2: mfn.Identity, __: None, nghp: None = None
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <m1(x_n)^T m2(x_n)>_p(x_n)
        - m1(.), m2(.) :: Identity mean functions

    :return: NxDxD
    """
    return p.cov + (p.mu[:, :, None] * p.mu[:, None, :])


@dispatch.expectation.register(Gaussian, mfn.Identity, NoneType, mfn.Linear, NoneType)
@check_shapes(
    "p: [N, D]",
    "return: [N, D, Q]",
)
def _expectation_gaussian_identity__linear(
    p: Gaussian, mean1: mfn.Identity, _: None, mean2: mfn.Linear, __: None, nghp: None = None
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <m1(x_n)^T m2(x_n)>_p(x_n)
        - m1(.) :: Identity mean function
        - m2(.) :: Linear mean function

    :return: NxDxQ
    """
    N = tf.shape(p.mu)[0]
    e_xxt = p.cov + (p.mu[:, :, None] * p.mu[:, None, :])  # NxDxD
    e_xxt_A = tf.linalg.matmul(e_xxt, tf.tile(mean2.A[None, ...], (N, 1, 1)))  # NxDxQ
    e_x_bt = p.mu[:, :, None] * mean2.b[None, None, :]  # NxDxQ

    return e_xxt_A + e_x_bt


@dispatch.expectation.register(Gaussian, mfn.Linear, NoneType, mfn.Identity, NoneType)
@check_shapes(
    "p: [N, D]",
    "return: [N, Q, D]",
)
def _expectation_gaussian_linear__identity(
    p: Gaussian, mean1: mfn.Linear, _: None, mean2: mfn.Identity, __: None, nghp: None = None
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <m1(x_n)^T m2(x_n)>_p(x_n)
        - m1(.) :: Linear mean function
        - m2(.) :: Identity mean function

    :return: NxQxD
    """
    N = tf.shape(p.mu)[0]
    e_xxt = p.cov + (p.mu[:, :, None] * p.mu[:, None, :])  # NxDxD
    e_A_xxt = tf.linalg.matmul(
        tf.tile(mean1.A[None, ...], (N, 1, 1)), e_xxt, transpose_a=True
    )  # NxQxD
    e_b_xt = mean1.b[None, :, None] * p.mu[:, None, :]  # NxQxD

    return e_A_xxt + e_b_xt


@dispatch.expectation.register(Gaussian, mfn.Linear, NoneType, mfn.Linear, NoneType)
@check_shapes(
    "p: [N, D]",
    "return: [N, Q1, Q2]",
)
def _expectation_gaussian_linear__linear(
    p: Gaussian, mean1: mfn.Linear, _: None, mean2: mfn.Linear, __: None, nghp: None = None
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <m1(x_n)^T m2(x_n)>_p(x_n)
        - m1(.), m2(.) :: Linear mean functions

    :return: NxQ1xQ2
    """
    e_xxt = p.cov + (p.mu[:, :, None] * p.mu[:, None, :])  # NxDxD
    e_A1t_xxt_A2 = tf.einsum("iq,nij,jz->nqz", mean1.A, e_xxt, mean2.A)  # NxQ1xQ2
    e_A1t_x_b2t = tf.einsum("iq,ni,z->nqz", mean1.A, p.mu, mean2.b)  # NxQ1xQ2
    e_b1_xt_A2 = tf.einsum("q,ni,iz->nqz", mean1.b, p.mu, mean2.A)  # NxQ1xQ2
    e_b1_b2t = mean1.b[:, None] * mean2.b[None, :]  # Q1xQ2

    return e_A1t_xxt_A2 + e_A1t_x_b2t + e_b1_xt_A2 + e_b1_b2t


# gpflow/expectations/misc.py

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Optional, Type, Union, cast

import tensorflow as tf

from .. import kernels
from .. import mean_functions as mfn
from ..experimental.check_shapes import check_shapes
from ..inducing_variables import InducingPoints, InducingVariables
from ..probability_distributions import DiagonalGaussian, Gaussian, MarkovGaussian
from . import dispatch
from .expectations import ExpectationObject, PackedExpectationObject, expectation

NoneType: Type[None] = type(None)

# ================ exKxz transpose and mean function handling =================


@dispatch.expectation.register(
    (Gaussian, MarkovGaussian), mfn.Identity, NoneType, kernels.Linear, InducingPoints
)
@check_shapes(
    "p: [N, D]",
    "inducing_variable: [M, D, P]",
    "return: [N, D, M]",
)
def _expectation_gaussian__linear_inducingpoints(
    p: Union[Gaussian, MarkovGaussian],
    mean: mfn.Identity,
    _: None,
    kernel: kernels.Linear,
    inducing_variable: InducingPoints,
    nghp: None = None,
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <x_n K_{x_n, Z}>_p(x_n)
        - K_{.,} :: Linear kernel
    or the equivalent for MarkovGaussian

    :return: NxDxM
    """
    return tf.linalg.adjoint(expectation(p, (kernel, inducing_variable), mean))


@dispatch.expectation.register(
    (Gaussian, MarkovGaussian), kernels.Kernel, InducingVariables, mfn.MeanFunction, NoneType
)
@check_shapes(
    "p: [N, D]",
    "inducing_variable: [M, D, P]",
    "return: [N, M, Q]",
)
def _expectation_gaussian_kernel_inducingvariables__meanfunction(
    p: Union[Gaussian, MarkovGaussian],
    kernel: kernels.Kernel,
    inducing_variable: InducingVariables,
    mean: mfn.MeanFunction,
    _: None,
    nghp: None = None,
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <K_{Z, x_n} m(x_n)>_p(x_n)
    or the equivalent for MarkovGaussian

    :return: NxMxQ
    """
    return tf.linalg.adjoint(expectation(p, mean, (kernel, inducing_variable), nghp=nghp))


@dispatch.expectation.register(Gaussian, mfn.Constant, NoneType, kernels.Kernel, InducingPoints)
@check_shapes(
    "p: [N, D]",
    "inducing_variable: [M, D, P]",
    "return: [N, Q, M]",
)
def _expectation_gaussian_constant__kernel_inducingpoints(
    p: Gaussian,
    constant_mean: mfn.Constant,
    _: None,
    kernel: kernels.Kernel,
    inducing_variable: InducingPoints,
    nghp: None = None,
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <m(x_n)^T K_{x_n, Z}>_p(x_n)
        - m(x_i) = c :: Constant function
        - K_{.,.}    :: Kernel function

    :return: NxQxM
    """
    c = constant_mean(p.mu)  # NxQ
    eKxz = expectation(p, (kernel, inducing_variable), nghp=nghp)  # NxM

    return c[..., None] * eKxz[:, None, :]


@dispatch.expectation.register(Gaussian, mfn.Linear, NoneType, kernels.Kernel, InducingPoints)
@check_shapes(
    "p: [N, D]",
    "inducing_variable: [M, D, P]",
    "return: [N, Q, M]",
)
def _expectation_gaussian_linear__kernel_inducingpoints(
    p: Gaussian,
    linear_mean: mfn.Linear,
    _: None,
    kernel: kernels.Kernel,
    inducing_variable: InducingPoints,
    nghp: None = None,
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <m(x_n)^T K_{x_n, Z}>_p(x_n)
        - m(x_i) = A x_i + b :: Linear mean function
        - K_{.,.}            :: Kernel function

    :return: NxQxM
    """
    N = tf.shape(p.mu)[0]
    D = tf.shape(p.mu)[1]
    exKxz = expectation(p, mfn.Identity(D), (kernel, inducing_variable), nghp=nghp)
    eKxz = expectation(p, (kernel, inducing_variable), nghp=nghp)
    eAxKxz = tf.linalg.matmul(
        tf.tile(linear_mean.A[None, :, :], (N, 1, 1)), exKxz, transpose_a=True
    )
    ebKxz = linear_mean.b[None, :, None] * eKxz[:, None, :]
    return eAxKxz + ebKxz


@dispatch.expectation.register(Gaussian, mfn.Identity, NoneType, kernels.Kernel, InducingPoints)
def _expectation_gaussian__kernel_inducingpoints(
    p: Gaussian,
    identity_mean: mfn.Identity,
    _: None,
    kernel: kernels.Kernel,
    inducing_variable: InducingPoints,
    nghp: None = None,
) -> tf.Tensor:
    """
    This prevents infinite recursion for kernels that don't have specific
    implementations of _expectation(p, identity_mean, None, kernel, inducing_variable).
    Recursion can arise because Identity is a subclass of Linear mean function
    so _expectation(p, linear_mean, none, kernel, inducing_variable) would call itself.
    More specific signatures (e.g. (p, identity_mean, None, RBF, inducing_variable)) will
    be found and used whenever available
    """
    raise NotImplementedError


# ============== Conversion to Gaussian from Diagonal or Markov ===============
# Catching missing DiagonalGaussian implementations by converting to full Gaussian:


@dispatch.expectation.register(
    DiagonalGaussian, object, (InducingVariables, NoneType), object, (InducingVariables, NoneType)
)
def _expectation_diagonal_generic(
    p: DiagonalGaussian,
    obj1: ExpectationObject,
    feat1: Optional[InducingVariables],
    obj2: ExpectationObject,
    feat2: Optional[InducingVariables],
    nghp: None = None,
) -> tf.Tensor:
    gaussian = Gaussian(p.mu, tf.linalg.diag(p.cov))
    return expectation(
        gaussian,
        cast(PackedExpectationObject, (obj1, feat1)),
        cast(PackedExpectationObject, (obj2, feat2)),
        nghp=nghp,
    )


# Catching missing MarkovGaussian implementations by converting to Gaussian (when indifferent):


@dispatch.expectation.register(
    MarkovGaussian, object, (InducingVariables, NoneType), object, (InducingVariables, NoneType)
)
def _expectation_markov_generic(
    p: MarkovGaussian,
    obj1: ExpectationObject,
    feat1: Optional[InducingVariables],
    obj2: ExpectationObject,
    feat2: Optional[InducingVariables],
    nghp: None = None,
) -> tf.Tensor:
    """
    Nota Bene: if only one object is passed, obj1 is
    associated with x_n, whereas obj2 with x_{n+1}

    """
    if obj2 is None:
        gaussian = Gaussian(p.mu[:-1], p.cov[0, :-1])
        return expectation(gaussian, cast(PackedExpectationObject, (obj1, feat1)), nghp=nghp)
    elif obj1 is None:
        gaussian = Gaussian(p.mu[1:], p.cov[0, 1:])
        return expectation(gaussian, cast(PackedExpectationObject, (obj2, feat2)), nghp=nghp)
    else:
        return expectation(
            p,
            cast(PackedExpectationObject, (obj1, feat1)),
            cast(PackedExpectationObject, (obj2, feat2)),
            nghp=nghp,
        )


# gpflow/expectations/products.py

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from functools import reduce
from typing import Type

import tensorflow as tf

from .. import kernels
from ..experimental.check_shapes import check_shapes
from ..inducing_variables import InducingPoints
from ..probability_distributions import DiagonalGaussian
from . import dispatch
from .expectations import expectation

NoneType: Type[None] = type(None)


@dispatch.expectation.register(DiagonalGaussian, kernels.Product, NoneType, NoneType, NoneType)
@check_shapes(
    "p: [N, D]",
    "return: [N]",
)
def _expectation_diagonal_product(
    p: DiagonalGaussian, kernel: kernels.Product, _: None, __: None, ___: None, nghp: None = None
) -> tf.Tensor:
    r"""
    Compute the expectation:
    <\HadamardProd_i diag(Ki_{X[:, active_dims_i], X[:, active_dims_i]})>_p(X)
        - \HadamardProd_i Ki_{.,.} :: Product kernel
        - p                        :: DiagonalGaussian distribution (p.cov NxD)

    :return: N
    """
    if not kernel.on_separate_dimensions:
        raise NotImplementedError(
            "Product currently needs to be defined on separate dimensions."
        )  # pragma: no cover

    exps = [expectation(p, k, nghp=nghp) for k in kernel.kernels]
    return reduce(tf.multiply, exps)


@dispatch.expectation.register(
    DiagonalGaussian, kernels.Product, InducingPoints, NoneType, NoneType
)
@check_shapes(
    "p: [N, D]",
    "inducing_variable: [M, D, P]",
    "return: [N, M]",
)
def _expectation_diagonal_product_inducingpoints(
    p: DiagonalGaussian,
    kernel: kernels.Product,
    inducing_variable: InducingPoints,
    __: None,
    ___: None,
    nghp: None = None,
) -> tf.Tensor:
    r"""
    Compute the expectation:
    <\HadamardProd_i Ki_{X[:, active_dims_i], Z[:, active_dims_i]}>_p(X)
        - \HadamardProd_i Ki_{.,.} :: Product kernel
        - p                        :: DiagonalGaussian distribution (p.cov NxD)

    :return: NxM
    """
    if not kernel.on_separate_dimensions:
        raise NotImplementedError(
            "Product currently needs to be defined on separate dimensions."
        )  # pragma: no cover

    exps = [expectation(p, (k, inducing_variable), nghp=nghp) for k in kernel.kernels]
    return reduce(tf.multiply, exps)


@dispatch.expectation.register(
    DiagonalGaussian, kernels.Product, InducingPoints, kernels.Product, InducingPoints
)
@check_shapes(
    "p: [N, D]",
    "feat1: [M, D, P]",
    "feat2: [M, D, P]",
    "return: [N, M, M]",
)
def _expectation_diagonal_product_inducingpoints__product_inducingpoints(
    p: DiagonalGaussian,
    kern1: kernels.Product,
    feat1: InducingPoints,
    kern2: kernels.Product,
    feat2: InducingPoints,
    nghp: None = None,
) -> tf.Tensor:
    r"""
    Compute the expectation:
    expectation[n] = < prodK_{Z, x_n} prodK_{x_n, Z} >_p(x_n)
                   = < (\HadamardProd_i Ki_{Z[:, active_dims_i], x[n, active_dims_i]})  <-- Mx1
               1xM -->  (\HadamardProd_j Kj_{x[n, active_dims_j], Z[:, active_dims_j]}) >_p(x_n)  (MxM)

        - \HadamardProd_i Ki_{.,.}, \HadamardProd_j Kj_{.,.} :: Product kernels
        - p                        :: DiagonalGaussian distribution (p.cov NxD)

    :return: NxMxM
    """
    if feat1 != feat2:
        raise NotImplementedError("Different inducing variables are not supported.")
    if kern1 != kern2:
        raise NotImplementedError(
            "Calculating the expectation over two " "different Product kernels is not supported."
        )

    kernel = kern1
    inducing_variable = feat1

    if not kernel.on_separate_dimensions:
        raise NotImplementedError(
            "Product currently needs to be defined on separate dimensions."
        )  # pragma: no cover

    exps = [
        expectation(p, (k, inducing_variable), (k, inducing_variable), nghp=nghp)
        for k in kernel.kernels
    ]
    return reduce(tf.multiply, exps)


# gpflow/expectations/quadratures.py

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Any, Callable, Optional, Type, Union, cast

import numpy as np
import tensorflow as tf

from .. import kernels
from .. import mean_functions as mfn
from ..base import TensorType
from ..covariances import Kuf
from ..experimental.check_shapes import check_shapes
from ..inducing_variables import InducingVariables
from ..probability_distributions import DiagonalGaussian, Gaussian, MarkovGaussian
from ..quadrature import mvnquad
from . import dispatch
from .expectations import ExpectationObject, PackedExpectationObject, quadrature_expectation

register = dispatch.quadrature_expectation.register


NoneType: Type[None] = type(None)
EllipsisType = Any


def get_eval_func(
    obj: ExpectationObject,
    inducing_variable: Optional[InducingVariables],
    slice: Union[slice, EllipsisType, None] = None,
) -> Callable[[TensorType], tf.Tensor]:
    """
    Return the function of interest (kernel or mean) for the expectation
    depending on the type of :obj: and whether any inducing are given
    """

    slice = ... if slice is None else slice
    if inducing_variable is not None:
        # kernel + inducing_variable combination
        if not isinstance(inducing_variable, InducingVariables) or not isinstance(
            obj, kernels.Kernel
        ):
            raise TypeError("If `inducing_variable` is supplied, `obj` must be a kernel.")
        return lambda x: tf.transpose(Kuf(inducing_variable, obj, x))[slice]
    elif isinstance(obj, mfn.MeanFunction):
        return lambda x: obj(x)[slice]  # type: ignore[misc]
    elif isinstance(obj, kernels.Kernel):
        return lambda x: obj(x, full_cov=False)  # type: ignore[call-arg, misc]

    raise NotImplementedError()


@dispatch.quadrature_expectation.register(
    (Gaussian, DiagonalGaussian),
    object,
    (InducingVariables, NoneType),
    object,
    (InducingVariables, NoneType),
)
@check_shapes(
    "p: [N, D]",
    "inducing_variable1: [M1, D, P]",
    "inducing_variable2: [M2, D, P]",
    "return: [N, ...]",
)
def _quadrature_expectation_gaussian(
    p: Union[Gaussian, DiagonalGaussian],
    obj1: ExpectationObject,
    inducing_variable1: Optional[InducingVariables],
    obj2: ExpectationObject,
    inducing_variable2: Optional[InducingVariables],
    nghp: Optional[int] = None,
) -> tf.Tensor:
    """
    General handling of quadrature expectations for Gaussians and DiagonalGaussians
    Fallback method for missing analytic expectations
    """
    nghp = 100 if nghp is None else nghp

    # logger.warning(
    #     "Quadrature is used to calculate the expectation. This means that "
    #     "an analytical implementations is not available for the given combination."
    # )

    if obj1 is None:
        raise NotImplementedError("First object cannot be None.")

    if not isinstance(p, DiagonalGaussian):
        cov = p.cov
    else:
        if (
            isinstance(obj1, kernels.Kernel)
            and isinstance(obj2, kernels.Kernel)
            and obj1.on_separate_dims(obj2)
        ):  # no joint expectations required
            eKxz1 = quadrature_expectation(
                p, cast(PackedExpectationObject, (obj1, inducing_variable1)), nghp=nghp
            )
            eKxz2 = quadrature_expectation(
                p, cast(PackedExpectationObject, (obj2, inducing_variable2)), nghp=nghp
            )
            return eKxz1[:, :, None] * eKxz2[:, None, :]
        cov = tf.linalg.diag(p.cov)

    if obj2 is None:

        def eval_func(x: TensorType) -> tf.Tensor:
            fn = get_eval_func(obj1, inducing_variable1)
            return fn(x)

    else:

        def eval_func(x: TensorType) -> tf.Tensor:
            fn1 = get_eval_func(obj1, inducing_variable1, np.s_[:, :, None])
            fn2 = get_eval_func(obj2, inducing_variable2, np.s_[:, None, :])
            return fn1(x) * fn2(x)

    return mvnquad(eval_func, p.mu, cov, nghp)


@dispatch.quadrature_expectation.register(
    MarkovGaussian, object, (InducingVariables, NoneType), object, (InducingVariables, NoneType)
)
@check_shapes(
    "p: [N, D]",
    "inducing_variable1: [M1, D, P]",
    "inducing_variable2: [M2, D, P]",
    "return: [N, ...]",
)
def _quadrature_expectation_markov(
    p: MarkovGaussian,
    obj1: ExpectationObject,
    inducing_variable1: Optional[InducingVariables],
    obj2: ExpectationObject,
    inducing_variable2: Optional[InducingVariables],
    nghp: Optional[int] = None,
) -> tf.Tensor:
    """
    Handling of quadrature expectations for Markov Gaussians (useful for time series)
    Fallback method for missing analytic expectations wrt Markov Gaussians
    Nota Bene: obj1 is always associated with x_n, whereas obj2 always with x_{n+1}
               if one requires e.g. <x_{n+1} K_{x_n, Z}>_p(x_{n:n+1}), compute the
               transpose and then transpose the result of the expectation
    """
    nghp = 40 if nghp is None else nghp

    # logger.warning(
    #     "Quadrature is used to calculate the expectation. This means that "
    #     "an analytical implementations is not available for the given combination."
    # )

    if obj2 is None:

        def eval_func(x: TensorType) -> tf.Tensor:
            return get_eval_func(obj1, inducing_variable1)(x)

        mu, cov = p.mu[:-1], p.cov[0, :-1]  # cross covariances are not needed
    elif obj1 is None:

        def eval_func(x: TensorType) -> tf.Tensor:
            return get_eval_func(obj2, inducing_variable2)(x)

        mu, cov = p.mu[1:], p.cov[0, 1:]  # cross covariances are not needed
    else:

        def eval_func(x: TensorType) -> tf.Tensor:
            x1 = tf.split(x, 2, 1)[0]
            x2 = tf.split(x, 2, 1)[1]
            res1 = get_eval_func(obj1, inducing_variable1, np.s_[:, :, None])(x1)
            res2 = get_eval_func(obj2, inducing_variable2, np.s_[:, None, :])(x2)
            return res1 * res2

        mu = tf.concat((p.mu[:-1, :], p.mu[1:, :]), 1)  # Nx2D
        cov_top = tf.concat((p.cov[0, :-1, :, :], p.cov[1, :-1, :, :]), 2)  # NxDx2D
        cov_bottom = tf.concat((tf.linalg.adjoint(p.cov[1, :-1, :, :]), p.cov[0, 1:, :, :]), 2)
        cov = tf.concat((cov_top, cov_bottom), 1)  # Nx2Dx2D

    return mvnquad(eval_func, mu, cov, nghp)


# gpflow/expectations/squared_exponentials.py

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Type, Union

import tensorflow as tf

from .. import kernels
from .. import mean_functions as mfn
from ..experimental.check_shapes import check_shapes
from ..inducing_variables import InducingPoints
from ..probability_distributions import DiagonalGaussian, Gaussian, MarkovGaussian
from ..utilities.ops import square_distance
from . import dispatch
from .expectations import expectation

NoneType: Type[None] = type(None)


@dispatch.expectation.register(Gaussian, kernels.SquaredExponential, NoneType, NoneType, NoneType)
@check_shapes(
    "p: [N, D]",
    "return: [N]",
)
def _expectation_gaussian_sqe(
    p: Gaussian, kernel: kernels.SquaredExponential, _: None, __: None, ___: None, nghp: None = None
) -> tf.Tensor:
    """
    Compute the expectation:
    <diag(K_{X, X})>_p(X)
        - K_{.,.} :: RBF kernel

    :return: N
    """
    return kernel(p.mu, full_cov=False)


@dispatch.expectation.register(
    Gaussian, kernels.SquaredExponential, InducingPoints, NoneType, NoneType
)
@check_shapes(
    "p: [N, D]",
    "inducing_variable: [M, D, P]",
    "return: [N, M]",
)
def _expectation_gaussian_sqe_inducingpoints(
    p: Gaussian,
    kernel: kernels.SquaredExponential,
    inducing_variable: InducingPoints,
    _: None,
    __: None,
    nghp: None = None,
) -> tf.Tensor:
    """
    Compute the expectation:
    <K_{X, Z}>_p(X)
        - K_{.,.} :: RBF kernel

    :return: NxM
    """
    # use only active dimensions
    Xcov = kernel.slice_cov(p.cov)
    Z, Xmu = kernel.slice(inducing_variable.Z, p.mu)
    D = tf.shape(Xmu)[1]

    lengthscales = kernel.lengthscales
    if not kernel.ard:
        lengthscales = tf.zeros((D,), dtype=lengthscales.dtype) + kernel.lengthscales

    chol_L_plus_Xcov = tf.linalg.cholesky(tf.linalg.diag(lengthscales ** 2) + Xcov)  # NxDxD

    all_diffs = tf.transpose(Z) - tf.expand_dims(Xmu, 2)  # NxDxM
    exponent_mahalanobis = tf.linalg.triangular_solve(
        chol_L_plus_Xcov, all_diffs, lower=True
    )  # NxDxM
    exponent_mahalanobis = tf.reduce_sum(tf.square(exponent_mahalanobis), 1)  # NxM
    exponent_mahalanobis = tf.exp(-0.5 * exponent_mahalanobis)  # NxM

    sqrt_det_L = tf.reduce_prod(lengthscales)
    sqrt_det_L_plus_Xcov = tf.exp(
        tf.reduce_sum(tf.math.log(tf.linalg.diag_part(chol_L_plus_Xcov)), axis=1)
    )
    determinants = sqrt_det_L / sqrt_det_L_plus_Xcov  # N

    return kernel.variance * (determinants[:, None] * exponent_mahalanobis)


@dispatch.expectation.register(
    Gaussian, mfn.Identity, NoneType, kernels.SquaredExponential, InducingPoints
)
@check_shapes(
    "p: [N, D]",
    "inducing_variable: [M, D, P]",
    "return: [N, D, M]",
)
def _expectation_gaussian__sqe_inducingpoints(
    p: Gaussian,
    mean: mfn.Identity,
    _: None,
    kernel: kernels.SquaredExponential,
    inducing_variable: InducingPoints,
    nghp: None = None,
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <x_n K_{x_n, Z}>_p(x_n)
        - K_{.,.} :: RBF kernel

    :return: NxDxM
    """
    Xmu, Xcov = p.mu, p.cov

    D = tf.shape(Xmu)[1]

    lengthscales = kernel.lengthscales
    if not kernel.ard:
        lengthscales = tf.zeros((D,), dtype=lengthscales.dtype) + lengthscales

    chol_L_plus_Xcov = tf.linalg.cholesky(tf.linalg.diag(lengthscales ** 2) + Xcov)  # NxDxD
    all_diffs = tf.transpose(inducing_variable.Z) - tf.expand_dims(Xmu, 2)  # NxDxM

    sqrt_det_L = tf.reduce_prod(lengthscales)
    sqrt_det_L_plus_Xcov = tf.exp(
        tf.reduce_sum(tf.math.log(tf.linalg.diag_part(chol_L_plus_Xcov)), axis=1)
    )
    determinants = sqrt_det_L / sqrt_det_L_plus_Xcov  # N

    exponent_mahalanobis = tf.linalg.cholesky_solve(chol_L_plus_Xcov, all_diffs)  # NxDxM
    non_exponent_term = tf.linalg.matmul(Xcov, exponent_mahalanobis, transpose_a=True)
    non_exponent_term = tf.expand_dims(Xmu, 2) + non_exponent_term  # NxDxM

    exponent_mahalanobis = tf.reduce_sum(all_diffs * exponent_mahalanobis, 1)  # NxM
    exponent_mahalanobis = tf.exp(-0.5 * exponent_mahalanobis)  # NxM

    return (
        kernel.variance
        * (determinants[:, None] * exponent_mahalanobis)[:, None, :]
        * non_exponent_term
    )


@dispatch.expectation.register(
    MarkovGaussian, mfn.Identity, NoneType, kernels.SquaredExponential, InducingPoints
)
@check_shapes(
    "p: [N, D]",
    "inducing_variable: [M, D, P]",
    "return: [N, D, M]",
)
def _expectation_markov__sqe_inducingpoints(
    p: MarkovGaussian,
    mean: mfn.Identity,
    _: None,
    kernel: kernels.SquaredExponential,
    inducing_variable: InducingPoints,
    nghp: None = None,
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <x_{n+1} K_{x_n, Z}>_p(x_{n:n+1})
        - K_{.,.} :: RBF kernel
        - p       :: MarkovGaussian distribution (p.cov 2x(N+1)xDxD)

    :return: NxDxM
    """
    Xmu, Xcov = p.mu, p.cov

    D = tf.shape(Xmu)[1]
    lengthscales = kernel.lengthscales
    if not kernel.ard:
        lengthscales = tf.zeros((D,), dtype=lengthscales.dtype) + lengthscales

    chol_L_plus_Xcov = tf.linalg.cholesky(tf.linalg.diag(lengthscales ** 2) + Xcov[0, :-1])  # NxDxD
    all_diffs = tf.transpose(inducing_variable.Z) - tf.expand_dims(Xmu[:-1], 2)  # NxDxM

    sqrt_det_L = tf.reduce_prod(lengthscales)
    sqrt_det_L_plus_Xcov = tf.exp(
        tf.reduce_sum(tf.math.log(tf.linalg.diag_part(chol_L_plus_Xcov)), axis=1)
    )
    determinants = sqrt_det_L / sqrt_det_L_plus_Xcov  # N

    exponent_mahalanobis = tf.linalg.cholesky_solve(chol_L_plus_Xcov, all_diffs)  # NxDxM
    non_exponent_term = tf.linalg.matmul(Xcov[1, :-1], exponent_mahalanobis, transpose_a=True)
    non_exponent_term = tf.expand_dims(Xmu[1:], 2) + non_exponent_term  # NxDxM

    exponent_mahalanobis = tf.reduce_sum(all_diffs * exponent_mahalanobis, 1)  # NxM
    exponent_mahalanobis = tf.exp(-0.5 * exponent_mahalanobis)  # NxM

    return (
        kernel.variance
        * (determinants[:, None] * exponent_mahalanobis)[:, None, :]
        * non_exponent_term
    )


@dispatch.expectation.register(
    (Gaussian, DiagonalGaussian),
    kernels.SquaredExponential,
    InducingPoints,
    kernels.SquaredExponential,
    InducingPoints,
)
@check_shapes(
    "p: [N, D]",
    "feat1: [M, D, P]",
    "feat2: [M, D, P]",
    "return: [N, M, M]",
)
def _expectation_gaussian_sqe_inducingpoints__sqe_inducingpoints(
    p: Union[Gaussian, DiagonalGaussian],
    kern1: kernels.SquaredExponential,
    feat1: InducingPoints,
    kern2: kernels.SquaredExponential,
    feat2: InducingPoints,
    nghp: None = None,
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <Ka_{Z1, x_n} Kb_{x_n, Z2}>_p(x_n)
        - Ka_{.,.}, Kb_{.,.} :: RBF kernels
    Ka and Kb as well as Z1 and Z2 can differ from each other, but this is supported
    only if the Gaussian p is Diagonal (p.cov NxD) and Ka, Kb have disjoint active_dims
    in which case the joint expectations simplify into a product of expectations

    :return: NxMxM
    """
    if kern1.on_separate_dims(kern2) and isinstance(
        p, DiagonalGaussian
    ):  # no joint expectations required
        eKxz1 = expectation(p, (kern1, feat1))
        eKxz2 = expectation(p, (kern2, feat2))
        return eKxz1[:, :, None] * eKxz2[:, None, :]

    if feat1 != feat2 or kern1 != kern2:
        raise NotImplementedError(
            "The expectation over two kernels has only an "
            "analytical implementation if both kernels are equal."
        )

    kernel = kern1
    inducing_variable = feat1

    # use only active dimensions
    Xcov = kernel.slice_cov(tf.linalg.diag(p.cov) if isinstance(p, DiagonalGaussian) else p.cov)
    Z, Xmu = kernel.slice(inducing_variable.Z, p.mu)

    N = tf.shape(Xmu)[0]
    D = tf.shape(Xmu)[1]

    squared_lengthscales = kernel.lengthscales ** 2
    if not kernel.ard:
        zero_lengthscales = tf.zeros((D,), dtype=squared_lengthscales.dtype)
        squared_lengthscales = squared_lengthscales + zero_lengthscales

    sqrt_det_L = tf.reduce_prod(0.5 * squared_lengthscales) ** 0.5
    C = tf.linalg.cholesky(0.5 * tf.linalg.diag(squared_lengthscales) + Xcov)  # NxDxD
    dets = sqrt_det_L / tf.exp(tf.reduce_sum(tf.math.log(tf.linalg.diag_part(C)), axis=1))  # N

    C_inv_mu = tf.linalg.triangular_solve(C, tf.expand_dims(Xmu, 2), lower=True)  # NxDx1
    C_inv_z = tf.linalg.triangular_solve(
        C, tf.tile(tf.expand_dims(0.5 * tf.transpose(Z), 0), [N, 1, 1]), lower=True
    )  # NxDxM
    mu_CC_inv_mu = tf.expand_dims(tf.reduce_sum(tf.square(C_inv_mu), 1), 2)  # Nx1x1
    z_CC_inv_z = tf.reduce_sum(tf.square(C_inv_z), 1)  # NxM
    zm_CC_inv_zn = tf.linalg.matmul(C_inv_z, C_inv_z, transpose_a=True)  # NxMxM
    two_z_CC_inv_mu = 2 * tf.linalg.matmul(C_inv_z, C_inv_mu, transpose_a=True)[:, :, 0]  # NxM
    # NxMxM
    exponent_mahalanobis = (
        mu_CC_inv_mu
        + tf.expand_dims(z_CC_inv_z, 1)
        + tf.expand_dims(z_CC_inv_z, 2)
        + 2 * zm_CC_inv_zn
        - tf.expand_dims(two_z_CC_inv_mu, 2)
        - tf.expand_dims(two_z_CC_inv_mu, 1)
    )
    exponent_mahalanobis = tf.exp(-0.5 * exponent_mahalanobis)  # NxMxM

    # Compute sqrt(self(Z)) explicitly to prevent automatic gradient from
    # being NaN sometimes, see pull request #615
    kernel_sqrt = tf.exp(-0.25 * square_distance(Z / kernel.lengthscales, None))
    return kernel.variance ** 2 * kernel_sqrt * tf.reshape(dets, [N, 1, 1]) * exponent_mahalanobis


# gpflow/expectations/sums.py

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import itertools
from functools import reduce
from typing import Type, Union

import tensorflow as tf

from .. import kernels
from .. import mean_functions as mfn
from ..experimental.check_shapes import check_shapes
from ..inducing_variables import InducingPoints
from ..probability_distributions import DiagonalGaussian, Gaussian, MarkovGaussian
from . import dispatch
from .expectations import expectation

NoneType: Type[None] = type(None)


@dispatch.expectation.register(Gaussian, kernels.Sum, NoneType, NoneType, NoneType)
@check_shapes(
    "p: [N, D]",
    "return: [N]",
)
def _expectation_gaussian_sum(
    p: Gaussian, kernel: kernels.Sum, _: None, __: None, ___: None, nghp: None = None
) -> tf.Tensor:
    r"""
    Compute the expectation:
    <\Sum_i diag(Ki_{X, X})>_p(X)
        - \Sum_i Ki_{.,.} :: Sum kernel

    :return: N
    """
    exps = [expectation(p, k, nghp=nghp) for k in kernel.kernels]
    return reduce(tf.add, exps)


@dispatch.expectation.register(Gaussian, kernels.Sum, InducingPoints, NoneType, NoneType)
@check_shapes(
    "p: [N, D]",
    "inducing_variable: [M, D, P]",
    "return: [N, M]",
)
def _expectation_gaussian_sum_inducingpoints(
    p: Gaussian,
    kernel: kernels.Sum,
    inducing_variable: InducingPoints,
    _: None,
    __: None,
    nghp: None = None,
) -> tf.Tensor:
    r"""
    Compute the expectation:
    <\Sum_i Ki_{X, Z}>_p(X)
        - \Sum_i Ki_{.,.} :: Sum kernel

    :return: NxM
    """
    exps = [expectation(p, (k, inducing_variable), nghp=nghp) for k in kernel.kernels]
    return reduce(tf.add, exps)


@dispatch.expectation.register(
    Gaussian, (mfn.Linear, mfn.Identity, mfn.Constant), NoneType, kernels.Sum, InducingPoints
)
@check_shapes(
    "p: [N, D]",
    "inducing_variable: [M, D, P]",
    "return: [N, Q, M]",
)
def _expectation_gaussian_linear__sum_inducingpoints(
    p: Gaussian,
    mean: Union[mfn.Linear, mfn.Identity, mfn.Constant],
    _: None,
    kernel: kernels.Sum,
    inducing_variable: InducingPoints,
    nghp: None = None,
) -> tf.Tensor:
    r"""
    Compute the expectation:
    expectation[n] = <m(x_n)^T (\Sum_i Ki_{x_n, Z})>_p(x_n)
        - \Sum_i Ki_{.,.} :: Sum kernel

    :return: NxQxM
    """
    exps = [expectation(p, mean, (k, inducing_variable), nghp=nghp) for k in kernel.kernels]
    return reduce(tf.add, exps)


@dispatch.expectation.register(MarkovGaussian, mfn.Identity, NoneType, kernels.Sum, InducingPoints)
@check_shapes(
    "p: [N, D]",
    "inducing_variable: [M, D, P]",
    "return: [N, D, M]",
)
def _expectation_markov__sum_inducingpoints(
    p: MarkovGaussian,
    mean: mfn.Identity,
    _: None,
    kernel: kernels.Sum,
    inducing_variable: InducingPoints,
    nghp: None = None,
) -> tf.Tensor:
    r"""
    Compute the expectation:
    expectation[n] = <x_{n+1} (\Sum_i Ki_{x_n, Z})>_p(x_{n:n+1})
        - \Sum_i Ki_{.,.} :: Sum kernel

    :return: NxDxM
    """
    exps = [expectation(p, mean, (k, inducing_variable), nghp=nghp) for k in kernel.kernels]
    return reduce(tf.add, exps)


@dispatch.expectation.register(
    (Gaussian, DiagonalGaussian), kernels.Sum, InducingPoints, kernels.Sum, InducingPoints
)
@check_shapes(
    "p: [N, D]",
    "feat1: [M1, D, P]",
    "feat2: [M2, D, P]",
    "return: [N, M1, M2]",
)
def _expectation_gaussian_sum_inducingpoints__sum_inducingpoints(
    p: Union[Gaussian, DiagonalGaussian],
    kern1: kernels.Sum,
    feat1: InducingPoints,
    kern2: kernels.Sum,
    feat2: InducingPoints,
    nghp: None = None,
) -> tf.Tensor:
    r"""
    Compute the expectation:
    expectation[n] = <(\Sum_i K1_i_{Z1, x_n}) (\Sum_j K2_j_{x_n, Z2})>_p(x_n)
        - \Sum_i K1_i_{.,.}, \Sum_j K2_j_{.,.} :: Sum kernels

    :return: NxM1xM2
    """
    crossexps = []

    if kern1 == kern2 and feat1 == feat2:  # avoid duplicate computation by using transposes
        for i, k1 in enumerate(kern1.kernels):
            crossexps.append(expectation(p, (k1, feat1), (k1, feat1), nghp=nghp))

            for k2 in kern1.kernels[:i]:
                eKK = expectation(p, (k1, feat1), (k2, feat2), nghp=nghp)
                eKK += tf.linalg.adjoint(eKK)
                crossexps.append(eKK)
    else:
        for k1, k2 in itertools.product(kern1.kernels, kern2.kernels):
            crossexps.append(expectation(p, (k1, feat1), (k2, feat2), nghp=nghp))

    return reduce(tf.add, crossexps)


# gpflow/experimental/__init__.py

# Copyright 2016-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
This package contains code that is still experimental. This means:

* Expect breaking changes.
* Expect poor documentation.
* Expect bugs.

When the code is considered mature it will be moved out of this package.
"""
from . import check_shapes, utils

__all__ = [
    "check_shapes",
    "utils",
]


# gpflow/experimental/utils.py

# Copyright 2016-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from functools import wraps
from typing import Any, Callable, TypeVar, cast
from warnings import warn

C = TypeVar("C", bound=Callable[..., Any])


def experimental(func: C) -> C:
    """
    Decorator that marks the decorated function as experimental.

    The first time an experimental function is called, a warning is printed.

    Example::

        @experimental
        def forty_two() -> int:
            return 42
    """

    has_warned = False

    @wraps(func)
    def wrap_experimental(*args: Any, **kwargs: Any) -> Any:
        nonlocal has_warned

        if not has_warned:
            name = f"{func.__module__}.{func.__qualname__}"
            warn(
                f"You're calling {name} which is considered *experimental*."
                " Expect: breaking changes, poor documentation, and bugs."
            )
            has_warned = True

        return func(*args, **kwargs)

    return cast(C, wrap_experimental)


# gpflow/inducing_variables/__init__.py

from . import multioutput
from .inducing_patch import InducingPatches
from .inducing_variables import InducingPoints, InducingVariables, Multiscale
from .multioutput import (
    FallbackSeparateIndependentInducingVariables,
    FallbackSharedIndependentInducingVariables,
    MultioutputInducingVariables,
    SeparateIndependentInducingVariables,
    SharedIndependentInducingVariables,
)

__all__ = [
    "FallbackSeparateIndependentInducingVariables",
    "FallbackSharedIndependentInducingVariables",
    "InducingPatches",
    "InducingPoints",
    "InducingVariables",
    "MultioutputInducingVariables",
    "Multiscale",
    "SeparateIndependentInducingVariables",
    "SharedIndependentInducingVariables",
    "inducing_patch",
    "inducing_variables",
    "multioutput",
]


# gpflow/inducing_variables/inducing_patch.py

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from .inducing_variables import InducingPoints


class InducingPatches(InducingPoints):
    pass


# gpflow/inducing_variables/inducing_variables.py

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import abc
from typing import Optional

import tensorflow as tf
import tensorflow_probability as tfp
from deprecated import deprecated

from ..base import Module, Parameter, TensorData, TensorType
from ..experimental.check_shapes import ErrorContext, Shape, check_shapes, register_get_shape
from ..utilities import positive


class InducingVariables(Module, abc.ABC):
    """
    Abstract base class for inducing variables.
    """

    @property
    @abc.abstractmethod
    def num_inducing(self) -> tf.Tensor:
        """
        Returns the number of inducing variables, relevant for example to determine the size of the
        variational distribution.
        """
        raise NotImplementedError

    @deprecated(
        reason="len(iv) should return an `int`, but this actually returns a `tf.Tensor`."
        " Use `iv.num_inducing` instead."
    )
    def __len__(self) -> tf.Tensor:
        return self.num_inducing

    @property
    @abc.abstractmethod
    def shape(self) -> Shape:
        """
        Return the shape of these inducing variables.

        Shape should be some variation of ``[M, D, P]``, where:

        * ``M`` is the number of inducing variables.
        * ``D`` is the number of input dimensions.
        * ``P`` is the number of output dimensions (1 if this is not a multi-output inducing
          variable).
        """


@register_get_shape(InducingVariables)
def get_scalar_shape(shaped: InducingVariables, context: ErrorContext) -> Shape:
    return shaped.shape


class InducingPointsBase(InducingVariables):
    @check_shapes(
        "Z: [M, D]",
    )
    def __init__(self, Z: TensorData, name: Optional[str] = None):
        """
        :param Z: The initial positions of the inducing points.
        """
        super().__init__(name=name)
        if not isinstance(Z, (tf.Variable, tfp.util.TransformedVariable)):
            Z = Parameter(Z)
        self.Z = Z

    @property  # type: ignore[misc]  # mypy doesn't like decorated properties.
    @check_shapes(
        "return: []",
    )
    def num_inducing(self) -> Optional[tf.Tensor]:
        return tf.shape(self.Z)[0]

    @property
    def shape(self) -> Shape:
        shape = self.Z.shape
        if not shape:
            return None
        return tuple(shape) + (1,)


class InducingPoints(InducingPointsBase):
    """
    Real-space inducing points
    """


class Multiscale(InducingPointsBase):
    r"""
    Multi-scale inducing variables

    Originally proposed in :cite:t:`NIPS2009_3876`.
    """

    @check_shapes(
        "Z: [M, D]",
        "scales: [M, D]",
    )
    def __init__(self, Z: TensorData, scales: TensorData):
        super().__init__(Z)
        # Multi-scale inducing_variable widths (std. dev. of Gaussian)
        self.scales = Parameter(scales, transform=positive())

    @staticmethod
    @check_shapes(
        "A: [N, D]",
        "B: [M, D]",
        "sc: [broadcast N, broadcast M, D]",
        "return: [N, M]",
    )
    def _cust_square_dist(A: TensorType, B: TensorType, sc: TensorType) -> tf.Tensor:
        """
        Custom version of _square_dist that allows sc to provide per-datapoint length
        scales.
        """
        return tf.reduce_sum(tf.square((tf.expand_dims(A, 1) - tf.expand_dims(B, 0)) / sc), 2)


# gpflow/kernels/__init__.py

from . import multioutput
from .base import Combination, Kernel, Product, Sum
from .changepoints import ChangePoints
from .convolutional import Convolutional
from .linears import Linear, Polynomial
from .misc import ArcCosine, Coregion
from .multioutput import (
    IndependentLatent,
    LinearCoregionalization,
    MultioutputKernel,
    SeparateIndependent,
    SharedIndependent,
)
from .periodic import Periodic
from .statics import Constant, Static, White

from .stationaries import (  # isort:skip
    # base classes:
    Stationary,
    IsotropicStationary,
    AnisotropicStationary,
    # actual kernel classes:
    Cosine,
    Exponential,
    Matern12,
    Matern32,
    Matern52,
    RationalQuadratic,
    SquaredExponential,
)

Bias = Constant
RBF = SquaredExponential

__all__ = [
    "AnisotropicStationary",
    "ArcCosine",
    "Bias",
    "ChangePoints",
    "Combination",
    "Constant",
    "Convolutional",
    "Coregion",
    "Cosine",
    "Exponential",
    "IndependentLatent",
    "IsotropicStationary",
    "Kernel",
    "Linear",
    "LinearCoregionalization",
    "Matern12",
    "Matern32",
    "Matern52",
    "MultioutputKernel",
    "Periodic",
    "Polynomial",
    "Product",
    "RBF",
    "RationalQuadratic",
    "SeparateIndependent",
    "SharedIndependent",
    "SquaredExponential",
    "Static",
    "Stationary",
    "Sum",
    "White",
    "base",
    "changepoints",
    "convolutional",
    "linears",
    "misc",
    "multioutput",
    "periodic",
    "statics",
    "stationaries",
]


# gpflow/kernels/base.py

# Copyright 2018-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

r"""
Kernels form a core component of GPflow models and allow prior information to
be encoded about a latent function of interest. The effect of choosing
different kernels, and how it is possible to combine multiple kernels is shown
in the `"Using kernels in GPflow" notebook <notebooks/kernels.html>`_.

Broadcasting over leading dimensions:
`kernel.K(X1, X2)` returns the kernel evaluated on every pair in X1 and X2.
E.g. if X1 has shape [S1, N1, D] and X2 has shape [S2, N2, D], kernel.K(X1, X2)
will return a tensor of shape [S1, N1, S2, N2]. Similarly, kernel.K(X1, X1)
returns a tensor of shape [S1, N1, S1, N1]. In contrast, the return shape of
kernel.K(X1) is [S1, N1, N1]. (Without leading dimensions, the behaviour of
kernel.K(X, None) is identical to kernel.K(X, X).)
"""

import abc
from functools import partial, reduce
from typing import Callable, List, Optional, Sequence, Tuple, Union, overload

import numpy as np
import tensorflow as tf

from ..base import AnyNDArray, Module, TensorType
from ..experimental.check_shapes import check_shapes

ActiveDims = Union[slice, Sequence[int]]
NormalizedActiveDims = Union[slice, AnyNDArray]


class Kernel(Module, metaclass=abc.ABCMeta):
    """
    The basic kernel class. Handles active dims.
    """

    def __init__(
        self, active_dims: Optional[ActiveDims] = None, name: Optional[str] = None
    ) -> None:
        """
        :param active_dims: active dimensions, either a slice or list of
            indices into the columns of X.
        :param name: optional kernel name.
        """
        super().__init__(name=name)
        self._active_dims = self._normalize_active_dims(active_dims)

    @staticmethod
    def _normalize_active_dims(value: Optional[ActiveDims]) -> NormalizedActiveDims:
        if value is None:
            return slice(None, None, None)
        if isinstance(value, slice):
            return value
        return np.array(value, dtype=int)

    @property
    def active_dims(self) -> NormalizedActiveDims:
        return self._active_dims

    @active_dims.setter
    def active_dims(self, value: ActiveDims) -> None:
        self._active_dims = self._normalize_active_dims(value)

    def on_separate_dims(self, other: "Kernel") -> bool:
        """
        Checks if the dimensions, over which the kernels are specified, overlap.
        Returns True if they are defined on different/separate dimensions and False otherwise.
        """
        if isinstance(self.active_dims, slice) or isinstance(other.active_dims, slice):
            # Be very conservative for kernels defined over slices of dimensions
            return False

        if self.active_dims is None or other.active_dims is None:
            return False

        this_dims = self.active_dims.reshape(-1, 1)
        other_dims = other.active_dims.reshape(1, -1)
        return not np.any(this_dims == other_dims)

    @overload
    def slice(self, X: TensorType, X2: TensorType) -> Tuple[tf.Tensor, tf.Tensor]:
        ...

    @overload
    def slice(self, X: TensorType, X2: None) -> Tuple[tf.Tensor, None]:
        ...

    @check_shapes(
        "X: [batch..., N, D]",
        "X2: [batch2..., N2, D]",
        "return[0]: [batch..., N, I]",
        "return[1]: [batch2..., N2, I]",
    )
    def slice(
        self, X: TensorType, X2: Optional[TensorType] = None
    ) -> Tuple[tf.Tensor, Optional[tf.Tensor]]:
        """
        Slice the correct dimensions for use in the kernel, as indicated by `self.active_dims`.

        :param X: Input 1.
        :param X2: Input 2, can be None.
        :return: Sliced X, X2.
        """
        dims = self.active_dims
        if isinstance(dims, slice):
            X = X[..., dims]
            if X2 is not None:
                X2 = X2[..., dims]
        elif dims is not None:
            X = tf.gather(X, dims, axis=-1)
            if X2 is not None:
                X2 = tf.gather(X2, dims, axis=-1)
        return X, X2

    @check_shapes(
        "cov: [N, *D_or_DD]",
        "return: [N, I, I]",
    )
    def slice_cov(self, cov: TensorType) -> tf.Tensor:
        """
        Slice the correct dimensions for use in the kernel, as indicated by
        `self.active_dims` for covariance matrices. This requires slicing the
        rows *and* columns. This will also turn flattened diagonal
        matrices into a tensor of full diagonal matrices.

        :param cov: Tensor of covariance matrices.
        :return: Sliced covariance matrices.
        """
        cov = tf.convert_to_tensor(cov)
        assert isinstance(cov, tf.Tensor)  # Hint for mypy.

        if cov.shape.ndims == 2:
            cov = tf.linalg.diag(cov)

        dims = self.active_dims

        if isinstance(dims, slice):
            return cov[..., dims, dims]
        elif dims is not None:
            nlast = tf.shape(cov)[-1]
            ndims = len(dims)

            cov_shape = tf.shape(cov)
            cov_reshaped = tf.reshape(cov, [-1, nlast, nlast])
            gather1 = tf.gather(tf.transpose(cov_reshaped, [2, 1, 0]), dims)
            gather2 = tf.gather(tf.transpose(gather1, [1, 0, 2]), dims)
            cov = tf.reshape(
                tf.transpose(gather2, [2, 0, 1]), tf.concat([cov_shape[:-2], [ndims, ndims]], 0)
            )

        return cov

    @check_shapes(
        "ard_parameter: [any...]",
    )
    def _validate_ard_active_dims(self, ard_parameter: TensorType) -> None:
        """
        Validate that ARD parameter matches the number of active_dims (provided active_dims
        has been specified as an array).
        """
        ard_parameter = tf.convert_to_tensor(ard_parameter)
        assert isinstance(ard_parameter, tf.Tensor)  # Hint for mypy.

        if self.active_dims is None or isinstance(self.active_dims, slice):
            # Can only validate parameter if active_dims is an array
            return

        if ard_parameter.shape.rank > 0 and ard_parameter.shape[0] != len(self.active_dims):
            raise ValueError(
                f"Size of `active_dims` {self.active_dims} does not match "
                f"size of ard parameter ({ard_parameter.shape[0]})"
            )

    @abc.abstractmethod
    @check_shapes(
        "X: [batch..., N, D]",
        "X2: [batch2..., N2, D]",
        "return: [batch..., N, batch2..., N2] if X2 is not None",
        "return: [batch..., N, N] if X2 is None",
    )
    def K(self, X: TensorType, X2: Optional[TensorType] = None) -> tf.Tensor:
        raise NotImplementedError

    @abc.abstractmethod
    @check_shapes(
        "X: [batch..., N, D]",
        "return: [batch..., N]",
    )
    def K_diag(self, X: TensorType) -> tf.Tensor:
        raise NotImplementedError

    @check_shapes(
        "X: [batch..., N, D]",
        "X2: [batch2..., N2, D]",
        "return: [batch..., N, batch2..., N2] if full_cov and (X2 is not None)",
        "return: [batch..., N, N] if full_cov and (X2 is None)",
        "return: [batch..., N] if not full_cov",
    )
    def __call__(
        self,
        X: TensorType,
        X2: Optional[TensorType] = None,
        *,
        full_cov: bool = True,
        presliced: bool = False,
    ) -> tf.Tensor:
        if (not full_cov) and (X2 is not None):
            raise ValueError("Ambiguous inputs: `not full_cov` and `X2` are not compatible.")

        if not presliced:
            X, X2 = self.slice(X, X2)

        if not full_cov:
            assert X2 is None
            return self.K_diag(X)

        else:
            return self.K(X, X2)

    def __add__(self, other: "Kernel") -> "Kernel":
        return Sum([self, other])

    def __mul__(self, other: "Kernel") -> "Kernel":
        return Product([self, other])


class Combination(Kernel):
    """
    Combine a list of kernels, e.g. by adding or multiplying (see inheriting
    classes).

    The names of the kernels to be combined are generated from their class
    names.
    """

    _reduction = None

    def __init__(self, kernels: Sequence[Kernel], name: Optional[str] = None) -> None:
        super().__init__(name=name)

        if not all(isinstance(k, Kernel) for k in kernels):
            raise TypeError("can only combine Kernel instances")  # pragma: no cover

        self.kernels: List[Kernel] = []
        self._set_kernels(kernels)

    def _set_kernels(self, kernels: Sequence[Kernel]) -> None:
        # add kernels to a list, flattening out instances of this class therein
        kernels_list: List[Kernel] = []
        for k in kernels:
            if isinstance(k, self.__class__):
                kernels_list.extend(k.kernels)
            else:
                kernels_list.append(k)
        self.kernels = kernels_list

    @property
    def on_separate_dimensions(self) -> bool:
        """
        Checks whether the kernels in the combination act on disjoint subsets
        of dimensions. Currently, it is hard to asses whether two slice objects
        will overlap, so this will always return False.

        :return: Boolean indicator.
        """
        if np.any([isinstance(k.active_dims, slice) for k in self.kernels]):
            # Be conservative in the case of a slice object
            return False
        else:
            dimlist = [k.active_dims for k in self.kernels]
            overlapping = False
            for i, dims_i in enumerate(dimlist):
                assert isinstance(dims_i, np.ndarray)
                for dims_j in dimlist[i + 1 :]:
                    assert isinstance(dims_j, np.ndarray)
                    if np.any(dims_i.reshape(-1, 1) == dims_j.reshape(1, -1)):
                        overlapping = True
            return not overlapping


class ReducingCombination(Combination):
    def __call__(
        self,
        X: TensorType,
        X2: Optional[TensorType] = None,
        *,
        full_cov: bool = True,
        presliced: bool = False,
    ) -> tf.Tensor:
        return self._reduce(
            [k(X, X2, full_cov=full_cov, presliced=presliced) for k in self.kernels]
        )

    def K(self, X: TensorType, X2: Optional[TensorType] = None) -> tf.Tensor:
        return self._reduce([k.K(X, X2) for k in self.kernels])

    def K_diag(self, X: TensorType) -> tf.Tensor:
        return self._reduce([k.K_diag(X) for k in self.kernels])

    @property
    @abc.abstractmethod
    def _reduce(self) -> Callable[[Sequence[TensorType]], TensorType]:
        pass


class Sum(ReducingCombination):
    @property
    def _reduce(self) -> Callable[[Sequence[TensorType]], TensorType]:
        return tf.add_n  # type: ignore[no-any-return]


class Product(ReducingCombination):
    @property
    def _reduce(self) -> Callable[[Sequence[TensorType]], TensorType]:
        return partial(reduce, tf.multiply)


# gpflow/kernels/changepoints.py

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Optional, Sequence

import tensorflow as tf

from ..base import Parameter, TensorType
from ..experimental.check_shapes import check_shape as cs
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..utilities import positive
from .base import Combination, Kernel


class ChangePoints(Combination):
    r"""
    The ChangePoints kernel defines a fixed number of change-points along a 1d
    input space where different kernels govern different parts of the space.

    The kernel is by multiplication and addition of the base kernels with
    sigmoid functions (σ). A single change-point kernel is defined as::

        K₁(x, x') * (1 - σ(x)) * (1 - σ(x')) + K₂(x, x') * σ(x) * σ(x')

    where K₁ is deactivated around the change-point and K₂ is activated. The
    single change-point version can be found in :cite:t:`lloyd2014`. Each sigmoid
    is a logistic function defined as::

        σ(x) = 1 / (1 + exp{-s(x - x₀)})

    parameterized by location "x₀" and steepness "s".

    The key reference is :cite:t:`lloyd2014`.
    """

    @check_shapes(
        "locations: [n_change_points]",
        "steepness: [broadcast n_change_points]",
    )
    def __init__(
        self,
        kernels: Sequence[Kernel],
        locations: TensorType,
        steepness: TensorType = 1.0,
        name: Optional[str] = None,
    ):
        """
        :param kernels: list of kernels defining the different regimes
        :param locations: list of change-point locations in the 1d input space
        :param steepness: the steepness parameter(s) of the sigmoids, this can be
            common between them or decoupled
        """
        if len(kernels) != len(locations) + 1:
            raise ValueError(
                "Number of kernels ({nk}) must be one more than the number of "
                "changepoint locations ({nl})".format(nk=len(kernels), nl=len(locations))
            )

        if isinstance(steepness, Sequence) and len(steepness) != len(locations):
            raise ValueError(
                "Dimension of steepness ({ns}) does not match number of changepoint "
                "locations ({nl})".format(ns=len(steepness), nl=len(locations))
            )

        super().__init__(kernels, name=name)

        self.locations = Parameter(locations)
        self.steepness = Parameter(steepness, transform=positive())

    def _set_kernels(self, kernels: Sequence[Kernel]) -> None:
        # it is not clear how to flatten out nested change-points
        self.kernels = list(kernels)

    @inherit_check_shapes
    def K(self, X: tf.Tensor, X2: Optional[tf.Tensor] = None) -> tf.Tensor:
        cs(X, "[batch..., N, 1]  # The `ChangePoints` kernel requires a 1D input space.")

        rank = tf.rank(X) - 2
        batch = tf.shape(X)[:-2]
        N = tf.shape(X)[-2]
        Ncp = tf.shape(self.locations)[0]
        sig_X = cs(self._sigmoids(X), "[batch..., N, 1, Ncp]")

        if X2 is None:
            rank2 = 0
            batch2 = tf.constant([], dtype=tf.int32)
            N2 = N
            sig_X2 = sig_X
            sig_X = cs(
                tf.reshape(sig_X, tf.concat([batch, [N, 1, Ncp]], 0)), "[batch..., N, 1, Ncp]"
            )
            sig_X2 = cs(
                tf.reshape(sig_X2, tf.concat([batch, [1, N, Ncp]], 0)), "[batch..., 1, N, Ncp]"
            )
        else:
            rank2 = tf.rank(X2) - 2
            batch2 = tf.shape(X2)[:-2]
            N2 = tf.shape(X2)[-2]

            sig_X2 = cs(self._sigmoids(X2), "[batch2..., N2, 1, Ncp]")

            ones = tf.ones((rank,), dtype=tf.int32)
            ones2 = tf.ones((rank2,), dtype=tf.int32)
            sig_X = cs(
                tf.reshape(sig_X, tf.concat([batch, [N], ones2, [1, Ncp]], 0)),
                "[batch..., N, ..., 1, Ncp]",
            )
            sig_X2 = cs(
                tf.reshape(sig_X2, tf.concat([ones, [1], batch2, [N2, Ncp]], 0)),
                "[..., 1, batch2..., N2, Ncp]",
            )

        # `starters` are the sigmoids going from 0 -> 1, whilst `stoppers` go
        # from 1 -> 0.
        starters = cs(sig_X * sig_X2, "[batch..., N, batch2..., N2, Ncp]")
        stoppers = cs((1 - sig_X) * (1 - sig_X2), "[batch..., N, batch2..., N2, Ncp]")

        # prepend `starters` with ones and append ones to `stoppers` since the
        # first kernel has no start and the last kernel has no end
        ones = tf.ones(tf.concat([batch, [N], batch2, [N2, 1]], 0), dtype=X.dtype)
        starters = cs(tf.concat([ones, starters], axis=-1), "[batch..., N, batch2..., N2, Nkern]")
        stoppers = cs(tf.concat([stoppers, ones], axis=-1), "[batch..., N, batch2..., N2, Nkern]")

        # now combine with the underlying kernels
        kernel_stack = cs(
            tf.stack([k(X, X2) for k in self.kernels], axis=-1),
            "[batch..., N, batch2..., N2, Nkern]",
        )
        return tf.reduce_sum(kernel_stack * starters * stoppers, axis=-1)

    @inherit_check_shapes
    def K_diag(self, X: tf.Tensor) -> tf.Tensor:
        cs(X, "[batch..., N, 1]  # The `ChangePoints` kernel requires a 1D input space.")

        batch = tf.shape(X)[:-2]
        N = tf.shape(X)[-2]
        Ncp = tf.shape(self.locations)[0]

        sig_X = cs(
            tf.reshape(self._sigmoids(X), tf.concat([batch, [N, Ncp]], 0)), "[batch..., N, Ncp]"
        )

        ones = tf.ones(tf.concat([batch, [N, 1]], 0), dtype=X.dtype)
        starters = cs(tf.concat([ones, sig_X * sig_X], axis=-1), "[batch..., N, Nkern]")
        stoppers = cs(tf.concat([(1 - sig_X) * (1 - sig_X), ones], axis=-1), "[batch..., N, Nkern]")

        kernel_stack = cs(
            tf.stack([k(X, full_cov=False) for k in self.kernels], axis=-1), "[batch..., N, Nkern]"
        )
        return tf.reduce_sum(kernel_stack * starters * stoppers, axis=-1)

    @check_shapes(
        "X: [batch...]",
        "return: [batch..., Ncp]",
    )
    def _sigmoids(self, X: tf.Tensor) -> tf.Tensor:
        locations = tf.sort(self.locations)  # ensure locations are ordered
        locations = tf.reshape(locations, (-1,))
        steepness = tf.reshape(self.steepness, (-1,))
        return tf.sigmoid(steepness * (X[..., None] - locations))


# gpflow/kernels/convolutional.py

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Optional, Sequence, cast

import numpy as np
import tensorflow as tf

from ..base import Parameter, TensorType
from ..config import default_float
from ..experimental.check_shapes import check_shape as cs
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..utilities import to_default_float
from .base import Kernel


class Convolutional(Kernel):
    r"""
    Plain convolutional kernel as described in :cite:t:`vdw2017convgp`. Defines
    a GP :math:`f()` that is constructed from a sum of responses of individual patches
    in an image:

    .. math::
       f(x) = \sum_p x^{[p]}

    where :math:`x^{[p]}` is the :math:`p`'th patch in the image.

    The key reference is :cite:t:`vdw2017convgp`.
    """

    @check_shapes(
        "image_shape: [2]",
        "patch_shape: [2]",
        "weights: [P]",
    )
    def __init__(
        self,
        base_kernel: Kernel,
        image_shape: Sequence[int],
        patch_shape: Sequence[int],
        weights: Optional[TensorType] = None,
        colour_channels: int = 1,
    ) -> None:
        super().__init__()
        self.image_shape = image_shape
        self.patch_shape = patch_shape
        self.base_kernel = base_kernel
        self.colour_channels = colour_channels
        self.weights = Parameter(
            np.ones(self.num_patches, dtype=default_float()) if weights is None else weights
        )

    @check_shapes(
        "X: [batch..., N, D]",
        "return: [batch..., N, P, S]",
    )
    def get_patches(self, X: TensorType) -> tf.Tensor:
        """
        Extracts patches from the images X. Patches are extracted separately for each of the colour
        channels.

        :param X: Images.
        :return: Patches.
        """
        # Roll the colour channel to the front, so it appears to
        # `tf.extract_image_patches()` as separate images. Then extract patches
        # and reshape to have the first axis the same as the number of images.
        # The separate patches will then be in the second axis.
        batch = tf.shape(X)[:-2]
        N = tf.shape(X)[-2]
        flat_batch = tf.reduce_prod(batch)
        num_data = flat_batch * N
        X = cs(
            tf.transpose(tf.reshape(X, [num_data, -1, self.colour_channels]), [0, 2, 1]),
            "[num_data, C, W_x_H]",
        )
        X = cs(
            tf.reshape(X, [-1, self.image_shape[0], self.image_shape[1], 1], name="rX"),
            "[num_data_x_C, W, H, 1]",
        )
        patches = cs(
            tf.image.extract_patches(
                X,
                [1, self.patch_shape[0], self.patch_shape[1], 1],
                [1, 1, 1, 1],
                [1, 1, 1, 1],
                "VALID",
            ),
            "[num_data_x_C, n_x_patches, n_y_patches, S]",
        )
        shp = tf.shape(patches)
        reshaped_patches = cs(
            tf.reshape(
                patches, tf.concat([batch, [N, self.colour_channels * shp[1] * shp[2], shp[3]]], 0)
            ),
            "[batch..., N, P, S]",
        )
        return to_default_float(reshaped_patches)

    @inherit_check_shapes
    def K(self, X: TensorType, X2: Optional[TensorType] = None) -> tf.Tensor:
        Xp = cs(self.get_patches(X), "[batch..., N, P, S]")
        W2 = cs(self.weights[:, None] * self.weights[None, :], "[P, P]")

        rank = tf.rank(Xp) - 3
        batch = tf.shape(Xp)[:-3]
        N = tf.shape(Xp)[-3]
        P = tf.shape(Xp)[-2]
        S = tf.shape(Xp)[-1]
        ones = tf.ones((rank,), dtype=tf.int32)

        if X2 is None:
            Xp = cs(tf.reshape(Xp, tf.concat([batch, [N * P, S]], 0)), "[batch..., N_x_P, S]")
            bigK = cs(self.base_kernel.K(Xp), "[batch..., N_x_P, N_x_P]")
            bigK = cs(
                tf.reshape(bigK, tf.concat([batch, [N, P, N, P]], 0)), "[batch..., N, P, N, P]"
            )
            W2 = cs(tf.reshape(W2, tf.concat([ones, [1, P, 1, P]], 0)), "[..., 1, P, 1, P]")
            W2bigK = cs(bigK * W2, "[batch..., N, P, N, P]")
            return cs(
                tf.reduce_sum(W2bigK, [rank + 1, rank + 3]) / self.num_patches ** 2.0,
                "[batch..., N, N]",
            )

        else:
            Xp2 = Xp if X2 is None else cs(self.get_patches(X2), "[batch2..., N2, P, S]")
            rank2 = tf.rank(Xp2) - 3
            ones2 = tf.ones((rank2,), dtype=tf.int32)
            bigK = cs(self.base_kernel.K(Xp, Xp2), "[batch..., N, P, batch2..., N2, P]")
            W2 = cs(
                tf.reshape(W2, tf.concat([ones, [1, P], ones2, [1, P]], 0)),
                "[..., 1, P, ..., 1, P]",
            )
            W2bigK = cs(bigK * W2, "[batch..., N, P, batch2..., N2, P]")
            return cs(
                tf.reduce_sum(W2bigK, [rank + 1, rank + rank2 + 3]) / self.num_patches ** 2.0,
                "[batch..., N, batch2..., N2]",
            )

    @inherit_check_shapes
    def K_diag(self, X: TensorType) -> tf.Tensor:
        Xp = cs(self.get_patches(X), "[batch..., N, P, S]")

        rank = tf.rank(Xp) - 3
        P = tf.shape(Xp)[-2]
        ones = tf.ones((rank,), dtype=tf.int32)

        W2 = cs(self.weights[:, None] * self.weights[None, :], "[P, P]")
        W2 = cs(tf.reshape(W2, tf.concat([ones, [1, P, P]], 0)), "[..., 1, P, P]")

        bigK = cs(self.base_kernel.K(Xp), "[batch..., N, P, P]")

        return tf.reduce_sum(bigK * W2, [rank + 1, rank + 2]) / self.num_patches ** 2.0

    @property
    def patch_len(self) -> int:
        return cast(int, np.prod(self.patch_shape))

    @property
    def num_patches(self) -> int:
        return (
            (self.image_shape[0] - self.patch_shape[0] + 1)
            * (self.image_shape[1] - self.patch_shape[1] + 1)
            * self.colour_channels
        )


# gpflow/kernels/linears.py

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Optional

import tensorflow as tf

from ..base import Parameter, TensorType
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..utilities import positive
from .base import ActiveDims, Kernel


class Linear(Kernel):
    """
    The linear kernel. Functions drawn from a GP with this kernel are linear, i.e. f(x) = cx.
    The kernel equation is

        k(x, y) = σ²xy

    where σ² is the variance parameter.
    """

    @check_shapes(
        "variance: [broadcast n_active_dims]",
    )
    def __init__(
        self, variance: TensorType = 1.0, active_dims: Optional[ActiveDims] = None
    ) -> None:
        """
        :param variance: the (initial) value for the variance parameter(s),
            to induce ARD behaviour this must be initialised as an array the same
            length as the the number of active dimensions e.g. [1., 1., 1.]
        :param active_dims: a slice or list specifying which columns of X are used
        """
        super().__init__(active_dims)
        self.variance = Parameter(variance, transform=positive())
        self._validate_ard_active_dims(self.variance)

    @property
    def ard(self) -> bool:
        """
        Whether ARD behaviour is active.
        """
        ndims: int = self.variance.shape.ndims
        return ndims > 0

    @inherit_check_shapes
    def K(self, X: TensorType, X2: Optional[TensorType] = None) -> tf.Tensor:
        if X2 is None:
            return tf.matmul(X * self.variance, X, transpose_b=True)
        else:
            return tf.tensordot(X * self.variance, X2, [[-1], [-1]])

    @inherit_check_shapes
    def K_diag(self, X: TensorType) -> tf.Tensor:
        return tf.reduce_sum(tf.square(X) * self.variance, axis=-1)


class Polynomial(Linear):
    """
    The Polynomial kernel. Functions drawn from a GP with this kernel are
    polynomials of degree `d`. The kernel equation is

        k(x, y) = (σ²xy + γ)ᵈ

    where:
    σ² is the variance parameter,
    γ is the offset parameter,
    d is the degree parameter.
    """

    @check_shapes(
        "variance: [broadcast n_active_dims]",
    )
    def __init__(
        self,
        degree: TensorType = 3.0,
        variance: TensorType = 1.0,
        offset: TensorType = 1.0,
        active_dims: Optional[ActiveDims] = None,
    ) -> None:
        """
        :param degree: the degree of the polynomial
        :param variance: the (initial) value for the variance parameter(s),
            to induce ARD behaviour this must be initialised as an array the same
            length as the the number of active dimensions e.g. [1., 1., 1.]
        :param offset: the offset of the polynomial
        :param active_dims: a slice or list specifying which columns of X are used
        """
        super().__init__(variance, active_dims)
        self.degree = degree
        self.offset = Parameter(offset, transform=positive())

    @inherit_check_shapes
    def K(self, X: TensorType, X2: Optional[TensorType] = None) -> tf.Tensor:
        return (super().K(X, X2) + self.offset) ** self.degree

    @inherit_check_shapes
    def K_diag(self, X: TensorType) -> tf.Tensor:
        return (super().K_diag(X) + self.offset) ** self.degree


# gpflow/kernels/misc.py

# Copyright 2018-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Optional

import numpy as np
import tensorflow as tf

from ..base import AnyNDArray, Parameter, TensorType
from ..experimental.check_shapes import check_shape as cs
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..utilities import positive, to_default_float
from .base import ActiveDims, Kernel


class ArcCosine(Kernel):
    """
    The Arc-cosine family of kernels which mimics the computation in neural
    networks. The order parameter specifies the assumed activation function.
    The Multi Layer Perceptron (MLP) kernel is closely related to the ArcCosine
    kernel of order 0.

    The key reference is :cite:t:`NIPS2009_3628`.
    """

    implemented_orders = {0, 1, 2}

    @check_shapes(
        "variance: []",
        "weight_variances: [broadcast n_active_dims]",
        "bias_variance: []",
    )
    def __init__(
        self,
        order: int = 0,
        variance: TensorType = 1.0,
        weight_variances: TensorType = 1.0,
        bias_variance: TensorType = 1.0,
        *,
        active_dims: Optional[ActiveDims] = None,
        name: Optional[str] = None,
    ) -> None:
        """
        :param order: specifies the activation function of the neural network
          the function is a rectified monomial of the chosen order
        :param variance: the (initial) value for the variance parameter
        :param weight_variances: the (initial) value for the weight_variances parameter,
            to induce ARD behaviour this must be initialised as an array the same
            length as the the number of active dimensions e.g. [1., 1., 1.]
        :param bias_variance: the (initial) value for the bias_variance parameter
            defaults to 1.0
        :param active_dims: a slice or list specifying which columns of X are used
        """
        super().__init__(active_dims=active_dims, name=name)

        if order not in self.implemented_orders:
            raise ValueError("Requested kernel order is not implemented.")
        self.order = order

        self.variance = Parameter(variance, transform=positive())
        self.bias_variance = Parameter(bias_variance, transform=positive())
        self.weight_variances = Parameter(weight_variances, transform=positive())
        self._validate_ard_active_dims(self.weight_variances)

    @property
    def ard(self) -> bool:
        """
        Whether ARD behaviour is active.
        """
        ndims: int = self.weight_variances.shape.ndims
        return ndims > 0

    @check_shapes(
        "X: [batch..., N, D]",
        "return: [batch..., N]",
    )
    def _diag_weighted_product(self, X: TensorType) -> tf.Tensor:
        return tf.reduce_sum(self.weight_variances * tf.square(X), axis=-1) + self.bias_variance

    @check_shapes(
        "X: [batch..., N, D]",
        "X2: [batch2..., N2, D]",
        "return: [batch..., N, batch2..., N2] if X2 is not None",
        "return: [batch..., N, N] if X2 is None",
    )
    def _full_weighted_product(self, X: TensorType, X2: Optional[TensorType]) -> tf.Tensor:
        if X2 is None:
            return (
                tf.linalg.matmul((self.weight_variances * X), X, transpose_b=True)
                + self.bias_variance
            )

        else:
            D = tf.shape(X)[-1]
            N = tf.shape(X)[-2]
            N2 = tf.shape(X2)[-2]
            batch = tf.shape(X)[:-2]
            batch2 = tf.shape(X2)[:-2]
            rank = tf.rank(X) - 2
            rank2 = tf.rank(X2) - 2
            ones = tf.ones((rank,), tf.int32)
            ones2 = tf.ones((rank2,), tf.int32)

            X = cs(
                tf.reshape(X, tf.concat([batch, ones2, [N, D]], 0)),
                "[batch..., broadcast batch2..., N, D]",
            )
            X2 = cs(
                tf.reshape(X2, tf.concat([ones, batch2, [N2, D]], 0)),
                "[broadcast batch..., batch2..., N2, D]",
            )
            result = cs(
                tf.linalg.matmul((self.weight_variances * X), X2, transpose_b=True)
                + self.bias_variance,
                "[batch..., batch2..., N, N2]",
            )

            indices = tf.concat(
                [
                    tf.range(rank),
                    [rank + rank2],
                    tf.range(rank2) + rank,
                    [rank + rank2 + 1],
                ],
                axis=0,
            )
            return tf.transpose(result, indices)

    @check_shapes(
        "theta: [any...]",
        "return: [any...]",
    )
    def _J(self, theta: TensorType) -> TensorType:
        """
        Implements the order dependent family of functions defined in equations
        4 to 7 in the reference paper.
        """
        if self.order == 0:
            return np.pi - theta
        elif self.order == 1:
            return tf.sin(theta) + (np.pi - theta) * tf.cos(theta)
        else:
            assert self.order == 2, f"Don't know how to handle order {self.order}."
            return 3.0 * tf.sin(theta) * tf.cos(theta) + (np.pi - theta) * (
                1.0 + 2.0 * tf.cos(theta) ** 2
            )

    @inherit_check_shapes
    def K(self, X: TensorType, X2: Optional[TensorType] = None) -> tf.Tensor:
        X_denominator = cs(tf.sqrt(self._diag_weighted_product(X)), "[batch..., N]")
        if X2 is None:
            X2_denominator = cs(X_denominator[..., None, :], "[batch..., 1, N]")
            X_denominator = cs(X_denominator[..., :, None], "[batch..., N, 1]")
            numerator = cs(self._full_weighted_product(X, None), "[batch..., N, N]")
        else:
            X2_denominator = cs(tf.sqrt(self._diag_weighted_product(X2)), "[batch2..., N2]")

            batch = tf.shape(X)[:-1]
            batch2 = tf.shape(X2)[:-1]
            ones = tf.ones((tf.rank(X) - 1,), tf.int32)
            ones2 = tf.ones((tf.rank(X2) - 1,), tf.int32)

            X_denominator = cs(
                tf.reshape(X_denominator, tf.concat([batch, ones2], 0)),
                "[batch..., N, broadcast batch2..., 1]",
            )
            X2_denominator = cs(
                tf.reshape(X2_denominator, tf.concat([ones, batch2], 0)),
                "[broadcast batch..., 1, batch2..., N2]",
            )
            numerator = cs(self._full_weighted_product(X, X2), "[batch..., N, batch2..., N2]")

        cos_theta = numerator / X_denominator / X2_denominator
        jitter = 1e-15
        theta = tf.acos(jitter + (1 - 2 * jitter) * cos_theta)

        return (
            self.variance
            * (1.0 / np.pi)
            * self._J(theta)
            * X_denominator ** self.order
            * X2_denominator ** self.order
        )

    @inherit_check_shapes
    def K_diag(self, X: TensorType) -> tf.Tensor:
        X_product = self._diag_weighted_product(X)
        const = (1.0 / np.pi) * self._J(to_default_float(0.0))
        return self.variance * const * X_product ** self.order


class Coregion(Kernel):
    """
    A Coregionalization kernel. The inputs to this kernel are _integers_ (we
    cast them from floats as needed) which usually specify the *outputs* of a
    Coregionalization model.

    The kernel function is an indexing of a positive-definite matrix:

      K(x, y) = B[x, y] .

    To ensure that B is positive-definite, it is specified by the two
    parameters of this kernel, W and kappa:

      B = W Wᵀ + diag(kappa) .

    We refer to the size of B as "output_dim x output_dim", since this is the
    number of outputs in a coregionalization model. We refer to the number of
    columns on W as 'rank': it is the number of degrees of correlation between
    the outputs.

    NB. There is a symmetry between the elements of W, which creates a local
    minimum at W=0. To avoid this, it is recommended to initialize the
    optimization (or MCMC chain) using a random W.
    """

    def __init__(
        self,
        output_dim: int,
        rank: int,
        *,
        active_dims: Optional[ActiveDims] = None,
        name: Optional[str] = None,
    ) -> None:
        """
        :param output_dim: number of outputs expected (0 <= X < output_dim)
        :param rank: number of degrees of correlation between outputs
        """

        # assert input_dim == 1, "Coregion kernel in 1D only"
        super().__init__(active_dims=active_dims, name=name)

        self.output_dim = output_dim
        self.rank = rank
        W: AnyNDArray = 0.1 * np.ones((self.output_dim, self.rank))
        kappa = np.ones(self.output_dim)
        self.W = Parameter(W)
        self.kappa = Parameter(kappa, transform=positive())

    @check_shapes(
        "return: [P, P]",
    )
    def output_covariance(self) -> tf.Tensor:
        B = tf.linalg.matmul(self.W, self.W, transpose_b=True) + tf.linalg.diag(self.kappa)
        return B

    @check_shapes(
        "return: [P]",
    )
    def output_variance(self) -> tf.Tensor:
        B_diag = tf.reduce_sum(tf.square(self.W), 1) + self.kappa
        return B_diag

    @inherit_check_shapes
    def K(self, X: TensorType, X2: Optional[TensorType] = None) -> tf.Tensor:
        cs(X, "[batch..., N, 1]  # The `Coregion` kernel requires a 1D input space.")

        B = cs(self.output_covariance(), "[O, O]")
        X = cs(tf.cast(X[..., 0], tf.int32), "[batch..., N]")
        if X2 is None:
            batch = tf.shape(X)[:-1]
            N = tf.shape(X)[-1]
            O = tf.shape(B)[-1]

            result = cs(tf.gather(B, X), "[batch..., N, O]")
            result = cs(tf.reshape(result, [-1, N, O]), "[flat_batch, N, O]")
            flat_X = cs(tf.reshape(X, [-1, N]), "[flat_batch, N]")
            result = cs(tf.gather(result, flat_X, axis=2, batch_dims=1), "[flat_batch, N, N]")
            result = cs(tf.reshape(result, tf.concat([batch, [N, N]], 0)), "[batch..., N, N]")
        else:
            X2 = cs(tf.cast(X2[..., 0], tf.int32), "[batch2..., N2]")

            rank2 = tf.rank(X2)

            result = cs(tf.gather(B, X2), "[batch2..., N2, O]")
            result = cs(
                tf.transpose(result, tf.concat([[rank2], tf.range(rank2)], 0)), "[O, batch2..., N2]"
            )
            result = cs(tf.gather(result, X), "[batch..., N, batch2..., N2]")

        return result

    @inherit_check_shapes
    def K_diag(self, X: TensorType) -> tf.Tensor:
        cs(X, "[batch..., N, 1]  # The `Coregion` kernel requires a 1D input space.")

        X = tf.cast(X[..., 0], tf.int32)
        B_diag = self.output_variance()
        return tf.gather(B_diag, X)


# gpflow/kernels/periodic.py

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Optional

import numpy as np
import tensorflow as tf

from ..base import Parameter, TensorType
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..utilities import positive
from ..utilities.ops import difference_matrix
from .base import ActiveDims, Kernel, NormalizedActiveDims
from .stationaries import IsotropicStationary


class Periodic(Kernel):
    """
    The periodic family of kernels. Can be used to wrap any Stationary kernel
    to transform it into a periodic version. The canonical form (based on the
    SquaredExponential kernel) can be found in Equation (47) of

    D.J.C.MacKay. Introduction to Gaussian processes. In C.M.Bishop, editor,
    Neural Networks and Machine Learning, pages 133--165. Springer, 1998.

    The derivation can be achieved by mapping the original inputs through the
    transformation u = (cos(x), sin(x)).

    For the SquaredExponential base kernel, the result can be expressed as:

        k(r) =  σ² exp{ -0.5 sin²(π r / γ) / ℓ²}

    where:
    r is the Euclidean distance between the input points
    ℓ is the lengthscales parameter,
    σ² is the variance parameter,
    γ is the period parameter.

    NOTE: usually we have a factor of 4 instead of 0.5 in front but this
        is absorbed into the lengthscales hyperparameter.
    NOTE: periodic kernel uses `active_dims` of a base kernel, therefore
        the constructor doesn't have it as an argument.
    """

    @check_shapes(
        "period: [broadcast n_active_dims]",
    )
    def __init__(self, base_kernel: IsotropicStationary, period: TensorType = 1.0) -> None:
        """
        :param base_kernel: the base kernel to make periodic; must inherit from Stationary
            Note that `active_dims` should be specified in the base kernel.
        :param period: the period; to induce a different period per active dimension
            this must be initialized with an array the same length as the number
            of active dimensions e.g. [1., 1., 1.]
        """
        if not isinstance(base_kernel, IsotropicStationary):
            raise TypeError("Periodic requires an IsotropicStationary kernel as the `base_kernel`")

        super().__init__()
        self.base_kernel = base_kernel
        self.period = Parameter(period, transform=positive())
        self.base_kernel._validate_ard_active_dims(self.period)

    @property
    def active_dims(self) -> NormalizedActiveDims:
        return self.base_kernel.active_dims

    @active_dims.setter
    def active_dims(self, value: ActiveDims) -> None:
        # type-ignore below is because mypy doesn't understand that getter and the setter of
        # `active_dims` have different types.
        self.base_kernel.active_dims = value  # type: ignore[assignment]

    @inherit_check_shapes
    def K_diag(self, X: TensorType) -> tf.Tensor:
        return self.base_kernel.K_diag(X)

    @inherit_check_shapes
    def K(self, X: TensorType, X2: Optional[TensorType] = None) -> tf.Tensor:
        r = np.pi * (difference_matrix(X, X2)) / self.period
        scaled_sine = tf.sin(r) / self.base_kernel.lengthscales
        if hasattr(self.base_kernel, "K_r"):
            sine_r = tf.reduce_sum(tf.abs(scaled_sine), -1)
            K = self.base_kernel.K_r(sine_r)
        else:
            sine_r2 = tf.reduce_sum(tf.square(scaled_sine), -1)
            K = self.base_kernel.K_r2(sine_r2)
        return K


# gpflow/kernels/statics.py

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Optional

import tensorflow as tf

from ..base import Parameter, TensorType
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..utilities import positive
from .base import ActiveDims, Kernel


class Static(Kernel):
    """
    Kernels who don't depend on the value of the inputs are 'Static'.  The only
    parameter is a variance, σ².
    """

    @check_shapes(
        "variance: []",
    )
    def __init__(
        self, variance: TensorType = 1.0, active_dims: Optional[ActiveDims] = None
    ) -> None:
        super().__init__(active_dims)
        self.variance = Parameter(variance, transform=positive())

    @inherit_check_shapes
    def K_diag(self, X: TensorType) -> tf.Tensor:
        return tf.fill(tf.shape(X)[:-1], tf.squeeze(self.variance))


class White(Static):
    """
    The White kernel: this kernel produces 'white noise'. The kernel equation is

        k(x_n, x_m) = δ(n, m) σ²

    where:
    δ(.,.) is the Kronecker delta,
    σ²  is the variance parameter.
    """

    @inherit_check_shapes
    def K(self, X: TensorType, X2: Optional[TensorType] = None) -> tf.Tensor:
        if X2 is None:
            d = tf.fill(tf.shape(X)[:-1], tf.squeeze(self.variance))
            return tf.linalg.diag(d)
        else:
            shape = tf.concat([tf.shape(X)[:-1], tf.shape(X2)[:-1]], axis=0)
            return tf.zeros(shape, dtype=X.dtype)


class Constant(Static):
    """
    The Constant (aka Bias) kernel. Functions drawn from a GP with this kernel
    are constant, i.e. f(x) = c, with c ~ N(0, σ^2). The kernel equation is

        k(x, y) = σ²

    where:
    σ²  is the variance parameter.
    """

    @inherit_check_shapes
    def K(self, X: TensorType, X2: Optional[TensorType] = None) -> tf.Tensor:
        if X2 is None:
            shape = tf.concat(
                [
                    tf.shape(X)[:-2],
                    tf.reshape(tf.shape(X)[-2], [1]),
                    tf.reshape(tf.shape(X)[-2], [1]),
                ],
                axis=0,
            )
        else:
            shape = tf.concat([tf.shape(X)[:-1], tf.shape(X2)[:-1]], axis=0)

        return tf.fill(shape, tf.squeeze(self.variance))


# gpflow/kernels/stationaries.py

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Any, Optional

import numpy as np
import tensorflow as tf

from ..base import Parameter, TensorType
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..utilities import positive
from ..utilities.ops import difference_matrix, square_distance
from .base import ActiveDims, Kernel


class Stationary(Kernel):
    """
    Base class for kernels that are stationary, that is, they only depend on

        d = x - x'

    This class handles 'ard' behaviour, which stands for 'Automatic Relevance
    Determination'. This means that the kernel has one lengthscale per
    dimension, otherwise the kernel is isotropic (has a single lengthscale).
    """

    @check_shapes(
        "variance: []",
        "lengthscales: [broadcast n_active_dims]",
    )
    def __init__(
        self, variance: TensorType = 1.0, lengthscales: TensorType = 1.0, **kwargs: Any
    ) -> None:
        """
        :param variance: the (initial) value for the variance parameter.
        :param lengthscales: the (initial) value for the lengthscale
            parameter(s), to induce ARD behaviour this must be initialised as
            an array the same length as the number of active dimensions
            e.g. [1., 1., 1.]. If only a single value is passed, this value
            is used as the lengthscale of each dimension.
        :param kwargs: accepts `name` and `active_dims`, which is a list or
            slice of indices which controls which columns of X are used (by
            default, all columns are used).
        """
        for kwarg in kwargs:
            if kwarg not in {"name", "active_dims"}:
                raise TypeError(f"Unknown keyword argument: {kwarg}")

        super().__init__(**kwargs)
        self.variance = Parameter(variance, transform=positive())
        self.lengthscales = Parameter(lengthscales, transform=positive())
        self._validate_ard_active_dims(self.lengthscales)

    @property
    def ard(self) -> bool:
        """
        Whether ARD behaviour is active.
        """
        ndims: int = self.lengthscales.shape.ndims
        return ndims > 0

    @check_shapes(
        "X: [broadcast any...]",
        "return: [any...]",
    )
    def scale(self, X: TensorType) -> TensorType:
        X_scaled = X / self.lengthscales if X is not None else X
        return X_scaled

    @inherit_check_shapes
    def K_diag(self, X: TensorType) -> tf.Tensor:
        return tf.fill(tf.shape(X)[:-1], tf.squeeze(self.variance))


class IsotropicStationary(Stationary):
    """
    Base class for isotropic stationary kernels, i.e. kernels that only
    depend on

        r = ‖x - x'‖

    Derived classes should implement one of:

        K_r2(self, r2): Returns the kernel evaluated on r² (r2), which is the
        squared scaled Euclidean distance Should operate element-wise on r2.

        K_r(self, r): Returns the kernel evaluated on r, which is the scaled
        Euclidean distance. Should operate element-wise on r.
    """

    @inherit_check_shapes
    def K(self, X: TensorType, X2: Optional[TensorType] = None) -> tf.Tensor:
        r2 = self.scaled_squared_euclid_dist(X, X2)
        return self.K_r2(r2)

    @check_shapes(
        "r2: [batch..., N]",
        "return: [batch..., N]",
    )
    def K_r2(self, r2: TensorType) -> tf.Tensor:
        if hasattr(self, "K_r"):
            # Clipping around the (single) float precision which is ~1e-45.
            r = tf.sqrt(tf.maximum(r2, 1e-36))
            return self.K_r(r)  # pylint: disable=no-member
        raise NotImplementedError

    @check_shapes(
        "X: [batch..., N, D]",
        "X2: [batch2..., N2, D]",
        "return: [batch..., N, batch2..., N2] if X2 is not None",
        "return: [batch..., N, N] if X2 is None",
    )
    def scaled_squared_euclid_dist(
        self, X: TensorType, X2: Optional[TensorType] = None
    ) -> tf.Tensor:
        """
        Returns ‖(X - X2ᵀ) / ℓ‖², i.e. the squared L₂-norm.
        """
        return square_distance(self.scale(X), self.scale(X2))


class AnisotropicStationary(Stationary):
    """
    Base class for anisotropic stationary kernels, i.e. kernels that only
    depend on

        d = x - x'

    Derived classes should implement K_d(self, d): Returns the kernel evaluated
    on d, which is the pairwise difference matrix, scaled by the lengthscale
    parameter ℓ (i.e. [(X - X2ᵀ) / ℓ]). The last axis corresponds to the
    input dimension.
    """

    @check_shapes(
        "variance: []",
        "lengthscales: [broadcast n_active_dims]",
    )
    def __init__(
        self, variance: TensorType = 1.0, lengthscales: TensorType = 1.0, **kwargs: Any
    ) -> None:
        """
        :param variance: the (initial) value for the variance parameter.
        :param lengthscales: the (initial) value for the lengthscale
            parameter(s), to induce ARD behaviour this must be initialised as
            an array the same length as the number of active dimensions
            e.g. [1., 1., 1.]. Note that anisotropic kernels can possess
            negative lengthscales. If only a single value is passed, this
            value is used as the lengthscale of each dimension.
        :param kwargs: accepts `name` and `active_dims`, which is a list or
            slice of indices which controls which columns of X are used (by
            default, all columns are used).
        """
        super().__init__(variance, lengthscales, **kwargs)

        if self.ard:
            self.lengthscales = Parameter(self.lengthscales.numpy())

    @inherit_check_shapes
    def K(self, X: TensorType, X2: Optional[TensorType] = None) -> tf.Tensor:
        return self.K_d(self.scaled_difference_matrix(X, X2))

    @check_shapes(
        "X: [batch..., N, D]",
        "X2: [batch2..., N2, D]",
        "return: [batch..., N, batch2..., N2, D] if X2 is not None",
        "return: [batch..., N, N, D] if X2 is None",
    )
    def scaled_difference_matrix(self, X: TensorType, X2: Optional[TensorType] = None) -> tf.Tensor:
        """
        Returns [(X - X2ᵀ) / ℓ]. If X has shape [..., N, D] and
        X2 has shape [..., M, D], the output will have shape [..., N, M, D].
        """
        return difference_matrix(self.scale(X), self.scale(X2))

    @check_shapes(
        "d: [batch..., N, D]",
        "return: [batch..., N]",
    )
    def K_d(self, d: TensorType) -> tf.Tensor:
        raise NotImplementedError


class SquaredExponential(IsotropicStationary):
    """
    The radial basis function (RBF) or squared exponential kernel. The kernel equation is

        k(r) = σ² exp{-½ r²}

    where:
    r   is the Euclidean distance between the input points, scaled by the lengthscales parameter ℓ.
    σ²  is the variance parameter

    Functions drawn from a GP with this kernel are infinitely differentiable!
    """

    @inherit_check_shapes
    def K_r2(self, r2: TensorType) -> tf.Tensor:
        return self.variance * tf.exp(-0.5 * r2)


class RationalQuadratic(IsotropicStationary):
    """
    Rational Quadratic kernel,

    k(r) = σ² (1 + r² / 2αℓ²)^(-α)

    σ² : variance
    ℓ  : lengthscales
    α  : alpha, determines relative weighting of small-scale and large-scale fluctuations

    For α → ∞, the RQ kernel becomes equivalent to the squared exponential.
    """

    def __init__(
        self,
        variance: TensorType = 1.0,
        lengthscales: TensorType = 1.0,
        alpha: TensorType = 1.0,
        active_dims: Optional[ActiveDims] = None,
    ) -> None:
        super().__init__(variance=variance, lengthscales=lengthscales, active_dims=active_dims)
        self.alpha = Parameter(alpha, transform=positive())

    @inherit_check_shapes
    def K_r2(self, r2: TensorType) -> tf.Tensor:
        return self.variance * (1 + 0.5 * r2 / self.alpha) ** (-self.alpha)


class Exponential(IsotropicStationary):
    """
    The Exponential kernel. It is equivalent to a Matern12 kernel with doubled lengthscales.
    """

    @check_shapes(
        "r: [batch..., N]",
        "return: [batch..., N]",
    )
    def K_r(self, r: TensorType) -> tf.Tensor:
        return self.variance * tf.exp(-0.5 * r)


class Matern12(IsotropicStationary):
    """
    The Matern 1/2 kernel. Functions drawn from a GP with this kernel are not
    differentiable anywhere. The kernel equation is

    k(r) = σ² exp{-r}

    where:
    r  is the Euclidean distance between the input points, scaled by the lengthscales parameter ℓ.
    σ² is the variance parameter
    """

    @check_shapes(
        "r: [batch..., N]",
        "return: [batch..., N]",
    )
    def K_r(self, r: TensorType) -> tf.Tensor:
        return self.variance * tf.exp(-r)


class Matern32(IsotropicStationary):
    """
    The Matern 3/2 kernel. Functions drawn from a GP with this kernel are once
    differentiable. The kernel equation is

    k(r) = σ² (1 + √3r) exp{-√3 r}

    where:
    r  is the Euclidean distance between the input points, scaled by the lengthscales parameter ℓ,
    σ² is the variance parameter.
    """

    @check_shapes(
        "r: [batch..., N]",
        "return: [batch..., N]",
    )
    def K_r(self, r: TensorType) -> tf.Tensor:
        sqrt3 = np.sqrt(3.0)
        return self.variance * (1.0 + sqrt3 * r) * tf.exp(-sqrt3 * r)


class Matern52(IsotropicStationary):
    """
    The Matern 5/2 kernel. Functions drawn from a GP with this kernel are twice
    differentiable. The kernel equation is

    k(r) = σ² (1 + √5r + 5/3r²) exp{-√5 r}

    where:
    r  is the Euclidean distance between the input points, scaled by the lengthscales parameter ℓ,
    σ² is the variance parameter.
    """

    @check_shapes(
        "r: [batch..., N]",
        "return: [batch..., N]",
    )
    def K_r(self, r: TensorType) -> tf.Tensor:
        sqrt5 = np.sqrt(5.0)
        return self.variance * (1.0 + sqrt5 * r + 5.0 / 3.0 * tf.square(r)) * tf.exp(-sqrt5 * r)


class Cosine(AnisotropicStationary):
    """
    The Cosine kernel. Functions drawn from a GP with this kernel are sinusoids
    (with a random phase).  The kernel equation is

        k(r) = σ² cos{2πd}

    where:
    d  is the sum of the per-dimension differences between the input points, scaled by the
    lengthscale parameter ℓ (i.e. Σᵢ [(X - X2ᵀ) / ℓ]ᵢ),
    σ² is the variance parameter.
    """

    @inherit_check_shapes
    def K_d(self, d: TensorType) -> tf.Tensor:
        d = tf.reduce_sum(d, axis=-1)
        return self.variance * tf.cos(2 * np.pi * d)


# gpflow/likelihoods/__init__.py

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Likelihoods are another core component of GPflow. This describes how likely the
data is under the assumptions made about the underlying latent functions
p(Y|F). Different likelihoods make different
assumptions about the distribution of the data, as such different data-types
(continuous, binary, ordinal, count) are better modelled with different
likelihood assumptions.

Use of any likelihood other than Gaussian typically introduces the need to use
an approximation to perform inference, if one isn't already needed.
Variational inference and MCMC models are included in GPflow and allow
approximate inference with non-Gaussian likelihoods. An introduction to these
models can be found :ref:`here <implemented_models>`. Specific notebooks
illustrating non-Gaussian likelihood regressions are available for
`classification <notebooks/classification.html>`_ (binary data), `ordinal
<notebooks/ordinal.html>`_ and `multiclass <notebooks/multiclass.html>`_.

Creating new likelihoods
------------------------

Likelihoods are defined by their
log-likelihood. When creating new likelihoods, the
:func:`logp <gpflow.likelihoods.Likelihood.logp>` method (log p(Y|F)), the
:func:`conditional_mean <gpflow.likelihoods.Likelihood.conditional_mean>`,
:func:`conditional_variance
<gpflow.likelihoods.Likelihood.conditional_variance>`.

In order to perform variational inference with non-Gaussian likelihoods a term
called ``variational expectations``, ∫ q(F) log p(Y|F) dF, needs to
be computed under a Gaussian distribution q(F) ~ N(μ, Σ).

The :func:`variational_expectations <gpflow.likelihoods.Likelihood.variational_expectations>`
method can be overriden if this can be computed in closed form, otherwise; if
the new likelihood inherits
:class:`Likelihood <gpflow.likelihoods.Likelihood>` the default will use
Gauss-Hermite numerical integration (works well when F is 1D
or 2D), if the new likelihood inherits from
:class:`MonteCarloLikelihood <gpflow.likelihoods.MonteCarloLikelihood>` the
integration is done by sampling (can be more suitable when F is higher dimensional).
"""

from .base import (
    Likelihood,
    MonteCarloLikelihood,
    QuadratureLikelihood,
    ScalarLikelihood,
    SwitchedLikelihood,
)
from .misc import GaussianMC
from .multiclass import MultiClass, RobustMax, Softmax
from .multilatent import (
    HeteroskedasticTFPConditional,
    MultiLatentLikelihood,
    MultiLatentTFPConditional,
)
from .scalar_continuous import Beta, Exponential, Gamma, Gaussian, StudentT
from .scalar_discrete import Bernoulli, Ordinal, Poisson

__all__ = [
    "Bernoulli",
    "Beta",
    "Exponential",
    "Gamma",
    "Gaussian",
    "GaussianMC",
    "HeteroskedasticTFPConditional",
    "Likelihood",
    "MonteCarloLikelihood",
    "MultiClass",
    "MultiLatentLikelihood",
    "MultiLatentTFPConditional",
    "Ordinal",
    "Poisson",
    "QuadratureLikelihood",
    "RobustMax",
    "ScalarLikelihood",
    "Softmax",
    "StudentT",
    "SwitchedLikelihood",
    "base",
    "misc",
    "multiclass",
    "multilatent",
    "scalar_continuous",
    "scalar_discrete",
    "utils",
]


# gpflow/likelihoods/base.py

# Copyright 2016-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import abc
from typing import Any, Callable, Iterable, Optional, Sequence, Union

import tensorflow as tf

from ..base import MeanAndVariance, Module, TensorType
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..quadrature import GaussianQuadrature, NDiagGHQuadrature, ndiag_mc

DEFAULT_NUM_GAUSS_HERMITE_POINTS = 20
"""
The number of Gauss-Hermite points to use for quadrature (fallback when a
likelihood method does not have an analytic method) if quadrature object is not
explicitly passed to likelihood constructor.
"""


class Likelihood(Module, abc.ABC):
    def __init__(
        self, input_dim: Optional[int], latent_dim: Optional[int], observation_dim: Optional[int]
    ) -> None:
        """
        A base class for likelihoods, which specifies an observation model
        connecting the latent functions ('F') to the data ('Y').

        All of the members of this class are expected to obey some shape conventions, as specified
        by latent_dim and observation_dim.

        If we're operating on an array of function values 'F', then the last dimension represents
        multiple functions (preceding dimensions could represent different data points, or
        different random samples, for example). Similarly, the last dimension of Y represents a
        single data point. We check that the dimensions are as this object expects.

        The return shapes of all functions in this class is the broadcasted shape of the arguments,
        excluding the last dimension of each argument.

        :param input_dim: the dimension of the input vector X for a single data point
        :param latent_dim: the dimension of the vector F of latent functions for a single data point
        :param observation_dim: the dimension of the observation vector Y for a single data point
        """
        super().__init__()
        self.input_dim = input_dim
        self.latent_dim = latent_dim
        self.observation_dim = observation_dim

    @check_shapes(
        "X: [broadcast batch..., input_dim]",
        "F: [broadcast batch..., latent_dim]",
        "Y: [broadcast batch..., observation_dim]",
        "return: [batch...]",
    )
    def log_prob(self, X: TensorType, F: TensorType, Y: TensorType) -> tf.Tensor:
        """
        The log probability density log p(Y|X,F)

        :param X: input tensor
        :param F: function evaluation tensor
        :param Y: observation tensor
        :returns: log pdf
        """
        return self._log_prob(X, F, Y)

    @abc.abstractmethod
    @check_shapes(
        "X: [broadcast batch..., input_dim]",
        "F: [broadcast batch..., latent_dim]",
        "Y: [broadcast batch..., observation_dim]",
        "return: [batch...]",
    )
    def _log_prob(self, X: TensorType, F: TensorType, Y: TensorType) -> tf.Tensor:
        raise NotImplementedError

    @check_shapes(
        "X: [broadcast batch..., input_dim]",
        "F: [broadcast batch..., latent_dim]",
        "return: [batch..., observation_dim]",
    )
    def conditional_mean(self, X: TensorType, F: TensorType) -> tf.Tensor:
        """
        The conditional mean of Y|X,F: [E[Y₁|X,F], ..., E[Yₖ|X,F]]
        where K = observation_dim

        :param X: input tensor
        :param F: function evaluation tensor
        :returns: mean
        """
        return self._conditional_mean(X, F)

    @check_shapes(
        "X: [broadcast batch..., input_dim]",
        "F: [broadcast batch..., latent_dim]",
        "return: [batch..., observation_dim]",
    )
    def _conditional_mean(self, X: TensorType, F: TensorType) -> tf.Tensor:
        raise NotImplementedError

    @check_shapes(
        "X: [broadcast batch..., input_dim]",
        "F: [broadcast batch..., latent_dim]",
        "return: [batch..., observation_dim]",
    )
    def conditional_variance(self, X: TensorType, F: TensorType) -> tf.Tensor:
        """
        The conditional marginal variance of Y|X,F: [var(Y₁|X,F), ..., var(Yₖ|X,F)]
        where K = observation_dim

        :param X: input tensor
        :param F: function evaluation tensor
        :returns: variance
        """
        return self._conditional_variance(X, F)

    @check_shapes(
        "X: [broadcast batch..., input_dim]",
        "F: [broadcast batch..., latent_dim]",
        "return: [batch..., observation_dim]",
    )
    def _conditional_variance(self, X: TensorType, F: TensorType) -> tf.Tensor:
        raise NotImplementedError

    @check_shapes(
        "X: [broadcast batch..., input_dim]",
        "Fmu: [broadcast batch..., latent_dim]",
        "Fvar: [broadcast batch..., latent_dim]",
        "return[0]: [batch..., observation_dim]",
        "return[1]: [batch..., observation_dim]",
    )
    def predict_mean_and_var(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType
    ) -> MeanAndVariance:
        """
        Given a Normal distribution for the latent function,
        return the mean and marginal variance of Y,

        i.e. if
            q(f) = N(Fmu, Fvar)

        and this object represents

            p(y|f)

        then this method computes the predictive mean

           ∫∫ y p(y|f)q(f) df dy

        and the predictive variance

           ∫∫ y² p(y|f)q(f) df dy  - [ ∫∫ y p(y|f)q(f) df dy ]²

        :param X: input tensor
        :param Fmu: mean function evaluation tensor
        :param Fvar: variance of function evaluation tensor
        :returns: mean and variance
        """
        return self._predict_mean_and_var(X, Fmu, Fvar)

    @abc.abstractmethod
    @check_shapes(
        "X: [broadcast batch..., input_dim]",
        "Fmu: [broadcast batch..., latent_dim]",
        "Fvar: [broadcast batch..., latent_dim]",
        "return[0]: [batch..., observation_dim]",
        "return[1]: [batch..., observation_dim]",
    )
    def _predict_mean_and_var(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType
    ) -> MeanAndVariance:
        raise NotImplementedError

    @check_shapes(
        "X: [broadcast batch..., input_dim]",
        "Fmu: [broadcast batch..., latent_dim]",
        "Fvar: [broadcast batch..., latent_dim]",
        "Y: [broadcast batch..., observation_dim]",
        "return: [batch...]",
    )
    def predict_log_density(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        r"""
        Given a Normal distribution for the latent function, and a datum Y,
        compute the log predictive density of Y,

        i.e. if
            q(F) = N(Fmu, Fvar)

        and this object represents

            p(y|F)

        then this method computes the predictive density

            log ∫ p(y=Y|F)q(F) df

        :param X: input tensor
        :param Fmu: mean function evaluation tensor
        :param Fvar: variance of function evaluation tensor
        :param Y: observation tensor
        :returns: log predictive density
        """
        return self._predict_log_density(X, Fmu, Fvar, Y)

    @abc.abstractmethod
    @check_shapes(
        "X: [broadcast batch..., input_dim]",
        "Fmu: [broadcast batch..., latent_dim]",
        "Fvar: [broadcast batch..., latent_dim]",
        "Y: [broadcast batch..., observation_dim]",
        "return: [batch...]",
    )
    def _predict_log_density(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        raise NotImplementedError

    @check_shapes(
        "X: [broadcast batch..., input_dim]",
        "Fmu: [broadcast batch..., latent_dim]",
        "Fvar: [broadcast batch..., latent_dim]",
        "Y: [broadcast batch..., observation_dim]",
        "return: [batch...]",
    )
    def variational_expectations(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        r"""
        Compute the expected log density of the data, given a Gaussian
        distribution for the function values,

        i.e. if
            q(f) = N(Fmu, Fvar)

        and this object represents

            p(y|f)

        then this method computes

           ∫ log(p(y=Y|f)) q(f) df.

        This only works if the broadcasting dimension of the statistics of q(f) (mean and variance)
        are broadcastable with that of the data Y.

        :param X: input tensor
        :param Fmu: mean function evaluation tensor
        :param Fvar: variance of function evaluation tensor
        :param Y: observation tensor
        :returns: expected log density of the data given q(F)
        """
        return self._variational_expectations(X, Fmu, Fvar, Y)

    @abc.abstractmethod
    @check_shapes(
        "X: [broadcast batch..., input_dim]",
        "Fmu: [broadcast batch..., latent_dim]",
        "Fvar: [broadcast batch..., latent_dim]",
        "Y: [broadcast batch..., observation_dim]",
        "return: [batch...]",
    )
    def _variational_expectations(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        raise NotImplementedError


class QuadratureLikelihood(Likelihood, abc.ABC):
    def __init__(
        self,
        input_dim: Optional[int],
        latent_dim: Optional[int],
        observation_dim: Optional[int],
        *,
        quadrature: Optional[GaussianQuadrature] = None,
    ) -> None:
        super().__init__(
            input_dim=input_dim, latent_dim=latent_dim, observation_dim=observation_dim
        )
        if quadrature is None:
            with tf.init_scope():
                quadrature = NDiagGHQuadrature(
                    self._quadrature_dim, DEFAULT_NUM_GAUSS_HERMITE_POINTS
                )
        self.quadrature = quadrature

    @property
    def _quadrature_dim(self) -> int:
        """
        This defines the number of dimensions over which to evaluate the
        quadrature. Generally, this is equal to self.latent_dim. This exists
        as a separate property to allow the ScalarLikelihood subclass to
        override it with 1 (broadcasting over observation/latent dimensions
        instead).
        """
        assert self.latent_dim is not None
        return self.latent_dim

    @check_shapes(
        "F: [broadcast batch..., latent_dim]",
        "X: [broadcast batch..., input_dim]",
        "Y: [broadcast batch..., observation_dim]",
        "return: [batch..., d]",
    )
    def _quadrature_log_prob(self, F: TensorType, X: TensorType, Y: TensorType) -> tf.Tensor:
        """
        Returns the appropriate log prob integrand for quadrature.

        Quadrature expects f(X), here logp(F), to return shape [N_quad_points]
        + batch_shape + [d']. Here d'=1, but log_prob() only returns
        [N_quad_points] + batch_shape, so we add an extra dimension.

        Also see _quadrature_reduction.
        """
        return tf.expand_dims(self.log_prob(X, F, Y), axis=-1)

    @check_shapes(
        "quadrature_result: [batch..., d]",
        "return: [batch...]",
    )
    def _quadrature_reduction(self, quadrature_result: TensorType) -> tf.Tensor:
        """
        Converts the quadrature integral appropriately.

        The return shape of quadrature is batch_shape + [d']. Here, d'=1, but
        we want predict_log_density and variational_expectations to return just
        batch_shape, so we squeeze out the extra dimension.

        Also see _quadrature_log_prob.
        """
        return tf.squeeze(quadrature_result, axis=-1)

    @inherit_check_shapes
    def _predict_log_density(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        r"""
        Here, we implement a default Gauss-Hermite quadrature routine, but some
        likelihoods (Gaussian, Poisson) will implement specific cases.
        :param X: input tensor
        :param Fmu: mean function evaluation tensor
        :param Fvar: variance of function evaluation tensor
        :param Y: observation tensor
        :returns: log predictive density
        """
        return self._quadrature_reduction(
            self.quadrature.logspace(self._quadrature_log_prob, Fmu, Fvar, X=X, Y=Y)
        )

    @inherit_check_shapes
    def _variational_expectations(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        r"""
        Here, we implement a default Gauss-Hermite quadrature routine, but some
        likelihoods (Gaussian, Poisson) will implement specific cases.
        :param X: input tensor
        :param Fmu: mean function evaluation tensor
        :param Fvar: variance of function evaluation tensor
        :param Y: observation tensor
        :returns: variational expectations
        """
        return self._quadrature_reduction(
            self.quadrature(self._quadrature_log_prob, Fmu, Fvar, X=X, Y=Y)
        )

    @inherit_check_shapes
    def _predict_mean_and_var(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType
    ) -> MeanAndVariance:
        r"""
        Here, we implement a default Gauss-Hermite quadrature routine, but some
        likelihoods (e.g. Gaussian) will implement specific cases.

        :param X: input tensor
        :param Fmu: mean function evaluation tensor
        :param Fvar: variance of function evaluation tensor
        :returns: mean and variance of Y
        """

        def conditional_mean(F: TensorType, X_: TensorType) -> tf.Tensor:
            return self.conditional_mean(X_, F)

        def conditional_y_squared(F: TensorType, X_: TensorType) -> tf.Tensor:
            return self.conditional_variance(X_, F) + tf.square(self.conditional_mean(X_, F))

        E_y, E_y2 = self.quadrature([conditional_mean, conditional_y_squared], Fmu, Fvar, X_=X)
        V_y = E_y2 - E_y ** 2
        return E_y, V_y


class ScalarLikelihood(QuadratureLikelihood, abc.ABC):
    """
    A likelihood class that helps with scalar likelihood functions: likelihoods where
    each scalar latent function is associated with a single scalar observation variable.

    If there are multiple latent functions, then there must be a corresponding number of data: we
    check for this.

    The `Likelihood` class contains methods to compute marginal statistics of functions
    of the latents and the data ϕ(y,x,f):

    * variational_expectations:  ϕ(y,x,f) = log p(y|x,f)
    * predict_log_density: ϕ(y,x,f) = p(y|x,f)

    Those statistics are computed after having first marginalized the latent processes f
    under a multivariate normal distribution q(x,f) that is fully factorized.

    Some univariate integrals can be done by quadrature: we implement quadrature routines for 1D
    integrals in this class, though they may be overwritten by inheriting classes where those
    integrals are available in closed form.
    """

    def __init__(self, **kwargs: Any) -> None:
        super().__init__(input_dim=None, latent_dim=None, observation_dim=None, **kwargs)

    @inherit_check_shapes
    def _log_prob(self, X: TensorType, F: TensorType, Y: TensorType) -> tf.Tensor:
        r"""
        Compute log p(Y|X,F), where by convention we sum out the last axis as it represented
        independent latent functions and observations.
        :param F: function evaluation tensor
        :param Y: observation tensor
        """
        return tf.reduce_sum(self._scalar_log_prob(X, F, Y), axis=-1)

    @abc.abstractmethod
    @check_shapes(
        "X: [broadcast batch..., input_dim]",
        "F: [broadcast batch..., latent_dim]",
        "Y: [broadcast batch..., observation_dim]",
        "return: [batch..., latent_dim]",
    )
    def _scalar_log_prob(self, X: TensorType, F: TensorType, Y: TensorType) -> tf.Tensor:
        raise NotImplementedError

    @property
    def _quadrature_dim(self) -> int:
        """
        Quadrature is over the latent dimensions. Generally, this is equal to
        self.latent_dim. This separate property allows the ScalarLikelihood
        subclass to override it with 1 (broadcasting over observation/latent
        dimensions instead).
        """
        return 1

    @inherit_check_shapes
    def _quadrature_log_prob(self, F: TensorType, X: TensorType, Y: TensorType) -> tf.Tensor:
        """
        Returns the appropriate log prob integrand for quadrature.

        Quadrature expects f(X), here logp(F), to return shape [N_quad_points]
        + batch_shape + [d']. Here d' corresponds to the last dimension of both
        F and Y, and _scalar_log_prob simply broadcasts over this.

        Also see _quadrature_reduction.
        """
        return self._scalar_log_prob(X, F, Y)

    @inherit_check_shapes
    def _quadrature_reduction(self, quadrature_result: TensorType) -> tf.Tensor:
        """
        Converts the quadrature integral appropriately.

        The return shape of quadrature is batch_shape + [d']. Here, d'
        corresponds to the last dimension of both F and Y, and we want to sum
        over the observations to obtain the overall predict_log_density or
        variational_expectations.

        Also see _quadrature_log_prob.
        """
        return tf.reduce_sum(quadrature_result, axis=-1)


class SwitchedLikelihood(ScalarLikelihood):
    def __init__(self, likelihood_list: Iterable[ScalarLikelihood], **kwargs: Any) -> None:
        """
        In this likelihood, we assume at extra column of Y, which contains
        integers that specify a likelihood from the list of likelihoods.
        """
        super().__init__(**kwargs)
        self.likelihoods = list(likelihood_list)

    @check_shapes(
        "args[all]: [batch..., .]",
        "return: [batch..., ...]",
    )
    def _partition_and_stitch(self, args: Sequence[TensorType], func_name: str) -> tf.Tensor:
        """
        args is a list of tensors, to be passed to self.likelihoods.<func_name>

        args[-1] is the 'Y' argument, which contains the indexes to self.likelihoods.

        This function splits up the args using dynamic_partition, calls the
        relevant function on the likelihoods, and re-combines the result.
        """
        # get the index from Y
        args_list = list(args)
        Y = args_list[-1]
        ind = Y[..., -1]
        ind = tf.cast(ind, tf.int32)
        Y = Y[..., :-1]
        args_list[-1] = Y

        # split up the arguments into chunks corresponding to the relevant likelihoods
        args_chunks = zip(*[tf.dynamic_partition(X, ind, len(self.likelihoods)) for X in args_list])

        # apply the likelihood-function to each section of the data
        funcs = [getattr(lik, func_name) for lik in self.likelihoods]
        results = [f(*args_i) for f, args_i in zip(funcs, args_chunks)]

        # stitch the results back together
        partitions = tf.dynamic_partition(tf.range(0, tf.size(ind)), ind, len(self.likelihoods))
        results = tf.dynamic_stitch(partitions, results)

        return results

    @inherit_check_shapes
    def _scalar_log_prob(self, X: TensorType, F: TensorType, Y: TensorType) -> tf.Tensor:
        return self._partition_and_stitch([X, F, Y], "_scalar_log_prob")

    @inherit_check_shapes
    def _predict_log_density(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        return self._partition_and_stitch([X, Fmu, Fvar, Y], "predict_log_density")

    @inherit_check_shapes
    def _variational_expectations(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        return self._partition_and_stitch([X, Fmu, Fvar, Y], "variational_expectations")

    @inherit_check_shapes
    def _predict_mean_and_var(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType
    ) -> MeanAndVariance:
        mvs = [lik.predict_mean_and_var(X, Fmu, Fvar) for lik in self.likelihoods]
        mu_list, var_list = zip(*mvs)
        mu = tf.concat(mu_list, axis=1)
        var = tf.concat(var_list, axis=1)
        return mu, var

    @inherit_check_shapes
    def _conditional_mean(self, X: TensorType, F: TensorType) -> tf.Tensor:
        raise NotImplementedError

    @inherit_check_shapes
    def _conditional_variance(self, X: TensorType, F: TensorType) -> tf.Tensor:
        raise NotImplementedError


class MonteCarloLikelihood(Likelihood):
    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self.num_monte_carlo_points = 100

    @check_shapes(
        "Fmu: [batch..., latent_dim]",
        "Fvar: [batch..., latent_dim]",
        "Ys.values(): [batch..., .]",
        "return: [broadcast n_funcs, batch..., .]",
    )
    def _mc_quadrature(
        self,
        funcs: Union[Callable[..., tf.Tensor], Iterable[Callable[..., tf.Tensor]]],
        Fmu: TensorType,
        Fvar: TensorType,
        logspace: bool = False,
        epsilon: Optional[TensorType] = None,
        **Ys: TensorType,
    ) -> tf.Tensor:
        return ndiag_mc(funcs, self.num_monte_carlo_points, Fmu, Fvar, logspace, epsilon, **Ys)

    @inherit_check_shapes
    def _predict_mean_and_var(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, epsilon: Optional[TensorType] = None
    ) -> MeanAndVariance:
        r"""
        Given a Normal distribution for the latent function,
        return the mean of Y

        if
            q(f) = N(Fmu, Fvar)

        and this object represents

            p(y|f)

        then this method computes the predictive mean

           ∫∫ y p(y|f)q(f) df dy

        and the predictive variance

           ∫∫ y² p(y|f)q(f) df dy  - [ ∫∫ y p(y|f)q(f) df dy ]²

        Here, we implement a default Monte Carlo routine.
        """

        def conditional_mean(F: TensorType, X_: TensorType) -> TensorType:
            return self.conditional_mean(X_, F)

        def conditional_y_squared(F: TensorType, X_: TensorType) -> TensorType:
            return self.conditional_variance(X_, F) + tf.square(self.conditional_mean(X_, F))

        E_y, E_y2 = self._mc_quadrature(
            [conditional_mean, conditional_y_squared],
            Fmu,
            Fvar,
            epsilon=epsilon,
            X_=X,
        )
        V_y = E_y2 - tf.square(E_y)
        return E_y, V_y  # [N, D]

    @inherit_check_shapes
    def _predict_log_density(
        self,
        X: TensorType,
        Fmu: TensorType,
        Fvar: TensorType,
        Y: TensorType,
        epsilon: Optional[TensorType] = None,
    ) -> tf.Tensor:
        r"""
        Given a Normal distribution for the latent function, and a datum Y,
        compute the log predictive density of Y.

        i.e. if
            q(x, f) = N(Fmu, Fvar)

        and this object represents

            p(y|x,f)

        then this method computes the predictive density

            log ∫ p(y=Y|x,f)q(x,f) df

        Here, we implement a default Monte Carlo routine.
        """

        def log_prob(F: TensorType, X_: TensorType, Y_: TensorType) -> tf.Tensor:
            return self.log_prob(X_, F, Y_)

        return tf.reduce_sum(
            self._mc_quadrature(log_prob, Fmu, Fvar, logspace=True, epsilon=epsilon, X_=X, Y_=Y),
            axis=-1,
        )

    @inherit_check_shapes
    def _variational_expectations(
        self,
        X: TensorType,
        Fmu: TensorType,
        Fvar: TensorType,
        Y: TensorType,
        epsilon: Optional[TensorType] = None,
    ) -> tf.Tensor:
        r"""
        Compute the expected log density of the data, given a Gaussian
        distribution for the function values.

        if
            q(x,f) = N(Fmu, Fvar)  - Fmu: [N, D]  Fvar: [N, D]

        and this object represents

            p(y|x,f)  - Y: [N, 1]

        then this method computes

           ∫ (log p(y|x,f)) q(x,f) df.


        Here, we implement a default Monte Carlo quadrature routine.
        """

        def log_prob(F: TensorType, X_: TensorType, Y_: TensorType) -> tf.Tensor:
            return self.log_prob(X_, F, Y_)

        return tf.reduce_sum(
            self._mc_quadrature(log_prob, Fmu, Fvar, epsilon=epsilon, X_=X, Y_=Y), axis=-1
        )


# gpflow/likelihoods/misc.py

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from .base import MonteCarloLikelihood
from .scalar_continuous import Gaussian


class GaussianMC(MonteCarloLikelihood, Gaussian):  # type: ignore[misc]
    """
    Stochastic version of Gaussian likelihood for demonstration purposes only.
    """

    pass


# gpflow/likelihoods/multiclass.py

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Any, Optional

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp

from ..base import MeanAndVariance, Module, Parameter, TensorType
from ..config import default_float
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..quadrature import hermgauss
from ..utilities import to_default_float, to_default_int
from .base import Likelihood, MonteCarloLikelihood


class Softmax(MonteCarloLikelihood):
    """
    The soft-max multi-class likelihood.  It can only provide a stochastic
    Monte-Carlo estimate of the variational expectations term, but this
    added variance tends to be small compared to that due to mini-batching
    (when using the SVGP model).
    """

    def __init__(self, num_classes: int, **kwargs: Any) -> None:
        super().__init__(input_dim=None, latent_dim=num_classes, observation_dim=None, **kwargs)
        self.num_classes = self.latent_dim

    @inherit_check_shapes
    def _log_prob(self, X: TensorType, F: TensorType, Y: TensorType) -> tf.Tensor:
        return -tf.nn.sparse_softmax_cross_entropy_with_logits(logits=F, labels=Y[:, 0])

    @inherit_check_shapes
    def _conditional_mean(self, X: TensorType, F: TensorType) -> tf.Tensor:
        return tf.nn.softmax(F)

    @inherit_check_shapes
    def _conditional_variance(self, X: TensorType, F: TensorType) -> tf.Tensor:
        p = self.conditional_mean(X, F)
        return p - p ** 2


class RobustMax(Module):
    r"""
    This class represent a multi-class inverse-link function. Given a vector
    :math:`f=[f_1, f_2, ... f_k]`, the result of the mapping is

    .. math::

       y = [y_1 ... y_k]

    with

    .. math::

       y_i = \left\{
       \begin{array}{ll}
           (1-\varepsilon)   & \textrm{if} \ i = \textrm{argmax}(f) \\
           \varepsilon/(k-1) & \textrm{otherwise}
       \end{array}
       \right.

    where :math:`k` is the number of classes.
    """

    @check_shapes(
        "epsilon: []",
    )
    def __init__(self, num_classes: int, epsilon: float = 1e-3, **kwargs: Any) -> None:
        """
        `epsilon` represents the fraction of 'errors' in the labels of the
        dataset. This may be a hard parameter to optimize, so by default
        it is set un-trainable, at a small value.
        """
        super().__init__(**kwargs)
        transform = tfp.bijectors.Sigmoid()
        prior = tfp.distributions.Beta(to_default_float(0.2), to_default_float(5.0))
        self.epsilon = Parameter(epsilon, transform=transform, prior=prior, trainable=False)
        self.num_classes = num_classes
        self._squash = 1e-6

    @check_shapes(
        "F: [broadcast batch..., latent_dim]",
        "return: [batch..., latent_dim]",
    )
    def __call__(self, F: TensorType) -> tf.Tensor:
        i = tf.argmax(F, 1)
        return tf.one_hot(i, self.num_classes, 1.0 - self.epsilon, self.eps_k1)

    @property  # type: ignore[misc]  # Mypy doesn't like decorated properties.
    @check_shapes(
        "return: []",
    )
    def eps_k1(self) -> tf.Tensor:
        return self.epsilon / (self.num_classes - 1.0)

    @check_shapes(
        "val: [batch...]",
        "return: [batch...]",
    )
    def safe_sqrt(self, val: TensorType) -> tf.Tensor:
        return tf.sqrt(tf.maximum(val, 1e-10))

    @check_shapes(
        "Y: [broadcast batch..., observation_dim]",
        "mu: [broadcast batch..., latent_dim]",
        "var: [broadcast batch..., latent_dim]",
        "gh_x: [n_quad_points]",
        "gh_w: [n_quad_points]",
        "return: [batch..., observation_dim]",
    )
    def prob_is_largest(
        self, Y: TensorType, mu: TensorType, var: TensorType, gh_x: TensorType, gh_w: TensorType
    ) -> tf.Tensor:
        Y = to_default_int(Y)
        # work out what the mean and variance is of the indicated latent function.
        oh_on = tf.cast(
            tf.one_hot(tf.reshape(Y, (-1,)), self.num_classes, 1.0, 0.0), dtype=mu.dtype
        )
        mu_selected = tf.reduce_sum(oh_on * mu, 1)
        var_selected = tf.reduce_sum(oh_on * var, 1)

        # generate Gauss Hermite grid
        X = tf.reshape(mu_selected, (-1, 1)) + gh_x * tf.reshape(
            self.safe_sqrt(2.0 * var_selected), (-1, 1)
        )

        # compute the CDF of the Gaussian between the latent functions and the grid (including the selected function)
        dist = (tf.expand_dims(X, 1) - tf.expand_dims(mu, 2)) / tf.expand_dims(
            self.safe_sqrt(var), 2
        )
        cdfs = 0.5 * (1.0 + tf.math.erf(dist / np.sqrt(2.0)))

        cdfs = cdfs * (1 - 2 * self._squash) + self._squash

        # blank out all the distances on the selected latent function
        oh_off = tf.cast(
            tf.one_hot(tf.reshape(Y, (-1,)), self.num_classes, 0.0, 1.0), dtype=mu.dtype
        )
        cdfs = cdfs * tf.expand_dims(oh_off, 2) + tf.expand_dims(oh_on, 2)

        # take the product over the latent functions, and the sum over the GH grid.
        return tf.reduce_prod(cdfs, axis=[1]) @ tf.reshape(gh_w / np.sqrt(np.pi), (-1, 1))


class MultiClass(Likelihood):
    def __init__(
        self, num_classes: int, invlink: Optional[RobustMax] = None, **kwargs: Any
    ) -> None:
        """
        A likelihood for multi-way classification.  Currently the only valid
        choice of inverse-link function (invlink) is an instance of RobustMax.

        For most problems, the stochastic `Softmax` likelihood may be more
        appropriate (note that you then cannot use Scipy optimizer).
        """
        super().__init__(input_dim=None, latent_dim=num_classes, observation_dim=None, **kwargs)
        self.num_classes = num_classes
        self.num_gauss_hermite_points = 20

        if invlink is None:
            invlink = RobustMax(self.num_classes)

        if not isinstance(invlink, RobustMax):
            raise NotImplementedError

        self.invlink = invlink

    @inherit_check_shapes
    def _log_prob(self, X: TensorType, F: TensorType, Y: TensorType) -> tf.Tensor:
        hits = tf.equal(tf.expand_dims(tf.argmax(F, 1), 1), tf.cast(Y, tf.int64))
        yes = tf.ones(tf.shape(Y), dtype=default_float()) - self.invlink.epsilon
        no = tf.zeros(tf.shape(Y), dtype=default_float()) + self.invlink.eps_k1
        p = tf.where(hits, yes, no)
        return tf.reduce_sum(tf.math.log(p), axis=-1)

    @inherit_check_shapes
    def _variational_expectations(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        gh_x, gh_w = hermgauss(self.num_gauss_hermite_points)
        p = self.invlink.prob_is_largest(Y, Fmu, Fvar, gh_x, gh_w)
        ve = p * tf.math.log(1.0 - self.invlink.epsilon) + (1.0 - p) * tf.math.log(
            self.invlink.eps_k1
        )
        return tf.reduce_sum(ve, axis=-1)

    @inherit_check_shapes
    def _predict_mean_and_var(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType
    ) -> MeanAndVariance:
        possible_outputs = [
            tf.fill(tf.stack([tf.shape(Fmu)[0], 1]), np.array(i, dtype=np.int64))
            for i in range(self.num_classes)
        ]
        ps = [self._predict_non_logged_density(X, Fmu, Fvar, po) for po in possible_outputs]
        ps = tf.transpose(tf.stack([tf.reshape(p, (-1,)) for p in ps]))
        return ps, ps - tf.square(ps)

    @inherit_check_shapes
    def _predict_log_density(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        return tf.reduce_sum(
            tf.math.log(self._predict_non_logged_density(X, Fmu, Fvar, Y)), axis=-1
        )

    @check_shapes(
        "X: [broadcast batch..., input_dim]",
        "Fmu: [broadcast batch..., latent_dim]",
        "Fvar: [broadcast batch..., latent_dim]",
        "Y: [broadcast batch..., observation_dim]",
        "return: [batch..., observation_dim]",
    )
    def _predict_non_logged_density(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        gh_x, gh_w = hermgauss(self.num_gauss_hermite_points)
        p = self.invlink.prob_is_largest(Y, Fmu, Fvar, gh_x, gh_w)
        den = p * (1.0 - self.invlink.epsilon) + (1.0 - p) * (self.invlink.eps_k1)
        return den

    @inherit_check_shapes
    def _conditional_mean(self, X: TensorType, F: TensorType) -> tf.Tensor:
        return self.invlink(F)

    @inherit_check_shapes
    def _conditional_variance(self, X: TensorType, F: TensorType) -> tf.Tensor:
        p = self.conditional_mean(X, F)
        return p - tf.square(p)


# gpflow/likelihoods/multilatent.py

# Copyright 2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Any, Callable, Optional, Type

import tensorflow as tf
import tensorflow_probability as tfp

from ..base import TensorType
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..utilities import positive
from .base import QuadratureLikelihood


class MultiLatentLikelihood(QuadratureLikelihood):
    r"""
    A Likelihood which assumes that a single dimensional observation is driven
    by multiple latent GPs.

    Note that this implementation does not allow for taking into account
    covariance between outputs.
    """

    def __init__(self, latent_dim: int, **kwargs: Any) -> None:
        super().__init__(
            input_dim=None,
            latent_dim=latent_dim,
            observation_dim=1,
            **kwargs,
        )


class MultiLatentTFPConditional(MultiLatentLikelihood):
    """
    MultiLatent likelihood where the conditional distribution
    is given by a TensorFlow Probability Distribution.
    """

    def __init__(
        self,
        latent_dim: int,
        conditional_distribution: Callable[..., tfp.distributions.Distribution],
        **kwargs: Any,
    ):
        """
        :param latent_dim: number of arguments to the `conditional_distribution` callable
        :param conditional_distribution: function from F to a tfp Distribution,
            where F has shape [..., latent_dim]
        """
        super().__init__(latent_dim, **kwargs)
        self.conditional_distribution = conditional_distribution

    @inherit_check_shapes
    def _log_prob(self, X: TensorType, F: TensorType, Y: TensorType) -> tf.Tensor:
        """
        The log probability density log p(Y|F)

        :param F: function evaluation Tensor, with shape [..., latent_dim]
        :param Y: observation Tensor, with shape [..., 1]:
        :returns: log pdf, with shape [...]
        """
        return tf.squeeze(self.conditional_distribution(F).log_prob(Y), -1)

    @inherit_check_shapes
    def _conditional_mean(self, X: TensorType, F: TensorType) -> tf.Tensor:
        """
        The conditional marginal mean of Y|F: [E(Y₁|F)]

        :param F: function evaluation Tensor, with shape [..., latent_dim]
        :returns: mean [..., 1]
        """
        return self.conditional_distribution(F).mean()

    @inherit_check_shapes
    def _conditional_variance(self, X: TensorType, F: TensorType) -> tf.Tensor:
        """
        The conditional marginal variance of Y|F: [Var(Y₁|F)]

        :param F: function evaluation Tensor, with shape [..., latent_dim]
        :returns: variance [..., 1]
        """
        return self.conditional_distribution(F).variance()


class HeteroskedasticTFPConditional(MultiLatentTFPConditional):
    """
    Heteroskedastic Likelihood where the conditional distribution
    is given by a TensorFlow Probability Distribution.
    The `loc` and `scale` of the distribution are given by a
    two-dimensional multi-output GP.
    """

    def __init__(
        self,
        distribution_class: Type[tfp.distributions.Distribution] = tfp.distributions.Normal,
        scale_transform: Optional[tfp.bijectors.Bijector] = None,
        **kwargs: Any,
    ) -> None:
        """
        :param distribution_class: distribution class parameterized by `loc` and `scale`
            as first and second argument, respectively.
        :param scale_transform: callable/bijector applied to the latent
            function modelling the scale to ensure its positivity.
            Typically, `tf.exp` or `tf.softplus`, but can be any function f: R -> R^+. Defaults to exp if not explicitly specified.
        """
        if scale_transform is None:
            scale_transform = positive(base="exp")
        self.scale_transform = scale_transform

        @check_shapes(
            "F: [batch..., 2]",
        )
        def conditional_distribution(F: TensorType) -> tfp.distributions.Distribution:
            loc = F[..., :1]
            scale = self.scale_transform(F[..., 1:])
            return distribution_class(loc, scale)

        super().__init__(
            latent_dim=2,
            conditional_distribution=conditional_distribution,
            **kwargs,
        )


# gpflow/likelihoods/scalar_continuous.py

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from math import sqrt
from typing import Any, Callable, Optional

import numpy as np
import tensorflow as tf

from .. import logdensities
from ..base import MeanAndVariance, TensorType
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..utilities.parameter_or_function import (
    ConstantOrFunction,
    ParameterOrFunction,
    evaluate_parameter_or_function,
    prepare_parameter_or_function,
)
from .base import ScalarLikelihood
from .utils import inv_probit

DEFAULT_LOWER_BOUND = 1e-6


class Gaussian(ScalarLikelihood):
    r"""
    The Gaussian likelihood is appropriate where uncertainties associated with
    the data are believed to follow a normal distribution, with constant
    variance.

    Very small uncertainties can lead to numerical instability during the
    optimization process. A lower bound of 1e-6 is therefore imposed on the
    likelihood variance by default.
    """

    def __init__(
        self,
        variance: Optional[ConstantOrFunction] = None,
        *,
        scale: Optional[ConstantOrFunction] = None,
        variance_lower_bound: float = DEFAULT_LOWER_BOUND,
        **kwargs: Any,
    ) -> None:
        """
        :param variance: The noise variance; must be greater than
            ``variance_lower_bound``. This is mutually exclusive with `scale`.
        :param scale: The noise scale; must be greater than
            ``sqrt(variance_lower_bound)``. This is mutually exclusive with `variance`.
        :param variance_lower_bound: The lower (exclusive) bound of ``variance``.
        :param kwargs: Keyword arguments forwarded to :class:`ScalarLikelihood`.
        """
        super().__init__(**kwargs)

        self.variance_lower_bound = variance_lower_bound
        self.scale_lower_bound = sqrt(variance_lower_bound)
        if scale is None:
            if variance is None:
                variance = 1.0
            self.variance: Optional[ParameterOrFunction] = prepare_parameter_or_function(
                variance, lower_bound=self.variance_lower_bound
            )
            self.scale: Optional[ParameterOrFunction] = None
        else:
            if variance is None:
                self.variance = None
                self.scale = prepare_parameter_or_function(
                    scale, lower_bound=self.scale_lower_bound
                )
            else:
                assert False, "Cannot set both `variance` and `scale`."

    @check_shapes(
        "X: [batch..., N, D]",
        "return: [broadcast batch..., broadcast N, broadcast P]",
    )
    def _variance(self, X: TensorType) -> tf.Tensor:
        if self.variance is not None:
            return evaluate_parameter_or_function(
                self.variance, X, lower_bound=self.variance_lower_bound
            )
        else:
            assert self.scale is not None  # For mypy.
            return (
                evaluate_parameter_or_function(self.scale, X, lower_bound=self.scale_lower_bound)
                ** 2
            )

    @check_shapes(
        "X: [batch..., N, D]",
        "return: [batch..., N, 1]",
    )
    def variance_at(self, X: TensorType) -> tf.Tensor:
        variance = self._variance(X)
        shape = tf.concat([tf.shape(X)[:-1], [1]], 0)
        return tf.broadcast_to(variance, shape)

    @inherit_check_shapes
    def _scalar_log_prob(self, X: TensorType, F: TensorType, Y: TensorType) -> tf.Tensor:
        return logdensities.gaussian(Y, F, self._variance(X))

    @inherit_check_shapes
    def _conditional_mean(self, X: TensorType, F: TensorType) -> tf.Tensor:  # pylint: disable=R0201
        return tf.identity(F)

    @inherit_check_shapes
    def _conditional_variance(self, X: TensorType, F: TensorType) -> tf.Tensor:
        shape = tf.shape(F)
        return tf.broadcast_to(self._variance(X), shape)

    @inherit_check_shapes
    def _predict_mean_and_var(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType
    ) -> MeanAndVariance:
        return tf.identity(Fmu), Fvar + self._variance(X)

    @inherit_check_shapes
    def _predict_log_density(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        return tf.reduce_sum(logdensities.gaussian(Y, Fmu, Fvar + self._variance(X)), axis=-1)

    @inherit_check_shapes
    def _variational_expectations(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        variance = self._variance(X)
        return tf.reduce_sum(
            -0.5 * np.log(2 * np.pi)
            - 0.5 * tf.math.log(variance)
            - 0.5 * ((Y - Fmu) ** 2 + Fvar) / variance,
            axis=-1,
        )


class Exponential(ScalarLikelihood):
    def __init__(self, invlink: Callable[[tf.Tensor], tf.Tensor] = tf.exp, **kwargs: Any) -> None:
        super().__init__(**kwargs)
        self.invlink = invlink

    @inherit_check_shapes
    def _scalar_log_prob(self, X: TensorType, F: TensorType, Y: TensorType) -> tf.Tensor:
        return logdensities.exponential(Y, self.invlink(F))

    @inherit_check_shapes
    def _conditional_mean(self, X: TensorType, F: TensorType) -> tf.Tensor:
        return self.invlink(F)

    @inherit_check_shapes
    def _conditional_variance(self, X: TensorType, F: TensorType) -> tf.Tensor:
        return tf.square(self.invlink(F))

    @inherit_check_shapes
    def _variational_expectations(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        if self.invlink is tf.exp:
            return tf.reduce_sum(-tf.exp(-Fmu + Fvar / 2) * Y - Fmu, axis=-1)
        return super()._variational_expectations(X, Fmu, Fvar, Y)


class StudentT(ScalarLikelihood):
    def __init__(
        self,
        scale: ConstantOrFunction = 1.0,
        df: float = 3.0,
        scale_lower_bound: float = DEFAULT_LOWER_BOUND,
        **kwargs: Any,
    ) -> None:
        """
        :param scale float: scale parameter
        :param df float: degrees of freedom
        """
        super().__init__(**kwargs)
        self.df = df
        self.scale_lower_bound = scale_lower_bound
        self.scale = prepare_parameter_or_function(scale, lower_bound=self.scale_lower_bound)

    @check_shapes(
        "X: [batch..., N, D]",
        "return: [broadcast batch..., broadcast N, broadcast P]",
    )
    def _scale(self, X: TensorType) -> tf.Tensor:
        return evaluate_parameter_or_function(self.scale, X, lower_bound=self.scale_lower_bound)

    @inherit_check_shapes
    def _scalar_log_prob(self, X: TensorType, F: TensorType, Y: TensorType) -> tf.Tensor:
        return logdensities.student_t(Y, F, self._scale(X), self.df)

    @inherit_check_shapes
    def _conditional_mean(self, X: TensorType, F: TensorType) -> tf.Tensor:
        return F

    @inherit_check_shapes
    def _conditional_variance(self, X: TensorType, F: TensorType) -> tf.Tensor:
        shape = tf.shape(F)
        var = (self._scale(X) ** 2) * (self.df / (self.df - 2.0))
        return tf.broadcast_to(var, shape)


class Gamma(ScalarLikelihood):
    """
    Use the transformed GP to give the *scale* (inverse rate) of the Gamma
    """

    def __init__(
        self,
        invlink: Callable[[tf.Tensor], tf.Tensor] = tf.exp,
        shape: ConstantOrFunction = 1.0,
        shape_lower_bound: float = DEFAULT_LOWER_BOUND,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        self.invlink = invlink
        self.shape_lower_bound = shape_lower_bound
        self.shape = prepare_parameter_or_function(shape, lower_bound=self.shape_lower_bound)

    @check_shapes(
        "X: [batch..., N, D]",
        "return: [broadcast batch..., broadcast N, broadcast P]",
    )
    def _shape(self, X: TensorType) -> tf.Tensor:
        return evaluate_parameter_or_function(self.shape, X, lower_bound=self.shape_lower_bound)

    @inherit_check_shapes
    def _scalar_log_prob(self, X: TensorType, F: TensorType, Y: TensorType) -> tf.Tensor:
        return logdensities.gamma(Y, self._shape(X), self.invlink(F))

    @inherit_check_shapes
    def _conditional_mean(self, X: TensorType, F: TensorType) -> tf.Tensor:
        return self._shape(X) * self.invlink(F)

    @inherit_check_shapes
    def _conditional_variance(self, X: TensorType, F: TensorType) -> tf.Tensor:
        scale = self.invlink(F)
        return self._shape(X) * (scale ** 2)

    @inherit_check_shapes
    def _variational_expectations(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        if self.invlink is tf.exp:
            shape = self._shape(X)
            return tf.reduce_sum(
                -shape * Fmu
                - tf.math.lgamma(shape)
                + (shape - 1.0) * tf.math.log(Y)
                - Y * tf.exp(-Fmu + Fvar / 2.0),
                axis=-1,
            )
        else:
            return super()._variational_expectations(X, Fmu, Fvar, Y)


class Beta(ScalarLikelihood):
    """
    This uses a reparameterisation of the Beta density. We have the mean of the
    Beta distribution given by the transformed process:

        m = invlink(f)

    and a scale parameter. The familiar α, β parameters are given by

        m     = α / (α + β)
        scale = α + β

    so:
        α = scale * m
        β  = scale * (1-m)
    """

    def __init__(
        self,
        invlink: Callable[[tf.Tensor], tf.Tensor] = inv_probit,
        scale: ConstantOrFunction = 1.0,
        scale_lower_bound: float = DEFAULT_LOWER_BOUND,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        self.scale_lower_bound = DEFAULT_LOWER_BOUND
        self.scale = prepare_parameter_or_function(scale, lower_bound=self.scale_lower_bound)
        self.invlink = invlink

    @check_shapes(
        "X: [batch..., N, D]",
        "return: [broadcast batch..., broadcast N, broadcast P]",
    )
    def _scale(self, X: TensorType) -> tf.Tensor:
        return evaluate_parameter_or_function(self.scale, X, lower_bound=self.scale_lower_bound)

    @inherit_check_shapes
    def _scalar_log_prob(self, X: TensorType, F: TensorType, Y: TensorType) -> tf.Tensor:
        mean = self.invlink(F)
        scale = self._scale(X)
        alpha = mean * scale
        beta = scale - alpha
        return logdensities.beta(Y, alpha, beta)

    @inherit_check_shapes
    def _conditional_mean(self, X: TensorType, F: TensorType) -> tf.Tensor:
        return self.invlink(F)

    @inherit_check_shapes
    def _conditional_variance(self, X: TensorType, F: TensorType) -> tf.Tensor:
        mean = self.invlink(F)
        var = (mean - tf.square(mean)) / (self._scale(X) + 1.0)
        shape = tf.shape(F)
        return tf.broadcast_to(var, shape)


# gpflow/likelihoods/scalar_discrete.py

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Any, Callable

import numpy as np
import tensorflow as tf

from .. import logdensities
from ..base import AnyNDArray, MeanAndVariance, Parameter, TensorType
from ..config import default_float
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..utilities import positive, to_default_int
from .base import ScalarLikelihood
from .utils import inv_probit


class Poisson(ScalarLikelihood):
    r"""
    Poisson likelihood for use with count data, where the rate is given by the (transformed) GP.

    let g(.) be the inverse-link function, then this likelihood represents

    p(yᵢ | fᵢ) = Poisson(yᵢ | g(fᵢ) * binsize)

    Note:binsize
    For use in a Log Gaussian Cox process (doubly stochastic model) where the
    rate function of an inhomogeneous Poisson process is given by a GP.  The
    intractable likelihood can be approximated via a Riemann sum (with bins
    of size 'binsize') and using this Poisson likelihood.
    """

    def __init__(
        self,
        invlink: Callable[[tf.Tensor], tf.Tensor] = tf.exp,
        binsize: float = 1.0,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        self.invlink = invlink
        self.binsize: AnyNDArray = np.array(binsize, dtype=default_float())

    @inherit_check_shapes
    def _scalar_log_prob(self, X: TensorType, F: TensorType, Y: TensorType) -> tf.Tensor:
        return logdensities.poisson(Y, self.invlink(F) * self.binsize)

    @inherit_check_shapes
    def _conditional_variance(self, X: TensorType, F: TensorType) -> tf.Tensor:
        return self.invlink(F) * self.binsize

    @inherit_check_shapes
    def _conditional_mean(self, X: TensorType, F: TensorType) -> tf.Tensor:
        return self.invlink(F) * self.binsize

    @inherit_check_shapes
    def _variational_expectations(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        if self.invlink is tf.exp:
            return tf.reduce_sum(
                Y * Fmu
                - tf.exp(Fmu + Fvar / 2) * self.binsize
                - tf.math.lgamma(Y + 1)
                + Y * tf.math.log(self.binsize),
                axis=-1,
            )
        return super()._variational_expectations(X, Fmu, Fvar, Y)


class Bernoulli(ScalarLikelihood):
    def __init__(
        self, invlink: Callable[[tf.Tensor], tf.Tensor] = inv_probit, **kwargs: Any
    ) -> None:
        super().__init__(**kwargs)
        self.invlink = invlink

    @inherit_check_shapes
    def _scalar_log_prob(self, X: TensorType, F: TensorType, Y: TensorType) -> tf.Tensor:
        return logdensities.bernoulli(Y, self.invlink(F))

    @inherit_check_shapes
    def _predict_mean_and_var(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType
    ) -> MeanAndVariance:
        if self.invlink is inv_probit:
            p = inv_probit(Fmu / tf.sqrt(1 + Fvar))
            return p, p - tf.square(p)
        else:
            # for other invlink, use quadrature
            return super()._predict_mean_and_var(X, Fmu, Fvar)

    @inherit_check_shapes
    def _predict_log_density(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        p = self.predict_mean_and_var(X, Fmu, Fvar)[0]
        return tf.reduce_sum(logdensities.bernoulli(Y, p), axis=-1)

    @inherit_check_shapes
    def _conditional_mean(self, X: TensorType, F: TensorType) -> tf.Tensor:
        return self.invlink(F)

    @inherit_check_shapes
    def _conditional_variance(self, X: TensorType, F: TensorType) -> tf.Tensor:
        p = self.conditional_mean(X, F)
        return p - (p ** 2)


class Ordinal(ScalarLikelihood):
    """
    A likelihood for doing ordinal regression.

    The data are integer values from 0 to k, and the user must specify (k-1)
    'bin edges' which define the points at which the labels switch. Let the bin
    edges be [a₀, a₁, ... aₖ₋₁], then the likelihood is

    p(Y=0|F) = ɸ((a₀ - F) / σ)
    p(Y=1|F) = ɸ((a₁ - F) / σ) - ɸ((a₀ - F) / σ)
    p(Y=2|F) = ɸ((a₂ - F) / σ) - ɸ((a₁ - F) / σ)
    ...
    p(Y=K|F) = 1 - ɸ((aₖ₋₁ - F) / σ)

    where ɸ is the cumulative density function of a Gaussian (the inverse probit
    function) and σ is a parameter to be learned.

    A reference is :cite:t:`chu2005gaussian`.
    """

    @check_shapes(
        "bin_edges: [num_bins_minus_1]",
    )
    def __init__(self, bin_edges: AnyNDArray, **kwargs: Any) -> None:
        """
        bin_edges is a numpy array specifying at which function value the
        output label should switch. If the possible Y values are 0...K, then
        the size of bin_edges should be (K-1).
        """
        super().__init__(**kwargs)
        self.bin_edges = bin_edges
        self.num_bins = bin_edges.size + 1
        self.sigma = Parameter(1.0, transform=positive())

    @inherit_check_shapes
    def _scalar_log_prob(self, X: TensorType, F: TensorType, Y: TensorType) -> tf.Tensor:
        Y = to_default_int(Y)
        scaled_bins_left = tf.concat([self.bin_edges / self.sigma, np.array([np.inf])], 0)
        scaled_bins_right = tf.concat([np.array([-np.inf]), self.bin_edges / self.sigma], 0)
        selected_bins_left = tf.gather(scaled_bins_left, Y)
        selected_bins_right = tf.gather(scaled_bins_right, Y)

        return tf.math.log(
            inv_probit(selected_bins_left - F / self.sigma)
            - inv_probit(selected_bins_right - F / self.sigma)
            + 1e-6
        )

    @check_shapes(
        "F: [batch..., latent_dim]",
        "return: [batch_and_latent_dim, num_bins]",
    )
    def _make_phi(self, F: TensorType) -> tf.Tensor:
        """
        A helper function for making predictions. Constructs a probability
        matrix where each row output the probability of the corresponding
        label, and the rows match the entries of F.

        Note that a matrix of F values is flattened.
        """
        scaled_bins_left = tf.concat([self.bin_edges / self.sigma, np.array([np.inf])], 0)
        scaled_bins_right = tf.concat([np.array([-np.inf]), self.bin_edges / self.sigma], 0)
        return inv_probit(scaled_bins_left - tf.reshape(F, (-1, 1)) / self.sigma) - inv_probit(
            scaled_bins_right - tf.reshape(F, (-1, 1)) / self.sigma
        )

    @inherit_check_shapes
    def _conditional_mean(self, X: TensorType, F: TensorType) -> tf.Tensor:
        phi = self._make_phi(F)
        Ys = tf.reshape(np.arange(self.num_bins, dtype=default_float()), (-1, 1))
        return tf.reshape(tf.linalg.matmul(phi, Ys), tf.shape(F))

    @inherit_check_shapes
    def _conditional_variance(self, X: TensorType, F: TensorType) -> tf.Tensor:
        phi = self._make_phi(F)
        Ys = tf.reshape(np.arange(self.num_bins, dtype=default_float()), (-1, 1))
        E_y = phi @ Ys
        E_y2 = phi @ (Ys ** 2)
        return tf.reshape(E_y2 - E_y ** 2, tf.shape(F))


# gpflow/likelihoods/utils.py

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
import tensorflow as tf

from ..experimental.check_shapes import check_shapes


@check_shapes(
    "x: [batch...]",
    "return: [batch...]",
)
def inv_probit(x: tf.Tensor) -> tf.Tensor:
    jitter = 1e-3  # ensures output is strictly between 0 and 1
    return 0.5 * (1.0 + tf.math.erf(x / np.sqrt(2.0))) * (1 - 2 * jitter) + jitter


# gpflow/models/__init__.py

from .cglb import CGLB
from .gplvm import GPLVM, BayesianGPLVM
from .gpmc import GPMC
from .gpr import GPR
from .model import BayesianModel, GPModel
from .sgpmc import SGPMC
from .sgpr import GPRFITC, SGPR
from .svgp import SVGP
from .training_mixins import ExternalDataTrainingLossMixin, InternalDataTrainingLossMixin
from .util import maximum_log_likelihood_objective, training_loss, training_loss_closure
from .vgp import VGP, VGPOpperArchambeau

__all__ = [
    "BayesianGPLVM",
    "BayesianModel",
    "CGLB",
    "ExternalDataTrainingLossMixin",
    "GPLVM",
    "GPMC",
    "GPModel",
    "GPR",
    "GPRFITC",
    "InternalDataTrainingLossMixin",
    "SGPMC",
    "SGPR",
    "SVGP",
    "VGP",
    "VGPOpperArchambeau",
    "cglb",
    "gplvm",
    "gpmc",
    "gpr",
    "maximum_log_likelihood_objective",
    "model",
    "sgpmc",
    "sgpr",
    "svgp",
    "training_loss",
    "training_loss_closure",
    "training_mixins",
    "util",
    "vgp",
]


# gpflow/models/cglb.py

# Copyright 2021 the GPflow authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Any, List, NamedTuple, Optional, Tuple

import tensorflow as tf

from ..base import InputData, MeanAndVariance, Parameter, RegressionData, TensorType
from ..config import default_float, default_int
from ..covariances import Kuf
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..utilities import add_noise_cov, assert_params_false, to_default_float
from .sgpr import SGPR


class CGLB(SGPR):
    """
    Conjugate Gradient Lower Bound.

    The key reference is :cite:t:`pmlr-v139-artemev21a`.

    :param cg_tolerance: Determines accuracy to which conjugate
        gradient is run when evaluating the elbo. Running more
        iterations of CG would increase the ELBO by at most
        `cg_tolerance`.

    :param max_cg_iters: Maximum number of iterations of CG to run
        per evaluation of the ELBO (or mean prediction).

    :param restart_cg_iters: How frequently to restart the CG iteration.
        Can be useful to avoid build up of numerical errors when
        many steps of CG are run.

    :param v_grad_optimization: If False, in every evaluation of the
        ELBO, CG is run to select a new auxilary vector `v`. If
        False, no CG is run when evaluating the ELBO but
        gradients with respect to `v` are tracked so that it can
        be optimized jointly with other parameters.

    """

    @check_shapes(
        "data[0]: [N, D]",
        "data[1]: [N, P]",
    )
    def __init__(
        self,
        data: RegressionData,
        *args: Any,
        cg_tolerance: float = 1.0,
        max_cg_iters: int = 100,
        restart_cg_iters: int = 40,
        v_grad_optimization: bool = False,
        **kwargs: Any,
    ) -> None:
        super().__init__(data, *args, **kwargs)
        n, b = self.data[1].shape
        self._v = Parameter(tf.zeros((b, n), dtype=default_float()), trainable=v_grad_optimization)
        self._cg_tolerance = cg_tolerance
        self._max_cg_iters = max_cg_iters
        self._restart_cg_iters = restart_cg_iters

    @property  # type: ignore[misc]  # Mypy doesn't like decorated properties.
    @check_shapes(
        "return: [N, P]",
    )
    def aux_vec(self) -> Parameter:
        return self._v

    @inherit_check_shapes
    def logdet_term(self, common: SGPR.CommonTensors) -> tf.Tensor:
        r"""
        Compute a lower bound on :math:`-0.5 * \log |K + σ²I|` based on a
        low-rank approximation to K.

        .. math::

           \log |K + σ²I| <= \log |Q + σ²I| + n * \log(1 + \textrm{tr}(K - Q)/(σ²n)).

        This bound is at least as tight as

        .. math::

           \log |K + σ²I| <=  \log |Q + σ²I| + \textrm{tr}(K - Q)/σ²,

        which appears in SGPR.
        """
        LB = common.LB
        AAT = common.AAT
        x, y = self.data
        num_data = to_default_float(tf.shape(y)[0])
        output_dim = to_default_float(tf.shape(y)[1])
        sigma_sq = self.likelihood.variance

        kdiag = self.kernel(x, full_cov=False)
        # t / σ²
        trace = tf.reduce_sum(kdiag) / sigma_sq - tf.reduce_sum(tf.linalg.diag_part(AAT))
        logdet_b = tf.reduce_sum(tf.math.log(tf.linalg.diag_part(LB)))
        logsigma_sq = num_data * tf.math.log(sigma_sq)

        # Correction term from Jensen's inequality
        logtrace = num_data * tf.math.log(1 + trace / num_data)
        return -output_dim * (logdet_b + 0.5 * logsigma_sq + 0.5 * logtrace)

    @inherit_check_shapes
    def quad_term(self, common: SGPR.CommonTensors) -> tf.Tensor:
        """
        Computes a lower bound on the quadratic term in the log
        marginal likelihood of conjugate GPR. The bound is based on
        an auxiliary vector, v. For :math:`Q ≺ K` and :math:`r=y - Kv`

        .. math::
            -0.5 * (rᵀQ⁻¹r + 2yᵀv - vᵀ K v ) <= -0.5 * yᵀK⁻¹y <= -0.5 * (2yᵀv - vᵀKv).

        Equality holds if :math:`r=0`, i.e. :math:`v = K⁻¹y`.

        If `self.aux_vec` is trainable, gradients are computed with
        respect to :math:`v` as well and :math:`v` can be optimized
        using gradient based methods.

        Otherwise, :math:`v` is updated with the method of conjugate
        gradients (CG). CG is run until :math:`0.5 * rᵀQ⁻¹r <= ϵ`,
        which ensures that the maximum bias due to this term is not
        more than :math:`ϵ`. The :math:`ϵ` is the CG tolerance.
        """
        x, y = self.data
        err = y - self.mean_function(x)
        sigma_sq = self.likelihood.variance
        K = add_noise_cov(self.kernel.K(x), sigma_sq)

        A = common.A
        LB = common.LB
        preconditioner = NystromPreconditioner(A, LB, sigma_sq)
        err_t = tf.transpose(err)

        v_init = self.aux_vec
        if not v_init.trainable:
            v = cglb_conjugate_gradient(
                K,
                err_t,
                v_init,
                preconditioner,
                self._cg_tolerance,
                self._max_cg_iters,
                self._restart_cg_iters,
            )
        else:
            v = v_init

        Kv = v @ K
        r = err_t - Kv
        _, error_bound = preconditioner(r)
        lb = tf.reduce_sum(v * (r + 0.5 * Kv))
        ub = lb + 0.5 * error_bound

        if not v_init.trainable:
            v_init.assign(v)

        return -ub

    @inherit_check_shapes
    def predict_f(
        self,
        Xnew: InputData,
        full_cov: bool = False,
        full_output_cov: bool = False,
        cg_tolerance: Optional[float] = 1e-3,
    ) -> MeanAndVariance:
        """
        The posterior mean for CGLB model is given by

        .. :math::
            m(xs) = K_{sf}v + Q_{ff}Q⁻¹r

        where :math:`r = y - K v`  is the residual from CG.

        Note that when :math:`v=0`, this agree with the SGPR mean,
        while if :math:`v = K⁻¹ y`, then :math:`r=0`, and the exact
        GP mean is recovered.

        :param cg_tolerance: float or None: If None, the cached value of
            :math:`v` is used. If float, conjugate gradient is run until :math:`rᵀQ⁻¹r < ϵ`.
        """
        assert_params_false(self.predict_f, full_output_cov=full_output_cov)

        x, y = self.data
        err = y - self.mean_function(x)
        kxx = self.kernel(x, x)
        ksf = self.kernel(Xnew, x)
        sigma_sq = self.likelihood.variance
        sigma = tf.sqrt(sigma_sq)
        iv = self.inducing_variable
        kernel = self.kernel
        matmul = tf.linalg.matmul
        trisolve = tf.linalg.triangular_solve

        kmat = add_noise_cov(kxx, sigma_sq)

        common = self._common_calculation()
        A, LB, L = common.A, common.LB, common.L

        v = self.aux_vec
        if cg_tolerance is not None:
            preconditioner = NystromPreconditioner(A, LB, sigma_sq)
            err_t = tf.transpose(err)
            v = cglb_conjugate_gradient(
                kmat,
                err_t,
                v,
                preconditioner,
                cg_tolerance,
                self._max_cg_iters,
                self._restart_cg_iters,
            )

        cg_mean = matmul(ksf, v, transpose_b=True)
        res = err - matmul(kmat, v, transpose_b=True)

        Kus = Kuf(iv, kernel, Xnew)
        Ares = matmul(A, res)  # The god of war!
        c = trisolve(LB, Ares, lower=True) / sigma
        tmp1 = trisolve(L, Kus, lower=True)
        tmp2 = trisolve(LB, tmp1, lower=True)
        sgpr_mean = matmul(tmp2, c, transpose_a=True)

        if full_cov:
            var = (
                kernel(Xnew)
                + matmul(tmp2, tmp2, transpose_a=True)
                - matmul(tmp1, tmp1, transpose_a=True)
            )
            var = tf.tile(var[None, ...], [self.num_latent_gps, 1, 1])  # [P, N, N]
        else:
            var = (
                kernel(Xnew, full_cov=False)
                + tf.reduce_sum(tf.square(tmp2), 0)
                - tf.reduce_sum(tf.square(tmp1), 0)
            )
            var = tf.tile(var[:, None], [1, self.num_latent_gps])

        mean = sgpr_mean + cg_mean + self.mean_function(Xnew)
        return mean, var

    @inherit_check_shapes
    def predict_y(
        self,
        Xnew: InputData,
        full_cov: bool = False,
        full_output_cov: bool = False,
        cg_tolerance: Optional[float] = 1e-3,
    ) -> MeanAndVariance:
        """
        Compute the mean and variance of the held-out data at the
        input points.
        """
        assert_params_false(self.predict_y, full_cov=full_cov, full_output_cov=full_output_cov)

        f_mean, f_var = self.predict_f(
            Xnew, full_cov=full_cov, full_output_cov=full_output_cov, cg_tolerance=cg_tolerance
        )
        return self.likelihood.predict_mean_and_var(Xnew, f_mean, f_var)

    @inherit_check_shapes
    def predict_log_density(
        self,
        data: RegressionData,
        full_cov: bool = False,
        full_output_cov: bool = False,
        cg_tolerance: Optional[float] = 1e-3,
    ) -> tf.Tensor:
        """
        Compute the log density of the data at the new data points.
        """
        assert_params_false(
            self.predict_log_density, full_cov=full_cov, full_output_cov=full_output_cov
        )

        x, y = data
        f_mean, f_var = self.predict_f(
            x, full_cov=full_cov, full_output_cov=full_output_cov, cg_tolerance=cg_tolerance
        )
        return self.likelihood.predict_log_density(x, f_mean, f_var, y)


class NystromPreconditioner:
    """
    Preconditioner of the form :math:`Q=(Q_ff + σ²I)⁻¹`,
    where L is lower triangular with :math: `LLᵀ = Kᵤᵤ`
    :math:`A = σ⁻²L⁻¹Kᵤₓ` and :math:`B = AAᵀ + I = LᵦLᵦᵀ`
    """

    @check_shapes(
        "A: [M, N]",
        "LB: [M, M]",
    )
    def __init__(self, A: tf.Tensor, LB: tf.Tensor, sigma_sq: float) -> None:
        self.A = A
        self.LB = LB
        self.sigma_sq = sigma_sq

    @check_shapes(
        "v: [B, N]",
        "return[0]: [B, N]",
        "return[1]: []",
    )
    def __call__(self, v: TensorType) -> Tuple[tf.Tensor, tf.Tensor]:
        """
        Computes :math:`vᵀQ^{-1}` and `vᵀQ^{-1}v`. Note that this is
        implemented as multipication of a row vector on the right.

        :param v: Vector we want to backsolve.
        """
        sigma_sq = self.sigma_sq
        A = self.A
        LB = self.LB

        trans = tf.transpose
        trisolve = tf.linalg.triangular_solve
        matmul = tf.linalg.matmul

        v = trans(v)
        Av = matmul(A, v)
        LBinvAv = trisolve(LB, Av)
        LBinvtLBinvAv = trisolve(trans(LB), LBinvAv, lower=False)

        rv = v - matmul(A, LBinvtLBinvAv, transpose_a=True)
        vtrv = tf.reduce_sum(rv * v)
        return trans(rv) / sigma_sq, vtrv / sigma_sq


@check_shapes(
    "K: [N, N]",
    "b: [B, N]",
    "initial: [P, N]",
    "return: [P, N]",
)
def cglb_conjugate_gradient(
    K: TensorType,
    b: TensorType,
    initial: TensorType,
    preconditioner: NystromPreconditioner,
    cg_tolerance: float,
    max_steps: int,
    restart_cg_step: int,
) -> tf.Tensor:
    """
    Conjugate gradient algorithm used in CGLB model. The method of
    conjugate gradient (Hestenes and Stiefel, 1952) produces a
    sequence of vectors :math:`v_0, v_1, v_2, ..., v_N` such that
    :math:`v_0` = initial, and (in exact arithmetic)
    :math:`Kv_n = b`. In practice, the v_i often converge quickly to
    approximate :math:`K^{-1}b`, and the algorithm can be stopped
    without running N iterations.

    We assume the preconditioner, :math:`Q`, satisfies :math:`Q ≺ K`,
    and stop the algorithm when :math:`r_i = b - Kv_i` satisfies
    :math:`||rᵢᵀ||_{Q⁻¹r}^2 = rᵢᵀQ⁻¹rᵢ <= ϵ`.

    :param K: Matrix we want to backsolve from. Must be PSD.
    :param b: Vector we want to backsolve.
    :param initial: Initial vector solution.
    :param preconditioner: Preconditioner function.
    :param cg_tolerance: Expected maximum error. This value is used
        as a decision boundary against stopping criteria.
    :param max_steps: Maximum number of CG iterations.
    :param restart_cg_step: Restart step at which the CG resets the
        internal state to the initial position using the currect
        solution vector :math:`v`. Can help avoid build up of
        numerical errors.

    :return: `v` where `v` approximately satisfies :math:`Kv = b`.
    """

    class CGState(NamedTuple):
        i: tf.Tensor
        v: tf.Tensor
        r: tf.Tensor
        p: tf.Tensor
        rz: tf.Tensor

    def stopping_criterion(state: CGState) -> tf.Tensor:
        return (0.5 * state.rz > cg_tolerance) and (state.i < max_steps)

    def cg_step(state: CGState) -> List[CGState]:
        Ap = state.p @ K
        denom = tf.reduce_sum(state.p * Ap, axis=-1)
        gamma = state.rz / denom
        v = state.v + gamma * state.p
        i = state.i + 1
        r = tf.cond(
            state.i % restart_cg_step == restart_cg_step - 1,
            lambda: b - v @ K,
            lambda: state.r - gamma * Ap,
        )
        z, new_rz = preconditioner(r)
        p = tf.cond(
            state.i % restart_cg_step == restart_cg_step - 1,
            lambda: z,
            lambda: z + state.p * new_rz / state.rz,
        )
        return [CGState(i, v, r, p, new_rz)]

    Kv = initial @ K
    r = b - Kv
    z, rz = preconditioner(r)
    p = z
    i = tf.constant(0, dtype=default_int())
    initial_state = CGState(i, initial, r, p, rz)
    final_state = tf.while_loop(stopping_criterion, cg_step, [initial_state])
    final_state = tf.nest.map_structure(tf.stop_gradient, final_state)
    return final_state[0].v


# gpflow/models/gplvm.py

# Copyright 2016-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Optional

import numpy as np
import tensorflow as tf

from .. import covariances, kernels, likelihoods
from ..base import InputData, MeanAndVariance, OutputData, Parameter, RegressionData, TensorType
from ..config import default_float, default_jitter
from ..expectations import expectation
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..inducing_variables import InducingPoints
from ..kernels import Kernel
from ..mean_functions import MeanFunction, Zero
from ..probability_distributions import DiagonalGaussian
from ..utilities import assert_params_false, positive, to_default_float
from ..utilities.ops import pca_reduce
from .gpr import GPR
from .model import GPModel
from .training_mixins import InternalDataTrainingLossMixin
from .util import InducingVariablesLike, data_input_to_tensor, inducingpoint_wrapper


class GPLVM(GPR):
    """
    Standard GPLVM where the likelihood can be optimised with respect to the latent X.
    """

    @check_shapes(
        "data: [N, P]",
        "X_data_mean: [N, Q]",
    )
    def __init__(
        self,
        data: OutputData,
        latent_dim: int,
        X_data_mean: Optional[tf.Tensor] = None,
        kernel: Optional[Kernel] = None,
        mean_function: Optional[MeanFunction] = None,
    ):
        """
        Initialise GPLVM object. This method only works with a Gaussian likelihood.

        :param data: y data matrix.
        :param latent_dim: the number of latent dimensions (Q)
        :param X_data_mean: latent positions, for the initialisation of the latent space.
        :param kernel: kernel specification, by default Squared Exponential
        :param mean_function: mean function, by default None.
        """
        if X_data_mean is None:
            X_data_mean = pca_reduce(data, latent_dim)

        num_latent_gps = X_data_mean.shape[1]
        if num_latent_gps != latent_dim:
            msg = "Passed in number of latent {0} does not match initial X {1}."
            raise ValueError(msg.format(latent_dim, num_latent_gps))

        if mean_function is None:
            mean_function = Zero()

        if kernel is None:
            kernel = kernels.SquaredExponential(lengthscales=tf.ones((latent_dim,)))

        if data.shape[1] < num_latent_gps:
            raise ValueError("More latent dimensions than observed.")

        gpr_data = (Parameter(X_data_mean), data_input_to_tensor(data))
        super().__init__(gpr_data, kernel, mean_function=mean_function)


class BayesianGPLVM(GPModel, InternalDataTrainingLossMixin):
    @check_shapes(
        "data: [N, P]",
        "X_data_mean: [N, Q]",
        "X_data_var: [N, Q]",
        "X_prior_mean: [N, Q]",
        "X_prior_var: [N, Q]",
    )
    def __init__(
        self,
        data: OutputData,
        X_data_mean: tf.Tensor,
        X_data_var: tf.Tensor,
        kernel: Kernel,
        num_inducing_variables: Optional[int] = None,
        inducing_variable: Optional[InducingVariablesLike] = None,
        X_prior_mean: Optional[TensorType] = None,
        X_prior_var: Optional[TensorType] = None,
    ):
        """
        Initialise Bayesian GPLVM object. This method only works with a Gaussian likelihood.

        :param data: data matrix, size N (number of points) x P (dimensions)
        :param X_data_mean: initial latent positions, size N (number of points) x
            Q (latent dimensions).
        :param X_data_var: variance of latent positions ([N, Q]), for the initialisation of the
            latent space.
        :param kernel: kernel specification, by default Squared Exponential
        :param num_inducing_variables: number of inducing points, M
        :param inducing_variable: matrix of inducing points, size M (inducing points) x
            Q (latent dimensions). By default random permutation of X_data_mean.
        :param X_prior_mean: prior mean used in KL term of bound. By default 0.
            Same size as X_data_mean.
        :param X_prior_var: prior variance used in KL term of bound. By default 1.
        """
        num_data, num_latent_gps = X_data_mean.shape
        super().__init__(kernel, likelihoods.Gaussian(), num_latent_gps=num_latent_gps)
        self.data = data_input_to_tensor(data)

        self.X_data_mean = Parameter(X_data_mean)
        self.X_data_var = Parameter(X_data_var, transform=positive())

        self.num_data = num_data
        self.output_dim = self.data.shape[-1]

        if (inducing_variable is None) == (num_inducing_variables is None):
            raise ValueError(
                "BayesianGPLVM needs exactly one of `inducing_variable` and"
                " `num_inducing_variables`"
            )

        if inducing_variable is None:
            # By default we initialize by subset of initial latent points
            # Note that tf.random.shuffle returns a copy, it does not shuffle in-place
            Z = tf.random.shuffle(X_data_mean)[:num_inducing_variables]
            inducing_variable = InducingPoints(Z)

        self.inducing_variable = inducingpoint_wrapper(inducing_variable)

        assert X_data_mean.shape[1] == self.num_latent_gps

        # deal with parameters for the prior mean variance of X
        if X_prior_mean is None:
            X_prior_mean = tf.zeros((self.num_data, self.num_latent_gps), dtype=default_float())
        if X_prior_var is None:
            X_prior_var = tf.ones((self.num_data, self.num_latent_gps))

        self.X_prior_mean = tf.convert_to_tensor(np.atleast_1d(X_prior_mean), dtype=default_float())
        self.X_prior_var = tf.convert_to_tensor(np.atleast_1d(X_prior_var), dtype=default_float())

    # type-ignore is because of changed method signature:
    @inherit_check_shapes
    def maximum_log_likelihood_objective(self) -> tf.Tensor:  # type: ignore[override]
        return self.elbo()

    @check_shapes(
        "return: []",
    )
    def elbo(self) -> tf.Tensor:
        """
        Construct a tensorflow function to compute the bound on the marginal
        likelihood.
        """
        Y_data = self.data

        pX = DiagonalGaussian(self.X_data_mean, self.X_data_var)

        num_inducing = self.inducing_variable.num_inducing
        psi0 = tf.reduce_sum(expectation(pX, self.kernel))
        psi1 = expectation(pX, (self.kernel, self.inducing_variable))
        psi2 = tf.reduce_sum(
            expectation(
                pX, (self.kernel, self.inducing_variable), (self.kernel, self.inducing_variable)
            ),
            axis=0,
        )
        cov_uu = covariances.Kuu(self.inducing_variable, self.kernel, jitter=default_jitter())
        L = tf.linalg.cholesky(cov_uu)
        sigma2 = self.likelihood.variance

        # Compute intermediate matrices
        A = tf.linalg.triangular_solve(L, tf.transpose(psi1), lower=True)
        tmp = tf.linalg.triangular_solve(L, psi2, lower=True)
        AAT = tf.linalg.triangular_solve(L, tf.transpose(tmp), lower=True) / sigma2
        B = AAT + tf.eye(num_inducing, dtype=default_float())
        LB = tf.linalg.cholesky(B)
        log_det_B = 2.0 * tf.reduce_sum(tf.math.log(tf.linalg.diag_part(LB)))
        c = tf.linalg.triangular_solve(LB, tf.linalg.matmul(A, Y_data), lower=True) / sigma2

        # KL[q(x) || p(x)]
        dX_data_var = (
            self.X_data_var
            if self.X_data_var.shape.ndims == 2
            else tf.linalg.diag_part(self.X_data_var)
        )
        NQ = to_default_float(tf.size(self.X_data_mean))
        D = to_default_float(tf.shape(Y_data)[1])
        KL = -0.5 * tf.reduce_sum(tf.math.log(dX_data_var))
        KL += 0.5 * tf.reduce_sum(tf.math.log(self.X_prior_var))
        KL -= 0.5 * NQ
        KL += 0.5 * tf.reduce_sum(
            (tf.square(self.X_data_mean - self.X_prior_mean) + dX_data_var) / self.X_prior_var
        )

        # compute log marginal bound
        ND = to_default_float(tf.size(Y_data))
        bound = -0.5 * ND * tf.math.log(2 * np.pi * sigma2)
        bound += -0.5 * D * log_det_B
        bound += -0.5 * tf.reduce_sum(tf.square(Y_data)) / sigma2
        bound += 0.5 * tf.reduce_sum(tf.square(c))
        bound += -0.5 * D * (tf.reduce_sum(psi0) / sigma2 - tf.reduce_sum(tf.linalg.diag_part(AAT)))
        bound -= KL
        return bound

    @inherit_check_shapes
    def predict_f(
        self, Xnew: InputData, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        Compute the mean and variance of the latent function at some new points.
        Note that this is very similar to the SGPR prediction, for which
        there are notes in the SGPR notebook.

        Note: This model does not allow full output covariances.

        :param Xnew: points at which to predict
        """
        assert_params_false(self.predict_f, full_output_cov=full_output_cov)

        pX = DiagonalGaussian(self.X_data_mean, self.X_data_var)

        Y_data = self.data
        num_inducing = self.inducing_variable.num_inducing
        psi1 = expectation(pX, (self.kernel, self.inducing_variable))
        psi2 = tf.reduce_sum(
            expectation(
                pX, (self.kernel, self.inducing_variable), (self.kernel, self.inducing_variable)
            ),
            axis=0,
        )
        jitter = default_jitter()
        Kus = covariances.Kuf(self.inducing_variable, self.kernel, Xnew)
        sigma2 = self.likelihood.variance
        L = tf.linalg.cholesky(covariances.Kuu(self.inducing_variable, self.kernel, jitter=jitter))

        A = tf.linalg.triangular_solve(L, tf.transpose(psi1), lower=True)
        tmp = tf.linalg.triangular_solve(L, psi2, lower=True)
        AAT = tf.linalg.triangular_solve(L, tf.transpose(tmp), lower=True) / sigma2
        B = AAT + tf.eye(num_inducing, dtype=default_float())
        LB = tf.linalg.cholesky(B)
        c = tf.linalg.triangular_solve(LB, tf.linalg.matmul(A, Y_data), lower=True) / sigma2
        tmp1 = tf.linalg.triangular_solve(L, Kus, lower=True)
        tmp2 = tf.linalg.triangular_solve(LB, tmp1, lower=True)
        mean = tf.linalg.matmul(tmp2, c, transpose_a=True)
        if full_cov:
            var = (
                self.kernel(Xnew)
                + tf.linalg.matmul(tmp2, tmp2, transpose_a=True)
                - tf.linalg.matmul(tmp1, tmp1, transpose_a=True)
            )
            shape = tf.stack([tf.shape(Y_data)[1], 1, 1])
            var = tf.tile(tf.expand_dims(var, 0), shape)
        else:
            var = (
                self.kernel(Xnew, full_cov=False)
                + tf.reduce_sum(tf.square(tmp2), axis=0)
                - tf.reduce_sum(tf.square(tmp1), axis=0)
            )
            shape = tf.stack([1, tf.shape(Y_data)[1]])
            var = tf.tile(tf.expand_dims(var, 1), shape)
        return mean + self.mean_function(Xnew), var

    @inherit_check_shapes
    def predict_log_density(
        self, data: RegressionData, full_cov: bool = False, full_output_cov: bool = False
    ) -> tf.Tensor:
        raise NotImplementedError


# gpflow/models/gpmc.py

# Copyright 2016-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Optional

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp

from ..base import InputData, MeanAndVariance, Parameter, RegressionData
from ..conditionals import conditional
from ..config import default_float, default_jitter
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..kernels import Kernel
from ..likelihoods import Likelihood
from ..mean_functions import MeanFunction
from ..utilities import assert_params_false, to_default_float
from .model import GPModel
from .training_mixins import InternalDataTrainingLossMixin
from .util import data_input_to_tensor


class GPMC(GPModel, InternalDataTrainingLossMixin):
    @check_shapes(
        "data[0]: [N, D]",
        "data[1]: [N, P]",
    )
    def __init__(
        self,
        data: RegressionData,
        kernel: Kernel,
        likelihood: Likelihood,
        mean_function: Optional[MeanFunction] = None,
        num_latent_gps: Optional[int] = None,
    ):
        """
        data is a tuple of X, Y with X, a data matrix, size [N, D] and Y, a data matrix, size [N, R]
        kernel, likelihood, mean_function are appropriate GPflow objects

        This is a vanilla implementation of a GP with a non-Gaussian
        likelihood. The latent function values are represented by centered
        (whitened) variables, so

            v ~ N(0, I)
            f = Lv + m(x)

        with

            L L^T = K

        """
        if num_latent_gps is None:
            num_latent_gps = self.calc_num_latent_gps_from_data(data, kernel, likelihood)
        super().__init__(kernel, likelihood, mean_function, num_latent_gps)
        self.data = data_input_to_tensor(data)
        self.num_data = self.data[0].shape[0]
        self.V = Parameter(np.zeros((self.num_data, self.num_latent_gps)))
        self.V.prior = tfp.distributions.Normal(
            loc=to_default_float(0.0), scale=to_default_float(1.0)
        )

    # type-ignore is because of changed method signature:
    @inherit_check_shapes
    def log_posterior_density(self) -> tf.Tensor:  # type: ignore[override]
        return self.log_likelihood() + self.log_prior_density()

    # type-ignore is because of changed method signature:
    @inherit_check_shapes
    def _training_loss(self) -> tf.Tensor:  # type: ignore[override]
        return -self.log_posterior_density()

    # type-ignore is because of changed method signature:
    @inherit_check_shapes
    def maximum_log_likelihood_objective(self) -> tf.Tensor:  # type: ignore[override]
        return self.log_likelihood()

    @check_shapes(
        "return: []",
    )
    def log_likelihood(self) -> tf.Tensor:
        r"""
        Construct a tf function to compute the likelihood of a general GP
        model.

            \log p(Y | V, theta).

        """
        X_data, Y_data = self.data
        K = self.kernel(X_data)
        L = tf.linalg.cholesky(
            K + tf.eye(tf.shape(X_data)[0], dtype=default_float()) * default_jitter()
        )
        F = tf.linalg.matmul(L, self.V) + self.mean_function(X_data)

        return tf.reduce_sum(self.likelihood.log_prob(X_data, F, Y_data))

    @inherit_check_shapes
    def predict_f(
        self, Xnew: InputData, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        Xnew is a data matrix, point at which we want to predict

        This method computes

            p(F* | (F=LV) )

        where F* are points on the GP at Xnew, F=LV are points on the GP at X.

        """
        assert_params_false(self.predict_f, full_output_cov=full_output_cov)

        X_data, _Y_data = self.data
        mu, var = conditional(
            Xnew, X_data, self.kernel, self.V, full_cov=full_cov, q_sqrt=None, white=True
        )
        return mu + self.mean_function(Xnew), var


# gpflow/models/gpr.py

# Copyright 2016-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Optional

import tensorflow as tf

import gpflow

from .. import posteriors
from ..base import InputData, MeanAndVariance, RegressionData, TensorData
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..kernels import Kernel
from ..likelihoods import Gaussian
from ..logdensities import multivariate_normal
from ..mean_functions import MeanFunction
from ..utilities import add_likelihood_noise_cov, assert_params_false
from .model import GPModel
from .training_mixins import InternalDataTrainingLossMixin
from .util import data_input_to_tensor


class GPR_deprecated(GPModel, InternalDataTrainingLossMixin):
    r"""
    Gaussian Process Regression.

    This is a vanilla implementation of GP regression with a Gaussian
    likelihood.  Multiple columns of Y are treated independently.

    The log likelihood of this model is given by

    .. math::
       \log p(Y \,|\, \mathbf f) =
            \mathcal N(Y \,|\, 0, \sigma_n^2 \mathbf{I})

    To train the model, we maximise the log _marginal_ likelihood
    w.r.t. the likelihood variance and kernel hyperparameters theta.
    The marginal likelihood is found by integrating the likelihood
    over the prior, and has the form

    .. math::
       \log p(Y \,|\, \sigma_n, \theta) =
            \mathcal N(Y \,|\, 0, \mathbf{K} + \sigma_n^2 \mathbf{I})
    """

    @check_shapes(
        "data[0]: [N, D]",
        "data[1]: [N, P]",
        "noise_variance: []",
    )
    def __init__(
        self,
        data: RegressionData,
        kernel: Kernel,
        mean_function: Optional[MeanFunction] = None,
        noise_variance: Optional[TensorData] = None,
        likelihood: Optional[Gaussian] = None,
    ):
        assert (noise_variance is None) or (
            likelihood is None
        ), "Cannot set both `noise_variance` and `likelihood`."
        if likelihood is None:
            if noise_variance is None:
                noise_variance = 1.0
            likelihood = gpflow.likelihoods.Gaussian(noise_variance)
        _, Y_data = data
        super().__init__(kernel, likelihood, mean_function, num_latent_gps=Y_data.shape[-1])
        self.data = data_input_to_tensor(data)

    # type-ignore is because of changed method signature:
    @inherit_check_shapes
    def maximum_log_likelihood_objective(self) -> tf.Tensor:  # type: ignore[override]
        return self.log_marginal_likelihood()

    @check_shapes(
        "return: []",
    )
    def log_marginal_likelihood(self) -> tf.Tensor:
        r"""
        Computes the log marginal likelihood.

        .. math::
            \log p(Y | \theta).

        """
        X, Y = self.data
        K = self.kernel(X)
        ks = add_likelihood_noise_cov(K, self.likelihood, X)
        L = tf.linalg.cholesky(ks)
        m = self.mean_function(X)

        # [R,] log-likelihoods for each independent dimension of Y
        log_prob = multivariate_normal(Y, m, L)
        return tf.reduce_sum(log_prob)

    @inherit_check_shapes
    def predict_f(
        self, Xnew: InputData, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        r"""
        This method computes predictions at X \in R^{N \x D} input points

        .. math::
            p(F* | Y)

        where F* are points on the GP at new data points, Y are noisy observations at training data
        points.
        """
        assert_params_false(self.predict_f, full_output_cov=full_output_cov)

        X, Y = self.data
        err = Y - self.mean_function(X)

        kmm = self.kernel(X)
        knn = self.kernel(Xnew, full_cov=full_cov)
        kmn = self.kernel(X, Xnew)
        kmm_plus_s = add_likelihood_noise_cov(kmm, self.likelihood, X)

        conditional = gpflow.conditionals.base_conditional
        f_mean_zero, f_var = conditional(
            kmn, kmm_plus_s, knn, err, full_cov=full_cov, white=False
        )  # [N, P], [N, P] or [P, N, N]
        f_mean = f_mean_zero + self.mean_function(Xnew)
        return f_mean, f_var


class GPR_with_posterior(GPR_deprecated):
    """
    This is an implementation of GPR that provides a posterior() method that
    enables caching for faster subsequent predictions.
    """

    def posterior(
        self,
        precompute_cache: posteriors.PrecomputeCacheType = posteriors.PrecomputeCacheType.TENSOR,
    ) -> posteriors.GPRPosterior:
        """
        Create the Posterior object which contains precomputed matrices for
        faster prediction.

        precompute_cache has three settings:

        - `PrecomputeCacheType.TENSOR` (or `"tensor"`): Precomputes the cached
          quantities and stores them as tensors (which allows differentiating
          through the prediction). This is the default.
        - `PrecomputeCacheType.VARIABLE` (or `"variable"`): Precomputes the cached
          quantities and stores them as variables, which allows for updating
          their values without changing the compute graph (relevant for AOT
          compilation).
        - `PrecomputeCacheType.NOCACHE` (or `"nocache"` or `None`): Avoids
          immediate cache computation. This is useful for avoiding extraneous
          computations when you only want to call the posterior's
          `fused_predict_f` method.
        """

        return posteriors.GPRPosterior(
            kernel=self.kernel,
            data=self.data,
            likelihood=self.likelihood,
            mean_function=self.mean_function,
            precompute_cache=precompute_cache,
        )

    @inherit_check_shapes
    def predict_f(
        self, Xnew: InputData, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        For backwards compatibility, GPR's predict_f uses the fused (no-cache)
        computation, which is more efficient during training.

        For faster (cached) prediction, predict directly from the posterior object, i.e.,:
            model.posterior().predict_f(Xnew, ...)
        """
        return self.posterior(posteriors.PrecomputeCacheType.NOCACHE).fused_predict_f(
            Xnew, full_cov=full_cov, full_output_cov=full_output_cov
        )


class GPR(GPR_with_posterior):
    # subclassed to ensure __class__ == "GPR"
    pass


# gpflow/models/model.py

# Copyright 2016-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import abc
from typing import Any, Optional

import tensorflow as tf

from ..base import InputData, MeanAndVariance, Module, RegressionData
from ..conditionals.util import sample_mvn
from ..experimental.check_shapes import check_shapes
from ..kernels import Kernel, MultioutputKernel
from ..likelihoods import Likelihood, SwitchedLikelihood
from ..mean_functions import MeanFunction, Zero
from ..utilities import assert_params_false, to_default_float


class BayesianModel(Module, metaclass=abc.ABCMeta):
    """ Bayesian model. """

    @check_shapes(
        "return: []",
    )
    def log_prior_density(self) -> tf.Tensor:
        """
        Sum of the log prior probability densities of all (constrained) variables in this model.
        """
        if self.trainable_parameters:
            return tf.add_n([p.log_prior_density() for p in self.trainable_parameters])
        else:
            return to_default_float(0.0)

    @check_shapes(
        "return: []",
    )
    def log_posterior_density(self, *args: Any, **kwargs: Any) -> tf.Tensor:
        """
        This may be the posterior with respect to the hyperparameters (e.g. for
        GPR) or the posterior with respect to the function (e.g. for GPMC and
        SGPMC). It assumes that maximum_log_likelihood_objective() is defined
        sensibly.
        """
        return self.maximum_log_likelihood_objective(*args, **kwargs) + self.log_prior_density()

    @check_shapes(
        "return: []",
    )
    def _training_loss(self, *args: Any, **kwargs: Any) -> tf.Tensor:
        """
        Training loss definition. To allow MAP (maximum a-posteriori) estimation,
        adds the log density of all priors to maximum_log_likelihood_objective().
        """
        return -(self.maximum_log_likelihood_objective(*args, **kwargs) + self.log_prior_density())

    @abc.abstractmethod
    @check_shapes(
        "return: []",
    )
    def maximum_log_likelihood_objective(self, *args: Any, **kwargs: Any) -> tf.Tensor:
        """
        Objective for maximum likelihood estimation. Should be maximized. E.g.
        log-marginal likelihood (hyperparameter likelihood) for GPR, or lower
        bound to the log-marginal likelihood (ELBO) for sparse and variational
        GPs.
        """
        raise NotImplementedError


class GPModel(BayesianModel):
    r"""
    A stateless base class for Gaussian process models, that is, those of the
    form

    .. math::
       :nowrap:

       \begin{align}
           \theta        & \sim p(\theta) \\
           f             & \sim \mathcal{GP}(m(x), k(x, x'; \theta)) \\
           f_i           & = f(x_i) \\
           y_i \,|\, f_i & \sim p(y_i|f_i)
       \end{align}

    This class mostly adds functionality for predictions. To use it, inheriting
    classes must define a predict_f function, which computes the means and
    variances of the latent function.

    These predictions are then pushed through the likelihood to obtain means
    and variances of held out data, self.predict_y.

    The predictions can also be used to compute the (log) density of held-out
    data via self.predict_log_density.

    It is also possible to draw samples from the latent GPs using
    self.predict_f_samples.
    """

    def __init__(
        self,
        kernel: Kernel,
        likelihood: Likelihood,
        mean_function: Optional[MeanFunction] = None,
        num_latent_gps: Optional[int] = None,
    ):
        super().__init__()
        assert num_latent_gps is not None, "GPModel requires specification of num_latent_gps"
        self.num_latent_gps = num_latent_gps
        if mean_function is None:
            mean_function = Zero()
        self.mean_function = mean_function
        self.kernel = kernel
        self.likelihood = likelihood

    @staticmethod
    @check_shapes(
        "data[0]: [batch..., N, D]",
        "data[1]: [batch..., N, P]",
    )
    def calc_num_latent_gps_from_data(
        data: RegressionData, kernel: Kernel, likelihood: Likelihood
    ) -> int:
        """
        Calculates the number of latent GPs required based on the data as well
        as the type of kernel and likelihood.
        """
        _, Y = data
        output_dim = Y.shape[-1]
        return GPModel.calc_num_latent_gps(kernel, likelihood, output_dim)

    @staticmethod
    def calc_num_latent_gps(kernel: Kernel, likelihood: Likelihood, output_dim: int) -> int:
        """
        Calculates the number of latent GPs required given the number of
        outputs `output_dim` and the type of likelihood and kernel.

        Note: It's not nice for `GPModel` to need to be aware of specific
        likelihoods as here. However, `num_latent_gps` is a bit more broken in
        general, we should fix this in the future. There are also some slightly
        problematic assumptions re the output dimensions of mean_function.
        See https://github.com/GPflow/GPflow/issues/1343
        """
        if isinstance(kernel, MultioutputKernel):
            # MultioutputKernels already have num_latent_gps attributes
            num_latent_gps: int = kernel.num_latent_gps
        elif isinstance(likelihood, SwitchedLikelihood):
            # the SwitchedLikelihood partitions/stitches based on the last
            # column in Y, but we should not add a separate latent GP for this!
            # hence decrement by 1
            num_latent_gps = output_dim - 1
            assert num_latent_gps > 0
        else:
            num_latent_gps = output_dim

        return num_latent_gps

    @abc.abstractmethod
    @check_shapes(
        "Xnew: [batch..., N, D]",
        "return[0]: [batch..., N, P]",
        "return[1]: [batch..., N, P, N, P] if full_cov and full_output_cov",
        "return[1]: [batch..., P, N, N] if full_cov and (not full_output_cov)",
        "return[1]: [batch..., N, P, P] if (not full_cov) and full_output_cov",
        "return[1]: [batch..., N, P] if (not full_cov) and (not full_output_cov)",
    )
    def predict_f(
        self, Xnew: InputData, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        raise NotImplementedError

    @check_shapes(
        "Xnew: [batch..., N, D]",
        "return: [batch..., N, P] if (num_samples is None)",
        "return: [batch..., S, N, P] if (num_samples is not None)",
    )
    def predict_f_samples(
        self,
        Xnew: InputData,
        num_samples: Optional[int] = None,
        full_cov: bool = True,
        full_output_cov: bool = False,
    ) -> tf.Tensor:
        """
        Produce samples from the posterior latent function(s) at the input points.

        Currently, the method does not support `full_output_cov=True` and `full_cov=True`.

        :param Xnew:
            Input locations at which to draw samples.
        :param num_samples:
            Number of samples to draw.
            If `None`, a single sample is drawn and the return shape is [..., N, P],
            for any positive integer the return shape contains an extra batch
            dimension, [..., S, N, P], with S = num_samples and P is the number of outputs.
        :param full_cov:
            If True, draw correlated samples over the inputs. Computes the Cholesky over the
            dense covariance matrix of size [num_data, num_data].
            If False, draw samples that are uncorrelated over the inputs.
        :param full_output_cov:
            If True, draw correlated samples over the outputs.
            If False, draw samples that are uncorrelated over the outputs.
        """
        if full_cov and full_output_cov:
            raise NotImplementedError(
                "The combination of both `full_cov` and `full_output_cov` is not supported."
            )

        # check below for shape info
        mean, cov = self.predict_f(Xnew, full_cov=full_cov, full_output_cov=full_output_cov)
        if full_cov:
            # mean: [..., N, P]
            # cov: [..., P, N, N]
            mean_for_sample = tf.linalg.adjoint(mean)  # [..., P, N]
            samples = sample_mvn(
                mean_for_sample, cov, full_cov, num_samples=num_samples
            )  # [..., (S), P, N]
            samples = tf.linalg.adjoint(samples)  # [..., (S), N, P]
        else:
            # mean: [..., N, P]
            # cov: [..., N, P] or [..., N, P, P]
            samples = sample_mvn(
                mean, cov, full_output_cov, num_samples=num_samples
            )  # [..., (S), N, P]
        return samples  # [..., (S), N, P]

    @check_shapes(
        "Xnew: [batch..., N, D]",
        "return[0]: [batch..., N, P]",
        "return[1]: [batch..., N, P, N, P] if full_cov and full_output_cov",
        "return[1]: [batch..., P, N, N] if full_cov and (not full_output_cov)",
        "return[1]: [batch..., N, P, P] if (not full_cov) and full_output_cov",
        "return[1]: [batch..., N, P] if (not full_cov) and (not full_output_cov)",
    )
    def predict_y(
        self, Xnew: InputData, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        Compute the mean and variance of the held-out data at the input points.
        """
        # See https://github.com/GPflow/GPflow/issues/1461
        assert_params_false(self.predict_y, full_cov=full_cov, full_output_cov=full_output_cov)

        f_mean, f_var = self.predict_f(Xnew, full_cov=full_cov, full_output_cov=full_output_cov)
        return self.likelihood.predict_mean_and_var(Xnew, f_mean, f_var)

    @check_shapes(
        "data[0]: [batch..., N, D]",
        "data[1]: [batch..., N, P]",
        "return: [batch..., N]",
    )
    def predict_log_density(
        self, data: RegressionData, full_cov: bool = False, full_output_cov: bool = False
    ) -> tf.Tensor:
        """
        Compute the log density of the data at the new data points.
        """
        # See https://github.com/GPflow/GPflow/issues/1461
        assert_params_false(self.predict_y, full_cov=full_cov, full_output_cov=full_output_cov)

        X, Y = data
        f_mean, f_var = self.predict_f(X, full_cov=full_cov, full_output_cov=full_output_cov)
        return self.likelihood.predict_log_density(X, f_mean, f_var, Y)


# gpflow/models/sgpmc.py

# Copyright 2016-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Optional

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp

from ..base import InputData, MeanAndVariance, Parameter, RegressionData
from ..conditionals import conditional
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..kernels import Kernel
from ..likelihoods import Likelihood
from ..mean_functions import MeanFunction
from ..utilities import to_default_float
from .model import GPModel
from .training_mixins import InternalDataTrainingLossMixin
from .util import InducingPointsLike, data_input_to_tensor, inducingpoint_wrapper


class SGPMC(GPModel, InternalDataTrainingLossMixin):
    r"""
    This is the Sparse Variational GP using MCMC (SGPMC).

    The key reference is :cite:t:`hensman2015mcmc`.

    The latent function values are represented by centered
    (whitened) variables, so

    .. math::
       :nowrap:

       \begin{align}
       \mathbf v & \sim N(0, \mathbf I) \\
       \mathbf u &= \mathbf L\mathbf v
       \end{align}

    with

    .. math::
        \mathbf L \mathbf L^\top = \mathbf K


    """

    @check_shapes(
        "data[0]: [N, D]",
        "data[1]: [N, P]",
    )
    def __init__(
        self,
        data: RegressionData,
        kernel: Kernel,
        likelihood: Likelihood,
        mean_function: Optional[MeanFunction] = None,
        num_latent_gps: Optional[int] = None,
        inducing_variable: Optional[InducingPointsLike] = None,
    ):
        """
        data is a tuple of X, Y with X, a data matrix, size [N, D] and Y, a data matrix, size [N, R]
        Z is a data matrix, of inducing inputs, size [M, D]
        kernel, likelihood, mean_function are appropriate GPflow objects
        """
        if num_latent_gps is None:
            num_latent_gps = self.calc_num_latent_gps_from_data(data, kernel, likelihood)
        super().__init__(kernel, likelihood, mean_function, num_latent_gps=num_latent_gps)
        self.data = data_input_to_tensor(data)
        self.num_data = data[0].shape[0]
        self.inducing_variable = inducingpoint_wrapper(inducing_variable)
        self.V = Parameter(np.zeros((self.inducing_variable.num_inducing, self.num_latent_gps)))
        self.V.prior = tfp.distributions.Normal(
            loc=to_default_float(0.0), scale=to_default_float(1.0)
        )

    # type-ignore is because of changed method signature:
    @inherit_check_shapes
    def log_posterior_density(self) -> tf.Tensor:  # type: ignore[override]
        return self.log_likelihood_lower_bound() + self.log_prior_density()

    # type-ignore is because of changed method signature:
    @inherit_check_shapes
    def _training_loss(self) -> tf.Tensor:  # type: ignore[override]
        return -self.log_posterior_density()

    # type-ignore is because of changed method signature:
    @inherit_check_shapes
    def maximum_log_likelihood_objective(self) -> tf.Tensor:  # type: ignore[override]
        return self.log_likelihood_lower_bound()

    def log_likelihood_lower_bound(self) -> tf.Tensor:
        """
        This function computes the optimal density for v, q*(v), up to a constant
        """
        # get the (marginals of) q(f): exactly predicting!
        X_data, Y_data = self.data
        fmean, fvar = self.predict_f(X_data, full_cov=False)
        return tf.reduce_sum(self.likelihood.variational_expectations(X_data, fmean, fvar, Y_data))

    @inherit_check_shapes
    def predict_f(
        self, Xnew: InputData, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        Xnew is a data matrix of the points at which we want to predict

        This method computes

            p(F* | (U=LV) )

        where F* are points on the GP at Xnew, F=LV are points on the GP at Z,

        """
        mu, var = conditional(
            Xnew,
            self.inducing_variable,
            self.kernel,
            self.V,
            full_cov=full_cov,
            q_sqrt=None,
            white=True,
            full_output_cov=full_output_cov,
        )
        return mu + self.mean_function(Xnew), var


# gpflow/models/sgpr.py

# Copyright 2016-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import NamedTuple, Optional, Tuple

import numpy as np
import tensorflow as tf

from .. import posteriors
from ..base import InputData, MeanAndVariance, RegressionData, TensorData
from ..config import default_float, default_jitter
from ..covariances.dispatch import Kuf, Kuu
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..inducing_variables import InducingPoints
from ..kernels import Kernel
from ..likelihoods import Gaussian
from ..mean_functions import MeanFunction
from ..utilities import add_noise_cov, assert_params_false, to_default_float
from .model import GPModel
from .training_mixins import InternalDataTrainingLossMixin
from .util import InducingPointsLike, data_input_to_tensor, inducingpoint_wrapper


class SGPRBase_deprecated(GPModel, InternalDataTrainingLossMixin):
    """
    Common base class for SGPR and GPRFITC that provides the common __init__
    and upper_bound() methods.
    """

    @check_shapes(
        "data[0]: [N, D]",
        "data[1]: [N, P]",
        "noise_variance: []",
    )
    def __init__(
        self,
        data: RegressionData,
        kernel: Kernel,
        inducing_variable: InducingPointsLike,
        *,
        mean_function: Optional[MeanFunction] = None,
        num_latent_gps: Optional[int] = None,
        noise_variance: Optional[TensorData] = None,
        likelihood: Optional[Gaussian] = None,
    ):
        """
        This method only works with a Gaussian likelihood, its variance is
        initialized to `noise_variance`.

        :param data: a tuple of (X, Y), where the inputs X has shape [N, D]
            and the outputs Y has shape [N, R].
        :param inducing_variable:  an InducingPoints instance or a matrix of
            the pseudo inputs Z, of shape [M, D].
        :param kernel: An appropriate GPflow kernel object.
        :param mean_function: An appropriate GPflow mean function object.
        """
        assert (noise_variance is None) or (
            likelihood is None
        ), "Cannot set both `noise_variance` and `likelihood`."
        if likelihood is None:
            if noise_variance is None:
                noise_variance = 1.0
            likelihood = Gaussian(noise_variance)
        X_data, Y_data = data_input_to_tensor(data)
        num_latent_gps = Y_data.shape[-1] if num_latent_gps is None else num_latent_gps
        super().__init__(kernel, likelihood, mean_function, num_latent_gps=num_latent_gps)

        self.data = X_data, Y_data
        self.num_data = X_data.shape[0]

        self.inducing_variable: InducingPoints = inducingpoint_wrapper(inducing_variable)

    @check_shapes(
        "return: []",
    )
    def upper_bound(self) -> tf.Tensor:
        """
        Upper bound for the sparse GP regression marginal likelihood.  Note that
        the same inducing points are used for calculating the upper bound, as are
        used for computing the likelihood approximation. This may not lead to the
        best upper bound. The upper bound can be tightened by optimising Z, just
        like the lower bound. This is especially important in FITC, as FITC is
        known to produce poor inducing point locations. An optimisable upper bound
        can be found in https://github.com/markvdw/gp_upper.

        The key reference is :cite:t:`titsias_2014`.

        The key quantity, the trace term, can be computed via

        >>> _, v = conditionals.conditional(X, model.inducing_variable.Z, model.kernel,
        ...                                 np.zeros((model.inducing_variable.num_inducing, 1)))

        which computes each individual element of the trace term.
        """
        X_data, Y_data = self.data

        sigma_sq = tf.squeeze(self.likelihood.variance_at(X_data), axis=-1)  # [N]
        sigma = tf.sqrt(sigma_sq)  # [N]

        Kdiag = self.kernel(X_data, full_cov=False)
        kuu = Kuu(self.inducing_variable, self.kernel, jitter=default_jitter())
        kuf = Kuf(self.inducing_variable, self.kernel, X_data)

        I = tf.eye(tf.shape(kuu)[0], dtype=default_float())

        L = tf.linalg.cholesky(kuu)
        A = tf.linalg.triangular_solve(L, kuf, lower=True)

        A_sigma = tf.linalg.triangular_solve(L, kuf / sigma, lower=True)
        AAT_sigma = tf.linalg.matmul(A_sigma, A_sigma, transpose_b=True)
        B = I + AAT_sigma
        LB = tf.linalg.cholesky(B)

        # Using the Trace bound, from Titsias' presentation
        c = tf.reduce_sum(Kdiag) - tf.reduce_sum(tf.square(A))

        # Alternative bound on max eigenval:
        cn_var = sigma_sq + c
        cn_std = tf.sqrt(cn_var)

        const = -0.5 * tf.reduce_sum(tf.math.log(2 * np.pi * sigma_sq))
        logdet = -tf.reduce_sum(tf.math.log(tf.linalg.diag_part(LB)))

        A_cn = tf.linalg.triangular_solve(L, kuf / cn_std, lower=True)
        AAT_cn = tf.linalg.matmul(A_cn, A_cn, transpose_b=True)

        err = Y_data - self.mean_function(X_data)
        LC = tf.linalg.cholesky(I + AAT_cn)
        v = tf.linalg.triangular_solve(
            LC, tf.linalg.matmul(A_cn, err / cn_std[:, None]), lower=True
        )
        quad = -0.5 * tf.reduce_sum(tf.square(err / cn_std[:, None])) + 0.5 * tf.reduce_sum(
            tf.square(v)
        )

        return const + logdet + quad


class SGPR_deprecated(SGPRBase_deprecated):
    """
    Sparse Variational GP regression.

    The key reference is :cite:t:`titsias2009variational`.
    """

    class CommonTensors(NamedTuple):
        sigma_sq: tf.Tensor
        sigma: tf.Tensor
        A: tf.Tensor
        B: tf.Tensor
        LB: tf.Tensor
        AAT: tf.Tensor
        L: tf.Tensor

    # type-ignore is because of changed method signature:
    @inherit_check_shapes
    def maximum_log_likelihood_objective(self) -> tf.Tensor:  # type: ignore[override]
        return self.elbo()

    @check_shapes(
        "return.sigma_sq: [N]",
        "return.sigma: [N]",
        "return.A: [M, N]",
        "return.B: [M, M]",
        "return.LB: [M, M]",
        "return.AAT: [M, M]",
    )
    def _common_calculation(self) -> "SGPR.CommonTensors":
        """
        Matrices used in log-det calculation

        :return:
            * :math:`σ²`,
            * :math:`σ`,
            * :math:`A = L⁻¹K_{uf}/σ`, where :math:`LLᵀ = Kᵤᵤ`,
            * :math:`B = AAT+I`,
            * :math:`LB` where :math`LBLBᵀ = B`,
            * :math:`AAT = AAᵀ`,
        """
        x, _ = self.data  # [N]
        iv = self.inducing_variable  # [M]

        sigma_sq = tf.squeeze(self.likelihood.variance_at(x), axis=-1)  # [N]
        sigma = tf.sqrt(sigma_sq)  # [N]

        kuf = Kuf(iv, self.kernel, x)  # [M, N]
        kuu = Kuu(iv, self.kernel, jitter=default_jitter())  # [M, M]
        L = tf.linalg.cholesky(kuu)  # [M, M]

        # Compute intermediate matrices
        A = tf.linalg.triangular_solve(L, kuf / sigma, lower=True)
        AAT = tf.linalg.matmul(A, A, transpose_b=True)
        B = add_noise_cov(AAT, tf.cast(1.0, AAT.dtype))
        LB = tf.linalg.cholesky(B)

        return self.CommonTensors(sigma_sq, sigma, A, B, LB, AAT, L)

    @check_shapes(
        "return: []",
    )
    def logdet_term(self, common: "SGPR.CommonTensors") -> tf.Tensor:
        r"""
        Bound from Jensen's Inequality:

        .. math::
            \log |K + σ²I| <= \log |Q + σ²I| + N * \log (1 + \textrm{tr}(K - Q)/(σ²N))

        :param common: A named tuple containing matrices that will be used
        :return: log_det, lower bound on :math:`-.5 * \textrm{output_dim} * \log |K + σ²I|`
        """
        sigma_sq = common.sigma_sq
        LB = common.LB
        AAT = common.AAT

        x, y = self.data
        outdim = to_default_float(tf.shape(y)[1])
        kdiag = self.kernel(x, full_cov=False)

        # tr(K) / σ²
        trace_k = tf.reduce_sum(kdiag / sigma_sq)
        # tr(Q) / σ²
        trace_q = tf.reduce_sum(tf.linalg.diag_part(AAT))
        # tr(K - Q) / σ²
        trace = trace_k - trace_q

        # 0.5 * log(det(B))
        half_logdet_b = tf.reduce_sum(tf.math.log(tf.linalg.diag_part(LB)))

        # sum log(σ²)
        log_sigma_sq = tf.reduce_sum(tf.math.log(sigma_sq))

        logdet_k = -outdim * (half_logdet_b + 0.5 * log_sigma_sq + 0.5 * trace)
        return logdet_k

    @check_shapes(
        "return: []",
    )
    def quad_term(self, common: "SGPR.CommonTensors") -> tf.Tensor:
        """
        :param common: A named tuple containing matrices that will be used
        :return: Lower bound on -.5 yᵀ(K + σ²I)⁻¹y
        """
        sigma = common.sigma
        A = common.A
        LB = common.LB

        x, y = self.data
        err = (y - self.mean_function(x)) / sigma[..., None]

        Aerr = tf.linalg.matmul(A, err)
        c = tf.linalg.triangular_solve(LB, Aerr, lower=True)

        # σ⁻² yᵀy
        err_inner_prod = tf.reduce_sum(tf.square(err))
        c_inner_prod = tf.reduce_sum(tf.square(c))

        quad = -0.5 * (err_inner_prod - c_inner_prod)
        return quad

    @check_shapes(
        "return: []",
    )
    def elbo(self) -> tf.Tensor:
        """
        Construct a tensorflow function to compute the bound on the marginal
        likelihood. For a derivation of the terms in here, see the associated
        SGPR notebook.
        """
        common = self._common_calculation()
        output_shape = tf.shape(self.data[-1])
        num_data = to_default_float(output_shape[0])
        output_dim = to_default_float(output_shape[1])
        const = -0.5 * num_data * output_dim * np.log(2 * np.pi)
        logdet = self.logdet_term(common)
        quad = self.quad_term(common)
        return const + logdet + quad

    @inherit_check_shapes
    def predict_f(
        self, Xnew: InputData, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        Compute the mean and variance of the latent function at some new points
        Xnew. For a derivation of the terms in here, see the associated SGPR
        notebook.
        """
        # could copy into posterior into a fused version

        assert_params_false(self.predict_f, full_output_cov=full_output_cov)

        X_data, Y_data = self.data
        num_inducing = self.inducing_variable.num_inducing
        err = Y_data - self.mean_function(X_data)
        kuf = Kuf(self.inducing_variable, self.kernel, X_data)
        kuu = Kuu(self.inducing_variable, self.kernel, jitter=default_jitter())
        Kus = Kuf(self.inducing_variable, self.kernel, Xnew)

        sigma_sq = tf.squeeze(self.likelihood.variance_at(X_data), axis=-1)
        sigma = tf.sqrt(sigma_sq)

        L = tf.linalg.cholesky(kuu)  # cache alpha, qinv
        A = tf.linalg.triangular_solve(L, kuf / sigma, lower=True)
        B = tf.linalg.matmul(A, A, transpose_b=True) + tf.eye(
            num_inducing, dtype=default_float()
        )  # cache qinv
        LB = tf.linalg.cholesky(B)  # cache alpha
        Aerr = tf.linalg.matmul(A, err / sigma[..., None])
        c = tf.linalg.triangular_solve(LB, Aerr, lower=True)
        tmp1 = tf.linalg.triangular_solve(L, Kus, lower=True)
        tmp2 = tf.linalg.triangular_solve(LB, tmp1, lower=True)
        mean = tf.linalg.matmul(tmp2, c, transpose_a=True)
        if full_cov:
            var = (
                self.kernel(Xnew)
                + tf.linalg.matmul(tmp2, tmp2, transpose_a=True)
                - tf.linalg.matmul(tmp1, tmp1, transpose_a=True)
            )
            var = tf.tile(var[None, ...], [self.num_latent_gps, 1, 1])  # [P, N, N]
        else:
            var = (
                self.kernel(Xnew, full_cov=False)
                + tf.reduce_sum(tf.square(tmp2), 0)
                - tf.reduce_sum(tf.square(tmp1), 0)
            )
            var = tf.tile(var[:, None], [1, self.num_latent_gps])

        return mean + self.mean_function(Xnew), var

    @check_shapes(
        "return[0]: [M, P]",
        "return[1]: [M, M]",
    )
    def compute_qu(self) -> Tuple[tf.Tensor, tf.Tensor]:
        """
        Computes the mean and variance of q(u) = N(mu, cov), the variational distribution on
        inducing outputs.

        SVGP with this q(u) should predict identically to SGPR.

        :return: mu, cov
        """
        X_data, Y_data = self.data

        kuf = Kuf(self.inducing_variable, self.kernel, X_data)
        kuu = Kuu(self.inducing_variable, self.kernel, jitter=default_jitter())

        var = tf.squeeze(self.likelihood.variance_at(X_data), axis=-1)
        std = tf.sqrt(var)
        scaled_kuf = kuf / std
        sig = kuu + tf.matmul(scaled_kuf, scaled_kuf, transpose_b=True)
        sig_sqrt = tf.linalg.cholesky(sig)

        sig_sqrt_kuu = tf.linalg.triangular_solve(sig_sqrt, kuu)

        cov = tf.linalg.matmul(sig_sqrt_kuu, sig_sqrt_kuu, transpose_a=True)
        err = Y_data - self.mean_function(X_data)
        scaled_err = err / std[..., None]
        mu = tf.linalg.matmul(
            sig_sqrt_kuu,
            tf.linalg.triangular_solve(sig_sqrt, tf.linalg.matmul(scaled_kuf, scaled_err)),
            transpose_a=True,
        )

        return mu, cov


class GPRFITC(SGPRBase_deprecated):
    """
    This implements GP regression with the FITC approximation.

    The key reference is :cite:t:`Snelson06sparsegaussian`.

    Implementation loosely based on code from GPML matlab library although
    obviously gradients are automatic in GPflow.
    """

    @check_shapes(
        "return[0]: [N, R]",
        "return[1]: [N]",
        "return[2]: [M, M]",
        "return[3]: [M, M]",
        "return[4]: [M, R]",
        "return[5]: [N, R]",
        "return[6]: [M, R]",
    )
    def common_terms(
        self,
    ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor]:
        X_data, Y_data = self.data
        num_inducing = self.inducing_variable.num_inducing
        err = Y_data - self.mean_function(X_data)  # size [N, R]
        Kdiag = self.kernel(X_data, full_cov=False)
        kuf = Kuf(self.inducing_variable, self.kernel, X_data)
        kuu = Kuu(self.inducing_variable, self.kernel, jitter=default_jitter())

        sigma_sq = tf.squeeze(self.likelihood.variance_at(X_data), axis=-1)

        Luu = tf.linalg.cholesky(kuu)  # => Luu Luu^T = kuu
        V = tf.linalg.triangular_solve(Luu, kuf)  # => V^T V = Qff = kuf^T kuu^-1 kuf

        diagQff = tf.reduce_sum(tf.square(V), 0)
        nu = Kdiag - diagQff + sigma_sq

        B = tf.eye(num_inducing, dtype=default_float()) + tf.linalg.matmul(
            V / nu, V, transpose_b=True
        )
        L = tf.linalg.cholesky(B)
        beta = err / tf.expand_dims(nu, 1)  # size [N, R]
        alpha = tf.linalg.matmul(V, beta)  # size [M, R]

        gamma = tf.linalg.triangular_solve(L, alpha, lower=True)  # size [M, R]

        return err, nu, Luu, L, alpha, beta, gamma

    # type-ignore is because of changed method signature:
    @inherit_check_shapes
    def maximum_log_likelihood_objective(self) -> tf.Tensor:  # type: ignore[override]
        return self.fitc_log_marginal_likelihood()

    @check_shapes(
        "return: []",
    )
    def fitc_log_marginal_likelihood(self) -> tf.Tensor:
        """
        Construct a tensorflow function to compute the bound on the marginal
        likelihood.
        """

        # FITC approximation to the log marginal likelihood is
        # log ( normal( y | mean, K_fitc ) )
        # where K_fitc = Qff + diag( \nu )
        # where Qff = Kfu kuu^{-1} kuf
        # with \nu_i = Kff_{i,i} - Qff_{i,i} + \sigma^2

        # We need to compute the Mahalanobis term -0.5* err^T K_fitc^{-1} err
        # (summed over functions).

        # We need to deal with the matrix inverse term.
        # K_fitc^{-1} = ( Qff + \diag( \nu ) )^{-1}
        #             = ( V^T V + \diag( \nu ) )^{-1}
        # Applying the Woodbury identity we obtain
        #             = \diag( \nu^{-1} )
        #                 - \diag( \nu^{-1} ) V^T ( I + V \diag( \nu^{-1} ) V^T )^{-1}
        #                     V \diag(\nu^{-1} )
        # Let \beta =  \diag( \nu^{-1} ) err
        # and let \alpha = V \beta
        # then Mahalanobis term = -0.5* (
        #    \beta^T err - \alpha^T Solve( I + V \diag( \nu^{-1} ) V^T, alpha )
        # )

        err, nu, _Luu, L, _alpha, _beta, gamma = self.common_terms()

        mahalanobisTerm = -0.5 * tf.reduce_sum(
            tf.square(err) / tf.expand_dims(nu, 1)
        ) + 0.5 * tf.reduce_sum(tf.square(gamma))

        # We need to compute the log normalizing term -N/2 \log 2 pi - 0.5 \log \det( K_fitc )

        # We need to deal with the log determinant term.
        # \log \det( K_fitc ) = \log \det( Qff + \diag( \nu ) )
        #                     = \log \det( V^T V + \diag( \nu ) )
        # Applying the determinant lemma we obtain
        #                     = \log [ \det \diag( \nu ) \det( I + V \diag( \nu^{-1} ) V^T ) ]
        #                     = \log [
        #                        \det \diag( \nu ) ] + \log [ \det( I + V \diag( \nu^{-1} ) V^T )
        #                     ]

        constantTerm = -0.5 * self.num_data * tf.math.log(tf.constant(2.0 * np.pi, default_float()))
        logDeterminantTerm = -0.5 * tf.reduce_sum(tf.math.log(nu)) - tf.reduce_sum(
            tf.math.log(tf.linalg.diag_part(L))
        )
        logNormalizingTerm = constantTerm + logDeterminantTerm

        return mahalanobisTerm + logNormalizingTerm * self.num_latent_gps

    @inherit_check_shapes
    def predict_f(
        self, Xnew: InputData, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        Compute the mean and variance of the latent function at some new points
        Xnew.
        """
        assert_params_false(self.predict_f, full_output_cov=full_output_cov)

        _, _, Luu, L, _, _, gamma = self.common_terms()
        Kus = Kuf(self.inducing_variable, self.kernel, Xnew)  # [M, N]

        w = tf.linalg.triangular_solve(Luu, Kus, lower=True)  # [M, N]

        tmp = tf.linalg.triangular_solve(tf.transpose(L), gamma, lower=False)
        mean = tf.linalg.matmul(w, tmp, transpose_a=True) + self.mean_function(Xnew)
        intermediateA = tf.linalg.triangular_solve(L, w, lower=True)

        if full_cov:
            var = (
                self.kernel(Xnew)
                - tf.linalg.matmul(w, w, transpose_a=True)
                + tf.linalg.matmul(intermediateA, intermediateA, transpose_a=True)
            )
            var = tf.tile(var[None, ...], [self.num_latent_gps, 1, 1])  # [P, N, N]
        else:
            var = (
                self.kernel(Xnew, full_cov=False)
                - tf.reduce_sum(tf.square(w), 0)
                + tf.reduce_sum(tf.square(intermediateA), 0)
            )  # [N, P]
            var = tf.tile(var[:, None], [1, self.num_latent_gps])

        return mean, var


class SGPR_with_posterior(SGPR_deprecated):
    """
    Sparse Variational GP regression.
    The key reference is :cite:t:`titsias2009variational`.

    This is an implementation of SGPR that provides a posterior() method that
    enables caching for faster subsequent predictions.
    """

    def posterior(
        self,
        precompute_cache: posteriors.PrecomputeCacheType = posteriors.PrecomputeCacheType.TENSOR,
    ) -> posteriors.SGPRPosterior:
        """
        Create the Posterior object which contains precomputed matrices for
        faster prediction.

        precompute_cache has three settings:

        - `PrecomputeCacheType.TENSOR` (or `"tensor"`): Precomputes the cached
          quantities and stores them as tensors (which allows differentiating
          through the prediction). This is the default.
        - `PrecomputeCacheType.VARIABLE` (or `"variable"`): Precomputes the cached
          quantities and stores them as variables, which allows for updating
          their values without changing the compute graph (relevant for AOT
          compilation).
        - `PrecomputeCacheType.NOCACHE` (or `"nocache"` or `None`): Avoids
          immediate cache computation. This is useful for avoiding extraneous
          computations when you only want to call the posterior's
          `fused_predict_f` method.
        """

        return posteriors.SGPRPosterior(
            kernel=self.kernel,
            data=self.data,
            inducing_variable=self.inducing_variable,
            likelihood=self.likelihood,
            num_latent_gps=self.num_latent_gps,
            mean_function=self.mean_function,
            precompute_cache=precompute_cache,
        )

    @inherit_check_shapes
    def predict_f(
        self, Xnew: InputData, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        For backwards compatibility, GPR's predict_f uses the fused (no-cache)
        computation, which is more efficient during training.

        For faster (cached) prediction, predict directly from the posterior object, i.e.,:
            model.posterior().predict_f(Xnew, ...)
        """
        return self.posterior(posteriors.PrecomputeCacheType.NOCACHE).fused_predict_f(
            Xnew, full_cov=full_cov, full_output_cov=full_output_cov
        )


class SGPR(SGPR_with_posterior):
    # subclassed to ensure __class__ == "SGPR"
    pass


# gpflow/models/svgp.py

# Copyright 2016-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Optional

import numpy as np
import tensorflow as tf

from .. import kullback_leiblers, posteriors
from ..base import AnyNDArray, InputData, MeanAndVariance, Parameter, RegressionData
from ..conditionals import conditional
from ..config import default_float
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..inducing_variables import InducingVariables
from ..kernels import Kernel
from ..likelihoods import Likelihood
from ..mean_functions import MeanFunction
from ..utilities import positive, triangular
from .model import GPModel
from .training_mixins import ExternalDataTrainingLossMixin
from .util import InducingVariablesLike, inducingpoint_wrapper


class SVGP_deprecated(GPModel, ExternalDataTrainingLossMixin):
    """
    This is the Sparse Variational GP (SVGP).

    The key reference is :cite:t:`hensman2014scalable`.
    """

    @check_shapes(
        "q_mu: [M, P]",
        "q_sqrt: [M, P] if q_diag",
        "q_sqrt: [P, M, M] if (not q_diag)",
    )
    def __init__(
        self,
        kernel: Kernel,
        likelihood: Likelihood,
        inducing_variable: InducingVariablesLike,
        *,
        mean_function: Optional[MeanFunction] = None,
        num_latent_gps: int = 1,
        q_diag: bool = False,
        q_mu: Optional[tf.Tensor] = None,
        q_sqrt: Optional[tf.Tensor] = None,
        whiten: bool = True,
        num_data: Optional[tf.Tensor] = None,
    ):
        """
        - kernel, likelihood, inducing_variables, mean_function are appropriate
          GPflow objects
        - num_latent_gps is the number of latent processes to use, defaults to 1
        - q_diag is a boolean. If True, the covariance is approximated by a
          diagonal matrix.
        - whiten is a boolean. If True, we use the whitened representation of
          the inducing points.
        - num_data is the total number of observations, defaults to X.shape[0]
          (relevant when feeding in external minibatches)
        """
        # init the super class, accept args
        super().__init__(kernel, likelihood, mean_function, num_latent_gps)
        self.num_data = num_data
        self.whiten = whiten
        self.inducing_variable: InducingVariables = inducingpoint_wrapper(inducing_variable)

        # init variational parameters
        num_inducing = self.inducing_variable.num_inducing
        self._init_variational_parameters(num_inducing, q_mu, q_sqrt, q_diag)

    @check_shapes(
        "q_mu: [M, P]",
        "q_sqrt: [M, P] if q_diag",
        "q_sqrt: [P, M, M] if (not q_diag)",
    )
    def _init_variational_parameters(
        self,
        num_inducing: int,
        q_mu: Optional[AnyNDArray],
        q_sqrt: Optional[AnyNDArray],
        q_diag: bool,
    ) -> None:
        """
        Constructs the mean and cholesky of the covariance of the variational Gaussian posterior.
        If a user passes values for `q_mu` and `q_sqrt` the routine checks if they have consistent
        and correct shapes. If a user does not specify any values for `q_mu` and `q_sqrt`, the
        routine initializes them, their shape depends on `num_inducing` and `q_diag`.

        Note: most often the comments refer to the number of observations (=output dimensions) with
        P, number of latent GPs with L, and number of inducing points M. Typically P equals L,
        but when certain multioutput kernels are used, this can change.

        Parameters
        ----------
        :param num_inducing:
            Number of inducing variables, typically refered to as M.
        :param q_mu:
            Mean of the variational Gaussian posterior. If None the function will initialise
            the mean with zeros. If not None, the shape of `q_mu` is checked.
        :param q_sqrt:
            Cholesky of the covariance of the variational Gaussian posterior.
            If None the function will initialise `q_sqrt` with identity matrix.
            If not None, the shape of `q_sqrt` is checked, depending on `q_diag`.
        :param q_diag:
            Used to check if `q_mu` and `q_sqrt` have the correct shape or to
            construct them with the correct shape. If `q_diag` is true,
            `q_sqrt` is two dimensional and only holds the square root of the
            covariance diagonal elements. If False, `q_sqrt` is three dimensional.
        """
        q_mu = np.zeros((num_inducing, self.num_latent_gps)) if q_mu is None else q_mu
        self.q_mu = Parameter(q_mu, dtype=default_float())  # [M, P]

        if q_sqrt is None:
            if q_diag:
                ones: AnyNDArray = np.ones(
                    (num_inducing, self.num_latent_gps), dtype=default_float()
                )
                self.q_sqrt = Parameter(ones, transform=positive())  # [M, P]
            else:
                np_q_sqrt: AnyNDArray = np.array(
                    [
                        np.eye(num_inducing, dtype=default_float())
                        for _ in range(self.num_latent_gps)
                    ]
                )
                self.q_sqrt = Parameter(np_q_sqrt, transform=triangular())  # [P, M, M]
        else:
            if q_diag:
                self.num_latent_gps = q_sqrt.shape[1]
                self.q_sqrt = Parameter(q_sqrt, transform=positive())  # [M, L|P]
            else:
                self.num_latent_gps = q_sqrt.shape[0]
                num_inducing = q_sqrt.shape[1]
                self.q_sqrt = Parameter(q_sqrt, transform=triangular())  # [L|P, M, M]

    @check_shapes(
        "return: []",
    )
    def prior_kl(self) -> tf.Tensor:
        return kullback_leiblers.prior_kl(
            self.inducing_variable, self.kernel, self.q_mu, self.q_sqrt, whiten=self.whiten
        )

    # type-ignore is because of changed method signature:
    @inherit_check_shapes
    def maximum_log_likelihood_objective(self, data: RegressionData) -> tf.Tensor:  # type: ignore[override]
        return self.elbo(data)

    @check_shapes(
        "return: []",
    )
    def elbo(self, data: RegressionData) -> tf.Tensor:
        """
        This gives a variational bound (the evidence lower bound or ELBO) on
        the log marginal likelihood of the model.
        """
        X, Y = data
        kl = self.prior_kl()
        f_mean, f_var = self.predict_f(X, full_cov=False, full_output_cov=False)
        var_exp = self.likelihood.variational_expectations(X, f_mean, f_var, Y)
        if self.num_data is not None:
            num_data = tf.cast(self.num_data, kl.dtype)
            minibatch_size = tf.cast(tf.shape(X)[0], kl.dtype)
            scale = num_data / minibatch_size
        else:
            scale = tf.cast(1.0, kl.dtype)
        return tf.reduce_sum(var_exp) * scale - kl

    @inherit_check_shapes
    def predict_f(
        self, Xnew: InputData, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        mu, var = conditional(
            Xnew,
            self.inducing_variable,
            self.kernel,
            self.q_mu,
            q_sqrt=self.q_sqrt,
            full_cov=full_cov,
            white=self.whiten,
            full_output_cov=full_output_cov,
        )
        # tf.debugging.assert_positive(var)  # We really should make the tests pass with this here
        return mu + self.mean_function(Xnew), var


class SVGP_with_posterior(SVGP_deprecated):
    """
    This is the Sparse Variational GP (SVGP).

    The key reference is :cite:t:`hensman2014scalable`.

    This class provides a posterior() method that enables caching for faster subsequent predictions.
    """

    def posterior(
        self,
        precompute_cache: posteriors.PrecomputeCacheType = posteriors.PrecomputeCacheType.TENSOR,
    ) -> posteriors.BasePosterior:
        """
        Create the Posterior object which contains precomputed matrices for
        faster prediction.

        precompute_cache has three settings:

        - `PrecomputeCacheType.TENSOR` (or `"tensor"`): Precomputes the cached
          quantities and stores them as tensors (which allows differentiating
          through the prediction). This is the default.
        - `PrecomputeCacheType.VARIABLE` (or `"variable"`): Precomputes the cached
          quantities and stores them as variables, which allows for updating
          their values without changing the compute graph (relevant for AOT
          compilation).
        - `PrecomputeCacheType.NOCACHE` (or `"nocache"` or `None`): Avoids
          immediate cache computation. This is useful for avoiding extraneous
          computations when you only want to call the posterior's
          `fused_predict_f` method.
        """
        return posteriors.create_posterior(
            self.kernel,
            self.inducing_variable,
            self.q_mu,
            self.q_sqrt,
            whiten=self.whiten,
            mean_function=self.mean_function,
            precompute_cache=precompute_cache,
        )

    @inherit_check_shapes
    def predict_f(
        self, Xnew: InputData, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        For backwards compatibility, SVGP's predict_f uses the fused (no-cache)
        computation, which is more efficient during training.

        For faster (cached) prediction, predict directly from the posterior object, i.e.,:
            model.posterior().predict_f(Xnew, ...)
        """
        return self.posterior(posteriors.PrecomputeCacheType.NOCACHE).fused_predict_f(
            Xnew, full_cov=full_cov, full_output_cov=full_output_cov
        )


class SVGP(SVGP_with_posterior):
    # subclassed to ensure __class__ == "SVGP"
    pass


# gpflow/models/training_mixins.py

# Copyright 2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
This module provides mixin classes to be used in conjunction with inheriting
from gpflow.models.BayesianModel (or its subclass gpflow.models.GPModel).

They provide a unified interface to obtain closures that return the training
loss, to be passed as the first argument to the minimize() method of the
optimizers defined in TensorFlow and GPflow.

All TrainingLossMixin classes assume that self._training_loss()
(which is provided by the BayesianModel base class), will be available. Note
that new models only need to implement the maximum_log_likelihood_objective
method that is defined as abstract in BayesianModel.

There are different mixins depending on whether the model already contains the
training data (InternalDataTrainingLossMixin), or requires it to be passed in
to the objective function (ExternalDataTrainingLossMixin).
"""
from typing import Callable, TypeVar, Union

import tensorflow as tf
from tensorflow.python.data.ops.iterator_ops import OwnedIterator as DatasetOwnedIterator

from ..base import InputData, OutputData, RegressionData
from ..experimental.check_shapes import check_shapes

Data = TypeVar("Data", RegressionData, InputData, OutputData)


class InternalDataTrainingLossMixin:
    """
    Mixin utility for training loss methods for models that own their own data. It provides

      - a uniform API for the training loss :meth:`training_loss`
      - a convenience method :meth:`training_loss_closure` for constructing the closure expected by
        various optimizers, namely :class:`gpflow.optimizers.Scipy` and subclasses of
        `tf.optimizers.Optimizer`.

    See :class:`ExternalDataTrainingLossMixin` for an equivalent mixin for models that do **not**
    own their own data.
    """

    @check_shapes(
        "return: []",
    )
    def training_loss(self) -> tf.Tensor:
        """
        Returns the training loss for this model.
        """
        # Type-ignore is because _training_loss should be added by implementing class.
        return self._training_loss()  # type: ignore[attr-defined]

    def training_loss_closure(self, *, compile: bool = True) -> Callable[[], tf.Tensor]:
        """
        Convenience method. Returns a closure which itself returns the training loss. This closure
        can be passed to the minimize methods on :class:`gpflow.optimizers.Scipy` and subclasses of
        `tf.optimizers.Optimizer`.

        :param compile: If `True` (default), compile the training loss function in a TensorFlow
            graph by wrapping it in tf.function()
        """
        closure = self.training_loss
        if compile:
            closure = tf.function(closure)
        return closure


class ExternalDataTrainingLossMixin:
    """
    Mixin utility for training loss methods for models that do **not** own their own data.
    It provides

      - a uniform API for the training loss :meth:`training_loss`
      - a convenience method :meth:`training_loss_closure` for constructing the closure expected by
        various optimizers, namely :class:`gpflow.optimizers.Scipy` and subclasses of
        `tf.optimizers.Optimizer`.

    See :class:`InternalDataTrainingLossMixin` for an equivalent mixin for models that **do** own
    their own data.
    """

    @check_shapes(
        "data[0]: [N, D]",
        "data[1]: [N, P]",
        "return: []",
    )
    def training_loss(self, data: Data) -> tf.Tensor:
        """
        Returns the training loss for this model.

        :param data: the data to be used for computing the model objective.
        """
        # Type-ignore is because _training_loss should be added by implementing class.
        return self._training_loss(data)  # type: ignore[attr-defined]

    def training_loss_closure(
        self,
        data: Union[Data, DatasetOwnedIterator],
        *,
        compile: bool = True,
    ) -> Callable[[], tf.Tensor]:
        """
        Returns a closure that computes the training loss, which by default is
        wrapped in tf.function(). This can be disabled by passing `compile=False`.

        :param data: the data to be used by the closure for computing the model
            objective. Can be the full dataset or an iterator, e.g.
            `iter(dataset.batch(batch_size))`, where dataset is an instance of
            tf.data.Dataset.
        :param compile: if True, wrap training loss in tf.function()
        """
        training_loss = self.training_loss

        if isinstance(data, DatasetOwnedIterator):
            if compile:
                # lambda because: https://github.com/GPflow/GPflow/issues/1929
                training_loss_lambda = lambda d: self.training_loss(d)
                input_signature = [data.element_spec]
                training_loss = tf.function(training_loss_lambda, input_signature=input_signature)

            def closure() -> tf.Tensor:
                assert isinstance(data, DatasetOwnedIterator)  # Hint for mypy.
                batch = next(data)
                return training_loss(batch)

        else:

            def closure() -> tf.Tensor:
                return training_loss(data)

            if compile:
                closure = tf.function(closure)

        return closure


# gpflow/models/util.py

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Any, Callable, Sequence, Union

import numpy as np
import tensorflow as tf

from ..base import AnyNDArray
from ..config import default_float
from ..experimental.check_shapes import check_shapes
from ..inducing_variables import InducingPoints, InducingVariables
from .model import BayesianModel
from .training_mixins import Data, ExternalDataTrainingLossMixin

InducingVariablesLike = Union[InducingVariables, tf.Tensor, AnyNDArray]
InducingPointsLike = Union[InducingPoints, tf.Tensor, AnyNDArray]


def inducingpoint_wrapper(inducing_variable: InducingVariablesLike) -> InducingVariables:
    """
    This wrapper allows transparently passing either an InducingVariables
    object or an array specifying InducingPoints positions.
    """
    if not isinstance(inducing_variable, InducingVariables):
        inducing_variable = InducingPoints(inducing_variable)
    return inducing_variable


def _assert_equal_data(
    data1: Union[tf.Tensor, Sequence[tf.Tensor]], data2: Union[tf.Tensor, Sequence[tf.Tensor]]
) -> None:
    if isinstance(data1, tf.Tensor) and isinstance(data2, tf.Tensor):
        tf.debugging.assert_equal(data1, data2)
    else:
        for v1, v2 in zip(data1, data2):
            tf.debugging.assert_equal(v1, v2)


@check_shapes(
    "data[0]: [N, D]",
    "data[1]: [N, P]",
)
def training_loss_closure(
    model: BayesianModel, data: Data, **closure_kwargs: Any
) -> Callable[[], tf.Tensor]:
    if isinstance(model, ExternalDataTrainingLossMixin):
        return model.training_loss_closure(data, **closure_kwargs)  # type: ignore[no-any-return]
    else:
        _assert_equal_data(model.data, data)
        return model.training_loss_closure(**closure_kwargs)  # type: ignore[no-any-return]


@check_shapes(
    "data[0]: [N, D]",
    "data[1]: [N, P]",
    "return: []",
)
def training_loss(model: BayesianModel, data: Data) -> tf.Tensor:
    if isinstance(model, ExternalDataTrainingLossMixin):
        return model.training_loss(data)
    else:
        _assert_equal_data(model.data, data)
        return model.training_loss()


@check_shapes(
    "data[0]: [N, D]",
    "data[1]: [N, P]",
    "return: []",
)
def maximum_log_likelihood_objective(model: BayesianModel, data: Data) -> tf.Tensor:
    if isinstance(model, ExternalDataTrainingLossMixin):
        return model.maximum_log_likelihood_objective(data)
    else:
        _assert_equal_data(model.data, data)
        return model.maximum_log_likelihood_objective()


def data_input_to_tensor(structure: Any) -> Any:
    """
    Converts non-tensor elements of a structure to TensorFlow tensors retaining the structure
    itself.

    The function doesn't keep original element's dtype and forcefully converts
    them to GPflow's default float type.
    """

    def convert_to_tensor(elem: Any) -> tf.Tensor:
        if tf.is_tensor(elem):
            return elem
        elif isinstance(elem, np.ndarray):
            return tf.convert_to_tensor(elem)
        return tf.convert_to_tensor(elem, dtype=default_float())

    return tf.nest.map_structure(convert_to_tensor, structure)


# gpflow/models/vgp.py

# Copyright 2016-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Optional

import numpy as np
import tensorflow as tf

import gpflow

from .. import posteriors
from ..base import InputData, MeanAndVariance, Parameter, RegressionData
from ..conditionals import conditional
from ..config import default_float, default_jitter
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..kernels import Kernel
from ..kullback_leiblers import gauss_kl
from ..likelihoods import Likelihood
from ..mean_functions import MeanFunction
from ..utilities import assert_params_false, is_variable, triangular, triangular_size
from .model import GPModel
from .training_mixins import InternalDataTrainingLossMixin
from .util import data_input_to_tensor


class VGP_deprecated(GPModel, InternalDataTrainingLossMixin):
    r"""
    This method approximates the Gaussian process posterior using a multivariate Gaussian.

    The idea is that the posterior over the function-value vector F is
    approximated by a Gaussian, and the KL divergence is minimised between
    the approximation and the posterior.

    This implementation is equivalent to SVGP with X=Z, but is more efficient.
    The whitened representation is used to aid optimization.

    The posterior approximation is

    .. math::

       q(\mathbf f) = N(\mathbf f \,|\, \boldsymbol \mu, \boldsymbol \Sigma)

    """

    @check_shapes(
        "data[0]: [N, D]",
        "data[1]: [N, P]",
    )
    def __init__(
        self,
        data: RegressionData,
        kernel: Kernel,
        likelihood: Likelihood,
        mean_function: Optional[MeanFunction] = None,
        num_latent_gps: Optional[int] = None,
    ):
        """
        data = (X, Y) contains the input points [N, D] and the observations [N, P]
        kernel, likelihood, mean_function are appropriate GPflow objects
        """
        if num_latent_gps is None:
            num_latent_gps = self.calc_num_latent_gps_from_data(data, kernel, likelihood)
        super().__init__(kernel, likelihood, mean_function, num_latent_gps)

        self.data = data_input_to_tensor(data)
        X_data, _Y_data = self.data

        static_num_data = X_data.shape[0]
        if static_num_data is None:
            q_sqrt_unconstrained_shape = (self.num_latent_gps, None)
        else:
            q_sqrt_unconstrained_shape = (self.num_latent_gps, triangular_size(static_num_data))
        self.num_data = Parameter(tf.shape(X_data)[0], shape=[], dtype=tf.int32, trainable=False)

        # Many functions below don't like `Parameter`s:
        dynamic_num_data = tf.convert_to_tensor(self.num_data)

        self.q_mu = Parameter(
            tf.zeros((dynamic_num_data, self.num_latent_gps)),
            shape=(static_num_data, num_latent_gps),
        )
        q_sqrt = tf.eye(dynamic_num_data, batch_shape=[self.num_latent_gps])
        self.q_sqrt = Parameter(
            q_sqrt,
            transform=triangular(),
            unconstrained_shape=q_sqrt_unconstrained_shape,
            constrained_shape=(num_latent_gps, static_num_data, static_num_data),
        )

    # type-ignore is because of changed method signature:
    @inherit_check_shapes
    def maximum_log_likelihood_objective(self) -> tf.Tensor:  # type: ignore[override]
        return self.elbo()

    @check_shapes(
        "return: []",
    )
    def elbo(self) -> tf.Tensor:
        r"""
        This method computes the variational lower bound on the likelihood,
        which is:

            E_{q(F)} [ \log p(Y|F) ] - KL[ q(F) || p(F)]

        with

            q(\mathbf f) = N(\mathbf f \,|\, \boldsymbol \mu, \boldsymbol \Sigma)

        """
        X_data, Y_data = self.data
        num_data = tf.convert_to_tensor(self.num_data)

        # Get prior KL.
        KL = gauss_kl(self.q_mu, self.q_sqrt)

        # Get conditionals
        K = self.kernel(X_data) + tf.eye(num_data, dtype=default_float()) * default_jitter()
        L = tf.linalg.cholesky(K)
        fmean = tf.linalg.matmul(L, self.q_mu) + self.mean_function(X_data)  # [NN, ND] -> ND
        q_sqrt_dnn = tf.linalg.band_part(self.q_sqrt, -1, 0)  # [D, N, N]
        L_tiled = tf.tile(tf.expand_dims(L, 0), tf.stack([self.num_latent_gps, 1, 1]))
        LTA = tf.linalg.matmul(L_tiled, q_sqrt_dnn)  # [D, N, N]
        fvar = tf.reduce_sum(tf.square(LTA), 2)

        fvar = tf.transpose(fvar)

        # Get variational expectations.
        var_exp = self.likelihood.variational_expectations(X_data, fmean, fvar, Y_data)

        return tf.reduce_sum(var_exp) - KL

    @inherit_check_shapes
    def predict_f(
        self, Xnew: InputData, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        assert_params_false(self.predict_f, full_output_cov=full_output_cov)

        X_data, _Y_data = self.data
        mu, var = conditional(
            Xnew,
            X_data,
            self.kernel,
            self.q_mu,
            q_sqrt=self.q_sqrt,
            full_cov=full_cov,
            white=True,
        )
        return mu + self.mean_function(Xnew), var


class VGP_with_posterior(VGP_deprecated):
    """
    This is an implementation of VGP that provides a posterior() method that
    enables caching for faster subsequent predictions.
    """

    def posterior(
        self,
        precompute_cache: posteriors.PrecomputeCacheType = posteriors.PrecomputeCacheType.TENSOR,
    ) -> posteriors.VGPPosterior:
        """
        Create the Posterior object which contains precomputed matrices for
        faster prediction.

        precompute_cache has three settings:

        - `PrecomputeCacheType.TENSOR` (or `"tensor"`): Precomputes the cached
          quantities and stores them as tensors (which allows differentiating
          through the prediction). This is the default.
        - `PrecomputeCacheType.VARIABLE` (or `"variable"`): Precomputes the cached
          quantities and stores them as variables, which allows for updating
          their values without changing the compute graph (relevant for AOT
          compilation).
        - `PrecomputeCacheType.NOCACHE` (or `"nocache"` or `None`): Avoids
          immediate cache computation. This is useful for avoiding extraneous
          computations when you only want to call the posterior's
          `fused_predict_f` method.
        """
        X_data, _Y_data = self.data
        return posteriors.VGPPosterior(
            self.kernel,
            X_data,
            self.q_mu,
            self.q_sqrt,
            mean_function=self.mean_function,
            precompute_cache=precompute_cache,
        )

    @inherit_check_shapes
    def predict_f(
        self, Xnew: InputData, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        For backwards compatibility, VGP's predict_f uses the fused (no-cache)
        computation, which is more efficient during training.

        For faster (cached) prediction, predict directly from the posterior object, i.e.,:
            model.posterior().predict_f(Xnew, ...)
        """
        return self.posterior(posteriors.PrecomputeCacheType.NOCACHE).fused_predict_f(
            Xnew, full_cov=full_cov, full_output_cov=full_output_cov
        )


class VGP(VGP_with_posterior):
    # subclassed to ensure __class__ == "VGP"
    pass


@check_shapes(
    "new_data[0]: [N, D]",
    "new_data[1]: [N, P]",
)
def update_vgp_data(vgp: VGP_deprecated, new_data: RegressionData) -> None:
    """
    Set the data on the given VGP model, and update its variational parameters.

    As opposed to many of the other models the VGP has internal parameters whose shape depends on
    the shape of the data. This functions updates the internal data of the given vgp, and updates
    the variational parameters to fit.

    This function requires that the input `vgp` were create with :class:`tf.Variable`s for
    `data`.
    """
    old_X_data, old_Y_data = vgp.data
    assert is_variable(old_X_data) and is_variable(
        old_Y_data
    ), "update_vgp_data requires the model to have been created with variable data."

    new_X_data, new_Y_data = new_data
    new_num_data = tf.shape(new_X_data)[0]
    f_mu, f_cov = vgp.predict_f(new_X_data, full_cov=True)  # [N, L], [L, N, N]

    # This model is hard-coded to use the whitened representation, i.e.  q_mu and q_sqrt
    # parametrize q(v), and u = f(X) = L v, where L = cholesky(K(X, X)) Hence we need to
    # back-transform from f_mu and f_cov to obtain the updated new_q_mu and new_q_sqrt:
    Knn = vgp.kernel(new_X_data, full_cov=True)  # [N, N]
    jitter_mat = default_jitter() * tf.eye(new_num_data, dtype=Knn.dtype)
    Lnn = tf.linalg.cholesky(Knn + jitter_mat)  # [N, N]
    new_q_mu = tf.linalg.triangular_solve(Lnn, f_mu)  # [N, L]
    tmp = tf.linalg.triangular_solve(Lnn[None], f_cov)  # [L, N, N], L⁻¹ f_cov
    S_v = tf.linalg.triangular_solve(Lnn[None], tf.linalg.matrix_transpose(tmp))  # [L, N, N]
    new_q_sqrt = tf.linalg.cholesky(S_v + jitter_mat)  # [L, N, N]

    old_X_data.assign(new_X_data)
    old_Y_data.assign(new_Y_data)
    vgp.num_data.assign(new_num_data)
    vgp.q_mu.assign(new_q_mu)
    vgp.q_sqrt.assign(new_q_sqrt)


class VGPOpperArchambeau(GPModel, InternalDataTrainingLossMixin):
    r"""
    This method approximates the Gaussian process posterior using a multivariate Gaussian.

    The key reference is :cite:t:`Opper:2009`.

    The idea is that the posterior over the function-value vector F is
    approximated by a Gaussian, and the KL divergence is minimised between
    the approximation and the posterior. It turns out that the optimal
    posterior precision shares off-diagonal elements with the prior, so
    only the diagonal elements of the precision need be adjusted.
    The posterior approximation is

    .. math::

       q(\mathbf f) = N(\mathbf f \,|\, \mathbf K \boldsymbol \alpha,
                         [\mathbf K^{-1} + \textrm{diag}(\boldsymbol \lambda))^2]^{-1})

    This approach has only 2ND parameters, rather than the N + N^2 of vgp,
    but the optimization is non-convex and in practice may cause difficulty.

    """

    @check_shapes(
        "data[0]: [N, D]",
        "data[1]: [N, P]",
    )
    def __init__(
        self,
        data: RegressionData,
        kernel: Kernel,
        likelihood: Likelihood,
        mean_function: Optional[MeanFunction] = None,
        num_latent_gps: Optional[int] = None,
    ):
        """
        data = (X, Y) contains the input points [N, D] and the observations [N, P]
        kernel, likelihood, mean_function are appropriate GPflow objects
        """
        if num_latent_gps is None:
            num_latent_gps = self.calc_num_latent_gps_from_data(data, kernel, likelihood)
        super().__init__(kernel, likelihood, mean_function, num_latent_gps)

        self.data = data_input_to_tensor(data)
        X_data, _Y_data = self.data
        self.num_data = X_data.shape[0]
        self.q_alpha = Parameter(np.zeros((self.num_data, self.num_latent_gps)))
        self.q_lambda = Parameter(
            np.ones((self.num_data, self.num_latent_gps)), transform=gpflow.utilities.positive()
        )

    # type-ignore is because of changed method signature:
    @inherit_check_shapes
    def maximum_log_likelihood_objective(self) -> tf.Tensor:  # type: ignore[override]
        return self.elbo()

    @check_shapes(
        "return: []",
    )
    def elbo(self) -> tf.Tensor:
        r"""
        q_alpha, q_lambda are variational parameters, size [N, R]
        This method computes the variational lower bound on the likelihood,
        which is:

        .. math::

           E_{q(F)} [ \log p(Y|F) ] - KL[ q(F) || p(F)]

        with

        .. math::

           q(f) = N(f |
               K \alpha + \textrm{mean},
               [K^-1 + \textrm{diag}(\textrm{square}(\lambda))]^-1) .
        """
        X_data, Y_data = self.data

        K = self.kernel(X_data)
        K_alpha = tf.linalg.matmul(K, self.q_alpha)
        f_mean = K_alpha + self.mean_function(X_data)

        # compute the variance for each of the outputs
        I = tf.tile(
            tf.eye(self.num_data, dtype=default_float())[None, ...], [self.num_latent_gps, 1, 1]
        )
        A = (
            I
            + tf.transpose(self.q_lambda)[:, None, ...]
            * tf.transpose(self.q_lambda)[:, :, None, ...]
            * K
        )
        L = tf.linalg.cholesky(A)
        Li = tf.linalg.triangular_solve(L, I)
        tmp = Li / tf.transpose(self.q_lambda)[:, None, ...]
        f_var = 1.0 / tf.square(self.q_lambda) - tf.transpose(tf.reduce_sum(tf.square(tmp), 1))

        # some statistics about A are used in the KL
        A_logdet = 2.0 * tf.reduce_sum(tf.math.log(tf.linalg.diag_part(L)))
        trAi = tf.reduce_sum(tf.square(Li))

        KL = 0.5 * (
            A_logdet
            + trAi
            - self.num_data * self.num_latent_gps
            + tf.reduce_sum(K_alpha * self.q_alpha)
        )

        v_exp = self.likelihood.variational_expectations(X_data, f_mean, f_var, Y_data)
        return tf.reduce_sum(v_exp) - KL

    @inherit_check_shapes
    def predict_f(
        self, Xnew: InputData, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        r"""
        The posterior variance of F is given by

        .. math::

           q(f) = N(f |
               K \alpha + \textrm{mean}, [K^-1 + \textrm{diag}(\lambda**2)]^-1)

        Here we project this to F*, the values of the GP at Xnew which is given
        by

        .. math::

           q(F*) = N ( F* | K_{*F} \alpha + \textrm{mean}, K_{**} - K_{*f}[K_{ff} +
                                           \textrm{diag}(\lambda**-2)]^-1 K_{f*} )

        Note: This model currently does not allow full output covariances
        """
        assert_params_false(self.predict_f, full_output_cov=full_output_cov)

        X_data, _ = self.data
        # compute kernel things
        Kx = self.kernel(X_data, Xnew)
        K = self.kernel(X_data)

        # predictive mean
        f_mean = tf.linalg.matmul(Kx, self.q_alpha, transpose_a=True) + self.mean_function(Xnew)

        # predictive var
        A = K + tf.linalg.diag(tf.transpose(1.0 / tf.square(self.q_lambda)))
        L = tf.linalg.cholesky(A)
        Kx_tiled = tf.tile(Kx[None, ...], [self.num_latent_gps, 1, 1])
        LiKx = tf.linalg.triangular_solve(L, Kx_tiled)
        if full_cov:
            f_var = self.kernel(Xnew) - tf.linalg.matmul(LiKx, LiKx, transpose_a=True)
        else:
            f_var = self.kernel(Xnew, full_cov=False) - tf.reduce_sum(tf.square(LiKx), axis=1)
        return f_mean, tf.transpose(f_var)


# gpflow/monitor/__init__.py

# Copyright 2020 GPflow authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
""" Provides basic functionality to monitor optimisation runs """

from .base import *
from .tensorboard import *

__all__ = [
    "ExecuteCallback",
    "ImageToTensorBoard",
    "ModelToTensorBoard",
    "Monitor",
    "MonitorTask",
    "MonitorTaskGroup",
    "ScalarToTensorBoard",
    "ToTensorBoard",
    "base",
    "tensorboard",
]


# gpflow/monitor/base.py

# Copyright 2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""" MonitorTask base classes """

from abc import ABC, abstractmethod
from typing import Any, Callable, Collection, Union

import tensorflow as tf

__all__ = [
    "ExecuteCallback",
    "Monitor",
    "MonitorTask",
    "MonitorTaskGroup",
]


class MonitorTask(ABC):
    """
    A base class for a monitoring task.

    All monitoring tasks are callable objects.
    A descendant class must implement the `run` method, which is the body of the monitoring task.
    """

    def __call__(self, step: int, **kwargs: Any) -> None:
        """
        It calls the 'run' function and sets the current step.

        :param step: current step in the optimisation.
        :param kwargs: additional keyword arguments that can be passed
            to the `run` method of the task. This is in particular handy for
            passing keyword argument to the callback of `ScalarToTensorBoard`.
        """
        self.current_step = tf.cast(step, tf.int64)
        self.run(**kwargs)

    @abstractmethod
    def run(self, **kwargs: Any) -> None:
        """
        Implements the task to be executed on __call__.
        The current step is available through `self.current_step`.

        :param kwargs: keyword arguments available to the run method.
        """
        raise NotImplementedError


class ExecuteCallback(MonitorTask):
    """ Executes a callback as task """

    def __init__(self, callback: Callable[..., None]) -> None:
        """
        :param callback: callable to be executed during the task.
            Arguments can be passed using keyword arguments.
        """
        super().__init__()
        self.callback = callback

    def run(self, **kwargs: Any) -> None:
        self.callback(**kwargs)


class MonitorTaskGroup:
    """
    Class for grouping `MonitorTask` instances. A group defines
    all the tasks that are run at the same frequency, given by `period`.

    A `MonitorTaskGroup` can exist of a single instance or a list of
    `MonitorTask` instances.
    """

    def __init__(
        self, task_or_tasks: Union[Collection[MonitorTask], MonitorTask], period: int = 1
    ) -> None:
        """
        :param task_or_tasks: a single instance or a list of `MonitorTask` instances.
            Each `MonitorTask` in the list will be run with the given `period`.
        :param period: defines how often to run the tasks; they will execute every `period`th step.
            For large values of `period` the tasks will be less frequently run. Defaults to
            running at every step (`period = 1`).
        """
        self._tasks: Collection[MonitorTask] = []
        self.tasks = task_or_tasks  # type: ignore[assignment]
        self._period = period

    @property
    def tasks(self) -> Collection[MonitorTask]:
        return self._tasks

    @tasks.setter
    def tasks(self, task_or_tasks: Union[Collection[MonitorTask], MonitorTask]) -> None:
        """Ensures the tasks are stored as a list. Even if there is only a single task."""
        if isinstance(task_or_tasks, MonitorTask):
            self._tasks = [task_or_tasks]
        else:
            assert isinstance(task_or_tasks, Collection)
            self._tasks = list(task_or_tasks)

    def __call__(self, step: int, **kwargs: Any) -> None:
        """Call each task in the group."""
        if step % self._period == 0:
            for task in self.tasks:
                task(step, **kwargs)


class Monitor:
    r"""
    Accepts any number of of `MonitorTaskGroup` instances, and runs them
    according to their specified periodicity.

    Example use-case::

        # Create some monitor tasks
        log_dir = "logs"
        model_task = ModelToTensorBoard(log_dir, model)
        image_task = ImageToTensorBoard(log_dir, plot_prediction, "image_samples")
        lml_task = ScalarToTensorBoard(log_dir, lambda: model.log_marginal_likelihood(), "lml")

        # Plotting tasks can be quite slow, so we want to run them less frequently.
        # We group them in a `MonitorTaskGroup` and set the period to 5.
        slow_tasks = MonitorTaskGroup(image_task, period=5)

        # The other tasks are fast. We run them at each iteration of the optimisation.
        fast_tasks = MonitorTaskGroup([model_task, lml_task], period=1)

        # We pass both groups to the `Monitor`
        monitor = Monitor(fast_tasks, slow_tasks)
    """

    def __init__(self, *task_groups: MonitorTaskGroup) -> None:
        """
        :param task_groups: a list of `MonitorTaskGroup`s to be executed.
        """
        self.task_groups = task_groups

    def __call__(self, step: int, **kwargs: Any) -> None:
        for group in self.task_groups:
            group(step, **kwargs)


# gpflow/monitor/tensorboard.py

# Copyright 2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""" Tasks that write to TensorBoard """

from io import BytesIO
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union

import numpy as np
import tensorflow as tf

from ..base import Parameter
from ..models import BayesianModel
from ..utilities import parameter_dict
from .base import MonitorTask

if TYPE_CHECKING:  # pragma: no cover
    import matplotlib


__all__ = [
    "ImageToTensorBoard",
    "ModelToTensorBoard",
    "ScalarToTensorBoard",
    "ToTensorBoard",
]


class ToTensorBoard(MonitorTask):
    writers: Dict[str, tf.summary.SummaryWriter] = {}

    def __init__(self, log_dir: str) -> None:
        """
        :param log_dir: directory in which to store the tensorboard files.
            Can be nested, e.g. ./logs/my_run/
        """
        super().__init__()
        if log_dir not in self.writers:
            self.writers[log_dir] = tf.summary.create_file_writer(log_dir)
        self.file_writer = self.writers[log_dir]

    def __call__(self, step: int, **kwargs: Any) -> None:
        with self.file_writer.as_default():
            super().__call__(step, **kwargs)
        self.file_writer.flush()


class ModelToTensorBoard(ToTensorBoard):
    """
    Monitoring task that creates a sensible TensorBoard for a model.

    Monitors all the model's parameters for which their name matches with `keywords_to_monitor`.
    By default, "kernel" and "likelihood" are elements of `keywords_to_monitor`.

    Example::

        keyword = "kernel", parameter = "kernel.lengthscale" => match
        keyword = "variational", parameter = "kernel.lengthscale" => no match
    """

    def __init__(
        self,
        log_dir: str,
        model: BayesianModel,
        *,
        max_size: int = 3,
        keywords_to_monitor: List[str] = ["kernel", "likelihood"],
        left_strip_character: str = ".",
    ) -> None:
        """
        :param log_dir: directory in which to store the tensorboard files.
            Can be a nested: for example, './logs/my_run/'.
        :param model: model to be monitord.
        :param max_size: maximum size of arrays (incl.) to store each
            element of the array independently as a scalar in the TensorBoard.
            Setting max_size to -1 will write all values. Use with care.
        :param keywords_to_monitor: specifies keywords to be monitored.
            If the parameter's name includes any of the keywords specified it
            will be monitored. By default, parameters that match the `kernel` or
            `likelihood` keyword are monitored.
            Adding a "*" to the list will match with all parameters,
            i.e. no parameters or variables will be filtered out.
        :param left_strip_character: certain frameworks prepend their variables with
            a character. GPflow adds a '.' and Keras add a '_', for example.
            When a `left_strip_character` is specified it will be stripped from the
            parameter's name. By default the '.' is left stripped, for example:
            ".likelihood.variance" becomes "likelihood.variance".
        """
        super().__init__(log_dir)
        self.model = model
        self.max_size = max_size
        self.keywords_to_monitor = keywords_to_monitor
        self.summarize_all = "*" in self.keywords_to_monitor
        self.left_strip_character = left_strip_character

    def run(self, **unused_kwargs: Any) -> None:
        for name, parameter in parameter_dict(self.model).items():
            # check if the parameter name matches any of the specified keywords
            if self.summarize_all or any(keyword in name for keyword in self.keywords_to_monitor):
                # keys are sometimes prepended with a character, which we strip
                name = name.lstrip(self.left_strip_character)
                self._summarize_parameter(name, parameter)

    def _summarize_parameter(self, name: str, param: Union[Parameter, tf.Variable]) -> None:
        """
        :param name: identifier used in tensorboard
        :param param: parameter to be stored in tensorboard
        """
        param = tf.reshape(param, (-1,))
        size = param.shape[0]

        if not isinstance(size, int):
            raise ValueError(
                f"The monitoring can not be autographed as the size of a parameter {param} "
                "is unknown at compile time. If compiling the monitor task is important, "
                "make sure the shape of all parameters is known beforehand. Otherwise, "
                "run the monitor outside the `tf.function`."
            )
        if size == 1:
            # if there's only one element do not add a numbered suffix
            tf.summary.scalar(name, param[0], step=self.current_step)
        else:
            it = range(size) if self.max_size == -1 else range(min(size, self.max_size))
            for i in it:
                tf.summary.scalar(f"{name}[{i}]", param[i], step=self.current_step)


class ScalarToTensorBoard(ToTensorBoard):
    """Stores the return value of a callback in a TensorBoard."""

    def __init__(self, log_dir: str, callback: Callable[[], float], name: str) -> None:
        """
        :param log_dir: directory in which to store the tensorboard files.
            For example, './logs/my_run/'.
        :param callback: callback to be executed and result written to TensorBoard.
            A callback can have arguments (e.g. data) passed to the function using
            keyword arguments.
            For example:
            ```
            lambda cb(x=None): 2 * x
            task = ScalarToTensorBoard(logdir, cb, "callback")

            # specify the argument of the function using kwargs, the names need to match.
            task(step, x=1)
            ```
        :param name: name used in TensorBoard.
        """
        super().__init__(log_dir)
        self.name = name
        self.callback = callback

    def run(self, **kwargs: Any) -> None:
        tf.summary.scalar(self.name, self.callback(**kwargs), step=self.current_step)


class ImageToTensorBoard(ToTensorBoard):
    def __init__(
        self,
        log_dir: str,
        plotting_function: Callable[
            ["matplotlib.figure.Figure", "matplotlib.figure.Axes"], "matplotlib.figure.Figure"
        ],
        name: Optional[str] = None,
        *,
        fig_kw: Optional[Dict[str, Any]] = None,
        subplots_kw: Optional[Dict[str, Any]] = None,
    ) -> None:
        """
        :param log_dir: directory in which to store the tensorboard files.
            Can be nested: for example, './logs/my_run/'.
        :param plotting_function: function performing the plotting.
        :param name: name used in TensorBoard.
        :params fig_kw: keyword arguments to be passed to Figure constructor, e.g. `figsize`.
        :params subplots_kw: keyword arguments to be passed to figure.subplots constructor, e.g.
            `nrows`, `ncols`, `sharex`, `sharey`. By default the default values
            from matplotlib.pyplot are used.
        """
        super().__init__(log_dir)
        self.plotting_function = plotting_function
        self.name = name
        self.fig_kw = fig_kw or {}
        self.subplots_kw = subplots_kw or {}

        try:
            from matplotlib.figure import Figure
        except ImportError:
            raise RuntimeError("ImageToTensorBoard requires the matplotlib package to be installed")

        self.fig = Figure(**self.fig_kw)
        if self.subplots_kw != {}:
            self.axes = self.fig.subplots(**self.subplots_kw)
        else:
            self.axes = self.fig.add_subplot(111)

    def _clear_axes(self) -> None:
        if isinstance(self.axes, np.ndarray):
            for ax in self.axes.flatten():
                ax.clear()
        else:
            self.axes.clear()

    def run(self, **unused_kwargs: Any) -> None:
        from matplotlib.backends.backend_agg import FigureCanvasAgg

        self._clear_axes()
        self.plotting_function(self.fig, self.axes)
        canvas = FigureCanvasAgg(self.fig)
        canvas.draw()

        # get PNG data from the figure
        png_buffer = BytesIO()
        canvas.print_png(png_buffer)
        png_encoded = png_buffer.getvalue()
        png_buffer.close()

        image_tensor = tf.io.decode_png(png_encoded)[None]

        # Write to TensorBoard
        tf.summary.image(self.name, image_tensor, step=self.current_step)


# gpflow/optimizers/__init__.py

from . import natgrad
from .mcmc import SamplingHelper
from .natgrad import *
from .scipy import Scipy

__all__ = [
    "NaturalGradient",
    "SamplingHelper",
    "Scipy",
    "XiNat",
    "XiSqrtMeanVar",
    "XiTransform",
    "mcmc",
    "natgrad",
    "scipy",
]


# gpflow/optimizers/mcmc.py

# Copyright 2019-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Callable, Optional, Sequence, Tuple

import tensorflow as tf

from gpflow.base import Parameter

__all__ = ["SamplingHelper"]


class SamplingHelper:
    """
    This helper makes it easy to read from variables being set with a prior and
    writes values back to the same variables.

    Example::

        model = ...  # Create a GPflow model
        hmc_helper = SamplingHelper(model.log_posterior_density, model.trainable_parameters)

        target_log_prob_fn = hmc_helper.target_log_prob_fn
        current_state = hmc_helper.current_state

        hmc = tfp.mcmc.HamiltonianMonteCarlo(target_log_prob_fn=target_log_prob_fn, ...)
        adaptive_hmc = tfp.mcmc.SimpleStepSizeAdaptation(hmc, ...)

        @tf.function
        def run_chain_fn():
            return mcmc.sample_chain(
                num_samples, num_burnin_steps, current_state, kernel=adaptive_hmc)

        hmc_samples = run_chain_fn()
        parameter_samples = hmc_helper.convert_to_constrained_values(hmc_samples)
    """

    def __init__(
        self, target_log_prob_fn: Callable[[], tf.Tensor], parameters: Sequence[Parameter]
    ) -> None:
        """
        :param target_log_prob_fn: a callable which returns the log-density of the model
            under the target distribution; needs to implicitly depend on the `parameters`.
            E.g. `model.log_posterior_density`.
        :param parameters: List of :class:`gpflow.Parameter` used as a state of the Markov chain.
            E.g. `model.trainable_parameters`
            Note that each parameter must have been given a prior.
        """
        if not all(isinstance(p, Parameter) and p.prior is not None for p in parameters):
            raise ValueError(
                "`parameters` should only contain gpflow.Parameter objects with priors"
            )

        self._parameters = parameters
        self._target_log_prob_fn = target_log_prob_fn
        self._variables = [p.unconstrained_variable for p in parameters]

    @property
    def current_state(self) -> Sequence[tf.Variable]:
        """Return the current state of the unconstrained variables, used in HMC."""

        return self._variables

    @property
    def target_log_prob_fn(
        self,
    ) -> Callable[..., Tuple[tf.Tensor, Callable[..., Tuple[tf.Tensor, Sequence[None]]]]]:
        """
        The target log probability, adjusted to allow for optimisation to occur on the tracked
        unconstrained underlying variables.
        """
        variables_list = self.current_state

        @tf.custom_gradient
        def _target_log_prob_fn_closure(
            *variables: tf.Variable,
        ) -> Tuple[tf.Tensor, Callable[..., Tuple[tf.Tensor, Sequence[None]]]]:
            for v_old, v_new in zip(variables_list, variables):
                v_old.assign(v_new)

            with tf.GradientTape(watch_accessed_variables=False) as tape:
                tape.watch(variables_list)
                log_prob = self._target_log_prob_fn()
                # Now need to correct for the fact that the prob fn is evaluated on the
                # constrained space while we wish to evaluate it in the unconstrained space
                for param in self._parameters:
                    if param.transform is not None:
                        x = param.unconstrained_variable
                        log_det_jacobian = param.transform.forward_log_det_jacobian(
                            x, x.shape.ndims
                        )
                        log_prob += tf.reduce_sum(log_det_jacobian)

            @tf.function
            def grad_fn(
                dy: tf.Tensor, variables: Optional[tf.Tensor] = None
            ) -> Tuple[tf.Tensor, Sequence[None]]:
                grad = tape.gradient(log_prob, variables_list)
                return grad, [None] * len(variables_list)

            return log_prob, grad_fn

        return _target_log_prob_fn_closure  # type: ignore[no-any-return]

    def convert_to_constrained_values(
        self, hmc_samples: Sequence[tf.Tensor]
    ) -> Sequence[tf.Tensor]:
        """
        Converts list of unconstrained values in `hmc_samples` to constrained
        versions. Each value in the list corresponds to an entry in parameters
        passed to the constructor; for parameters that have a transform, the
        constrained representation is returned.
        """
        values = []
        for hmc_value, param in zip(hmc_samples, self._parameters):
            if param.transform is not None:
                value = param.transform.forward(hmc_value)
            else:
                value = hmc_value
            values.append(value)
        return values


# gpflow/optimizers/natgrad.py

# Copyright 2018-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import abc
import functools
from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Union

import numpy as np
import tensorflow as tf

from ..base import Parameter, _to_constrained
from ..experimental.check_shapes import check_shapes

Scalar = Union[float, tf.Tensor, np.ndarray]
LossClosure = Callable[[], tf.Tensor]
NatGradParameters = Union[Tuple[Parameter, Parameter], Tuple[Parameter, Parameter, "XiTransform"]]

__all__ = [
    "NaturalGradient",
    "XiNat",
    "XiSqrtMeanVar",
    "XiTransform",
]


#
# Xi transformations necessary for natural gradient optimizer.
# Abstract class and two implementations: XiNat and XiSqrtMeanVar.
#


class XiTransform(metaclass=abc.ABCMeta):
    """
    XiTransform is the base class that implements three transformations necessary
    for the natural gradient calculation wrt any parameterization.
    """

    @staticmethod
    @abc.abstractmethod
    @check_shapes(
        "mean: [N, D]",
        "varsqrt: [D, N, N]",
        "return[0]: [N, D]",
        "return[1]: [D, N, N]",
    )
    def meanvarsqrt_to_xi(mean: tf.Tensor, varsqrt: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
        """
        Transforms the parameter `mean` and `varsqrt` to `xi1`, `xi2`

        :param mean: the mean parameter
        :param varsqrt: the varsqrt parameter
        :return: tuple (xi1, xi2), the xi parameters
        """

    @staticmethod
    @abc.abstractmethod
    @check_shapes(
        "xi1: [N, D]",
        "xi2: [D, N, N]",
        "return[0]: [N, D]",
        "return[1]: [D, N, N]",
    )
    def xi_to_meanvarsqrt(xi1: tf.Tensor, xi2: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
        """
        Transforms the parameter `xi1`, `xi2` to `mean`, `varsqrt`

        :param xi1: the ξ₁ parameter
        :param xi2: the ξ₂ parameter
        :return: tuple (mean, varsqrt), the meanvarsqrt parameters
        """

    @staticmethod
    @abc.abstractmethod
    @check_shapes(
        "nat1: [N, D]",
        "nat2: [D, N, N]",
        "return[0]: [N, D]",
        "return[1]: [D, N, N]",
    )
    def naturals_to_xi(nat1: tf.Tensor, nat2: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
        """
        Applies the transform so that `nat1`, `nat2` is mapped to `xi1`, `xi2`

        :param nat1: the θ₁ parameter
        :param nat2: the θ₂ parameter
        :return: tuple `xi1`, `xi2`
        """


class XiNat(XiTransform):
    """
    This is the default transform. Using the natural directly saves the forward mode
    gradient, and also gives the analytic optimal solution for gamma=1 in the case
    of Gaussian likelihood.
    """

    @staticmethod
    @check_shapes(
        "mean: [N, D]",
        "varsqrt: [D, N, N]",
        "return[0]: [N, D]",
        "return[1]: [D, N, N]",
    )
    def meanvarsqrt_to_xi(mean: tf.Tensor, varsqrt: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
        return meanvarsqrt_to_natural(mean, varsqrt)

    @staticmethod
    @check_shapes(
        "xi1: [N, D]",
        "xi2: [D, N, N]",
        "return[0]: [N, D]",
        "return[1]: [D, N, N]",
    )
    def xi_to_meanvarsqrt(xi1: tf.Tensor, xi2: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
        return natural_to_meanvarsqrt(xi1, xi2)

    @staticmethod
    @check_shapes(
        "nat1: [N, D]",
        "nat2: [D, N, N]",
        "return[0]: [N, D]",
        "return[1]: [D, N, N]",
    )
    def naturals_to_xi(nat1: tf.Tensor, nat2: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
        return nat1, nat2


class XiSqrtMeanVar(XiTransform):
    """
    This transformation will perform natural gradient descent on the model parameters,
    so saves the conversion to and from Xi.
    """

    @staticmethod
    @check_shapes(
        "mean: [N, D]",
        "varsqrt: [D, N, N]",
        "return[0]: [N, D]",
        "return[1]: [D, N, N]",
    )
    def meanvarsqrt_to_xi(mean: tf.Tensor, varsqrt: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
        return mean, varsqrt

    @staticmethod
    @check_shapes(
        "xi1: [N, D]",
        "xi2: [D, N, N]",
        "return[0]: [N, D]",
        "return[1]: [D, N, N]",
    )
    def xi_to_meanvarsqrt(xi1: tf.Tensor, xi2: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
        return xi1, xi2

    @staticmethod
    @check_shapes(
        "nat1: [N, D]",
        "nat2: [D, N, N]",
        "return[0]: [N, D]",
        "return[1]: [D, N, N]",
    )
    def naturals_to_xi(nat1: tf.Tensor, nat2: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
        return natural_to_meanvarsqrt(nat1, nat2)


class NaturalGradient(tf.optimizers.Optimizer):
    """
    Implements a natural gradient descent optimizer for variational models
    that are based on a distribution q(u) = N(q_mu, q_sqrt q_sqrtᵀ) that is
    parameterized by mean q_mu and lower-triangular Cholesky factor q_sqrt
    of the covariance.

    Note that this optimizer does not implement the standard API of
    tf.optimizers.Optimizer. Its only public method is minimize(), which has
    a custom signature (var_list needs to be a list of (q_mu, q_sqrt) tuples,
    where q_mu and q_sqrt are gpflow.Parameter instances, not tf.Variable).

    Note furthermore that the natural gradients are implemented only for the
    full covariance case (i.e., q_diag=True is NOT supported).

    When using in your work, please cite :cite:t:`salimbeni18`.
    """

    def __init__(
        self, gamma: Scalar, xi_transform: XiTransform = XiNat(), name: Optional[str] = None
    ) -> None:
        """
        :param gamma: natgrad step length
        :param xi_transform: default ξ transform (can be overridden in the call to minimize())
            The XiNat default choice works well in general.
        """
        name = self.__class__.__name__ if name is None else name
        super().__init__(name)
        self.gamma = gamma
        self.xi_transform = xi_transform

    @check_shapes(
        "var_list[all][0]: [N, D]",
        "var_list[all][1]: [D, N, N]",
    )
    def minimize(
        self,
        loss_fn: LossClosure,
        var_list: Sequence[NatGradParameters],
    ) -> None:
        """
        Minimizes objective function of the model.
        Natural Gradient optimizer works with variational parameters only.

        GPflow implements the `XiNat` (default) and `XiSqrtMeanVar` transformations
        for parameters. Custom transformations that implement the `XiTransform`
        interface are also possible.

        :param loss_fn: Loss function.
        :param var_list: List of pair tuples of variational parameters or
            triplet tuple with variational parameters and ξ transformation.
            If ξ is not specified, will use self.xi_transform.
            For example, `var_list` could be::

                var_list = [
                    (q_mu1, q_sqrt1),
                    (q_mu2, q_sqrt2, XiSqrtMeanVar())
                ]
        """
        parameters = [(v[0], v[1], (v[2] if len(v) > 2 else None)) for v in var_list]  # type: ignore[misc]
        self._natgrad_steps(loss_fn, parameters)

    @check_shapes(
        "parameters[all][0]: [N, D]",
        "parameters[all][1]: [D, N, N]",
    )
    def _natgrad_steps(
        self,
        loss_fn: LossClosure,
        parameters: Sequence[Tuple[Parameter, Parameter, Optional[XiTransform]]],
    ) -> None:
        """
        Computes gradients of loss_fn() w.r.t. q_mu and q_sqrt, and updates
        these parameters using the natgrad backwards step, for all sets of
        variational parameters passed in.

        :param loss_fn: Loss function.
        :param parameters: List of tuples (q_mu, q_sqrt, xi_transform)
        """
        q_mus, q_sqrts, xis = zip(*parameters)
        q_mu_vars = [p.unconstrained_variable for p in q_mus]
        q_sqrt_vars = [p.unconstrained_variable for p in q_sqrts]

        with tf.GradientTape(watch_accessed_variables=False) as tape:
            tape.watch(q_mu_vars + q_sqrt_vars)
            loss = loss_fn()

        q_mu_grads, q_sqrt_grads = tape.gradient(loss, [q_mu_vars, q_sqrt_vars])
        # NOTE that these are the gradients in *unconstrained* space

        with tf.name_scope(f"{self._name}/natural_gradient_steps"):
            for q_mu_grad, q_sqrt_grad, q_mu, q_sqrt, xi_transform in zip(
                q_mu_grads, q_sqrt_grads, q_mus, q_sqrts, xis
            ):
                self._natgrad_apply_gradients(q_mu_grad, q_sqrt_grad, q_mu, q_sqrt, xi_transform)

    @check_shapes(
        "q_mu_grad: [N, D]",
        "q_sqrt_grad: [D, N_N_transformed...]",
        "q_mu: [N, D]",
        "q_sqrt: [D, N, N]",
    )
    def _natgrad_apply_gradients(
        self,
        q_mu_grad: tf.Tensor,
        q_sqrt_grad: tf.Tensor,
        q_mu: Parameter,
        q_sqrt: Parameter,
        xi_transform: Optional[XiTransform] = None,
    ) -> None:
        """
        This function does the backward step on the q_mu and q_sqrt parameters,
        given the gradients of the loss function with respect to their unconstrained
        variables. I.e., it expects the arguments to come from

            with tf.GradientTape() as tape:
                loss = loss_function()
            q_mu_grad, q_mu_sqrt = tape.gradient(loss, [q_mu, q_sqrt])

        (Note that tape.gradient() returns the gradients in *unconstrained* space!)

        Implements equation [10] from :cite:t:`salimbeni18`.

        In addition, for convenience with the rest of GPflow, this code computes ∂L/∂η using
        the chain rule (the following assumes a numerator layout where the gradient is a row
        vector; note that TensorFlow actually returns a column vector), where L is the loss:

        ∂L/∂η = (∂L / ∂[q_mu, q_sqrt])(∂[q_mu, q_sqrt] / ∂η)

        In total there are three derivative calculations:
        natgrad of L w.r.t ξ  = (∂ξ / ∂θ) [(∂L / ∂[q_mu, q_sqrt]) (∂[q_mu, q_sqrt] / ∂η)]ᵀ

        Note that if ξ = θ (i.e. [q_mu, q_sqrt]) some of these calculations are the identity.
        In the code η = eta, ξ = xi, θ = nat.

        :param q_mu_grad: gradient of loss w.r.t. q_mu (in unconstrained space)
        :param q_sqrt_grad: gradient of loss w.r.t. q_sqrt (in unconstrained space)
        :param q_mu: parameter for the mean of q(u) with shape [M, L]
        :param q_sqrt: parameter for the square root of the covariance of q(u)
            with shape [L, M, M] (the diagonal parametrization, q_diag=True, is NOT supported)
        :param xi_transform: the ξ transform to use (self.xi_transform if not specified)
        """
        if xi_transform is None:
            xi_transform = self.xi_transform

        # 1) the ordinary gpflow gradient
        dL_dmean = _to_constrained(q_mu_grad, q_mu.transform)
        dL_dvarsqrt = _to_constrained(q_sqrt_grad, q_sqrt.transform)

        with tf.GradientTape(persistent=True, watch_accessed_variables=False) as tape:
            tape.watch([q_mu.unconstrained_variable, q_sqrt.unconstrained_variable])

            # the three parameterizations as functions of [q_mu, q_sqrt]
            eta1, eta2 = meanvarsqrt_to_expectation(q_mu, q_sqrt)
            # we need these to calculate the relevant gradients
            meanvarsqrt = expectation_to_meanvarsqrt(eta1, eta2)

            if not isinstance(xi_transform, XiNat):
                nat1, nat2 = meanvarsqrt_to_natural(q_mu, q_sqrt)
                xi1_nat, xi2_nat = xi_transform.naturals_to_xi(nat1, nat2)
                dummy_tensors = tf.ones_like(xi1_nat), tf.ones_like(xi2_nat)
                with tf.GradientTape(watch_accessed_variables=False) as forward_tape:
                    forward_tape.watch(dummy_tensors)
                    dummy_gradients = tape.gradient(
                        [xi1_nat, xi2_nat], [nat1, nat2], output_gradients=dummy_tensors
                    )

        # 2) the chain rule to get ∂L/∂η, where η (eta) are the expectation parameters
        dL_deta1, dL_deta2 = tape.gradient(
            meanvarsqrt, [eta1, eta2], output_gradients=[dL_dmean, dL_dvarsqrt]
        )

        if not isinstance(xi_transform, XiNat):
            nat_dL_xi1, nat_dL_xi2 = forward_tape.gradient(
                dummy_gradients, dummy_tensors, output_gradients=[dL_deta1, dL_deta2]
            )
        else:
            nat_dL_xi1, nat_dL_xi2 = dL_deta1, dL_deta2

        del tape  # Remove "persistent" tape

        xi1, xi2 = xi_transform.meanvarsqrt_to_xi(q_mu, q_sqrt)
        xi1_new = xi1 - self.gamma * nat_dL_xi1
        xi2_new = xi2 - self.gamma * nat_dL_xi2

        # Transform back to the model parameters [q_mu, q_sqrt]
        mean_new, varsqrt_new = xi_transform.xi_to_meanvarsqrt(xi1_new, xi2_new)

        q_mu.assign(mean_new)
        q_sqrt.assign(varsqrt_new)

    def get_config(self) -> Dict[str, Any]:
        config: Dict[str, Any] = super().get_config()
        config.update({"gamma": self._serialize_hyperparameter("gamma")})
        return config


#
# Auxiliary gaussian parameter conversion functions.
#
# The following functions expect their first and second inputs to have shape
# [D, N, 1] and [D, N, N], respectively. Return values are also of shapes [D, N, 1] and [D, N, N].


def swap_dimensions(
    method: Callable[[tf.Tensor, tf.Tensor], Tuple[tf.Tensor, tf.Tensor]]
) -> Callable[..., Tuple[tf.Tensor, tf.Tensor]]:
    """
    Converts between GPflow indexing and tensorflow indexing
    `method` is a function that broadcasts over the first dimension (i.e. like all tensorflow matrix
    ops):

    * `method` inputs [D, N, 1], [D, N, N]
    * `method` outputs [D, N, 1], [D, N, N]

    :return: Function that broadcasts over the final dimension (i.e. compatible with GPflow):

        * inputs: [N, D], [D, N, N]
        * outputs: [N, D], [D, N, N]
    """

    @functools.wraps(method)
    @check_shapes(
        "a_nd: [N, D] if swap",
        "a_nd: [D, N, 1] if not swap",
        "b_dnn: [D, N, N]",
        "return[0]: [N, D] if swap",
        "return[0]: [D, N, 1] if not swap",
        "return[1]: [D, N, N]",
    )
    def wrapper(
        a_nd: tf.Tensor, b_dnn: tf.Tensor, swap: bool = True
    ) -> Tuple[tf.Tensor, tf.Tensor]:
        if swap:
            a_dn1 = tf.linalg.adjoint(a_nd)[:, :, None]
            A_dn1, B_dnn = method(a_dn1, b_dnn)
            A_nd = tf.linalg.adjoint(A_dn1[:, :, 0])
            return A_nd, B_dnn
        else:
            return method(a_nd, b_dnn)

    return wrapper


@swap_dimensions
@check_shapes(
    "nat1: [D, N, 1]",
    "nat2: [D, N, N]",
    "return[0]: [D, N, 1]",
    "return[1]: [D, N, N]",
)
def natural_to_meanvarsqrt(nat1: tf.Tensor, nat2: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
    var_sqrt_inv = tf.linalg.cholesky(-2 * nat2)
    var_sqrt = _inverse_lower_triangular(var_sqrt_inv)
    S = tf.linalg.matmul(var_sqrt, var_sqrt, transpose_a=True)
    mu = tf.linalg.matmul(S, nat1)
    # We need the decomposition of S as L L^T, not as L^T L,
    # hence we need another cholesky.
    return mu, tf.linalg.cholesky(S)


@swap_dimensions
@check_shapes(
    "mu: [D, N, 1]",
    "s_sqrt: [D, N, N]",
    "return[0]: [D, N, 1]",
    "return[1]: [D, N, N]",
)
def meanvarsqrt_to_natural(mu: tf.Tensor, s_sqrt: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
    s_sqrt_inv = _inverse_lower_triangular(s_sqrt)
    s_inv = tf.linalg.matmul(s_sqrt_inv, s_sqrt_inv, transpose_a=True)
    return tf.linalg.matmul(s_inv, mu), -0.5 * s_inv


@swap_dimensions
@check_shapes(
    "nat1: [D, N, 1]",
    "nat2: [D, N, N]",
    "return[0]: [D, N, 1]",
    "return[1]: [D, N, N]",
)
def natural_to_expectation(nat1: tf.Tensor, nat2: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
    args = natural_to_meanvarsqrt(nat1, nat2, swap=False)
    return meanvarsqrt_to_expectation(*args, swap=False)


@swap_dimensions
@check_shapes(
    "eta1: [D, N, 1]",
    "eta2: [D, N, N]",
    "return[0]: [D, N, 1]",
    "return[1]: [D, N, N]",
)
def expectation_to_natural(eta1: tf.Tensor, eta2: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
    args = expectation_to_meanvarsqrt(eta1, eta2, swap=False)
    return meanvarsqrt_to_natural(*args, swap=False)


@swap_dimensions
@check_shapes(
    "eta1: [D, N, 1]",
    "eta2: [D, N, N]",
    "return[0]: [D, N, 1]",
    "return[1]: [D, N, N]",
)
def expectation_to_meanvarsqrt(eta1: tf.Tensor, eta2: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
    var = eta2 - tf.linalg.matmul(eta1, eta1, transpose_b=True)
    return eta1, tf.linalg.cholesky(var)


@swap_dimensions
@check_shapes(
    "m: [D, N, 1]",
    "v_sqrt: [D, N, N]",
    "return[0]: [D, N, 1]",
    "return[1]: [D, N, N]",
)
def meanvarsqrt_to_expectation(m: tf.Tensor, v_sqrt: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
    v = tf.linalg.matmul(v_sqrt, v_sqrt, transpose_b=True)
    return m, v + tf.linalg.matmul(m, m, transpose_b=True)


@check_shapes(
    "M: [D, N, N]",
    "return: [D, N, N]",
)
def _inverse_lower_triangular(M: tf.Tensor) -> tf.Tensor:
    """
    Take inverse of lower triangular (e.g. Cholesky) matrix. This function
    broadcasts over the first index.

    :param M: Tensor with lower triangular structure of shape [D, N, N]
    :return: The inverse of the Cholesky decomposition. Same shape as input.
    """
    if M.shape.ndims != 3:  # pragma: no cover
        raise ValueError("Number of dimensions for input is required to be 3.")
    D, N = tf.shape(M)[0], tf.shape(M)[1]
    I_dnn = tf.eye(N, dtype=M.dtype)[None, :, :] * tf.ones((D, 1, 1), dtype=M.dtype)
    return tf.linalg.triangular_solve(M, I_dnn)


# gpflow/optimizers/scipy.py

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import warnings
from typing import Any, Callable, Iterable, List, Optional, Sequence, Tuple, Union

import numpy as np
import scipy.optimize
import tensorflow as tf
from scipy.optimize import OptimizeResult

from ..base import AnyNDArray
from ..monitor.base import Monitor

__all__ = ["Scipy"]

Variables = Iterable[tf.Variable]  # deprecated
StepCallback = Union[Callable[[int, Sequence[tf.Variable], Sequence[tf.Tensor]], None], Monitor]
LossClosure = Callable[[], tf.Tensor]


class Scipy:
    def minimize(
        self,
        closure: LossClosure,
        variables: Sequence[tf.Variable],
        method: Optional[str] = "L-BFGS-B",
        step_callback: Optional[StepCallback] = None,
        compile: bool = True,
        allow_unused_variables: bool = False,
        **scipy_kwargs: Any,
    ) -> OptimizeResult:
        """
        Minimize `closure`.

        Minimize is a wrapper around the `scipy.optimize.minimize` function handling the packing and
        unpacking of a list of shaped variables on the TensorFlow side vs. the flat numpy array
        required on the Scipy side.

        :param closure: A closure that re-evaluates the model, returning the loss to be minimized.
        :param variables: The list (tuple) of variables to be optimized
            (typically `model.trainable_variables`)
        :param method: The type of solver to use in SciPy. Defaults to "L-BFGS-B".
        :param step_callback: If not None, a callable that gets called once after each optimisation
            step. The callable is passed the arguments `step`, `variables`, and `values`. `step` is
            the optimisation step counter, `variables` is the list of trainable variables as above,
            and `values` is the corresponding list of tensors of matching shape that contains their
            value at this optimisation step.
        :param compile: If True, wraps the evaluation function (the passed `closure` as well as its
            gradient computation) inside a `tf.function()`, which will improve optimization speed in
            most cases.
        :param allow_unused_variables: Whether to allow variables that are not actually used in the
            closure.
        :param scipy_kwargs: Arguments passed through to `scipy.optimize.minimize`.
            Note that Scipy's minimize() takes a `callback` argument, but you probably want to use
            our wrapper and pass in `step_callback`.
        :returns:
            The optimization result represented as a Scipy ``OptimizeResult`` object.
            See the Scipy documentation for description of attributes.
        """
        if not callable(closure):
            raise TypeError(
                "The 'closure' argument is expected to be a callable object."
            )  # pragma: no cover
        variables = tuple(variables)
        if not all(isinstance(v, tf.Variable) for v in variables):
            raise TypeError(
                "The 'variables' argument is expected to only contain tf.Variable instances"
                " (use model.trainable_variables, not model.trainable_parameters)"
            )  # pragma: no cover
        initial_params = self.initial_parameters(variables)

        func = self.eval_func(
            closure, variables, compile=compile, allow_unused_variables=allow_unused_variables
        )
        if step_callback is not None:
            if "callback" in scipy_kwargs:
                raise ValueError("Callback passed both via `step_callback` and `callback`")

            callback = self.callback_func(variables, step_callback)
            scipy_kwargs.update(dict(callback=callback))

        return scipy.optimize.minimize(
            func, initial_params, jac=True, method=method, **scipy_kwargs
        )

    @classmethod
    def initial_parameters(cls, variables: Sequence[tf.Variable]) -> tf.Tensor:
        return cls.pack_tensors(variables)

    @classmethod
    def eval_func(
        cls,
        closure: LossClosure,
        variables: Sequence[tf.Variable],
        compile: bool = True,
        allow_unused_variables: bool = False,
    ) -> Callable[[AnyNDArray], Tuple[AnyNDArray, AnyNDArray]]:
        first_call = True

        def _tf_eval(x: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
            nonlocal first_call

            values = cls.unpack_tensors(variables, x)
            cls.assign_tensors(variables, values)

            if first_call:
                # Only check for unconnected gradients on the first function evaluation.
                loss, grads = _compute_loss_and_gradients(
                    closure, variables, tf.UnconnectedGradients.NONE
                )
                grads = cls._filter_unused_variables(variables, grads, allow_unused_variables)
                first_call = False
            else:
                loss, grads = _compute_loss_and_gradients(
                    closure, variables, tf.UnconnectedGradients.ZERO
                )

            return loss, cls.pack_tensors(grads)

        if compile:
            _tf_eval = tf.function(_tf_eval)

        def _eval(x: AnyNDArray) -> Tuple[AnyNDArray, AnyNDArray]:
            loss, grad = _tf_eval(tf.convert_to_tensor(x))
            return loss.numpy().astype(np.float64), grad.numpy().astype(np.float64)

        return _eval

    @staticmethod
    def _filter_unused_variables(
        variables: Sequence[tf.Variable], grads: Sequence[tf.Tensor], allow_unused_variables: bool
    ) -> Sequence[tf.Tensor]:
        filtered_grads = []
        unused_variables = []
        for i, grad in enumerate(grads):
            if grad is None:
                variable = variables[i]
                filtered_grads.append(tf.zeros_like(variable))
                unused_variables.append(variable.name)
            else:
                filtered_grads.append(grad)

        if unused_variables:
            msg = (
                "Some variables does not have a gradient, and appear unused in / not connected to"
                f" the loss closure: {unused_variables}."
            )
            if allow_unused_variables:
                warnings.warn(msg)
            else:
                raise ValueError(msg)

        return filtered_grads

    @classmethod
    def callback_func(
        cls, variables: Sequence[tf.Variable], step_callback: StepCallback
    ) -> Callable[[AnyNDArray], None]:
        step = 0  # type: int

        def _callback(x: AnyNDArray) -> None:
            nonlocal step

            if isinstance(step_callback, Monitor):
                step_callback(step)
            else:
                values = cls.unpack_tensors(variables, x)
                step_callback(step, variables, values)

            step += 1

        return _callback

    @staticmethod
    def pack_tensors(tensors: Sequence[Union[tf.Tensor, tf.Variable]]) -> tf.Tensor:
        flats = [tf.reshape(tensor, (-1,)) for tensor in tensors]
        tensors_vector = tf.concat(flats, axis=0)
        return tensors_vector

    @staticmethod
    def unpack_tensors(
        to_tensors: Sequence[Union[tf.Tensor, tf.Variable]], from_vector: tf.Tensor
    ) -> List[tf.Tensor]:
        s = 0
        values = []
        for target_tensor in to_tensors:
            shape = tf.shape(target_tensor)
            dtype = target_tensor.dtype
            tensor_size = tf.reduce_prod(shape)
            tensor_vector = from_vector[s : s + tensor_size]
            tensor = tf.reshape(tf.cast(tensor_vector, dtype), shape)
            values.append(tensor)
            s += tensor_size
        return values

    @staticmethod
    def assign_tensors(to_tensors: Sequence[tf.Variable], values: Sequence[tf.Tensor]) -> None:
        if len(to_tensors) != len(values):
            raise ValueError("to_tensors and values should have same length")
        for target, value in zip(to_tensors, values):
            target.assign(value)


def _compute_loss_and_gradients(
    loss_closure: LossClosure,
    variables: Sequence[tf.Variable],
    unconnected_gradients: tf.UnconnectedGradients,
) -> Tuple[tf.Tensor, Sequence[tf.Tensor]]:
    with tf.GradientTape(watch_accessed_variables=False) as tape:
        tape.watch(variables)
        loss = loss_closure()
    grads = tape.gradient(loss, variables, unconnected_gradients=unconnected_gradients)
    return loss, grads


# gpflow/quadrature/__init__.py

from .base import GaussianQuadrature
from .deprecated import hermgauss, mvhermgauss, mvnquad, ndiag_mc, ndiagquad
from .gauss_hermite import NDiagGHQuadrature

__all__ = [
    "GaussianQuadrature",
    "NDiagGHQuadrature",
    "base",
    "deprecated",
    "gauss_hermite",
    "hermgauss",
    "mvhermgauss",
    "mvnquad",
    "ndiag_mc",
    "ndiagquad",
]


# gpflow/quadrature/base.py

# Copyright 2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import abc
from typing import Any, Callable, Iterable, Tuple, Union

import tensorflow as tf

from ..base import TensorType
from ..experimental.check_shapes import check_shapes


class GaussianQuadrature:
    """
    Abstract class implementing quadrature methods to compute Gaussian Expectations.
    Inheriting classes must provide the method _build_X_W to create points and weights
    to be used for quadrature.
    """

    @abc.abstractmethod
    @check_shapes(
        "mean: [batch..., D]",
        "var: [batch..., D]",
        "return[0]: [N_quad_points, batch..., D]",
        "return[1]: [N_quad_points, broadcast batch..., 1]",
    )
    def _build_X_W(self, mean: TensorType, var: TensorType) -> Tuple[tf.Tensor, tf.Tensor]:
        raise NotImplementedError

    @check_shapes(
        "mean: [in_batch..., D]",
        "var: [in_batch..., D]",
        "return: [n_funs..., out_batch..., broadcast D]",
    )
    def __call__(
        self,
        fun: Union[Callable[..., tf.Tensor], Iterable[Callable[..., tf.Tensor]]],
        mean: TensorType,
        var: TensorType,
        *args: Any,
        **kwargs: Any,
    ) -> tf.Tensor:
        r"""
        Compute the Gaussian Expectation of a function f::

            X ~ N(mean, var)
            E[f(X)] = ∫f(x, *args, **kwargs)p(x)dx

        Using the formula::

            E[f(X)] = sum_{i=1}^{N_quad_points} f(x_i) * w_i

        where x_i, w_i must be provided by the inheriting class through self._build_X_W.

        :param fun: Callable or Iterable of Callables that operates elementwise, with
            signature f(X, \*args, \*\*kwargs). Moreover, it must satisfy the shape-mapping::

                X shape: [N_quad_points, batch..., d].
                f(X) shape: [N_quad_points, batch..., broadcast d].

            In most cases, f should only operate over the last dimension of X
        :param mean: Array/Tensor representing the mean of a d-Variate Gaussian distribution
        :param var: Array/Tensor representing the variance of a d-Variate Gaussian distribution
        :param args: Passed to fun
        :param kargs: Passed to fun
        :return: Gaussian expectation of fun
        """

        X, W = self._build_X_W(mean, var)
        if isinstance(fun, Iterable):
            return [tf.reduce_sum(f(X, *args, **kwargs) * W, axis=0) for f in fun]
        return tf.reduce_sum(fun(X, *args, **kwargs) * W, axis=0)

    @check_shapes(
        "mean: [in_batch..., D]",
        "var: [in_batch..., D]",
        "return: [n_fun..., out_batch..., broadcast D]",
    )
    def logspace(
        self,
        fun: Union[Callable[..., tf.Tensor], Iterable[Callable[..., tf.Tensor]]],
        mean: TensorType,
        var: TensorType,
        *args: Any,
        **kwargs: Any,
    ) -> tf.Tensor:
        r"""
        Compute the Gaussian log-Expectation of a the exponential of a function f::

            X ~ N(mean, var)
            log E[exp[f(X)]] = log ∫exp[f(x, *args, **kwargs)]p(x)dx

        Using the formula::

            log E[exp[f(X)]] = log sum_{i=1}^{N_quad_points} exp[f(x_i) + log w_i]

        where x_i, w_i must be provided by the inheriting class through self._build_X_W.
        The computations broadcast along batch-dimensions, represented by [batch...].

        :param fun: Callable or Iterable of Callables that operates elementwise, with
            signature f(X, \*args, \*\*kwargs). Moreover, it must satisfy the shape-mapping::

                X shape: [N_quad_points, batch..., d].
                f(X) shape: [N_quad_points, batch..., broadcast d].

            In most cases, f should only operate over the last dimension of X
        :param mean: Array/Tensor representing the mean of a d-Variate Gaussian distribution
        :param var: Array/Tensor representing the variance of a d-Variate Gaussian distribution
        :param args: Passed to fun
        :param kwargs: Passed to fun
        :return: Gaussian log-expectation of the exponential of a function f
        """

        X, W = self._build_X_W(mean, var)
        logW = tf.math.log(W)
        if isinstance(fun, Iterable):
            return [tf.reduce_logsumexp(f(X, *args, **kwargs) + logW, axis=0) for f in fun]
        return tf.reduce_logsumexp(fun(X, *args, **kwargs) + logW, axis=0)


# gpflow/quadrature/deprecated.py

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import itertools
import warnings
from functools import wraps
from typing import Callable, Iterable, List, Optional, Tuple, Union

import numpy as np
import tensorflow as tf

from ..base import AnyNDArray, TensorType
from ..config import default_float
from ..experimental.check_shapes import check_shapes
from ..utilities import to_default_float
from .gauss_hermite import NDiagGHQuadrature


@check_shapes(
    "return[0]: [n_quad_points]",
    "return[1]: [n_quad_points]",
)
def hermgauss(n: int) -> Tuple[AnyNDArray, AnyNDArray]:
    # Type-ignore is because for some versions mypy can't find np.polynomial.hermite
    x, w = np.polynomial.hermite.hermgauss(n)
    x, w = x.astype(default_float()), w.astype(default_float())
    return x, w


@check_shapes(
    "return[0]: [n_quad_points, D]",
    "return[1]: [n_quad_points]",
)
def mvhermgauss(H: int, D: int) -> Tuple[AnyNDArray, AnyNDArray]:
    """
    Return the evaluation locations 'xn', and weights 'wn' for a multivariate
    Gauss-Hermite quadrature.

    The outputs can be used to approximate the following type of integral:
    int exp(-x)*f(x) dx ~ sum_i w[i,:]*f(x[i,:])

    :param H: Number of Gauss-Hermite evaluation points.
    :param D: Number of input dimensions. Needs to be known at call-time.
    :return: eval_locations 'x' (H**DxD), weights 'w' (H**D)
    """
    gh_x, gh_w = hermgauss(H)
    x: AnyNDArray = np.array(list(itertools.product(*(gh_x,) * D)))  # H**DxD
    w = np.prod(np.array(list(itertools.product(*(gh_w,) * D))), 1)  # H**D
    return x, w


@check_shapes(
    "means: [N, Din]",
    "covs: [N, Din, Din]",
    "return: [N, Dout...]",
)
def mvnquad(
    func: Callable[[tf.Tensor], tf.Tensor],
    means: TensorType,
    covs: TensorType,
    H: int,
    Din: Optional[int] = None,
    Dout: Optional[Tuple[int, ...]] = None,
) -> tf.Tensor:
    """
    Computes N Gaussian expectation integrals of a single function 'f'
    using Gauss-Hermite quadrature.

    :param f: integrand function. Takes one input of shape ?xD.
    :param H: Number of Gauss-Hermite evaluation points.
    :param Din: Number of input dimensions. Needs to be known at call-time.
    :param Dout: Number of output dimensions. Defaults to (). Dout is assumed
        to leave out the item index, i.e. f actually maps (?xD)->(?x*Dout).
    :return: quadratures
    """
    # Figure out input shape information
    if Din is None:
        Din = means.shape[1]

    if Din is None:
        raise ValueError(
            "If `Din` is passed as `None`, `means` must have a known shape. "
            "Running mvnquad in `autoflow` without specifying `Din` and `Dout` "
            "is problematic. Consider using your own session."
        )  # pragma: no cover

    xn, wn = mvhermgauss(H, Din)
    N = means.shape[0]

    # transform points based on Gaussian parameters
    cholXcov = tf.linalg.cholesky(covs)  # NxDxD
    Xt = tf.linalg.matmul(
        cholXcov, tf.tile(xn[None, :, :], (N, 1, 1)), transpose_b=True
    )  # NxDxH**D
    X = 2.0 ** 0.5 * Xt + tf.expand_dims(means, 2)  # NxDxH**D
    Xr = tf.reshape(tf.transpose(X, [2, 0, 1]), (-1, Din))  # (H**D*N)xD

    # perform quadrature
    fevals = func(Xr)
    if Dout is None:
        Dout = tuple((d if type(d) is int else d.value) for d in fevals.shape[1:])

    if any([d is None for d in Dout]):
        raise ValueError(
            "If `Dout` is passed as `None`, the output of `func` must have known "
            "shape. Running mvnquad in `autoflow` without specifying `Din` and `Dout` "
            "is problematic. Consider using your own session."
        )  # pragma: no cover
    fX = tf.reshape(
        fevals,
        (
            H ** Din,
            N,
        )
        + Dout,
    )
    wr = np.reshape(wn * np.pi ** (-Din * 0.5), (-1,) + (1,) * (1 + len(Dout)))
    return tf.reduce_sum(fX * wr, 0)


@check_shapes(
    "Fmu: [broadcast Din, N...]",
    "Fvar: [broadcast Din, N...]",
    "Ys.values(): [N...]",
    "return: [broadcast Dout, N...]",
)
def ndiagquad(
    funcs: Union[Callable[..., tf.Tensor], Iterable[Callable[..., tf.Tensor]]],
    H: int,
    Fmu: Union[TensorType, Tuple[TensorType, ...], List[TensorType]],
    Fvar: Union[TensorType, Tuple[TensorType, ...], List[TensorType]],
    logspace: bool = False,
    **Ys: TensorType,
) -> tf.Tensor:
    """
    Computes N Gaussian expectation integrals of one or more functions
    using Gauss-Hermite quadrature. The Gaussians must be independent.

    The means and variances of the Gaussians are specified by Fmu and Fvar.
    The N-integrals are assumed to be taken wrt the last dimensions of Fmu, Fvar.

    `Fmu`, `Fvar`, `Ys` should all have same shape, with overall size `N`.

    :param funcs: the integrand(s):
        Callable or Iterable of Callables that operates elementwise
    :param H: number of Gauss-Hermite quadrature points
    :param Fmu: array/tensor or `Din`-tuple/list thereof
    :param Fvar: array/tensor or `Din`-tuple/list thereof
    :param logspace: if True, funcs are the log-integrands and this calculates
        the log-expectation of exp(funcs)
    :param Ys: arrays/tensors; deterministic arguments to be passed by name
    :return: shape is the same as that of the first Fmu
    """
    warnings.warn(
        "Please use gpflow.quadrature.NDiagGHQuadrature instead "
        "(note the changed convention of how multi-dimensional quadrature is handled)",
        DeprecationWarning,
    )

    n_gh = H
    if isinstance(Fmu, (tuple, list)):
        dim = len(Fmu)
        shape = tf.shape(Fmu[0])
        Fmu = tf.stack(Fmu, axis=-1)
        Fvar = tf.stack(Fvar, axis=-1)
    else:
        dim = 1
        shape = tf.shape(Fmu)

    Fmu = tf.reshape(Fmu, (-1, dim))
    Fvar = tf.reshape(Fvar, (-1, dim))

    Ys = {Yname: tf.reshape(Y, (-1, 1)) for Yname, Y in Ys.items()}

    def wrapper(old_fun: Callable[..., tf.Tensor]) -> Callable[..., tf.Tensor]:
        @wraps(old_fun)
        def new_fun(X: TensorType, **Ys: TensorType) -> tf.Tensor:
            Xs = tf.unstack(tf.expand_dims(X, axis=-2), axis=-1)
            fun_eval = old_fun(*Xs, **Ys)
            return tf.cond(
                pred=tf.less(tf.rank(fun_eval), tf.rank(X)),
                true_fn=lambda: fun_eval[..., tf.newaxis],
                false_fn=lambda: fun_eval,
            )

        return new_fun

    if isinstance(funcs, Iterable):
        funcs = [wrapper(f) for f in funcs]
    else:
        funcs = wrapper(funcs)

    quadrature = NDiagGHQuadrature(dim, n_gh)
    if logspace:
        result = quadrature.logspace(funcs, Fmu, Fvar, **Ys)
    else:
        result = quadrature(funcs, Fmu, Fvar, **Ys)

    if isinstance(result, list):
        result = [tf.reshape(r, shape) for r in result]
    else:
        result = tf.reshape(result, shape)

    return result


@check_shapes(
    "Fmu: [N, Din]",
    "Fvar: [N, Din]",
    "Ys.values(): [broadcast N, .]",
    "return: [broadcast n_funs, N, P]",
)
def ndiag_mc(
    funcs: Union[Callable[..., tf.Tensor], Iterable[Callable[..., tf.Tensor]]],
    S: int,
    Fmu: TensorType,
    Fvar: TensorType,
    logspace: bool = False,
    epsilon: Optional[TensorType] = None,
    **Ys: TensorType,
) -> tf.Tensor:
    """
    Computes N Gaussian expectation integrals of one or more functions
    using Monte Carlo samples. The Gaussians must be independent.

    `Fmu`, `Fvar`, `Ys` should all have same shape, with overall size `N`.

    :param funcs: the integrand(s):
        Callable or Iterable of Callables that operates elementwise
    :param S: number of Monte Carlo sampling points
    :param Fmu: array/tensor
    :param Fvar: array/tensor
    :param logspace: if True, funcs are the log-integrands and this calculates
        the log-expectation of exp(funcs)
    :param Ys: arrays/tensors; deterministic arguments to be passed by name
    :return: shape is the same as that of the first Fmu
    """
    N, D = tf.shape(Fmu)[0], tf.shape(Fvar)[1]

    if epsilon is None:
        epsilon = tf.random.normal(shape=[S, N, D], dtype=default_float())

    mc_x = Fmu[None, :, :] + tf.sqrt(Fvar[None, :, :]) * epsilon
    mc_Xr = tf.reshape(mc_x, (S * N, D))

    for name, Y in Ys.items():
        D_out = Y.shape[1]
        # we can't rely on broadcasting and need tiling
        mc_Yr = tf.tile(Y[None, ...], [S, 1, 1])  # [S, N, D]_out
        Ys[name] = tf.reshape(mc_Yr, (S * N, D_out))  # S * [N, _]out

    def eval_func(func: Callable[..., tf.Tensor]) -> tf.Tensor:
        feval = func(mc_Xr, **Ys)
        feval = tf.reshape(feval, (S, N, -1))
        if logspace:
            log_S = tf.math.log(to_default_float(S))
            return tf.reduce_logsumexp(feval, axis=0) - log_S  # [N, D]
        else:
            return tf.reduce_mean(feval, axis=0)

    if isinstance(funcs, Iterable):
        return [eval_func(f) for f in funcs]
    else:
        return eval_func(funcs)


# gpflow/quadrature/gauss_hermite.py

# Copyright 2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Sequence, Tuple

import numpy as np
import tensorflow as tf

from ..base import TensorType
from ..config import default_float
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from .base import GaussianQuadrature


@check_shapes(
    "return[0]: [N]",
    "return[1]: [N]",
)
def gh_points_and_weights(n_gh: int) -> Tuple[tf.Tensor, tf.Tensor]:
    r"""
    Given the number of Gauss-Hermite points n_gh,
    returns the points z and the weights dz to perform the following
    uni-dimensional gaussian quadrature:

    X ~ N(mean, stddev²)
    E[f(X)] = ∫ f(x) p(x) dx = \sum_{i=1}^{n_gh} f(mean + stddev*z_i) dz_i

    :param n_gh: Number of Gauss-Hermite points
    :returns: Points z and weights dz to compute uni-dimensional gaussian expectation
    """
    z, dz = np.polynomial.hermite.hermgauss(n_gh)
    z = z * np.sqrt(2)
    dz = dz / np.sqrt(np.pi)
    z, dz = z.astype(default_float()), dz.astype(default_float())
    return tf.convert_to_tensor(z), tf.convert_to_tensor(dz)


@check_shapes(
    "xs[all]: [.]",
    "return: [N_product, D]",
)
def list_to_flat_grid(xs: Sequence[TensorType]) -> tf.Tensor:
    """
    :param xs: List with d rank-1 Tensors, with shapes N1, N2, ..., Nd
    :return: Tensor with shape [N1*N2*...*Nd, d] representing the flattened
        d-dimensional grid built from the input tensors xs
    """
    return tf.reshape(tf.stack(tf.meshgrid(*xs), axis=-1), (-1, len(xs)))


@check_shapes(
    "zs[all]: [.]",
    "dzs[all]: [.]",
    "return[0]: [N_product, D]",
    "return[1]: [N_product, 1]",
)
def reshape_Z_dZ(
    zs: Sequence[TensorType], dzs: Sequence[TensorType]
) -> Tuple[tf.Tensor, tf.Tensor]:
    """
    :param zs: List with d rank-1 Tensors, with shapes N1, N2, ..., Nd
    :param dzs: List with d rank-1 Tensors, with shapes N1, N2, ..., Nd
    :returns: points Z, Tensor with shape [N1*N2*...*Nd, D],
        and weights dZ, Tensor with shape [N1*N2*...*Nd, 1]
    """
    Z = list_to_flat_grid(zs)
    dZ = tf.reduce_prod(list_to_flat_grid(dzs), axis=-1, keepdims=True)
    return Z, dZ


@check_shapes(
    "x: [batch...]",
    "return: [n, batch...]",
)
def repeat_as_list(x: TensorType, n: int) -> Sequence[tf.Tensor]:
    """
    :param x: Array/Tensor to be repeated
    :param n: Integer with the number of repetitions
    :return: List of n repetitions of Tensor x
    """
    return [x for _ in range(n)]


@check_shapes(
    "return[0]: [n_quad_points, D]",
    "return[1]: [n_quad_points, 1]",
)
def ndgh_points_and_weights(dim: int, n_gh: int) -> Tuple[tf.Tensor, tf.Tensor]:
    r"""
    :param dim: dimension of the multivariate normal
    :param n_gh: number of Gauss-Hermite points per dimension
    :returns: points Z, Tensor with shape [n_gh**dim, D],
        and weights dZ, Tensor with shape [n_gh**dim, 1]
    """
    z, dz = gh_points_and_weights(n_gh)
    zs = repeat_as_list(z, dim)
    dzs = repeat_as_list(dz, dim)
    return reshape_Z_dZ(zs, dzs)


class NDiagGHQuadrature(GaussianQuadrature):
    def __init__(self, dim: int, n_gh: int) -> None:
        """
        :param dim: dimension of the multivariate normal
        :param n_gh: number of Gauss-Hermite points per dimension
        """
        self.dim = dim
        self.n_gh = n_gh
        self.n_gh_total = n_gh ** dim
        Z, dZ = ndgh_points_and_weights(self.dim, self.n_gh)
        self.Z = tf.ensure_shape(Z, (self.n_gh_total, self.dim))
        self.dZ = tf.ensure_shape(dZ, (self.n_gh_total, 1))

    @inherit_check_shapes
    def _build_X_W(self, mean: TensorType, var: TensorType) -> Tuple[tf.Tensor, tf.Tensor]:
        """
        :param mean: Array/Tensor with shape [b1, b2, ..., bX, dim], usually [N, dim],
            representing the mean of a dim-Variate Gaussian distribution
        :param var: Array/Tensor with shape b1, b2, ..., bX, dim], usually [N, dim],
            representing the variance of a dim-Variate Gaussian distribution
        :return: points X, Tensor with shape [n_gh_total, b1, b2, ..., bX, dim],
            usually [n_gh_total, N, dim],
            and weights W, a Tensor with shape [n_gh_total, b1, b2, ..., bX, 1],
            usually [n_gh_total, N, 1]
        """

        batch_shape_broadcast = tf.ones(tf.rank(mean) - 1, dtype=tf.int32)
        shape_aux = tf.concat([[self.n_gh_total], batch_shape_broadcast], axis=0)

        # mean, var: [b1, b2, ..., bX, dim], usually [N, dim]
        mean = tf.expand_dims(mean, 0)
        stddev = tf.expand_dims(tf.sqrt(var), 0)
        # mean, stddev: [1, b1, b2, ..., bX, dim], usually [1, N, dim]

        Z = tf.cast(tf.reshape(self.Z, tf.concat([shape_aux, [self.dim]], axis=0)), mean.dtype)
        dZ = tf.cast(tf.reshape(self.dZ, tf.concat([shape_aux, [1]], axis=0)), mean.dtype)

        X = mean + stddev * Z
        W = dZ
        # X: [n_gh_total, b1, b2, ..., bX, dim], usually [n_gh_total, N, dim]
        # W: [n_gh_total,  1,  1, ...,  1,   1], usually [n_gh_total, N,   1]

        return X, W


# gpflow/utilities/__init__.py

from .bijectors import positive, triangular, triangular_size
from .misc import is_variable, set_trainable, to_default_float, to_default_int, training_loop
from .model_utils import add_likelihood_noise_cov, add_noise_cov, assert_params_false
from .multipledispatch import Dispatcher
from .traversal import (
    deepcopy,
    freeze,
    leaf_components,
    multiple_assign,
    parameter_dict,
    print_summary,
    read_values,
    reset_cache_bijectors,
    select_dict_parameters_with_prior,
    tabulate_module_summary,
    traverse_module,
)

__all__ = [
    "Dispatcher",
    "add_likelihood_noise_cov",
    "add_noise_cov",
    "assert_params_false",
    "bijectors",
    "deepcopy",
    "freeze",
    "is_variable",
    "leaf_components",
    "misc",
    "model_utils",
    "multiple_assign",
    "multipledispatch",
    "ops",
    "parameter_dict",
    "parameter_or_function",
    "positive",
    "print_summary",
    "read_values",
    "reset_cache_bijectors",
    "select_dict_parameters_with_prior",
    "set_trainable",
    "tabulate_module_summary",
    "to_default_float",
    "to_default_int",
    "training_loop",
    "traversal",
    "traverse_module",
    "triangular",
    "triangular_size",
]


# gpflow/utilities/bijectors.py

# Copyright 2019-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Optional

import tensorflow as tf
import tensorflow_probability as tfp

from .. import config
from ..experimental.check_shapes import check_shapes
from .misc import to_default_float

__all__ = ["positive", "triangular", "triangular_size"]


def positive(lower: Optional[float] = None, base: Optional[str] = None) -> tfp.bijectors.Bijector:
    """
    Returns a positive bijector (a reversible transformation from real to positive numbers).

    :param lower: overrides default lower bound
        (if None, defaults to gpflow.config.default_positive_minimum())
    :param base: overrides base positive bijector
        (if None, defaults to gpflow.config.default_positive_bijector())
    :returns: a bijector instance
    """
    bijector = base if base is not None else config.default_positive_bijector()
    bijector = config.positive_bijector_type_map()[bijector.lower()]()

    lower_bound = lower if lower is not None else config.default_positive_minimum()

    if lower_bound != 0.0:
        shift = tfp.bijectors.Shift(to_default_float(lower_bound))
        bijector = tfp.bijectors.Chain([shift, bijector])  # from unconstrained to constrained
    return bijector


def triangular() -> tfp.bijectors.Bijector:
    """
    Returns instance of a (lower) triangular bijector.
    """
    return tfp.bijectors.FillTriangular()


@check_shapes(
    "n: []",
    "return: []",
)
def triangular_size(n: tf.Tensor) -> tf.Tensor:
    """
    Returns the number of non-zero elements in an `n` by `n` triangular matrix.
    """
    return n * (n + 1) // 2


# gpflow/utilities/misc.py

# Copyright 2017-2021 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Callable, Iterable, List, Optional, Union

import tensorflow as tf
import tensorflow_probability as tfp

from ..base import TensorData
from ..config import default_float, default_int
from ..experimental.check_shapes import check_shapes

__all__ = [
    "is_variable",
    "set_trainable",
    "to_default_float",
    "to_default_int",
    "training_loop",
]


@check_shapes(
    "x: [any...]",
    "return: [any...]",
)
def to_default_int(x: TensorData) -> tf.Tensor:
    return tf.cast(x, dtype=default_int())


@check_shapes(
    "x: [any...]",
    "return: [any...]",
)
def to_default_float(x: TensorData) -> tf.Tensor:
    return tf.cast(x, dtype=default_float())


def set_trainable(model: Union[tf.Module, Iterable[tf.Module]], flag: bool) -> None:
    """
    Set trainable flag for all :class:`tf.Variable`\ s and :class:`gpflow.Parameter`\ s in a
    :class:`tf.Module` or collection of :class:`tf.Module`\ s.
    """
    modules = [model] if isinstance(model, tf.Module) else model

    for mod in modules:
        for variable in mod.variables:
            variable._trainable = flag


def is_variable(t: TensorData) -> bool:
    """
    Returns whether the `t` is a TensorFlow variable.
    """
    return isinstance(t, (tf.Variable, tfp.util.TransformedVariable))


def training_loop(
    closure: Callable[[], tf.Tensor],
    optimizer: Optional[tf.optimizers.Optimizer] = None,
    var_list: Optional[List[tf.Variable]] = None,
    maxiter: int = 1_000,
    compile: bool = False,
) -> None:
    """
    Simple generic training loop. At each iteration uses a GradientTape to compute
    the gradients of a loss function with respect to a set of variables.

    :param closure: Callable that constructs a loss function based on data and model being trained
    :param optimizer: tf.optimizers or tf.keras.optimizers that updates variables by applying the
        corresponding loss gradients. Adam is a default optimizer with default settings.
    :param var_list: List of model variables to be learnt during training
    :param maxiter: Maximum number of
    :return:
    """

    safe_optimizer = tf.optimizers.Adam() if optimizer is None else optimizer
    safe_var_list = [] if var_list is None else var_list

    def optimization_step() -> None:
        with tf.GradientTape(watch_accessed_variables=False) as tape:
            tape.watch(safe_var_list)
            loss = closure()
        grads = tape.gradient(loss, safe_var_list)
        safe_optimizer.apply_gradients(zip(grads, safe_var_list))

    if compile:
        optimization_step = tf.function(optimization_step)

    for _ in range(maxiter):
        optimization_step()


# gpflow/utilities/model_utils.py

from typing import Any, Callable

import tensorflow as tf

from ..base import TensorType
from ..experimental.check_shapes import check_shapes
from ..likelihoods import Gaussian


def assert_params_false(
    called_method: Callable[..., Any],
    **kwargs: bool,
) -> None:
    """
    Asserts that parameters are ``False``.

    :param called_method: The method or function that is calling this. Used for nice error messages.
    :param kwargs: Parameters that must be ``False``.
    :raises NotImplementedError: If any ``kwargs`` are ``True``.
    """
    errors_str = ", ".join(f"{param}={value}" for param, value in kwargs.items() if value)
    if errors_str:
        raise NotImplementedError(
            f"{called_method.__qualname__} does not currently support: {errors_str}"
        )


@check_shapes(
    "K: [batch..., N, N]",
    "likelihood_variance: [broadcast batch..., broadcast N]",
    "return: [batch..., N, N]",
)
def add_noise_cov(K: tf.Tensor, likelihood_variance: TensorType) -> tf.Tensor:
    """
    Returns K + σ², where σ² is the diagonal likelihood noise variance.
    """
    k_diag = tf.linalg.diag_part(K)
    return tf.linalg.set_diag(K, k_diag + likelihood_variance)


@check_shapes(
    "K: [batch..., N, N]",
    "X: [batch..., N, D]",
    "return: [batch..., N, N]",
)
def add_likelihood_noise_cov(K: tf.Tensor, likelihood: Gaussian, X: TensorType) -> tf.Tensor:
    """
    Returns K + σ², where σ² is the likelihood noise variance.
    """
    return add_noise_cov(K, tf.squeeze(likelihood.variance_at(X), axis=-1))


# gpflow/utilities/multipledispatch.py

# Copyright 2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Any, Callable, Optional, Tuple, Type, TypeVar, Union

from multipledispatch import Dispatcher as GeneratorDispatcher
from multipledispatch.dispatcher import str_signature, variadic_signature_matches
from multipledispatch.variadic import isvariadic

__all__ = ["Dispatcher"]


AnyCallable = Callable[..., Any]
_C = TypeVar("_C", bound=AnyCallable)
Types = Union[Type[Any], Tuple[Type[Any], ...]]


class Dispatcher(GeneratorDispatcher):
    """
    multipledispatch.Dispatcher uses a generator to yield the
    desired function implementation, which is problematic as TensorFlow's
    autograph is not able to compile code that passes through generators.

    This class overwrites the problematic method in the original
    Dispatcher and solely makes use of simple for-loops, which are
    compilable by AutoGraph.
    """

    def register(self, *types: Types, **kwargs: Any) -> Callable[[_C], _C]:
        # Override to add type hints...
        result: Callable[[_C], _C] = super().register(*types, **kwargs)
        return result

    def dispatch(self, *types: Types) -> Optional[AnyCallable]:
        """
        Returns matching function for `types`; if not existing returns None.
        """
        if types in self.funcs:
            result: AnyCallable = self.funcs[types]
            return result

        return self.get_first_occurrence(*types)

    def dispatch_or_raise(self, *types: Types) -> AnyCallable:
        """
        Returns matching function for `types`; if not existing raises an error.
        """
        f = self.dispatch(*types)
        if f is None:
            raise NotImplementedError(
                f"Could not find signature for {self.name}: <{str_signature(types)}>"
            )
        return f

    def get_first_occurrence(self, *types: Types) -> Optional[AnyCallable]:
        """
        Returns the first occurrence of a matching function

        Based on `multipledispatch.Dispatcher.dispatch_iter`, which
        returns an iterator of matching functions. This method uses
        the same logic to select functions, but simply returns the first
        element of the iterator. If no matching functions are found,
        `None` is returned.
        """
        n = len(types)
        for signature in self.ordering:
            if len(signature) == n and all(map(issubclass, types, signature)):  # type: ignore[arg-type]
                result: AnyCallable = self.funcs[signature]
                return result
            elif len(signature) and isvariadic(signature[-1]):
                if variadic_signature_matches(types, signature):
                    result = self.funcs[signature]
                    return result
        return None


# gpflow/utilities/ops.py

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import copy
from typing import Any, Callable, List, Optional, Union

import tensorflow as tf
import tensorflow_probability as tfp

from ..base import AnyNDArray
from ..experimental.check_shapes import check_shapes


@check_shapes(
    "value: []",
    "return: [N, N]",
)
def eye(num: int, value: tf.Tensor, dtype: Optional[tf.DType] = None) -> tf.Tensor:
    if dtype is not None:
        value = tf.cast(value, dtype)
    return tf.linalg.diag(tf.fill([num], value))


@check_shapes(
    "tensor: [any...]",
    "return: [transposed_any...]",
)
def leading_transpose(tensor: tf.Tensor, perm: List[Any], leading_dim: int = 0) -> tf.Tensor:
    """
    Transposes tensors with leading dimensions.

    Leading dimensions in permutation list represented via ellipsis `...` and is of type
    List[Union[int, type(...)]  (please note, due to mypy issues, List[Any] is used instead).  When
    leading dimensions are found, `transpose` method considers them as a single grouped element
    indexed by 0 in `perm` list. So, passing `perm=[-2, ..., -1]`, you assume that your input tensor
    has [..., A, B] shape, and you want to move leading dims between A and B dimensions.  Dimension
    indices in permutation list can be negative or positive. Valid positive indices start from 1 up
    to the tensor rank, viewing leading dimensions `...` as zero index.

    Example::

        a = tf.random.normal((1, 2, 3, 4, 5, 6))
        # [..., A, B, C],
        # where A is 1st element,
        # B is 2nd element and
        # C is 3rd element in
        # permutation list,
        # leading dimensions are [1, 2, 3]
        # which are 0th element in permutation list
        b = leading_transpose(a, [3, -3, ..., -2])  # [C, A, ..., B]
        sess.run(b).shape

        output> (6, 4, 1, 2, 3, 5)

    :param tensor: TensorFlow tensor.
    :param perm: List of permutation indices.
    :returns: TensorFlow tensor.
    :raises ValueError: when `...` cannot be found.

    """
    perm = copy.copy(perm)
    idx = perm.index(...)
    perm[idx] = leading_dim

    rank = tf.rank(tensor)
    perm_tf = perm % rank

    leading_dims = tf.range(rank - len(perm) + 1)
    perm = tf.concat([perm_tf[:idx], leading_dims, perm_tf[idx + 1 :]], 0)
    return tf.transpose(tensor, perm)


@check_shapes(
    "a: [a_shape...]",
    "b: [b_shape...]",
    "return: [a_shape..., b_shape...]",
)
def broadcasting_elementwise(
    op: Callable[[tf.Tensor, tf.Tensor], tf.Tensor], a: tf.Tensor, b: tf.Tensor
) -> tf.Tensor:
    """
    Apply binary operation `op` to every pair in tensors `a` and `b`.

    :param op: binary operator on tensors, e.g. tf.add, tf.substract
    """
    flatres = op(tf.reshape(a, [-1, 1]), tf.reshape(b, [1, -1]))
    return tf.reshape(flatres, tf.concat([tf.shape(a), tf.shape(b)], 0))


@check_shapes(
    "X: [batch..., N, D]",
    "X2: [batch2..., N2, D]",
    "return: [batch..., N, batch2..., N2] if X2 is not None",
    "return: [batch..., N, N] if X2 is None",
)
def square_distance(X: tf.Tensor, X2: Optional[tf.Tensor]) -> tf.Tensor:
    """
    Returns ||X - X2ᵀ||²

    Due to the implementation and floating-point imprecision, the
    result may actually be very slightly negative for entries very
    close to each other.
    """
    if X2 is None:
        Xs = tf.reduce_sum(tf.square(X), axis=-1, keepdims=True)
        dist = -2 * tf.matmul(X, X, transpose_b=True)
        dist += Xs + tf.linalg.adjoint(Xs)
        return dist
    Xs = tf.reduce_sum(tf.square(X), axis=-1)
    X2s = tf.reduce_sum(tf.square(X2), axis=-1)
    dist = -2 * tf.tensordot(X, X2, [[-1], [-1]])
    dist += broadcasting_elementwise(tf.add, Xs, X2s)
    return dist


@check_shapes(
    "X: [batch..., N, D]",
    "X2: [batch2..., N2, D]",
    "return: [batch..., N, batch2..., N2, D] if X2 is not None",
    "return: [batch..., N, N, D] if X2 is None",
)
def difference_matrix(X: tf.Tensor, X2: Optional[tf.Tensor]) -> tf.Tensor:
    """
    Returns (X - X2ᵀ)
    """
    if X2 is None:
        X2 = X
        diff = X[..., :, tf.newaxis, :] - X2[..., tf.newaxis, :, :]
        return diff
    Xshape = tf.shape(X)
    X2shape = tf.shape(X2)
    X = tf.reshape(X, (-1, Xshape[-1]))
    X2 = tf.reshape(X2, (-1, X2shape[-1]))
    diff = X[:, tf.newaxis, :] - X2[tf.newaxis, :, :]
    diff = tf.reshape(diff, tf.concat((Xshape[:-1], X2shape[:-1], [Xshape[-1]]), 0))
    return diff


@check_shapes(
    "X: [N, D]",
    "latent_dim: []",
    "return: [N, Q]",
)
def pca_reduce(X: tf.Tensor, latent_dim: tf.Tensor) -> tf.Tensor:
    """
    Linearly reduce the dimensionality of the input points `X` to `latent_dim` dimensions.

    :param X: Data to reduce.
    :param latent_dim: Number of latent dimension, Q < D.
    :return: PCA projection array.

    """
    if latent_dim > X.shape[1]:  # pragma: no cover
        raise ValueError("Cannot have more latent dimensions than observed")
    X_cov = tfp.stats.covariance(X)
    evals, evecs = tf.linalg.eigh(X_cov)
    W = evecs[:, -latent_dim:]
    return (X - tf.reduce_mean(X, axis=0, keepdims=True)) @ W


# gpflow/utilities/parameter_or_function.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Optional, Union

import tensorflow as tf

from ..base import Parameter, TensorData, TensorType
from ..experimental.check_shapes import check_shapes
from ..functions import Function
from . import positive

ConstantOrFunction = Union[Function, TensorData]
ParameterOrFunction = Union[Function, Parameter]


def prepare_parameter_or_function(
    value: ConstantOrFunction,
    *,
    lower_bound: Optional[float] = None,
) -> ParameterOrFunction:
    if isinstance(value, Function):
        return value
    else:
        if lower_bound is None:
            return Parameter(value)
        else:
            return Parameter(value, transform=positive(lower_bound))


@check_shapes(
    "X: [batch..., N, D]",
    "return: [broadcast batch..., broadcast N, broadcast P]",
)
def evaluate_parameter_or_function(
    value: ParameterOrFunction,
    X: TensorType,
    *,
    lower_bound: Optional[float] = None,
) -> TensorType:
    if isinstance(value, Function):
        result = value(X)
        if lower_bound is not None:
            result = tf.maximum(result, lower_bound)
        return result
    else:
        return value


# gpflow/utilities/traversal.py

# Copyright 2017-2021 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import copy
import re
from functools import lru_cache
from typing import Any, Callable, Dict, Mapping, Optional, Pattern, Tuple, Type, TypeVar, Union

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
from packaging.version import Version
from tabulate import tabulate
from tensorflow.python.util.object_identity import Reference

from ..base import AnyNDArray, Parameter
from ..config import default_summary_fmt

__all__ = [
    "deepcopy",
    "freeze",
    "leaf_components",
    "multiple_assign",
    "parameter_dict",
    "print_summary",
    "read_values",
    "reset_cache_bijectors",
    "select_dict_parameters_with_prior",
    "tabulate_module_summary",
]
LeafComponent = Union[tf.Variable, tf.Tensor, Parameter]
LeafVariable = Union[tf.Variable, Parameter]
HashableTensor = Union[Reference, Parameter]
TraverseInput = TypeVar("TraverseInput", tf.Variable, tf.Module, Parameter)
State = TypeVar("State")
Path = str
Accumulator = Tuple[Path, State]
TraverseUpdateCallable = Callable[[TraverseInput, Path, State], State]


def multiple_assign(module: tf.Module, parameters: Mapping[Path, tf.Tensor]) -> None:
    """
    Multiple assign takes a dictionary with new values. Dictionary keys are paths to the
    `tf.Variable`s or `gpflow.Parameter` of the input module.

    :param module: `tf.Module`.
    :param parameters: a dictionary with keys of the form ".module.path.to.variable" and new value tensors.
    """
    reference_var_dict = parameter_dict(module)
    for path, value in parameters.items():
        reference_var_dict[path].assign(value)


def read_values(module: tf.Module) -> Dict[Path, AnyNDArray]:
    """Returns a dictionary of numpy values of the module parameters (variables)."""
    return {k: v.numpy() for k, v in parameter_dict(module).items()}


def parameter_dict(module: tf.Module) -> Dict[Path, LeafVariable]:
    """
    Returns a dictionary of parameters (variables) for the `tf.Module` component.
    Dictionary keys are relative paths to the attributes to which parameters (variables) assigned to.

        class SubModule(tf.Module):
            def __init__(self):
                self.parameter = gpflow.Parameter(1.0)
                self.variable = tf.Variable(1.0)

        class Module(tf.Module):
            def __init__(self):
                self.submodule = SubModule()

        m = Module()
        params = parameter_dict(m)
        # {
        #   ".submodule.parameter": <parameter object>,
        #   ".submodule.variable": <variable object>
        # }
    """
    param_dict = leaf_components(module)
    return {f".{key.split('.', 1)[-1]}": value for key, value in param_dict.items()}


def print_summary(module: tf.Module, fmt: Optional[str] = None) -> None:
    """
    Prints a summary of the parameters and variables contained in a tf.Module.
    """
    fmt = fmt if fmt is not None else default_summary_fmt()
    if fmt == "notebook":
        from IPython.core.display import HTML, display

        tab = tabulate_module_summary(module, "html")
        display(HTML(tab))
    else:
        print(tabulate_module_summary(module, fmt))


def tabulate_module_summary(module: tf.Module, tablefmt: Optional[str] = None) -> str:
    def get_transform(path: Path, var: LeafComponent) -> Optional[str]:
        if hasattr(var, "transform") and var.transform is not None:
            if isinstance(var.transform, tfp.bijectors.Chain):
                return " + ".join(b.__class__.__name__ for b in var.transform.bijectors[::-1])
            return var.transform.__class__.__name__  # type: ignore[no-any-return]
        return None

    def get_prior(path: Path, var: LeafComponent) -> Optional[str]:
        if hasattr(var, "prior") and var.prior is not None:
            return var.prior.name  # type: ignore[no-any-return]
        return None

    # list of (column_name: str, column_getter: Callable[[tf.Variable], str]) tuples:
    column_definition = [
        ("name", lambda path, var: path),
        ("class", lambda path, var: var.__class__.__name__),
        ("transform", get_transform),
        ("prior", get_prior),
        ("trainable", lambda path, var: var.trainable),
        ("shape", lambda path, var: var.shape),
        ("dtype", lambda path, var: var.dtype.name),
        ("value", lambda path, var: _str_tensor_value(var.numpy())),
    ]
    column_names, column_getters = zip(*column_definition)

    merged_leaf_components = _merge_leaf_components(leaf_components(module))

    column_values = [
        [getter(path, variable) for getter in column_getters]
        for path, variable in merged_leaf_components.items()
    ]
    # mypy claims it's wrong to pass a `None` tablefmt below. I think `tabulate` has bad type hints.
    return tabulate(column_values, headers=column_names, tablefmt=tablefmt)  # type: ignore[arg-type]


def leaf_components(input: tf.Module) -> Mapping[Path, LeafVariable]:
    return _get_leaf_components(input)


def _merge_leaf_components(input: Mapping[Path, LeafComponent]) -> Mapping[Path, LeafComponent]:

    ref_fn: Callable[[LeafComponent], HashableTensor] = lambda x: (
        x if isinstance(x, Parameter) else x.ref()
    )
    deref_fn: Callable[[HashableTensor], LeafComponent] = lambda x: (
        x if isinstance(x, Parameter) else x.deref()
    )

    input_values = {ref_fn(value) for value in input.values()}
    if len(input_values) == len(input):
        return input

    tmp_dict: Dict[HashableTensor, Path] = {}
    for key, value in input.items():
        ref = ref_fn(value)
        if ref in tmp_dict:
            tmp_dict[ref] = f"{tmp_dict[ref]}\n{key}"
        else:
            tmp_dict[ref] = key
    return {key: deref_fn(ref) for ref, key in tmp_dict.items()}


def _get_leaf_components(input_module: tf.Module) -> Mapping[Path, LeafVariable]:
    """
    Returns a list of tuples each corresponding to a gpflow.Parameter or tf.Variable in the each
    submodules of a given tf.Module. Each tuple consists of an specific Parameter (or Variable) and
    its relative path inside the module, which is constructed recursively by adding a prefix with
    the path to the current module. Designed to be used as a helper for the method 'print_summary'.

    :param input_module: tf.Module including keras.Model, keras.layers.Layer and gpflow.Module.
    :return:
    """
    target_types = (Parameter, tf.Variable)
    input_name = input_module.__class__.__name__
    state: Dict[Path, LeafVariable] = {}
    accumulator = (input_name, state)

    def update_state(
        parameter_or_variable: LeafVariable, path: Path, state: Dict[Path, LeafVariable]
    ) -> Dict[Path, LeafVariable]:
        state[path] = parameter_or_variable
        return state

    state = traverse_module(input_module, accumulator, update_state, target_types)
    return state


def reset_cache_bijectors(input_module: tf.Module) -> tf.Module:
    """
    Recursively finds tfp.bijectors.Bijector-s inside the components of the tf.Module using `traverse_component`.
    Resets the caches stored inside each tfp.bijectors.Bijector.

    :param input_module: tf.Module including keras.Model, keras.layers.Layer and gpflow.Module.
    :returns: same object but with all bijector caches reset
    """
    if Version(tfp.__version__) >= Version("0.11.0"):
        if hasattr(tfp.bijectors.Identity()._cache, "clear"):
            # implementation in `master` branch (checked 29 Sep 2020) provides clear():

            def _clear_bijector_cache(bijector: tfp.bijectors.Bijector) -> None:
                bijector._cache.clear()

        else:  # pragma: no cover
            # previous versions (including the versions 0.11.0 and 0.11.1 released as of 29 Sep 2020) provide reset(), but its implementation is broken

            def _clear_bijector_cache(bijector: tfp.bijectors.Bijector) -> None:
                # workaround for broken implementation of bijector._cache.reset():
                cache = bijector._cache
                cache_type = type(cache.forward)
                assert type(cache.inverse) == cache_type
                cache.__init__(cache.forward._func, cache.inverse._func, cache_type)

    else:  # pragma: no cover
        # fallback for backwards-compatibility with tensorflow_probability < 0.11.0

        def _clear_bijector_cache(bijector: tfp.bijectors.Bijector) -> None:
            # `_from_x` and `_from_y` are cache dictionaries for forward and inverse transformations
            bijector._from_x.clear()
            bijector._from_y.clear()

    target_types = (tfp.bijectors.Bijector,)
    accumulator = ("", None)

    def clear_bijector(bijector: tfp.bijectors.Bijector, _: Path, state: None) -> None:
        if not isinstance(bijector, tfp.bijectors.Bijector):
            return  # skip submodules that are not bijectors

        _clear_bijector_cache(bijector)

        if isinstance(bijector, tfp.bijectors.Chain):
            # recursively clear caches of sub-bijectors
            for m in bijector.submodules:
                if isinstance(m, tfp.bijectors.Bijector):
                    _clear_bijector_cache(m)

        return state

    traverse_module(input_module, accumulator, clear_bijector, target_types)
    return input_module


M = TypeVar("M", bound=tf.Module)


def deepcopy(input_module: M, memo: Optional[Dict[int, Any]] = None) -> M:
    """
    Returns a deepcopy of the input tf.Module. To do that first resets the caches stored inside each
    tfp.bijectors.Bijector to allow the deepcopy of the tf.Module.

    :param input_module: tf.Module including keras.Model, keras.layers.Layer and gpflow.Module.
    :param memo: passed through to func:`copy.deepcopy`
        (see https://docs.python.org/3/library/copy.html).
    :return: Returns a deepcopy of an input object.
    """
    return copy.deepcopy(reset_cache_bijectors(input_module), memo)  # type: ignore[no-any-return]


def freeze(input_module: M) -> M:
    """
    Returns a deepcopy of the input tf.Module with constants instead of variables and parameters.

    :param input_module: tf.Module or gpflow.Module.
    :return: Returns a frozen deepcopy of an input object.
    """
    objects_to_freeze = _get_leaf_components(input_module)
    memo_tensors = {id(v): tf.convert_to_tensor(v) for v in objects_to_freeze.values()}
    module_copy = deepcopy(input_module, memo_tensors)
    return module_copy


def traverse_module(
    m: TraverseInput,
    acc: Accumulator[State],
    update_cb: TraverseUpdateCallable[TraverseInput, State],
    target_types: Tuple[Type[Any], ...],
) -> State:
    """
    Recursively traverses `m`, accumulating in `acc` a path and a state until it finds an object of
    type in `target_types` to apply `update_cb` to update the accumulator `acc` and/or the object.

    :param m: tf.Module, tf.Variable or gpflow.Parameter
    :param acc: Tuple of path and state
    :param update_cb: Callable
    :param target_types: target class types
    :return:
    """
    path, state = acc

    new_state = state

    if isinstance(m, target_types):
        return update_cb(m, path, state)

    if isinstance(m, (list, tuple)):
        for term_idx, subterm in enumerate(m):
            new_acc = (f"{path}[{term_idx}]", new_state)
            new_state = traverse_module(subterm, new_acc, update_cb, target_types)
    elif isinstance(m, dict):
        for term_idx, subterm in m.items():
            new_acc = (f"{path}['{term_idx}']", new_state)
            new_state = traverse_module(subterm, new_acc, update_cb, target_types)
    elif isinstance(m, tf.Module):
        for name, submodule in vars(m).items():
            ignored_attributes = m._TF_MODULE_IGNORED_PROPERTIES
            # NOTE(awav): since tfp version 0.10.0, tfp.bijectors.Bijector instances have
            # `_parameters` dictionary with "self" references that cause
            # infinite recursive loop.
            if isinstance(m, tfp.bijectors.Bijector):
                ignored_attributes = ignored_attributes.union({"_parameters"})
            if name in ignored_attributes:
                continue
            new_acc = (f"{path}.{name}", new_state)
            new_state = traverse_module(submodule, new_acc, update_cb, target_types)
    return new_state


@lru_cache()
def _first_three_elements_regexp() -> Pattern[str]:
    num_re = r"[+\-]?(?:0|[1-9]\d*)(?:\.\d*)?(?:[eE][+\-]?\d+)?"
    pat_re = rf"^(?:(\[+)\s*)?({num_re})(?:\s+({num_re})(?:\s+({num_re}))?)?.*?"
    return re.compile(pat_re)


def _str_tensor_value(value: AnyNDArray) -> str:
    value_str = str(np.around(value, 5))
    if value.size <= 3:
        return value_str

    max_chars = 500
    value_str = value_str[:max_chars]
    regexp = _first_three_elements_regexp()
    match = regexp.match(value_str)
    assert match is not None
    brackets, elem1, elem2, elem3 = match.groups()

    out = f"{elem1}"
    if elem2 is not None:
        out = f"{out}{f', {elem2}'}"
        if elem3 is not None:
            out = f"{out}{f', {elem3}'}"
    if brackets is not None:
        out = f"{brackets}{out}..."

    return out


def select_dict_parameters_with_prior(model: tf.Module) -> Dict[Path, Parameter]:
    """Collects parameters with prior into a dictionary."""
    return {
        k: p
        for k, p in parameter_dict(model).items()
        if hasattr(p, "prior") and p.prior is not None
    }


# tests/benchmark/test_benchmark_api.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Collection
from unittest.mock import Mock

import pandas as pd
import pytest

import benchmark.datasets as ds
import benchmark.models as md
import benchmark.plotters as pl
from benchmark.benchmark_api import BenchmarkSet, BenchmarkSuite, BenchmarkTask
from benchmark.grouping import GroupingKey as GK
from benchmark.grouping import GroupingSpec


@pytest.mark.parametrize(
    "benchmark_set,expected_tasks",
    [
        (
            BenchmarkSet(
                name="metrics",
                datasets=[
                    ds.tiny_linear,
                    ds.tiny_sine,
                ],
                models=[
                    md.gpr,
                ],
                plots=[
                    pl.metrics_box_plot,
                ],
                do_compile=[False],
                do_optimise=[False],
                do_predict=True,
                do_posterior=True,
                file_by=GroupingSpec((GK.DATASET,), minimise=False),
                row_by=GroupingSpec((GK.METRIC,), minimise=False),
                column_by=GroupingSpec((GK.PLOTTER,), minimise=False),
                line_by=None,
                repetitions=2,
            ),
            [
                BenchmarkTask(
                    dataset_name="tiny_linear",
                    model_name="gpr",
                    do_compile=False,
                    do_optimise=False,
                    do_predict=True,
                    do_posterior=True,
                    repetitions=2,
                ),
                BenchmarkTask(
                    dataset_name="tiny_sine",
                    model_name="gpr",
                    do_compile=False,
                    do_optimise=False,
                    do_predict=True,
                    do_posterior=True,
                    repetitions=2,
                ),
            ],
        ),
        (
            BenchmarkSet(
                name="metrics",
                datasets=[
                    ds.tiny_linear,
                ],
                models=[
                    md.gpr,
                    md.svgp,
                ],
                plots=[
                    pl.metrics_box_plot,
                ],
                do_compile=[False],
                do_optimise=[False],
                do_predict=True,
                do_posterior=True,
                file_by=GroupingSpec((GK.DATASET,), minimise=False),
                row_by=GroupingSpec((GK.METRIC,), minimise=False),
                column_by=GroupingSpec((GK.PLOTTER,), minimise=False),
                line_by=None,
                repetitions=2,
            ),
            [
                BenchmarkTask(
                    dataset_name="tiny_linear",
                    model_name="gpr",
                    do_compile=False,
                    do_optimise=False,
                    do_predict=True,
                    do_posterior=True,
                    repetitions=2,
                ),
                BenchmarkTask(
                    dataset_name="tiny_linear",
                    model_name="svgp",
                    do_compile=False,
                    do_optimise=False,
                    do_predict=True,
                    do_posterior=True,
                    repetitions=2,
                ),
            ],
        ),
        (
            BenchmarkSet(
                name="metrics",
                datasets=[
                    ds.tiny_linear,
                ],
                models=[
                    md.gpr,
                ],
                plots=[
                    pl.metrics_box_plot,
                ],
                do_compile=[False, True],
                do_optimise=[False],
                do_predict=True,
                do_posterior=True,
                file_by=GroupingSpec((GK.DATASET,), minimise=False),
                row_by=GroupingSpec((GK.METRIC,), minimise=False),
                column_by=GroupingSpec((GK.PLOTTER,), minimise=False),
                line_by=None,
                repetitions=2,
            ),
            [
                BenchmarkTask(
                    dataset_name="tiny_linear",
                    model_name="gpr",
                    do_compile=False,
                    do_optimise=False,
                    do_predict=True,
                    do_posterior=True,
                    repetitions=2,
                ),
                BenchmarkTask(
                    dataset_name="tiny_linear",
                    model_name="gpr",
                    do_compile=True,
                    do_optimise=False,
                    do_predict=True,
                    do_posterior=True,
                    repetitions=2,
                ),
            ],
        ),
        (
            BenchmarkSet(
                name="metrics",
                datasets=[
                    ds.tiny_linear,
                ],
                models=[
                    md.gpr,
                ],
                plots=[
                    pl.metrics_box_plot,
                ],
                do_compile=[False],
                do_optimise=[False, True],
                do_predict=True,
                do_posterior=True,
                file_by=GroupingSpec((GK.DATASET,), minimise=False),
                row_by=GroupingSpec((GK.METRIC,), minimise=False),
                column_by=GroupingSpec((GK.PLOTTER,), minimise=False),
                line_by=None,
                repetitions=2,
            ),
            [
                BenchmarkTask(
                    dataset_name="tiny_linear",
                    model_name="gpr",
                    do_compile=False,
                    do_optimise=False,
                    do_predict=True,
                    do_posterior=True,
                    repetitions=2,
                ),
                BenchmarkTask(
                    dataset_name="tiny_linear",
                    model_name="gpr",
                    do_compile=False,
                    do_optimise=True,
                    do_predict=True,
                    do_posterior=True,
                    repetitions=2,
                ),
            ],
        ),
    ],
)
def test_benchmark_set__get_tasks(
    benchmark_set: BenchmarkSet, expected_tasks: Collection[BenchmarkTask]
) -> None:
    assert expected_tasks == benchmark_set.get_tasks()


@pytest.mark.parametrize(
    "benchmark_set,expected_metrics",
    [
        (
            BenchmarkSet(
                name="metrics",
                datasets=[
                    ds.tiny_linear,
                    ds.tiny_sine,
                ],
                models=[
                    md.gpr,
                ],
                plots=[
                    pl.metrics_box_plot,
                ],
                do_compile=[False],
                do_optimise=[False],
                do_predict=True,
                do_posterior=True,
                file_by=GroupingSpec((GK.DATASET,), minimise=False),
                row_by=GroupingSpec((GK.METRIC,), minimise=False),
                column_by=GroupingSpec((GK.PLOTTER,), minimise=False),
                line_by=None,
                repetitions=2,
            ),
            pd.DataFrame(
                [
                    (0, "tiny_linear", "gpr", False, False, 0),
                    (1, "tiny_sine", "gpr", False, False, 0),
                    (5, "tiny_linear", "gpr", False, False, 1),
                ],
                columns=["id", "dataset", "model", "do_compile", "do_optimise", "repetition"],
            ),
        ),
        (
            BenchmarkSet(
                name="metrics",
                datasets=[
                    ds.tiny_linear,
                ],
                models=[
                    md.gpr,
                    md.svgp,
                ],
                plots=[
                    pl.metrics_box_plot,
                ],
                do_compile=[False],
                do_optimise=[False],
                do_predict=True,
                do_posterior=True,
                file_by=GroupingSpec((GK.DATASET,), minimise=False),
                row_by=GroupingSpec((GK.METRIC,), minimise=False),
                column_by=GroupingSpec((GK.PLOTTER,), minimise=False),
                line_by=None,
                repetitions=2,
            ),
            pd.DataFrame(
                [
                    (0, "tiny_linear", "gpr", False, False, 0),
                    (2, "tiny_linear", "svgp", False, False, 0),
                    (5, "tiny_linear", "gpr", False, False, 1),
                ],
                columns=["id", "dataset", "model", "do_compile", "do_optimise", "repetition"],
            ),
        ),
        (
            BenchmarkSet(
                name="metrics",
                datasets=[
                    ds.tiny_linear,
                ],
                models=[
                    md.gpr,
                ],
                plots=[
                    pl.metrics_box_plot,
                ],
                do_compile=[False, True],
                do_optimise=[False],
                do_predict=True,
                do_posterior=True,
                file_by=GroupingSpec((GK.DATASET,), minimise=False),
                row_by=GroupingSpec((GK.METRIC,), minimise=False),
                column_by=GroupingSpec((GK.PLOTTER,), minimise=False),
                line_by=None,
                repetitions=2,
            ),
            pd.DataFrame(
                [
                    (0, "tiny_linear", "gpr", False, False, 0),
                    (3, "tiny_linear", "gpr", True, False, 0),
                    (5, "tiny_linear", "gpr", False, False, 1),
                ],
                columns=["id", "dataset", "model", "do_compile", "do_optimise", "repetition"],
            ),
        ),
        (
            BenchmarkSet(
                name="metrics",
                datasets=[
                    ds.tiny_linear,
                ],
                models=[
                    md.gpr,
                ],
                plots=[
                    pl.metrics_box_plot,
                ],
                do_compile=[False],
                do_optimise=[False, True],
                do_predict=True,
                do_posterior=True,
                file_by=GroupingSpec((GK.DATASET,), minimise=False),
                row_by=GroupingSpec((GK.METRIC,), minimise=False),
                column_by=GroupingSpec((GK.PLOTTER,), minimise=False),
                line_by=None,
                repetitions=2,
            ),
            pd.DataFrame(
                [
                    (0, "tiny_linear", "gpr", False, False, 0),
                    (4, "tiny_linear", "gpr", False, True, 0),
                    (5, "tiny_linear", "gpr", False, False, 1),
                ],
                columns=["id", "dataset", "model", "do_compile", "do_optimise", "repetition"],
            ),
        ),
        (
            BenchmarkSet(
                name="metrics",
                datasets=[
                    ds.tiny_linear,
                ],
                models=[
                    md.gpr,
                ],
                plots=[
                    pl.metrics_box_plot,
                ],
                do_compile=[False],
                do_optimise=[False],
                do_predict=True,
                do_posterior=True,
                file_by=GroupingSpec((GK.DATASET,), minimise=False),
                row_by=GroupingSpec((GK.METRIC,), minimise=False),
                column_by=GroupingSpec((GK.PLOTTER,), minimise=False),
                line_by=None,
                repetitions=5,
            ),
            pd.DataFrame(
                [
                    (0, "tiny_linear", "gpr", False, False, 0),
                    (5, "tiny_linear", "gpr", False, False, 1),
                    (6, "tiny_linear", "gpr", False, False, 2),
                    (7, "tiny_linear", "gpr", False, False, 3),
                ],
                columns=["id", "dataset", "model", "do_compile", "do_optimise", "repetition"],
            ),
        ),
    ],
)
def test_benchmark_set__filter_metrics(
    benchmark_set: BenchmarkSet,
    expected_metrics: pd.DataFrame,
) -> None:
    metrics = pd.DataFrame(
        [
            (0, "tiny_linear", "gpr", False, False, 0),
            (1, "tiny_sine", "gpr", False, False, 0),
            (2, "tiny_linear", "svgp", False, False, 0),
            (3, "tiny_linear", "gpr", True, False, 0),
            (4, "tiny_linear", "gpr", False, True, 0),
            (5, "tiny_linear", "gpr", False, False, 1),
            (6, "tiny_linear", "gpr", False, False, 2),
            (7, "tiny_linear", "gpr", False, False, 3),
        ],
        columns=["id", "dataset", "model", "do_compile", "do_optimise", "repetition"],
    )
    after = benchmark_set.filter_metrics(metrics).reset_index(drop=True)
    pd.testing.assert_frame_equal(expected_metrics, after)


def test_benchmark_suite__get_tasks() -> None:
    set1 = Mock(BenchmarkSet)
    set1.name = "set1"
    set1.get_tasks.return_value = [
        BenchmarkTask(
            dataset_name="ds",
            model_name="md",
            do_compile=False,
            do_optimise=False,
            do_predict=True,
            do_posterior=False,
            repetitions=2,
        ),
        BenchmarkTask(
            dataset_name="ds1",
            model_name="md1",
            do_compile=False,
            do_optimise=False,
            do_predict=True,
            do_posterior=True,
            repetitions=2,
        ),
    ]
    set2 = Mock(BenchmarkSet)
    set2.name = "set2"
    set2.get_tasks.return_value = [
        BenchmarkTask(
            dataset_name="ds",
            model_name="md",
            do_compile=False,
            do_optimise=False,
            do_predict=False,
            do_posterior=True,
            repetitions=2,
        ),
        BenchmarkTask(
            dataset_name="ds2",
            model_name="md1",
            do_compile=False,
            do_optimise=False,
            do_predict=True,
            do_posterior=True,
            repetitions=2,
        ),
    ]
    set3 = Mock(BenchmarkSet)
    set3.name = "set3"
    set3.get_tasks.return_value = [
        BenchmarkTask(
            dataset_name="ds",
            model_name="md",
            do_compile=False,
            do_optimise=False,
            do_predict=False,
            do_posterior=False,
            repetitions=5,
        ),
        BenchmarkTask(
            dataset_name="ds1",
            model_name="md3",
            do_compile=False,
            do_optimise=False,
            do_predict=True,
            do_posterior=True,
            repetitions=2,
        ),
    ]
    suite = BenchmarkSuite(
        name="test", description="Suite used in a test.", sets=[set1, set2, set3]
    )

    assert [
        BenchmarkTask(
            dataset_name="ds",
            model_name="md",
            do_compile=False,
            do_optimise=False,
            do_predict=True,
            do_posterior=True,
            repetitions=5,
        ),
        BenchmarkTask(
            dataset_name="ds1",
            model_name="md1",
            do_compile=False,
            do_optimise=False,
            do_predict=True,
            do_posterior=True,
            repetitions=2,
        ),
        BenchmarkTask(
            dataset_name="ds2",
            model_name="md1",
            do_compile=False,
            do_optimise=False,
            do_predict=True,
            do_posterior=True,
            repetitions=2,
        ),
        BenchmarkTask(
            dataset_name="ds1",
            model_name="md3",
            do_compile=False,
            do_optimise=False,
            do_predict=True,
            do_posterior=True,
            repetitions=2,
        ),
    ] == suite.get_tasks()


# tests/benchmark/test_grouping.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Collection, Sequence, Tuple
from unittest.mock import Mock

import pandas as pd
import pytest

from benchmark.grouping import GroupingKey as GK
from benchmark.grouping import GroupingSpec, _iter_by, group
from benchmark.metadata import BenchmarkMetadata
from benchmark.plotter_api import Plotter

METRICS_COLUMNS = ["run_id", "dataset", "model", "metric", "value"]
METRICS_DF = pd.DataFrame(
    [
        ("run_1", "ds1", "md1", "mt1", 1.0),
        ("run_2", "ds1", "md1", "mt1", 2.0),
        ("run_1", "ds2", "md1", "mt1", 3.0),
        ("run_1", "ds1", "md2", "mt1", 4.0),
        ("run_1", "ds1", "md1", "mt2", 5.0),
    ],
    columns=METRICS_COLUMNS,
)

PL1 = Mock(Plotter)
PL1.name = "pl1"
PL2 = Mock(Plotter)
PL2.name = "pl2"

BM1 = Mock(BenchmarkMetadata)
BM1.run_id = "run_1"
BM1.git_branch_name = "branch_1"
BM1.user = "user_1"
BM1.timestamp = "20220805_111111.111111"

BM2 = Mock(BenchmarkMetadata)
BM2.run_id = "run_2"
BM2.git_branch_name = "branch_2"
BM2.user = "user_2"
BM2.timestamp = "20220805_222222.222222"


def test_iter_by() -> None:
    assert [
        [],
        [GK.USER],
        [GK.RAM],
        [GK.HOSTNAME],
        [GK.USER, GK.RAM],
        [GK.TIMESTAMP],
        [GK.USER, GK.HOSTNAME],
        [GK.HOSTNAME, GK.RAM],
        [GK.USER, GK.TIMESTAMP],
        [GK.TIMESTAMP, GK.RAM],
        [GK.USER, GK.HOSTNAME, GK.RAM],
        [GK.HOSTNAME, GK.TIMESTAMP],
        [GK.USER, GK.TIMESTAMP, GK.RAM],
        [GK.USER, GK.HOSTNAME, GK.TIMESTAMP],
        [GK.HOSTNAME, GK.TIMESTAMP, GK.RAM],
        [GK.USER, GK.HOSTNAME, GK.TIMESTAMP, GK.RAM],
    ] == list(
        _iter_by(
            [
                GK.USER,
                GK.HOSTNAME,
                GK.TIMESTAMP,
                GK.RAM,
            ]
        )
    )


@pytest.mark.parametrize(
    "spec,expected",
    [
        # Group by no key:
        (
            GroupingSpec(by=[], minimise=False),
            [
                (
                    (),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md1", "mt1", 1.0),
                            ("run_2", "ds1", "md1", "mt1", 2.0),
                            ("run_1", "ds2", "md1", "mt1", 3.0),
                            ("run_1", "ds1", "md2", "mt1", 4.0),
                            ("run_1", "ds1", "md1", "mt2", 5.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL1, PL2],
                ),
            ],
        ),
        # Group by a metadata key:
        (
            GroupingSpec(by=[GK.GIT_BRANCH_NAME], minimise=False),
            [
                (
                    ("branch_1",),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md1", "mt1", 1.0),
                            ("run_1", "ds2", "md1", "mt1", 3.0),
                            ("run_1", "ds1", "md2", "mt1", 4.0),
                            ("run_1", "ds1", "md1", "mt2", 5.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL1, PL2],
                ),
                (
                    ("branch_2",),
                    pd.DataFrame(
                        [
                            ("run_2", "ds1", "md1", "mt1", 2.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL1, PL2],
                ),
            ],
        ),
        # Group by a metrics_df key:
        (
            GroupingSpec(by=[GK.DATASET], minimise=False),
            [
                (
                    ("ds1",),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md1", "mt1", 1.0),
                            ("run_2", "ds1", "md1", "mt1", 2.0),
                            ("run_1", "ds1", "md2", "mt1", 4.0),
                            ("run_1", "ds1", "md1", "mt2", 5.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL1, PL2],
                ),
                (
                    ("ds2",),
                    pd.DataFrame(
                        [
                            ("run_1", "ds2", "md1", "mt1", 3.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL1, PL2],
                ),
            ],
        ),
        # Group by a (the) plotter key:
        (
            GroupingSpec(by=[GK.PLOTTER], minimise=False),
            [
                (
                    (),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md1", "mt1", 1.0),
                            ("run_2", "ds1", "md1", "mt1", 2.0),
                            ("run_1", "ds2", "md1", "mt1", 3.0),
                            ("run_1", "ds1", "md2", "mt1", 4.0),
                            ("run_1", "ds1", "md1", "mt2", 5.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL1],
                ),
                (
                    (),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md1", "mt1", 1.0),
                            ("run_2", "ds1", "md1", "mt1", 2.0),
                            ("run_1", "ds2", "md1", "mt1", 3.0),
                            ("run_1", "ds1", "md2", "mt1", 4.0),
                            ("run_1", "ds1", "md1", "mt2", 5.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL2],
                ),
            ],
        ),
        # Group by all the keys:
        (
            GroupingSpec(
                by=[
                    GK.USER,
                    GK.GIT_BRANCH_NAME,
                    GK.TIMESTAMP,
                    GK.DATASET,
                    GK.MODEL,
                    GK.METRIC,
                    GK.PLOTTER,
                ],
                minimise=False,
            ),
            [
                (
                    ("user_1", "branch_1", "20220805_111111.111111", "ds1", "md1", "mt1"),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md1", "mt1", 1.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL1],
                ),
                (
                    ("user_1", "branch_1", "20220805_111111.111111", "ds1", "md1", "mt1"),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md1", "mt1", 1.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL2],
                ),
                (
                    ("user_1", "branch_1", "20220805_111111.111111", "ds1", "md1", "mt2"),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md1", "mt2", 5.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL1],
                ),
                (
                    ("user_1", "branch_1", "20220805_111111.111111", "ds1", "md1", "mt2"),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md1", "mt2", 5.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL2],
                ),
                (
                    ("user_1", "branch_1", "20220805_111111.111111", "ds1", "md2", "mt1"),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md2", "mt1", 4.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL1],
                ),
                (
                    ("user_1", "branch_1", "20220805_111111.111111", "ds1", "md2", "mt1"),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md2", "mt1", 4.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL2],
                ),
                (
                    ("user_1", "branch_1", "20220805_111111.111111", "ds2", "md1", "mt1"),
                    pd.DataFrame(
                        [
                            ("run_1", "ds2", "md1", "mt1", 3.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL1],
                ),
                (
                    ("user_1", "branch_1", "20220805_111111.111111", "ds2", "md1", "mt1"),
                    pd.DataFrame(
                        [
                            ("run_1", "ds2", "md1", "mt1", 3.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL2],
                ),
                (
                    ("user_2", "branch_2", "20220805_222222.222222", "ds1", "md1", "mt1"),
                    pd.DataFrame(
                        [
                            ("run_2", "ds1", "md1", "mt1", 2.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL1],
                ),
                (
                    ("user_2", "branch_2", "20220805_222222.222222", "ds1", "md1", "mt1"),
                    pd.DataFrame(
                        [
                            ("run_2", "ds1", "md1", "mt1", 2.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL2],
                ),
            ],
        ),
        # Group by all the keys, and minimise:
        (
            GroupingSpec(
                by=[
                    GK.USER,
                    GK.GIT_BRANCH_NAME,
                    GK.TIMESTAMP,
                    GK.DATASET,
                    GK.MODEL,
                    GK.METRIC,
                    GK.PLOTTER,
                ],
                minimise=True,
            ),
            [
                (
                    ("branch_1", "ds1", "md1", "mt1"),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md1", "mt1", 1.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL1],
                ),
                (
                    ("branch_1", "ds1", "md1", "mt1"),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md1", "mt1", 1.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL2],
                ),
                (
                    ("branch_1", "ds1", "md1", "mt2"),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md1", "mt2", 5.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL1],
                ),
                (
                    ("branch_1", "ds1", "md1", "mt2"),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md1", "mt2", 5.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL2],
                ),
                (
                    ("branch_1", "ds1", "md2", "mt1"),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md2", "mt1", 4.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL1],
                ),
                (
                    ("branch_1", "ds1", "md2", "mt1"),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md2", "mt1", 4.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL2],
                ),
                (
                    ("branch_1", "ds2", "md1", "mt1"),
                    pd.DataFrame(
                        [
                            ("run_1", "ds2", "md1", "mt1", 3.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL1],
                ),
                (
                    ("branch_1", "ds2", "md1", "mt1"),
                    pd.DataFrame(
                        [
                            ("run_1", "ds2", "md1", "mt1", 3.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL2],
                ),
                (
                    ("branch_2", "ds1", "md1", "mt1"),
                    pd.DataFrame(
                        [
                            ("run_2", "ds1", "md1", "mt1", 2.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL1],
                ),
                (
                    ("branch_2", "ds1", "md1", "mt1"),
                    pd.DataFrame(
                        [
                            ("run_2", "ds1", "md1", "mt1", 2.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL2],
                ),
            ],
        ),
    ],
)
def test_group(
    spec: GroupingSpec,
    expected: Sequence[Tuple[Tuple[str, ...], pd.DataFrame, Collection[Plotter]]],
) -> None:
    actual = group(METRICS_DF, [PL1, PL2], [BM1, BM2], spec)

    assert len(expected) == len(actual)
    for expected_gr, actual_gr in zip(expected, actual):
        expected_key, expected_df, expected_plotters = expected_gr
        actual_key, actual_df, actual_plotters = actual_gr

        assert expected_key == actual_key
        pd.testing.assert_frame_equal(expected_df, actual_df.reset_index(drop=True))
        assert expected_plotters == list(actual_plotters)


# tests/benchmark/test_metadata.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from benchmark.metadata import BenchmarkMetadata


def test_metadata_smoke() -> None:
    # Just test that we can create one without crashing:
    BenchmarkMetadata.create("test_suite")


# tests/benchmark/test_paths.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from pathlib import Path
from unittest.mock import Mock

import pytest

from benchmark.metadata import BenchmarkMetadata
from benchmark.paths import setup_dest


@pytest.mark.parametrize("exist", [False, True])
@pytest.mark.parametrize("no_subdir", [False, True])
def test_setup_dest(tmp_path: Path, exist: bool, no_subdir: bool) -> None:
    metadata = Mock(BenchmarkMetadata)
    metadata.run_id = "test_run_id"
    dest = tmp_path / "dest"
    expected = dest if no_subdir else (dest / "test_run_id")

    if exist:
        expected.mkdir(parents=True, exist_ok=True)
        with pytest.raises(FileExistsError):
            setup_dest(metadata, dest, no_subdir=no_subdir)
    else:
        expected.parent.mkdir(parents=True, exist_ok=True)
        assert expected == setup_dest(metadata, dest, no_subdir=no_subdir)


# tests/benchmark/test_plot.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from pathlib import Path

from benchmark.plot import _crawl_sources


def test_crawl_sources(tmp_path: Path) -> None:
    def setup_source(source_path: Path, test_id: int) -> None:
        source_path.mkdir(parents=True)
        (source_path / "metrics.csv").write_text(
            f"""test_col_1, test_col_2
{test_id}1, {test_id}2
{test_id}3, {test_id}4
"""
        )
        (source_path / "metadata.json").write_text(
            f"""{{
    "suite_name": "test",
    "argv": [],
    "user": "test_user",
    "hostname": "test_host",
    "timestamp": "20220808_085348.161698",
    "py_ver": "3.10.0",
    "tf_ver": "2.8.0",
    "np_ver": "1.20.0",
    "ram": 12,
    "cpu_name": "Test CPU",
    "cpu_count": 8,
    "cpu_frequency": 14,
    "gpu_names": ["Test GPU"],
    "git_branch_name": "develop",
    "git_commit": "1234",
    "run_id": "test_id_{test_id}"
}}
"""
        )

    source1 = tmp_path / "source1"
    source2 = tmp_path / "source2"
    source21 = source2 / "source21"
    source22 = source2 / "source22"

    setup_source(source1, 1)
    setup_source(source21, 21)
    setup_source(source22, 22)

    cycle = source22 / "cycle"
    cycle.symlink_to("../..")

    metrics_df, metadata = _crawl_sources([source1, source2])
    assert len(metrics_df) == 6
    assert {"test_id_1", "test_id_21", "test_id_22"} == set(m.run_id for m in metadata)


# tests/benchmark/test_registry.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import AbstractSet

import pytest

from benchmark.registry import Named, Registry, Tagged, TaggedRegistry
from benchmark.tag import Tag


def test_registry() -> None:
    class TestNamed(Named):
        def __init__(self, name: str) -> None:
            self.name = name

    o1 = TestNamed("o1")
    o2 = TestNamed("o2")

    reg: Registry[TestNamed] = Registry()

    assert () == reg.names()
    assert () == reg.all()

    assert o1 == reg.add(o1)
    assert o2 == reg.add(o2)

    assert ("o1", "o2") == reg.names()
    assert (o1, o2) == reg.all()
    assert o1 == reg.get("o1")
    assert o2 == reg.get("o2")

    o1_2 = TestNamed("o1")
    with pytest.raises(AssertionError):
        reg.add(o1_2)


def test_tagged_registry() -> None:
    class TestTag(Tag["TestTag"]):
        pass

    t1 = TestTag("t1")
    t2 = TestTag("t2")

    class TestTagged(Tagged[TestTag]):
        def __init__(self, name: str, tags: AbstractSet[TestTag]) -> None:
            self.name = name
            self.tags = tags

    o1 = TestTagged("o1", {t1})
    o2 = TestTagged("o2", {t1, t2})
    o3 = TestTagged("o3", {t2})

    reg: TaggedRegistry[TestTagged, TestTag] = TaggedRegistry()

    assert () == reg.names()
    assert () == reg.all()
    assert () == reg.where(t1)
    assert () == reg.where(t1 & ~t2)

    assert o1 == reg.add(o1)
    assert o2 == reg.add(o2)
    assert o3 == reg.add(o3)

    assert ("o1", "o2", "o3") == reg.names()
    assert (o1, o2, o3) == reg.all()
    assert (o1, o2) == reg.where(t1)
    assert (o1,) == reg.where(t1 & ~t2)
    assert o1 == reg.get("o1")
    assert o2 == reg.get("o2")
    assert o3 == reg.get("o3")

    o1_2 = TestTagged("o1", {t2})
    with pytest.raises(AssertionError):
        reg.add(o1_2)


# tests/benchmark/test_tag.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import AbstractSet

import pytest

from benchmark.tag import NO_REQ, Tag, TagReq


class TestTag(Tag["TestTag"]):
    pass


t1 = TestTag("t1")
t2 = TestTag("t2")
t3 = TestTag("t3")


@pytest.mark.parametrize(
    "req,tags,expected_satisfied,expected_repr",
    [
        # Test single operations:
        (NO_REQ, {}, True, "NO_REQ"),
        (t1, {t1, t2}, True, "t1"),
        (t3, {t1, t2}, False, "t3"),
        (t1 & t2, {t1, t2}, True, "t1 & t2"),
        (t1 & t3, {t1, t2}, False, "t1 & t3"),
        (t1 | t2, {t2}, True, "t1 | t2"),
        (t1 | t3, {t2}, False, "t1 | t3"),
        (~t1, {t1, t2}, False, "~t1"),
        (~t3, {t1, t2}, True, "~t3"),
        # Test composition:
        (t1 | (t2 & ~t3), {}, False, "t1 | (t2 & (~t3))"),
        (t1 | (t2 & ~t3), {t3}, False, "t1 | (t2 & (~t3))"),
        (t1 | (t2 & ~t3), {t2}, True, "t1 | (t2 & (~t3))"),
        (t1 | (t2 & ~t3), {t2, t3}, False, "t1 | (t2 & (~t3))"),
        (t1 | (t2 & ~t3), {t1}, True, "t1 | (t2 & (~t3))"),
        (t1 | (t2 & ~t3), {t1, t3}, True, "t1 | (t2 & (~t3))"),
        (t1 | (t2 & ~t3), {t1, t2}, True, "t1 | (t2 & (~t3))"),
        (t1 | (t2 & ~t3), {t1, t2, t3}, True, "t1 | (t2 & (~t3))"),
    ],
)
def test_tag_req(
    req: TagReq[TestTag], tags: AbstractSet[TestTag], expected_satisfied: bool, expected_repr: str
) -> None:
    assert expected_satisfied == req.satisfied(tags)
    assert expected_repr == repr(req)


# tests/examples/test_check_shapes_examples.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Code examples used in `check_shapes` docstrings.
"""


def test_example__basic() -> None:
    # This test / example is placed *above* the imports, to ensure that the example has the
    # necessary imports.

    # pylint: disable=import-outside-toplevel, redefined-outer-name, reimported

    # [basic]

    import tensorflow as tf

    from gpflow.experimental.check_shapes import check_shapes

    @tf.function
    @check_shapes(
        "features: [batch..., n_features]",
        "weights: [n_features]",
        "return: [batch...]",
    )
    def linear_model(features: tf.Tensor, weights: tf.Tensor) -> tf.Tensor:
        return tf.einsum("...i,i -> ...", features, weights)

    # [basic]

    w = tf.ones((3,))
    for batch_shape in [(), (2,), (2, 4)]:
        f = tf.ones(batch_shape + (3,))
        linear_model(f, w)


from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Iterable, Mapping, Optional, Sequence, Tuple
from unittest.mock import MagicMock

import numpy as np
import pytest

from gpflow.base import AnyNDArray
from gpflow.experimental.check_shapes import (
    DocstringFormat,
    ErrorContext,
    Shape,
    ShapeChecker,
    ShapeCheckingState,
    check_shape,
    check_shapes,
    disable_check_shapes,
    get_check_shapes,
    get_enable_check_shapes,
    get_enable_function_call_precompute,
    get_rewrite_docstrings,
    get_shape,
    inherit_check_shapes,
    register_get_shape,
    set_enable_check_shapes,
    set_enable_function_call_precompute,
    set_rewrite_docstrings,
)


def test_example__disable__manual() -> None:
    old_value = get_enable_check_shapes()

    try:

        # [disable__manual]

        set_enable_check_shapes(ShapeCheckingState.DISABLED)

        # [disable__manual]

    finally:
        set_enable_check_shapes(old_value)


def test_example__disable__context_manager() -> None:
    def performance_sensitive_function() -> None:
        pass

    # [disable__context_manager]

    with disable_check_shapes():
        performance_sensitive_function()

    # [disable__context_manager]


def test_example__pytest_fixture() -> None:
    # pylint: disable=unused-variable
    # [pytest_fixture]

    @pytest.fixture(autouse=True)
    def enable_shape_checks() -> Iterable[None]:
        old_enable = get_enable_check_shapes()
        old_rewrite_docstrings = get_rewrite_docstrings()
        old_function_call_precompute = get_enable_function_call_precompute()
        set_enable_check_shapes(ShapeCheckingState.ENABLED)
        set_rewrite_docstrings(DocstringFormat.SPHINX)
        set_enable_function_call_precompute(True)
        yield
        set_enable_function_call_precompute(old_function_call_precompute)
        set_rewrite_docstrings(old_rewrite_docstrings)
        set_enable_check_shapes(old_enable)

    # [pytest_fixture]


def test_example__argument_ref_attribute() -> None:
    # [argument_ref_attribute]

    @dataclass
    class Statistics:
        mean: AnyNDArray
        std: AnyNDArray

    @check_shapes(
        "data: [n_rows, n_columns]",
        "return.mean: [n_columns]",
        "return.std: [n_columns]",
    )
    def compute_statistics(data: AnyNDArray) -> Statistics:
        return Statistics(np.mean(data, axis=0), np.std(data, axis=0))

    # [argument_ref_attribute]

    compute_statistics(np.ones((4, 3)))


def test_example__argument_ref_index() -> None:
    # [argument_ref_index]

    @check_shapes(
        "data: [n_rows, n_columns]",
        "return[0]: [n_columns]",
        "return[1]: [n_columns]",
    )
    def compute_mean_and_std(data: AnyNDArray) -> Tuple[AnyNDArray, AnyNDArray]:
        return np.mean(data, axis=0), np.std(data, axis=0)

    # [argument_ref_index]

    compute_mean_and_std(np.ones((4, 3)))


def test_example__argument_ref_all() -> None:
    # [argument_ref_all]

    @check_shapes(
        "data[all]: [., n_columns]",
        "return: [., n_columns]",
    )
    def concat_rows(data: Sequence[AnyNDArray]) -> AnyNDArray:
        return np.concatenate(data, axis=0)

    concat_rows(
        [
            np.ones((1, 3)),
            np.ones((4, 3)),
        ]
    )

    # [argument_ref_all]


def test_example__argument_ref_keys() -> None:
    # [argument_ref_keys]

    @check_shapes(
        "data.keys(): [.]",
        "return: []",
    )
    def sum_key_lengths(data: Mapping[Tuple[int, ...], str]) -> int:
        return sum(len(k) for k in data)

    sum_key_lengths(
        {
            (3,): "foo",
            (1, 2): "bar",
        }
    )

    # [argument_ref_keys]


def test_example__argument_ref_values() -> None:
    # [argument_ref_values]

    @check_shapes(
        "data.values(): [., n_columns]",
        "return: [., n_columns]",
    )
    def concat_rows(data: Mapping[str, AnyNDArray]) -> AnyNDArray:
        return np.concatenate(list(data.values()), axis=0)

    concat_rows(
        {
            "foo": np.ones((1, 3)),
            "bar": np.ones((4, 3)),
        }
    )

    # [argument_ref_values]


def test_example__argument_ref_optional() -> None:
    # [argument_ref_optional]

    @check_shapes(
        "x1: [n_rows_1, n_inputs]",
        "x2: [n_rows_2, n_inputs]",
        "return: [n_rows_1, n_rows_2]",
    )
    def squared_exponential_kernel(
        variance: float, x1: AnyNDArray, x2: Optional[AnyNDArray] = None
    ) -> AnyNDArray:
        if x2 is None:
            x2 = x1
        cov: AnyNDArray = variance * np.exp(
            -0.5 * np.sum((x1[:, None, :] - x2[None, :, :]) ** 2, axis=2)
        )
        return cov

    squared_exponential_kernel(1.0, np.ones((3, 2)), np.ones((4, 2)))
    squared_exponential_kernel(3.2, np.ones((3, 2)))

    # [argument_ref_optional]


def test_example__dimension_spec_constant() -> None:
    # [dimension_spec_constant]

    @check_shapes(
        "v1: [2]",
        "v2: [2]",
    )
    def vector_2d_distance(v1: AnyNDArray, v2: AnyNDArray) -> float:
        return float(np.sqrt(np.sum((v1 - v2) ** 2)))

    # [dimension_spec_constant]

    vector_2d_distance(np.ones((2,)), np.ones((2,)))


def test_example__dimension_spec_variable() -> None:
    # [dimension_spec_variable]

    @check_shapes(
        "v1: [d]",
        "v2: [d]",
    )
    def vector_distance(v1: AnyNDArray, v2: AnyNDArray) -> float:
        return float(np.sqrt(np.sum((v1 - v2) ** 2)))

    # [dimension_spec_variable]

    vector_distance(np.ones((3,)), np.ones((3,)))


def test_example__dimension_spec_anonymous__dot() -> None:
    # [dimension_spec_anonymous__dot]

    @check_shapes(
        "v: [.]",
    )
    def vector_length(v: AnyNDArray) -> float:
        return float(np.sqrt(np.sum(v ** 2)))

    # [dimension_spec_anonymous__dot]

    vector_length(np.ones((3,)))


def test_example__dimension_spec_anonymous__none() -> None:
    # [dimension_spec_anonymous__none]

    @check_shapes(
        "v: [None]",
    )
    def vector_length(v: AnyNDArray) -> float:
        return float(np.sqrt(np.sum(v ** 2)))

    # [dimension_spec_anonymous__none]

    vector_length(np.ones((3,)))


def test_example__dimension_spec_variable_rank__star() -> None:
    # [dimension_spec_variable_rank__star]

    @check_shapes(
        "x: [*batch, n_columns]",
        "return: [*batch]",
    )
    def batch_mean(x: AnyNDArray) -> AnyNDArray:
        mean: AnyNDArray = np.mean(x, axis=-1)
        return mean

    # [dimension_spec_variable_rank__star]

    for batch_shape in [(), (2,), (2, 4)]:
        x = np.ones(batch_shape + (3,))
        batch_mean(x)


def test_example__dimension_spec_variable_rank__ellipsis() -> None:
    # [dimension_spec_variable_rank__ellipsis]

    @check_shapes(
        "x: [batch..., n_columns]",
        "return: [batch...]",
    )
    def batch_mean(x: AnyNDArray) -> AnyNDArray:
        mean: AnyNDArray = np.mean(x, axis=-1)
        return mean

    # [dimension_spec_variable_rank__ellipsis]

    for batch_shape in [(), (2,), (2, 4)]:
        x = np.ones(batch_shape + (3,))
        batch_mean(x)


def test_example__dimension_spec_anonymous_variable_rank__star() -> None:
    # [dimension_spec_anonymous_variable_rank__star]

    @check_shapes(
        "x: [*]",
    )
    def rank(x: AnyNDArray) -> int:
        return len(x.shape)

    # [dimension_spec_anonymous_variable_rank__star]

    for batch_shape in [(), (2,), (2, 4)]:
        x = np.ones(batch_shape)
        rank(x)


def test_example__dimension_spec_anonymous_variable_rank__ellipsis() -> None:
    # [dimension_spec_anonymous_variable_rank__ellipsis]

    @check_shapes(
        "x: [...]",
    )
    def rank(x: AnyNDArray) -> int:
        return len(x.shape)

    # [dimension_spec_anonymous_variable_rank__ellipsis]

    for batch_shape in [(), (2,), (2, 4)]:
        x = np.ones(batch_shape)
        rank(x)


def test_example__dimension_spec__scalar() -> None:
    # [dimension_spec__scalar]

    @check_shapes(
        "x: [...]",
        "return: []",
    )
    def mean(x: AnyNDArray) -> AnyNDArray:
        mean: AnyNDArray = np.sum(x) / x.size
        return mean

    # [dimension_spec__scalar]

    for batch_shape in [(), (2,), (2, 4)]:
        x = np.ones(batch_shape)
        mean(x)


def test_example__dimension_spec_broadcast() -> None:

    # [dimension_spec_broadcast]

    @check_shapes(
        "a: [broadcast batch...]",
        "b: [broadcast batch...]",
        "return: [batch...]",
    )
    def add(a: AnyNDArray, b: AnyNDArray) -> AnyNDArray:
        return a + b

    # [dimension_spec_broadcast]

    add(np.ones((3, 1)), np.ones((1, 4)))


def test_example__bool_spec_argument_ref() -> None:
    # pylint: disable=unused-argument

    # [bool_spec_argument_ref]

    @check_shapes(
        "a: [broadcast batch...] if check_a",
        "b: [broadcast batch...] if check_b",
        "return: [batch...]",
    )
    def add(a: AnyNDArray, b: AnyNDArray, check_a: bool = True, check_b: bool = True) -> AnyNDArray:
        return a + b

    add(np.ones((3, 1)), np.ones((1, 4)), check_b=False)

    # [bool_spec_argument_ref]


def test_example__bool_spec_argument_ref_is_none() -> None:
    # pylint: disable=unused-argument

    # [bool_spec_argument_ref_is_none]

    @check_shapes(
        "a: [n_a]",
        "b: [n_b]",
        "return: [n_a, n_a] if b is None",
        "return: [n_a, n_b] if b is not None",
    )
    def square(a: AnyNDArray, b: Optional[AnyNDArray] = None) -> AnyNDArray:
        if b is None:
            b = a
        result: AnyNDArray = a[:, None] * b[None, :]
        return result

    square(np.ones((3,)))
    square(np.ones((3,)), np.ones((4,)))

    # [bool_spec_argument_ref_is_none]


def test_example__bool_spec_or() -> None:
    # pylint: disable=unused-argument

    # [bool_spec_or]

    @check_shapes(
        "a: [broadcast batch...] if check_all or check_a",
        "b: [broadcast batch...] if check_all or check_b",
        "return: [batch...]",
    )
    def add(
        a: AnyNDArray,
        b: AnyNDArray,
        check_all: bool = False,
        check_a: bool = True,
        check_b: bool = True,
    ) -> AnyNDArray:
        return a + b

    add(np.ones((3, 1)), np.ones((1, 4)), check_b=False)

    # [bool_spec_or]


def test_example__bool_spec_and() -> None:
    # pylint: disable=unused-argument

    # [bool_spec_and]

    @check_shapes(
        "a: [broadcast batch...] if enable_checks and check_a",
        "b: [broadcast batch...] if enable_checks and check_b",
        "return: [batch...]",
    )
    def add(
        a: AnyNDArray,
        b: AnyNDArray,
        enable_checks: bool = True,
        check_a: bool = True,
        check_b: bool = True,
    ) -> AnyNDArray:
        return a + b

    add(np.ones((3, 1)), np.ones((1, 4)), check_b=False)

    # [bool_spec_and]


def test_example__bool_spec_not() -> None:
    # pylint: disable=unused-argument

    # [bool_spec_not]

    @check_shapes(
        "a: [broadcast batch...] if not disable_checks",
        "b: [broadcast batch...] if not disable_checks",
        "return: [batch...]",
    )
    def add(a: AnyNDArray, b: AnyNDArray, disable_checks: bool = False) -> AnyNDArray:
        return a + b

    add(np.ones((3, 1)), np.ones((1, 4)))

    # [bool_spec_not]


def test_example__bool_spec__composition() -> None:
    # pylint: disable=unused-argument

    # [bool_spec__composition]

    @check_shapes(
        "a: [j] if a_vector",
        "a: [i, j] if (not a_vector)",
        "b: [j] if b_vector",
        "b: [j, k] if (not b_vector)",
        "return: [1, 1] if a_vector and b_vector",
        "return: [1, k] if a_vector and (not b_vector)",
        "return: [i, 1] if (not a_vector) and b_vector",
        "return: [i, k] if (not a_vector) and (not b_vector)",
    )
    def multiply(a: AnyNDArray, b: AnyNDArray, a_vector: bool, b_vector: bool) -> AnyNDArray:
        if a_vector:
            a = a[None, :]
        if b_vector:
            b = b[:, None]

        return a @ b

    multiply(np.ones((4,)), np.ones((4, 5)), a_vector=True, b_vector=False)

    # [bool_spec__composition]

    multiply(np.ones((4,)), np.ones((4,)), a_vector=True, b_vector=True)
    multiply(np.ones((3, 4)), np.ones((4,)), a_vector=False, b_vector=True)
    multiply(np.ones((3, 4)), np.ones((4, 5)), a_vector=False, b_vector=False)


def test_example__note_spec__global() -> None:
    # [note_spec__global]

    @check_shapes(
        "features: [batch..., n_features]",
        "# linear_model currently only supports a single output.",
        "weights: [n_features]",
        "return: [batch...]",
    )
    def linear_model(features: AnyNDArray, weights: AnyNDArray) -> AnyNDArray:
        prediction: AnyNDArray = np.einsum("...i,i -> ...", features, weights)
        return prediction

    # [note_spec__global]

    linear_model(np.ones((4, 3)), np.ones((3,)))


def test_example__note_spec__local() -> None:
    # [note_spec__local]

    @check_shapes(
        "features: [batch..., n_features]",
        "weights: [n_features] # linear_model currently only supports a single output.",
        "return: [batch...]",
    )
    def linear_model(features: AnyNDArray, weights: AnyNDArray) -> AnyNDArray:
        prediction: AnyNDArray = np.einsum("...i,i -> ...", features, weights)
        return prediction

    # [note_spec__local]

    linear_model(np.ones((4, 3)), np.ones((3,)))


def test_example__reuse__inherit_check_shapes() -> None:
    # [reuse__inherit_check_shapes]

    class Model(ABC):
        @abstractmethod
        @check_shapes(
            "features: [batch..., n_features]",
            "return: [batch...]",
        )
        def predict(self, features: AnyNDArray) -> AnyNDArray:
            pass

    class LinearModel(Model):
        @check_shapes(
            "weights: [n_features]",
        )
        def __init__(self, weights: AnyNDArray) -> None:
            self._weights = weights

        @inherit_check_shapes
        def predict(self, features: AnyNDArray) -> AnyNDArray:
            prediction: AnyNDArray = np.einsum("...i,i -> ...", features, self._weights)
            return prediction

    # [reuse__inherit_check_shapes]

    model = LinearModel(np.ones((3,)))
    model.predict(np.ones((10, 3)))


def test_example__reuse__functional() -> None:
    # [reuse__functional]

    check_metric_shapes = check_shapes(
        "actual: [n_rows, n_labels]",
        "predicted: [n_rows, n_labels]",
        "return: []",
    )

    @check_metric_shapes
    def rmse(actual: AnyNDArray, predicted: AnyNDArray) -> float:
        return float(np.mean(np.sqrt(np.mean((predicted - actual) ** 2, axis=-1))))

    @check_metric_shapes
    def mape(actual: AnyNDArray, predicted: AnyNDArray) -> float:
        return float(np.mean(np.abs((predicted - actual) / actual)))

    # [reuse__functional]

    actual = np.ones((10, 3))
    predicted = np.ones((10, 3))
    rmse(actual, predicted)
    mape(actual, predicted)


def test_example__reuse__get_check_shapes() -> None:
    # [reuse__get_check_shapes]

    class Model(ABC):
        @abstractmethod
        @check_shapes(
            "features: [batch..., n_features]",
            "return: [batch...]",
        )
        def predict(self, features: AnyNDArray) -> AnyNDArray:
            pass

    @check_shapes(
        "test_features: [n_rows, n_features]",
        "test_labels: [n_rows]",
    )
    def evaluate_model(model: Model, test_features: AnyNDArray, test_labels: AnyNDArray) -> float:
        prediction = model.predict(test_features)
        return float(np.mean(np.sqrt(np.mean((prediction - test_labels) ** 2, axis=-1))))

    def test_evaluate_model() -> None:
        fake_features = np.ones((10, 3))
        fake_labels = np.ones((10,))
        fake_predictions = np.ones((10,))

        @get_check_shapes(Model.predict)
        def fake_predict(features: AnyNDArray) -> AnyNDArray:
            assert features is fake_features
            return fake_predictions

        fake_model = MagicMock(spec=Model, predict=fake_predict)

        assert pytest.approx(0.0) == evaluate_model(fake_model, fake_features, fake_labels)

    # [reuse__get_check_shapes]

    test_evaluate_model()


def test_example__intermediate_results() -> None:
    # [intermediate_results]

    @check_shapes(
        "weights: [n_features, n_labels]",
        "test_features: [n_rows, n_features]",
        "test_labels: [n_rows, n_labels]",
        "return: []",
    )
    def loss(weights: AnyNDArray, test_features: AnyNDArray, test_labels: AnyNDArray) -> AnyNDArray:
        prediction = check_shape(test_features @ weights, "[n_rows, n_labels]")
        error: AnyNDArray = check_shape(prediction - test_labels, "[n_rows, n_labels]")
        square_error = check_shape(error ** 2, "[n_rows, n_labels]")
        mean_square_error = check_shape(np.mean(square_error, axis=-1), "[n_rows]")
        root_mean_square_error = check_shape(np.sqrt(mean_square_error), "[n_rows]")
        loss: AnyNDArray = np.mean(root_mean_square_error)
        return loss

    # [intermediate_results]

    loss(
        np.ones((5, 3)),
        np.ones((10, 5)),
        np.ones((10, 3)),
    )


def test_example__shape_checker__raw() -> None:
    # [shape_checker__raw]

    def linear_model(features: AnyNDArray, weights: AnyNDArray) -> AnyNDArray:
        checker = ShapeChecker()
        checker.check_shape(features, "[batch..., n_features]")
        checker.check_shape(weights, "[n_features]")
        prediction: AnyNDArray = checker.check_shape(
            np.einsum("...i,i -> ...", features, weights), "[batch...]"
        )
        return prediction

    # [shape_checker__raw]

    linear_model(np.ones((4, 3)), np.ones((3,)))


def test_example__disable_function_call_precompute() -> None:
    def buggy_function() -> None:
        pass

    old_value = get_enable_function_call_precompute()

    try:

        # [disable_function_call_precompute]

        set_enable_function_call_precompute(True)

        buggy_function()

        # [disable_function_call_precompute]

    finally:
        set_enable_function_call_precompute(old_value)


def test_example__doc_rewrite() -> None:

    # [doc_rewrite__definition]

    @check_shapes(
        "features: [batch..., n_features]",
        "weights: [n_features]",
        "return: [batch...]",
    )
    def linear_model(features: AnyNDArray, weights: AnyNDArray) -> AnyNDArray:
        """
        Computes a prediction from a linear model.

        :param features: Data to make predictions from.
        :param weights: Model weights.
        :returns: Model predictions.
        """
        prediction: AnyNDArray = np.einsum("...i,i -> ...", features, weights)
        return prediction

    # [doc_rewrite__definition]

    expected_doc = (
        # [doc_rewrite__rewritten]
        """
        Computes a prediction from a linear model.

        :param features:
            * **features** has shape [*batch*..., *n_features*].

            Data to make predictions from.
        :param weights:
            * **weights** has shape [*n_features*].

            Model weights.
        :returns:
            * **return** has shape [*batch*...].

            Model predictions.
        """
        # [doc_rewrite__rewritten]
    )

    assert expected_doc == linear_model.__doc__


def test_example__doc_rewrite__disable() -> None:

    old_rewrite_docstrings = get_rewrite_docstrings()

    try:
        # [doc_rewrite__disable]
        set_rewrite_docstrings(None)
        # [doc_rewrite__disable]
    finally:
        set_rewrite_docstrings(old_rewrite_docstrings)


def test_example__custom_type() -> None:
    # pylint: disable=protected-access,unused_variable

    # [custom_type]

    class LinearModel:
        @check_shapes(
            "weights: [n_features]",
        )
        def __init__(self, weights: AnyNDArray) -> None:
            self._weights = weights

        @check_shapes(
            "self: [n_features]",
            "features: [batch..., n_features]",
            "return: [batch...]",
        )
        def predict(self, features: AnyNDArray) -> AnyNDArray:
            prediction: AnyNDArray = np.einsum("...i,i -> ...", features, self._weights)
            return prediction

    @register_get_shape(LinearModel)
    def get_linear_model_shape(model: LinearModel, context: ErrorContext) -> Shape:
        shape: Shape = model._weights.shape
        return shape

    @check_shapes(
        "model: [n_features]",
        "test_features: [n_rows, n_features]",
        "test_labels: [n_rows]",
        "return: []",
    )
    def loss(model: LinearModel, test_features: AnyNDArray, test_labels: AnyNDArray) -> float:
        prediction = model.predict(test_features)
        return float(np.mean(np.sqrt(np.mean((prediction - test_labels) ** 2, axis=-1))))

    # [custom_type]

    model = LinearModel(np.ones((3,)))
    loss(model, np.ones((10, 3)), np.ones((10,)))


# tests/gpflow/conftest.py

#  Copyright 2021 The GPflow Contributors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#  http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

from typing import cast

import numpy as np
import pytest
from _pytest.fixtures import SubRequest


@pytest.fixture(name="full_cov", params=[True, False])
def _full_cov_fixture(request: SubRequest) -> bool:
    return cast(bool, request.param)


@pytest.fixture(name="full_output_cov", params=[True, False])
def _full_output_cov_fixture(request: SubRequest) -> bool:
    return cast(bool, request.param)


@pytest.fixture(name="whiten", params=[True, False])
def _whiten_fixture(request: SubRequest) -> bool:
    return cast(bool, request.param)


@pytest.fixture(name="rng")
def _rng_fixture() -> np.random.Generator:
    return np.random.default_rng(20220523)


# tests/gpflow/test_all.py

# Copyright 2022 the GPflow authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import ast
import importlib
from pathlib import Path
from types import ModuleType
from typing import Any, List, Optional, Sequence

import pytest

import gpflow


def is_dunder(name: str) -> bool:
    return name.startswith("__") and name.endswith("__")


def get_module_path(module: ModuleType) -> Path:
    path_str = module.__file__
    assert path_str is not None
    return Path(path_str)


def is_package(module: ModuleType) -> bool:
    return get_module_path(module).name == "__init__.py"


def find_modules(root: ModuleType) -> Sequence[ModuleType]:
    root_path = get_module_path(root).parent
    parent = root_path.parent
    result: List[ModuleType] = []
    for path in root_path.glob("**/*.py"):
        relative_path = path.relative_to(parent)
        if relative_path.name == "__init__.py":
            module_name__slashes = str(relative_path.parent)
        else:
            module_name__slashes = str(relative_path)[:-3]
        module_name = module_name__slashes.replace("/", ".")
        result.append(importlib.import_module(module_name))
    return result


_MODULES = find_modules(gpflow)
_PACKAGES = [m for m in _MODULES if is_package(m)]
_MODULES_WITH_ALL = [m for m in _MODULES if hasattr(m, "__all__")]


@pytest.mark.parametrize("package", _PACKAGES)
def test_all_present_and_up_to_date(package: ModuleType) -> None:
    imported_sorted = sorted(attr for attr in dir(package) if not is_dunder(attr))
    all_list = getattr(package, "__all__", None)
    assert all_list is not None, f"Package {package} is missing an explicit __all__."
    all_sorted = sorted(attr for attr in all_list if not is_dunder(attr))
    assert imported_sorted == all_sorted, (
        f"{package}.__all__ is outdated."
        f" Imported values are {imported_sorted}, but exported values are {all_list}."
    )


@pytest.mark.parametrize("module", _MODULES_WITH_ALL)
def test_all_sorted(module: ModuleType) -> None:
    all_list = getattr(module, "__all__", None)
    assert all_list is not None  # Hint for mypy.
    all_sorted = sorted(all_list)
    assert (
        all_sorted == all_list
    ), f"{module}.__all__ is not sorted. Expected {all_sorted}, found {all_list}."


@pytest.mark.parametrize("module", _MODULES_WITH_ALL)
def test_all_static(module: ModuleType) -> None:
    module_path = get_module_path(module)
    tree = ast.parse(module_path.read_text(), str(module_path))
    assert isinstance(tree, ast.Module)

    def is_all_assignment(t: Any) -> bool:
        return (
            isinstance(t, ast.Assign)
            and len(t.targets) == 1
            and isinstance(t.targets[0], ast.Name)
            and t.targets[0].id == "__all__"
        )

    all_assignment: Optional[ast.Assign] = None
    for t in tree.body:
        if is_all_assignment(t):
            all_assignment = t  # type: ignore[assignment]
    assert all_assignment is not None, f"No `__all__` found in module {module}."

    error_message = "__all__ must be a static list of constant strings. Some tools expect this."

    all_value = all_assignment.value
    assert isinstance(all_value, ast.List), error_message

    assert all(
        # Actual type depends on Python version:
        (isinstance(t, ast.Constant) or isinstance(t, ast.Str))
        for t in all_value.elts
    ), error_message


# tests/gpflow/test_base.py

# Copyright 2020 the GPflow authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import tempfile
from typing import Sequence

import numpy as np
import pytest
import tensorflow as tf
import tensorflow_probability as tfp

import gpflow
from gpflow.base import AnyNDArray, PriorOn, TensorData
from gpflow.config import default_float
from gpflow.utilities import positive, triangular


def test_parameter_assign_validation() -> None:
    with pytest.raises(tf.errors.InvalidArgumentError):
        param = gpflow.Parameter(0.0, transform=positive())

    param = gpflow.Parameter(0.1, transform=positive())
    param.assign(0.2)
    with pytest.raises(tf.errors.InvalidArgumentError):
        param.assign(0.0)


def test_cast_to_dtype_precision_issue() -> None:
    """
    TensorFlow's tf.cast(value, dtype) implicitly does a tf.convert_to_tensor(value)
    *before* the cast when the value is not a tensor already. When value is a python float,
    this results in the following behaviour:

    >>> tf.cast(0.2, tf.float64)
    <tf.Tensor: id=37, shape=(), dtype=float64, numpy=0.20000000298023224>

    instead of the expected expansion of 0.2 to float64 precision that you get when
    passing in an object that already carries dtype information, such as a numpy array
    (which has float64 precision by default):

    >>> tf.cast(np.array(0.2), tf.float64)
    <tf.Tensor: id=40, shape=(), dtype=float64, numpy=0.2>

    This affected *all* gpflow.Parameter objects, resulting in numerical discrepancies
    between GPflow 1 and 2, due to the pass through _cast_to_dtype, which is now fixed.
    This is the corresponding regression test.
    """
    p = gpflow.Parameter(0.2, dtype=np.float64)
    actual_value = p.numpy()
    assert actual_value.dtype == np.float64
    expected_value = np.float64(0.2)
    assert actual_value == expected_value


def test_parameter_saved() -> None:
    dtype = tf.float64

    class Model(tf.Module):
        def __init__(self) -> None:
            self.p = gpflow.Parameter(0.1, dtype=dtype, transform=gpflow.utilities.positive())

        @tf.function(input_signature=[tf.TensorSpec([], dtype=dtype)])
        def exec(self, x: tf.Tensor) -> tf.Tensor:
            return tf.square(x * self.p)

    m0 = Model()
    x = tf.convert_to_tensor(2.0, dtype=dtype)
    expected = m0.exec(x)
    with tempfile.TemporaryDirectory() as dirname:
        tf.saved_model.save(m0, dirname)
        m1 = tf.saved_model.load(dirname)
        actual = m1.exec(x)
        np.testing.assert_equal(actual, expected)


@pytest.mark.parametrize("value", [0.0, [1.2, 1.1]])
def test_construct_parameter_from_existing_parameter_check_value(value: TensorData) -> None:
    initial_parameter = gpflow.Parameter(value)
    new_parameter = gpflow.Parameter(initial_parameter)

    np.testing.assert_equal(new_parameter.numpy(), value)


@pytest.mark.parametrize("value", [0.0, [1.2, 1.1]])
def test_construct_parameter_from_existing_parameter_override_value(value: TensorData) -> None:
    initial_parameter = gpflow.Parameter(value)
    new_parameter = gpflow.Parameter(initial_parameter + 1.0)

    np.testing.assert_equal(new_parameter.numpy(), np.array(value) + 1.0)


def test_construct_parameter_from_existing_parameter_check_transform() -> None:
    transform = tfp.bijectors.Sigmoid(
        tf.constant(0.0, dtype=tf.float64), tf.constant(2.0, dtype=tf.float64)
    )
    initial_parameter = gpflow.Parameter([1.2, 1.1], transform=transform)
    new_parameter = gpflow.Parameter(initial_parameter)

    assert new_parameter.transform == transform


def test_construct_parameter_from_existing_parameter_override_transform() -> None:
    initial_parameter = gpflow.Parameter([1.2, 1.1])

    transform = tfp.bijectors.Sigmoid(
        tf.constant(0.0, dtype=tf.float64), tf.constant(2.0, dtype=tf.float64)
    )
    new_parameter = gpflow.Parameter(initial_parameter, transform=transform)

    assert new_parameter.transform == transform


def test_construct_parameter_from_existing_parameter_check_prior() -> None:
    prior = tfp.distributions.Normal(0.0, 1.0)
    initial_parameter = gpflow.Parameter([1.2, 1.1], prior=prior)
    new_parameter = gpflow.Parameter(initial_parameter)

    assert new_parameter.prior == prior


def test_construct_parameter_from_existing_parameter_override_prior() -> None:
    initial_parameter = gpflow.Parameter([1.2, 1.1])

    prior = tfp.distributions.Normal(0.0, 1.0)
    new_parameter = gpflow.Parameter(initial_parameter, prior=prior)

    assert new_parameter.prior == prior


@pytest.mark.parametrize("prior_on", [PriorOn.CONSTRAINED, PriorOn.UNCONSTRAINED])
def test_construct_parameter_from_existing_parameter_check_prior_on(prior_on: PriorOn) -> None:
    initial_parameter = gpflow.Parameter([1.2, 1.1], prior_on=prior_on)
    new_parameter = gpflow.Parameter(initial_parameter)

    assert new_parameter.prior_on == prior_on


@pytest.mark.parametrize("prior_on", [PriorOn.CONSTRAINED, PriorOn.UNCONSTRAINED])
def test_construct_parameter_from_existing_parameter_override_prior_on(prior_on: PriorOn) -> None:
    initial_parameter = gpflow.Parameter([1.2, 1.1])
    new_parameter = gpflow.Parameter(initial_parameter, prior_on=prior_on)

    assert new_parameter.prior_on == prior_on


@pytest.mark.parametrize("trainable", [True, False])
def test_construct_parameter_from_existing_parameter_check_trainable(trainable: bool) -> None:
    initial_parameter = gpflow.Parameter([1.2, 1.1], trainable=trainable)
    new_parameter = gpflow.Parameter(initial_parameter)

    assert new_parameter.trainable == trainable


@pytest.mark.parametrize("trainable", [True, False])
def test_construct_parameter_from_existing_parameter_override_trainable(trainable: bool) -> None:
    initial_parameter = gpflow.Parameter([1.2, 1.1], trainable=trainable)
    new_parameter = gpflow.Parameter(initial_parameter, trainable=not trainable)

    assert new_parameter.trainable is not trainable


@pytest.mark.parametrize("dtype", [tf.float32, tf.float64])
def test_construct_parameter_from_existing_parameter_check_dtype(dtype: tf.DType) -> None:
    initial_parameter = gpflow.Parameter([1.1, 2.1], dtype=dtype)
    new_parameter = gpflow.Parameter(initial_parameter)

    assert new_parameter.dtype == dtype


@pytest.mark.parametrize("dtype", [tf.float32, tf.float64])
def test_construct_parameter_from_existing_parameter_override_dtype(dtype: tf.DType) -> None:
    initial_parameter = gpflow.Parameter([1.1, 2.1])
    new_parameter = gpflow.Parameter(initial_parameter, dtype=dtype)

    assert new_parameter.dtype == dtype


def test_construct_parameter_from_existing_parameter_check_name() -> None:
    transform = tfp.bijectors.Sigmoid(
        tf.constant(0.0, dtype=tf.float64), tf.constant(2.0, dtype=tf.float64)
    )
    initial_parameter = gpflow.Parameter([1.2, 1.1], transform=transform)
    new_parameter = gpflow.Parameter(initial_parameter)

    assert new_parameter.name == transform.name


def test_construct_parameter_from_existing_parameter_override_name() -> None:
    initial_parameter = gpflow.Parameter([1.2, 1.1])
    transform = tfp.bijectors.Sigmoid(
        tf.constant(0.0, dtype=tf.float64), tf.constant(2.0, dtype=tf.float64)
    )
    new_parameter = gpflow.Parameter(initial_parameter, transform=transform)

    assert new_parameter.name == transform.name


def test_construct_parameter_from_existing_parameter_value_becomes_invalid() -> None:
    initial_parameter = gpflow.Parameter(0.0)
    transform = tfp.bijectors.Reciprocal()

    with pytest.raises(tf.errors.InvalidArgumentError) as exc:
        gpflow.Parameter(initial_parameter, transform=transform)

    assert "gpflow.Parameter" in exc.value.message


def test_construct_parameter_with_variable_shape() -> None:
    parameter = gpflow.Parameter([[1, 2, 3]], shape=[None, None])

    values: Sequence[AnyNDArray] = [
        np.ones((0, 0), dtype=default_float()),
        np.ones((1, 3), dtype=default_float()),
        np.ones((3, 1), dtype=default_float()),
        np.ones((3, 4), dtype=default_float()),
    ]

    for value in values:
        parameter.assign(value)
        np.testing.assert_equal(value, parameter.numpy())


def test_construct_parameter_with_variable_shape__different_constrained_shape() -> None:
    parameter = gpflow.Parameter(
        [[1, 0], [2, 3]],
        transform=triangular(),
        unconstrained_shape=[None],
        constrained_shape=[None, None],
    )

    values: Sequence[AnyNDArray] = [
        # The triangular() transform doesn't appear to support 0x0 matrices.
        np.tril(np.ones((1, 1), dtype=default_float())),
        np.tril(np.ones((2, 2), dtype=default_float())),
        np.tril(np.ones((3, 3), dtype=default_float())),
    ]

    for value in values:
        parameter.assign(value)
        np.testing.assert_equal(value, parameter.numpy())


# tests/gpflow/test_base_prior.py

from typing import Any, Tuple, Type

import numpy as np
import pytest
import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow_probability.python.bijectors import Exp
from tensorflow_probability.python.distributions import Uniform

import gpflow
from gpflow.base import AnyNDArray, PriorOn
from gpflow.config import set_default_float
from gpflow.utilities import to_default_float

np.random.seed(1)


class Datum:
    X: AnyNDArray = 10 * np.random.randn(5, 1)
    Y: AnyNDArray = 10 * np.random.randn(5, 1)
    lengthscale = 3.3


def test_gpr_objective_equivalence() -> None:
    """
    In Maximum Likelihood Estimation (MLE), i.e. when there are no priors on
    the parameters, the objective should not depend on any transforms on the
    parameters.
    We use GPR as a simple model that has an objective.
    """
    data = (Datum.X, Datum.Y)
    l_value = Datum.lengthscale

    l_variable = tf.Variable(l_value, dtype=gpflow.default_float(), trainable=True)
    m1 = gpflow.models.GPR(data, kernel=gpflow.kernels.SquaredExponential(lengthscales=l_value))
    m2 = gpflow.models.GPR(data, kernel=gpflow.kernels.SquaredExponential())
    m2.kernel.lengthscales = gpflow.Parameter(l_variable, transform=None)
    assert np.allclose(
        m1.kernel.lengthscales.numpy(), m2.kernel.lengthscales.numpy()
    )  # consistency check

    assert np.allclose(
        m1.log_marginal_likelihood().numpy(), m2.log_marginal_likelihood().numpy()
    ), "MLE objective should not depend on Parameter transform"


def test_log_prior_with_no_prior() -> None:
    """
    A parameter without any prior should have zero log-prior,
    even if it has a transform to constrain it.
    """
    param = gpflow.Parameter(5.3, transform=gpflow.utilities.positive())
    assert param.log_prior_density().numpy() == 0.0


def test_log_prior_for_uniform_prior() -> None:
    """
    If we assign a Uniform prior to a parameter, we should not expect the value of the prior density
    to change with the parameter value, even if it has a transform associated with it.
    """

    uniform_prior = Uniform(low=np.float64(0), high=np.float64(100))
    param = gpflow.Parameter(1.0, transform=gpflow.utilities.positive(), prior=uniform_prior)
    low_value = param.log_prior_density().numpy()
    param.assign(10.0)
    high_value = param.log_prior_density().numpy()

    assert np.isclose(low_value, high_value)


def test_log_prior_on_unconstrained() -> None:
    """
    A parameter with an Exp transform, and a uniform prior on its unconstrained, should have a
    prior in the constrained space that scales as 1/value.
    """

    initial_value = 1.0
    scale_factor = 10.0
    uniform_prior = Uniform(low=np.float64(0), high=np.float64(100))
    param = gpflow.Parameter(
        initial_value,
        transform=Exp(),
        prior=uniform_prior,
        prior_on=PriorOn.UNCONSTRAINED,
    )
    low_value = param.log_prior_density().numpy()
    param.assign(scale_factor * initial_value)
    high_value = param.log_prior_density().numpy()

    assert np.isclose(low_value, high_value + np.log(scale_factor))


class DummyModel(gpflow.models.BayesianModel):
    value = 3.3
    log_scale = 0.4

    def __init__(self, with_transform: bool) -> None:
        super().__init__()

        prior = tfp.distributions.Normal(to_default_float(1.0), to_default_float(1.0))

        scale = np.exp(self.log_scale)
        if with_transform:
            transform = tfp.bijectors.Shift(to_default_float(0.0))(
                tfp.bijectors.Scale(to_default_float(scale))
            )
        else:
            transform = None

        self.theta = gpflow.Parameter(self.value, prior=prior, transform=transform)

    def maximum_log_likelihood_objective(self, *args: Any, **kwargs: Any) -> tf.Tensor:
        assert not args
        assert not kwargs
        return (self.theta + 5) ** 2


def test_map_invariance_to_transform() -> None:
    m1 = DummyModel(with_transform=True)
    m2 = DummyModel(with_transform=False)
    assert np.allclose(
        m1.log_posterior_density().numpy(), m2.log_posterior_density().numpy()
    ), "log posterior density should not be affected by a transform"


def get_gpmc_model_params() -> Tuple[Any, ...]:
    kernel = gpflow.kernels.Matern32()
    likelihood = gpflow.likelihoods.Gaussian()
    data = [np.random.randn(5, 1), np.random.randn(5, 1)]
    return data, kernel, likelihood


@pytest.mark.parametrize(
    "model_class, args",
    [
        (gpflow.models.GPMC, get_gpmc_model_params()),
        # (gpflow.models.SGPMC, get_SGPMC_model_params()) # Fails due to inducing_variable=None bug
    ],
)
def test_v_prior_dtypes(model_class: Type[Any], args: Tuple[Any, ...]) -> None:
    with gpflow.config.as_context():
        set_default_float(np.float32)
        m = model_class(*args)
        assert m.V.prior.dtype == np.float32
        set_default_float(np.float64)
        m = model_class(*args)
        assert m.V.prior.dtype == np.float64


# tests/gpflow/test_base_training.py

# Copyright 2018 the GPflow authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


from typing import Any, Mapping

import numpy as np
import pytest
import tensorflow as tf

import gpflow
from gpflow.models import SVGP
from gpflow.utilities import leaf_components, multiple_assign, read_values

rng = np.random.RandomState(0)


class Data:
    H0 = 5
    H1 = 2
    M = 10
    D = 1
    Z = rng.rand(M, 1)
    ls = 2.0
    ls_new = 1.5
    var = 1.0
    var_new = 10.0


# ------------------------------------------
# Fixtures
# ------------------------------------------


@pytest.fixture
def model() -> SVGP:
    return SVGP(
        kernel=gpflow.kernels.SquaredExponential(lengthscales=Data.ls, variance=Data.var),
        likelihood=gpflow.likelihoods.Gaussian(),
        inducing_variable=Data.Z,
        q_diag=True,
    )


# ------------------------------------------
# Reference
# ------------------------------------------

model_param_updates = {
    ".kernel.lengthscales": Data.ls_new,
    ".likelihood.variance": Data.var_new,
    ".inducing_variable.Z": np.zeros_like(Data.Z),
    ".q_sqrt": 0.5 * np.ones((Data.M, 1)),
}
model_wrong_path = [
    {"kernel.lengthscales": Data.ls_new},
    {".Gaussian.variance": Data.var_new},
    {"inducing_variable.Z": np.zeros_like(Data.Z)},
    {".q_std": 0.5 * np.ones((Data.M, 1))},
]

model_wrong_value = [
    {".likelihood.variance": np.ones((2, 1), dtype=np.int32)},
    {".inducing_variable.Z": [1, 2, 3]},
]


@pytest.mark.parametrize("var_update_dict", [model_param_updates])
def test_multiple_assign_updates_correct_values(
    model: SVGP, var_update_dict: Mapping[str, Any]
) -> None:
    old_value_dict = leaf_components(model)
    multiple_assign(model, var_update_dict)
    for path, variable in leaf_components(model).items():
        if path in var_update_dict.keys():
            np.testing.assert_almost_equal(variable.numpy(), var_update_dict[path], decimal=7)
        else:
            np.testing.assert_equal(variable.numpy(), old_value_dict[path].numpy())


@pytest.mark.parametrize("wrong_var_update_dict", model_wrong_path)
def test_multiple_assign_fails_with_invalid_path(
    model: SVGP, wrong_var_update_dict: Mapping[str, Any]
) -> None:
    with pytest.raises(KeyError):
        multiple_assign(model, wrong_var_update_dict)


@pytest.mark.parametrize("wrong_var_update_dict", model_wrong_value)
def test_multiple_assign_fails_with_invalid_values(
    model: SVGP, wrong_var_update_dict: Mapping[str, Any]
) -> None:
    with pytest.raises(ValueError):
        multiple_assign(model, wrong_var_update_dict)


def test_dict_utilities(model: SVGP) -> None:
    """
    Test both `parameter_dict()` and `read_values()`
    """

    class SubModule(tf.Module):
        def __init__(self) -> None:
            self.parameter = gpflow.Parameter(1.0)
            self.variable = tf.Variable(1.0)

    class Module(tf.Module):
        def __init__(self) -> None:
            self.submodule = SubModule()
            self.top_parameter = gpflow.Parameter(3.0)

    m = Module()
    params = gpflow.utilities.parameter_dict(m)
    # {
    #   ".submodule.parameter": <parameter object>,
    #   ".submodule.variable": <variable object>
    # }
    assert list(params.keys()) == [
        ".submodule.parameter",
        ".submodule.variable",
        ".top_parameter",
    ]
    assert list(params.values()) == [
        m.submodule.parameter,
        m.submodule.variable,
        m.top_parameter,
    ]

    for k, v in read_values(m).items():
        assert params[k].numpy() == v


# tests/gpflow/test_functions.py

# Copyright 2017 the GPflow authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Any, Sequence, Tuple, Type

import numpy as np
import pytest
from numpy.testing import assert_allclose

import gpflow
from gpflow.base import AnyNDArray, TensorType
from gpflow.config import default_int
from gpflow.experimental.check_shapes import check_shapes
from gpflow.functions import (
    Additive,
    Constant,
    Linear,
    MeanFunction,
    Polynomial,
    Product,
    SwitchedMeanFunction,
    Zero,
)
from gpflow.inducing_variables import InducingPoints

rng = np.random.RandomState(99021)


class Datum:
    input_dim, output_dim = 3, 2
    N, Ntest, M = 20, 30, 10


_mean_functions = [
    Zero(),
    Linear(
        A=rng.randn(Datum.input_dim, Datum.output_dim),
        b=rng.randn(Datum.output_dim, 1).reshape(-1),
    ),
    Constant(c=rng.randn(Datum.output_dim, 1).reshape(-1)),
    Polynomial(degree=2, input_dim=Datum.input_dim, output_dim=Datum.output_dim),
]


@pytest.mark.parametrize("mean_function_1", _mean_functions)
@pytest.mark.parametrize("mean_function_2", _mean_functions)
@pytest.mark.parametrize("operation", ["+", "*"])
def test_mean_functions_output_shape(
    mean_function_1: MeanFunction, mean_function_2: MeanFunction, operation: str
) -> None:
    """
    Test the output shape for basic and compositional mean functions, also
    check that the combination of mean functions returns the correct class
    """
    X = np.random.randn(Datum.N, Datum.input_dim)
    Y = mean_function_1(X)
    # basic output shape check
    assert Y.shape in [(Datum.N, Datum.output_dim), (Datum.N, 1)]

    # composed mean function output shape check
    if operation == "+":
        mean_composed = mean_function_1 + mean_function_2
    elif operation == "*":
        mean_composed = mean_function_1 * mean_function_2
    else:
        raise NotImplementedError

    Y_composed = mean_composed(X)
    assert Y_composed.shape in [(Datum.N, Datum.output_dim), (Datum.N, 1)]


@pytest.mark.parametrize("mean_function_1", _mean_functions)
@pytest.mark.parametrize("mean_function_2", _mean_functions)
@pytest.mark.parametrize("operation", ["+", "*"])
def test_mean_functions_composite_type(
    mean_function_1: MeanFunction, mean_function_2: MeanFunction, operation: str
) -> None:
    if operation == "+":
        mean_composed = mean_function_1 + mean_function_2
        assert isinstance(mean_composed, Additive)
    elif operation == "*":
        mean_composed = mean_function_1 * mean_function_2
        assert isinstance(mean_composed, Product)
    else:
        raise NotImplementedError


_linear_functions = [
    Linear(
        A=rng.randn(Datum.input_dim, Datum.output_dim),
        b=rng.randn(Datum.output_dim, 1).reshape(-1),
    )
    for _ in range(3)
]

# Append inverse of first Linear mean function in _linear_functions
_linear_functions.append(Linear(A=-1.0 * _linear_functions[0].A, b=-1.0 * _linear_functions[0].b))

_constant_functions = [Constant(c=rng.randn(Datum.output_dim, 1).reshape(-1)) for _ in range(3)]
# Append inverse of first Constant mean function in _constant_functions
_constant_functions.append(Constant(c=-1.0 * _constant_functions[0].c))


@check_shapes(
    "X: [N, D]",
    "Y: [N, Q]",
)
def _create_GPR_model_with_bias(
    X: TensorType, Y: TensorType, mean_function: MeanFunction
) -> gpflow.models.GPR:
    return gpflow.models.GPR(
        (X, Y), mean_function=mean_function, kernel=gpflow.kernels.Bias(Datum.input_dim)
    )


@pytest.mark.parametrize("mean_functions", [_linear_functions, _constant_functions])
def test_mean_functions_distributive_property(mean_functions: Sequence[MeanFunction]) -> None:
    """
    Tests that distributive property of addition and multiplication holds for mean functions
    (both Constant and Linear): A * (B + C) = A * B + A * C
    """
    X, Y = rng.randn(Datum.N, Datum.input_dim), rng.randn(Datum.N, Datum.output_dim)
    Xtest = rng.randn(30, Datum.input_dim)
    A, B, C = mean_functions[0], mean_functions[1], mean_functions[2]
    lhs = Product(A, Additive(B, C))  # A * (B + C)
    rhs = Additive(Product(A, B), Product(A, C))  # A * B + A * C

    model_lhs = _create_GPR_model_with_bias(X, Y, mean_function=lhs)
    model_rhs = _create_GPR_model_with_bias(X, Y, mean_function=rhs)

    mu_lhs, var_lhs = model_lhs.predict_f(Xtest)
    mu_rhs, var_rhs = model_rhs.predict_f(Xtest)

    assert_allclose(mu_lhs, mu_rhs)
    assert_allclose(var_lhs, var_rhs)


@pytest.mark.parametrize("mean_functions", [_linear_functions, _constant_functions])
def test_mean_functions_A_minus_A_equals_zero(mean_functions: Sequence[MeanFunction]) -> None:
    """
    Tests that the addition the inverse of a mean function to itself is equivalent to having a
    Zero mean function: A + (-A) = 0
    """
    X, Y = rng.randn(Datum.N, Datum.input_dim), rng.randn(Datum.N, Datum.output_dim)
    Xtest = rng.randn(30, Datum.input_dim)
    A, A_inverse = mean_functions[0], mean_functions[-1]
    lhs = Additive(A, A_inverse)  # A + (-A)
    rhs = Zero()  # 0

    model_lhs = _create_GPR_model_with_bias(X, Y, mean_function=lhs)
    model_rhs = _create_GPR_model_with_bias(X, Y, mean_function=rhs)

    mu_lhs, var_lhs = model_lhs.predict_f(Xtest)
    mu_rhs, var_rhs = model_rhs.predict_f(Xtest)

    assert_allclose(mu_lhs, mu_rhs)
    assert_allclose(var_lhs, var_rhs)


@pytest.mark.parametrize("mean_functions", [_linear_functions])
def test_linear_mean_functions_associative_property(mean_functions: Sequence[MeanFunction]) -> None:
    """
    Tests that associative property of addition holds for linear mean functions:
    A + (B + (-A)) = B = (A + B) + (-A)
    """
    X, Y = rng.randn(Datum.N, Datum.input_dim), rng.randn(Datum.N, Datum.output_dim)
    Xtest = rng.randn(30, Datum.input_dim)
    A, B, A_inverse = mean_functions[0], mean_functions[1], mean_functions[-1]

    lhs = Additive(A, Additive(B, A_inverse))  # A + (B + (-A))
    rhs = Additive(Additive(A, B), A_inverse)  # (A + B) + (-A)

    model_lhs = _create_GPR_model_with_bias(X, Y, mean_function=lhs)
    model_b = _create_GPR_model_with_bias(X, Y, mean_function=B)
    model_rhs = _create_GPR_model_with_bias(X, Y, mean_function=rhs)

    mu_lhs, var_lhs = model_lhs.predict_f(Xtest)
    mu_b, var_b = model_b.predict_f(Xtest)
    mu_rhs, var_rhs = model_rhs.predict_f(Xtest)

    assert_allclose(mu_lhs, mu_b)
    assert_allclose(var_lhs, var_b)
    assert_allclose(mu_b, mu_rhs)
    assert_allclose(var_b, var_rhs)


@pytest.mark.parametrize("batch", [(), (1, 2)])
@pytest.mark.parametrize("degree", [0, 1, 3])
@pytest.mark.parametrize("input_dim", [0, 1, 3])
@pytest.mark.parametrize("output_dim", [1, 2])
def test_polynomial__sanity(
    batch: Tuple[int, ...], degree: int, input_dim: int, output_dim: int
) -> None:
    p = Polynomial(degree, input_dim, output_dim)
    X = np.ones(batch + (input_dim,))
    Y = p(X)
    assert (batch) + (output_dim,) == Y.shape
    assert_allclose(1.0, Y)  # Polynomial is initialised to constant 1.0.


def test_polynomial__compute_powers() -> None:
    assert_allclose(
        [
            (0, 0, 0),
            (0, 0, 1),
            (0, 0, 2),
            (0, 1, 0),
            (0, 1, 1),
            (0, 2, 0),
            (1, 0, 0),
            (1, 0, 1),
            (1, 1, 0),
            (2, 0, 0),
        ],
        list(Polynomial.compute_powers(degree=2, input_dim=3)),
    )


def test_polynomial__1d() -> None:
    # Test on a 1D quadratic function, where we can easily work out the maths.
    p = Polynomial(degree=2, w=[1.0, 2.0, 3.0])
    X = np.array([[1.0], [2.0]])
    Y = p(X)
    assert_allclose(
        [
            [1.0 + 2.0 * 1.0 + 3.0 * (1.0 ** 2)],
            [1.0 + 2.0 * 2.0 + 3.0 * (2.0 ** 2)],
        ],
        Y,
    )


def test_polynomial__linear() -> None:
    # Test on a 3D linear function, where we can easily work out the maths.
    p = Polynomial(degree=1, input_dim=3, w=[1.0, 2.0, 3.0, 4.0])
    X = np.array([1.0, 2.0, 3.0])
    Y = p(X)
    assert_allclose([1.0 + 2.0 * 3.0 + 3.0 * 2.0 + 4.0 * 1.0], Y)


@pytest.mark.parametrize("N, D", [[10, 3]])
def test_switched_mean_function(N: int, D: int) -> None:
    """
    Test for the SwitchedMeanFunction.
    """
    X: AnyNDArray = np.hstack([rng.randn(N, D), 1.0 * rng.randint(0, 2, N).reshape(-1, 1)])
    zeros, ones = Constant(np.zeros(1)), Constant(np.ones(1))
    switched_mean = SwitchedMeanFunction([zeros, ones])

    np_list: AnyNDArray = np.array([0.0, 1.0])
    result_ref = (np_list[X[:, D].astype(default_int())]).reshape(-1, 1)
    result = switched_mean(X)

    assert_allclose(result, result_ref)


def test_bug_277_regression() -> None:
    """
    See github issue #277. This is a regression test.
    """
    model1, model2 = Linear(), Linear()
    assert model1.b.numpy() == model2.b.numpy()
    model2.b.assign([1.0])
    assert not model1.b.numpy() == model2.b.numpy()


_model_classes = [
    gpflow.models.GPR,
    gpflow.models.SGPR,
    gpflow.models.GPRFITC,
    gpflow.models.SVGP,
    gpflow.models.VGP,
    gpflow.models.GPMC,
    gpflow.models.SGPMC,
]


@pytest.mark.parametrize("model_class", _model_classes)
def test_models_with_mean_functions_changes(model_class: Type[Any]) -> None:
    """
    Simply check that all models have a higher prediction with a constant mean
    function than with a zero mean function.

    For compositions of mean functions check that multiplication/ addition of
    a constant results in a higher prediction, whereas addition of zero/
    mutliplication with one does not.
    """
    data = rng.randn(Datum.N, Datum.input_dim), rng.randn(Datum.N, 1)
    Xnew = rng.randn(Datum.Ntest, Datum.input_dim)
    inducing_variable = InducingPoints(Z=rng.randn(Datum.M, Datum.input_dim))
    kernel = gpflow.kernels.Matern32()
    likelihood = gpflow.likelihoods.Gaussian()
    zero_mean = Zero()
    non_zero_mean = Constant(c=np.ones(1) * 10)

    if model_class in [gpflow.models.GPR]:
        model_zero_mean = model_class(data, kernel=kernel, mean_function=zero_mean)
        model_non_zero_mean = model_class(data, kernel=kernel, mean_function=non_zero_mean)
    elif model_class in [gpflow.models.VGP]:
        model_zero_mean = model_class(
            data, likelihood=likelihood, kernel=kernel, mean_function=zero_mean
        )
        model_non_zero_mean = model_class(
            data, likelihood=likelihood, kernel=kernel, mean_function=non_zero_mean
        )
    elif model_class in [gpflow.models.SVGP]:
        model_zero_mean = model_class(
            kernel=kernel,
            likelihood=likelihood,
            inducing_variable=inducing_variable,
            mean_function=zero_mean,
        )
        model_non_zero_mean = model_class(
            kernel=kernel,
            likelihood=likelihood,
            inducing_variable=inducing_variable,
            mean_function=non_zero_mean,
        )
    elif model_class in [gpflow.models.SGPR, gpflow.models.GPRFITC]:
        model_zero_mean = model_class(
            data,
            kernel=kernel,
            inducing_variable=inducing_variable,
            mean_function=zero_mean,
        )
        model_non_zero_mean = model_class(
            data,
            kernel=kernel,
            inducing_variable=inducing_variable,
            mean_function=non_zero_mean,
        )
    elif model_class in [gpflow.models.SGPMC]:
        model_zero_mean = model_class(
            data,
            kernel=kernel,
            likelihood=likelihood,
            inducing_variable=inducing_variable,
            mean_function=zero_mean,
        )
        model_non_zero_mean = model_class(
            data,
            kernel=kernel,
            likelihood=likelihood,
            inducing_variable=inducing_variable,
            mean_function=non_zero_mean,
        )
    elif model_class in [gpflow.models.GPMC]:
        model_zero_mean = model_class(
            data, kernel=kernel, likelihood=likelihood, mean_function=zero_mean
        )
        model_non_zero_mean = model_class(
            data, kernel=kernel, likelihood=likelihood, mean_function=non_zero_mean
        )
    else:
        raise NotImplementedError

    mu_zero, var_zero = model_zero_mean.predict_f(Xnew)
    mu_non_zero, var_non_zero = model_non_zero_mean.predict_f(Xnew)
    # predictive variance remains unchanged after modifying mean function
    assert np.all(var_zero.numpy() == var_non_zero.numpy())
    # predictive mean changes after modifying mean function
    assert not np.all(mu_zero.numpy() == mu_non_zero.numpy())


# tests/gpflow/test_inducing_variables.py

# Copyright 2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Callable, Optional, Tuple

import numpy as np
import pytest
import tensorflow as tf

import gpflow
import gpflow.inducing_variables as giv
from gpflow.experimental.check_shapes import get_shape
from tests.gpflow.experimental.check_shapes.utils import TestContext


def test_inducing_points_with_variable_shape() -> None:
    N, M1, D, P = 50, 13, 3, 1
    X, Y = np.random.randn(N, D), np.random.randn(N, P)

    Z1 = np.random.randn(M1, D)

    # use explicit tf.Variable with None shape:
    iv = giv.InducingPoints(
        tf.Variable(Z1, trainable=False, dtype=gpflow.default_float(), shape=(None, D))
    )
    # Note that we cannot have Z be trainable if we want to be able to change its shape;
    # TensorFlow optimizers expect shape to be known at construction time.

    m = gpflow.models.SGPR(data=(X, Y), kernel=gpflow.kernels.Matern32(), inducing_variable=iv)

    # Check 1: that we can still optimize with None shape
    opt = tf.optimizers.Adam()

    @tf.function
    def optimization_step() -> None:
        opt.minimize(m.training_loss, m.trainable_variables)

    optimization_step()

    # Check 2: that we can successfully assign a new Z with different number of inducing points!
    Z2 = np.random.randn(M1 + 1, D)
    m.inducing_variable.Z.assign(Z2)

    # Check 3: that we can also optimize with changed Z tensor
    optimization_step()


@pytest.mark.parametrize(
    "iv_factory,expected_shape",
    [
        (lambda t: giv.InducingPoints(t), (7, 3, 1)),
        (lambda t: giv.Multiscale(t, t), (7, 3, 1)),
        (lambda t: giv.InducingPatches(t), (7, 3, 1)),
        (
            lambda t: giv.FallbackSharedIndependentInducingVariables(giv.InducingPoints(t)),
            (7, 3, None),
        ),
        (
            lambda t: giv.FallbackSeparateIndependentInducingVariables(
                [
                    giv.InducingPoints(t),
                    giv.InducingPoints(t),
                ]
            ),
            (7, 3, 2),
        ),
        (
            lambda t: giv.SharedIndependentInducingVariables(giv.InducingPoints(t)),
            (7, 3, None),
        ),
        (
            lambda t: giv.SeparateIndependentInducingVariables(
                [
                    giv.InducingPoints(t),
                    giv.InducingPoints(t),
                ]
            ),
            (7, 3, 2),
        ),
    ],
)
@pytest.mark.parametrize("none_shape", [False, True])
def test_shape(
    iv_factory: Callable[[tf.Tensor], giv.InducingVariables],
    none_shape: bool,
    expected_shape: Tuple[Optional[int], ...],
) -> None:
    M = expected_shape[0]
    ones = tf.ones((7, 3))
    kwargs = {"shape": tf.TensorShape(None)} if none_shape else {}
    iv = iv_factory(tf.Variable(ones, **kwargs))
    if none_shape:
        assert get_shape(iv, TestContext()) is None
    else:
        assert expected_shape == get_shape(iv, TestContext())
    assert M == iv.num_inducing
    assert M == len(iv)


@pytest.mark.parametrize("none_shape", [False, True])
def test_shape__inconsistent(none_shape: bool) -> None:
    kwargs = {"shape": tf.TensorShape(None)} if none_shape else {}
    iv = giv.SeparateIndependentInducingVariables(
        [
            giv.InducingPoints(tf.Variable(tf.ones((5, 3)), **kwargs)),
            giv.InducingPoints(tf.Variable(tf.ones((7, 3)), **kwargs)),
        ]
    )
    assert get_shape(iv, TestContext()) is None
    with pytest.raises(tf.errors.InvalidArgumentError):
        iv.num_inducing
    with pytest.raises(tf.errors.InvalidArgumentError):
        len(iv)


# tests/gpflow/test_initial_value.py

# Copyright 2022 the GPflow authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import ast
from pathlib import Path
from typing import List, Sequence, Union

import pytest

import gpflow


def find_py_files() -> Sequence[Path]:
    """
    Find all ``.py`` files of the GPflow user-facing code.
    """
    root_module = gpflow
    root_path_str = root_module.__file__
    assert root_path_str is not None
    root_path = Path(root_path_str).parent
    return tuple(root_path.glob("**/*.py"))


_PY_FILES = find_py_files()


@pytest.mark.parametrize("path", _PY_FILES, ids=str)
def test_assignment_value(path: Path) -> None:
    """
    Python allows assignments without a value, generally used for typing. For example::

        @tf.function
        def foo(a: tf.Tensor, b: tf.Tensor) -> tf.Tensor:
            result: tf.Tensor  # <-- Assignment without a value. Only for typing.
            result = a + b
            return result

    However, ``TensorFlow < 2.9.0`` cannot compile these.

    This test parses all our ``.py`` files and look for assignments without a value, to ensure all
    our code can be compiled by TensorFlow.

    See also: https://github.com/tensorflow/tensorflow/issues/56024
    """

    tree = ast.parse(path.read_text(), str(path))

    class RequireAssignmentValue(ast.NodeVisitor):
        """
        When ``.visit(tree)`` is called this will recursively visit all nodes of ``tree``.

        Visitor methods are called based on the types of the nodes of the tree. If a node has type
        ``Foo``, then ``visit_Foo(...)`` below will be called. If no such method exists
        ``generic_visit`` is called insted, which just visits the children of the node recursively.
        """

        def __init__(self) -> None:
            super().__init__()

            self.in_function_stack: List[bool] = []
            """
            Keeps track of whether we're inside a function.  We only raise errors for assignments in
            functions - I think variables on classes are handled differently, and doesn't cause
            these problems.
            """

            self.bad_assigns: List[ast.AST] = []
            """
            List of bad assignments, so that we can report multiple errors at the same time.
            """

        def set_in_function(self, node: ast.AST, in_function: bool) -> None:
            """
            Set whether we're inside a function, and visit children recursively.
            """
            self.in_function_stack.append(in_function)
            self.generic_visit(node)
            self.in_function_stack.pop()

        def visit_ClassDef(self, node: ast.ClassDef) -> None:
            """
            Called for class definitions.
            """
            self.set_in_function(node, False)

        def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
            """
            Called for function (and method) definitions.
            """
            self.set_in_function(node, True)

        def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
            """
            Called for async function definitions.
            """
            self.set_in_function(node, True)

        def check_assign(self, node: Union[ast.Assign, ast.AugAssign, ast.AnnAssign]) -> None:
            """
            Checks whether this is a valid assignment, and visit children recursively.
            """
            if self.in_function_stack and self.in_function_stack[-1]:
                if not node.value:
                    self.bad_assigns.append(node)
            self.generic_visit(node)

        def visit_Assign(self, node: ast.Assign) -> None:
            """
            Called for a regular assignment.
            """
            self.check_assign(node)

        def visit_AugAssign(self, node: ast.AugAssign) -> None:
            """
            Called for +=, -=, etc.
            """
            self.check_assign(node)

        def visit_AnnAssign(self, node: ast.AnnAssign) -> None:
            """
            Called for an assigment with a type annotation.
            """
            self.check_assign(node)

    visitor = RequireAssignmentValue()
    visitor.visit(tree)

    if visitor.bad_assigns:
        lines = [""]
        lines.extend(f"{path}:{n.lineno}" for n in visitor.bad_assigns)
        lines.append(
            "TensorFlow cannot compile assignments without a value,"
            " so for performance reasons we require all assgnments to have a value."
            " See: https://github.com/tensorflow/tensorflow/issues/56024 "
        )
        raise AssertionError("\n".join(lines))


# tests/gpflow/test_kullback_leiblers.py

# Copyright 2017 the GPflow authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# -*- coding: utf-8 -*-

from unittest.mock import patch

import numpy as np
import pytest
import tensorflow as tf
from numpy.testing import assert_allclose, assert_almost_equal

import gpflow
from gpflow import Parameter, default_float, default_jitter
from gpflow.base import AnyNDArray, TensorType
from gpflow.experimental.check_shapes import ShapeChecker, check_shapes
from gpflow.inducing_variables import InducingPoints
from gpflow.kernels import Kernel
from gpflow.kullback_leiblers import gauss_kl, prior_kl
from gpflow.utilities.bijectors import triangular

rng = np.random.RandomState(0)

# ------------------------------------------
# Fixtures
# ------------------------------------------

Ln = 2
Nn = 10
Mn = 50


@pytest.fixture(scope="module")
def kernel() -> Kernel:
    k = gpflow.kernels.Matern32() + gpflow.kernels.White()
    k.kernels[1].variance.assign(0.01)
    return k


@pytest.fixture(scope="module")
@check_shapes(
    "return: [N, 1, 1]",
)
def inducing_points() -> InducingPoints:
    return InducingPoints(rng.randn(Nn, 1))


@pytest.fixture(scope="module")
@check_shapes(
    "return: [N, L]",
)
def mu() -> Parameter:
    return Parameter(rng.randn(Nn, Ln))


# ------------------------------------------
# Helpers
# ------------------------------------------


@check_shapes(
    "return: [N, M, M]",
)
def make_sqrt(N: int, M: int) -> TensorType:
    return np.array([np.tril(rng.randn(M, M)) for _ in range(N)])


@check_shapes(
    "return: [N, M, M]",
)
def make_K_batch(N: int, M: int) -> TensorType:
    K_np = rng.randn(N, M, M)
    beye: AnyNDArray = np.array([np.eye(M) for _ in range(N)])
    return 0.1 * (K_np + np.transpose(K_np, (0, 2, 1))) + beye


@check_shapes(
    "q_mu: [broadcast batch...]",
    "q_sigma: [broadcast batch...]",
    "p_var: [broadcast batch...]",
    "return: []",
)
def compute_kl_1d(q_mu: TensorType, q_sigma: TensorType, p_var: TensorType = 1.0) -> TensorType:
    p_var = tf.ones_like(q_sigma) if p_var is None else p_var
    q_var = tf.square(q_sigma)
    kl = 0.5 * (q_var / p_var + tf.square(q_mu) / p_var - 1 + tf.math.log(p_var / q_var))
    return tf.reduce_sum(kl)


# ------------------------------------------
# Data classes: storing constants
# ------------------------------------------


class Datum:
    cs = ShapeChecker().check_shape

    M, N = 5, 4

    mu = cs(rng.randn(M, N), "[M, N]")
    A = cs(rng.randn(M, M), "[M, M]")
    I = cs(np.eye(M), "[M, M]")
    K = cs(A @ A.T + default_jitter() * I, "[M, M]")
    sqrt = cs(make_sqrt(N, M), "[N, M, M]")
    sqrt_diag = cs(rng.randn(M, N), "[M, N]")
    K_batch = cs(make_K_batch(N, M), "[N, M, M]")
    K_cholesky = cs(np.linalg.cholesky(K), "[M, M]")


@pytest.mark.parametrize("diag", [True, False])
def test_kl_k_cholesky(diag: bool) -> None:
    """
    Test that passing K or K_cholesky yield the same answer
    """
    q_mu = Datum.mu
    q_sqrt = Datum.sqrt_diag if diag else Datum.sqrt
    kl_K = gauss_kl(q_mu, q_sqrt, K=Datum.K)
    kl_K_chol = gauss_kl(q_mu, q_sqrt, K_cholesky=Datum.K_cholesky)

    np.testing.assert_allclose(kl_K.numpy(), kl_K_chol.numpy())


@pytest.mark.parametrize("white", [True, False])
def test_diags(white: bool) -> None:
    """
    The covariance of q(x) can be Cholesky matrices or diagonal matrices.
    Here we make sure the behaviours overlap.
    """
    # the chols are diagonal matrices, with the same entries as the diag representation.
    chol_from_diag = tf.stack(
        [tf.linalg.diag(Datum.sqrt_diag[:, i]) for i in range(Datum.N)]  # [N, M, M]
    )
    kl_diag = gauss_kl(Datum.mu, Datum.sqrt_diag, Datum.K if white else None)
    kl_dense = gauss_kl(Datum.mu, chol_from_diag, Datum.K if white else None)

    np.testing.assert_allclose(kl_diag, kl_dense)


@pytest.mark.parametrize("diag", [True, False])
def test_whitened(diag: bool) -> None:
    """
    Check that K=Identity and K=None give same answer
    """
    chol_from_diag = tf.stack(
        [tf.linalg.diag(Datum.sqrt_diag[:, i]) for i in range(Datum.N)]  # [N, M, M]
    )
    s = Datum.sqrt_diag if diag else chol_from_diag

    kl_white = gauss_kl(Datum.mu, s)
    kl_nonwhite = gauss_kl(Datum.mu, s, Datum.I)

    np.testing.assert_allclose(kl_white, kl_nonwhite)


@pytest.mark.parametrize("shared_k", [True, False])
@pytest.mark.parametrize("diag", [True, False])
def test_sumkl_equals_batchkl(shared_k: bool, diag: bool) -> None:
    """
    gauss_kl implicitely performs a sum of KL divergences
    This test checks that doing the sum outside of the function is equivalent
    For q(X)=prod q(x_l) and p(X)=prod p(x_l), check that sum KL(q(x_l)||p(x_l)) = KL(q(X)||p(X))
    Here, q(X) has covariance [L, M, M]
    p(X) has covariance [L, M, M] ( or [M, M] )
    Here, q(x_i) has covariance [1, M, M]
    p(x_i) has covariance [M, M]
    """
    s = Datum.sqrt_diag if diag else Datum.sqrt
    kl_batch = gauss_kl(Datum.mu, s, Datum.K if shared_k else Datum.K_batch)
    kl_sum = []
    for n in range(Datum.N):
        q_mu_n = Datum.mu[:, n][:, None]  # [M, 1]
        q_sqrt_n = (
            Datum.sqrt_diag[:, n][:, None] if diag else Datum.sqrt[n, :, :][None, :, :]
        )  # [1, M, M] or [M, 1]
        K_n = Datum.K if shared_k else Datum.K_batch[n, :, :][None, :, :]  # [1, M, M] or [M, M]
        kl_n = gauss_kl(q_mu_n, q_sqrt_n, K=K_n)
        kl_sum.append(kl_n)

    kl_sum = tf.reduce_sum(kl_sum)
    assert_almost_equal(kl_sum, kl_batch)


@patch("tensorflow.__version__", "2.1.0")
def test_sumkl_equals_batchkl_shared_k_not_diag_mocked_tf21() -> None:
    """
    Version of test_sumkl_equals_batchkl with shared_k=True and diag=False
    that tests the TensorFlow < 2.2 workaround with tiling still works.
    """
    kl_batch = gauss_kl(Datum.mu, Datum.sqrt, Datum.K)
    kl_sum = []
    for n in range(Datum.N):
        q_mu_n = Datum.mu[:, n][:, None]  # [M, 1]
        q_sqrt_n = Datum.sqrt[n, :, :][None, :, :]  # [1, M, M] or [M, 1]
        K_n = Datum.K  # [1, M, M] or [M, M]
        kl_n = gauss_kl(q_mu_n, q_sqrt_n, K=K_n)
        kl_sum.append(kl_n)

    kl_sum = tf.reduce_sum(kl_sum)
    assert_almost_equal(kl_sum, kl_batch)


@pytest.mark.parametrize("dim", [0, 1])
@pytest.mark.parametrize("white", [True, False])
def test_oned(white: bool, dim: bool) -> None:
    """
    Check that the KL divergence matches a 1D by-hand calculation.
    """
    mu1d = Datum.mu[dim, :][None, :]  # [1, N]
    s1d = Datum.sqrt[:, dim, dim][:, None, None]  # [N, 1, 1]
    K1d = Datum.K_batch[:, dim, dim][:, None, None]  # [N, 1, 1]

    kl = gauss_kl(mu1d, s1d, K1d if not white else None)
    kl_1d = compute_kl_1d(
        tf.reshape(mu1d, (-1,)),  # N
        tf.reshape(s1d, (-1,)),  # N
        None if white else tf.reshape(K1d, (-1,)),
    )  # N
    np.testing.assert_allclose(kl, kl_1d)


def test_unknown_size_inputs() -> None:
    """
    Test for #725 and #734. When the shape of the Gaussian's mean had at least
    one unknown parameter, `gauss_kl` would blow up. This happened because
    `tf.size` can only output types `tf.int32` or `tf.int64`.
    """
    mu: AnyNDArray = np.ones([1, 4], dtype=default_float())
    sqrt: AnyNDArray = np.ones([4, 1, 1], dtype=default_float())

    known_shape = gauss_kl(*map(tf.constant, [mu, sqrt]))
    unknown_shape = gauss_kl(mu, sqrt)

    np.testing.assert_allclose(known_shape, unknown_shape)


@pytest.mark.parametrize("white", [True, False])
def test_q_sqrt_constraints(
    inducing_points: InducingPoints, kernel: Kernel, mu: AnyNDArray, white: bool
) -> None:
    """Test that sending in an unconstrained q_sqrt returns the same conditional
    evaluation and gradients. This is important to match the behaviour of the KL, which
    enforces q_sqrt is triangular.
    """

    tril: AnyNDArray = np.tril(rng.randn(Ln, Nn, Nn))

    q_sqrt_constrained = Parameter(tril, transform=triangular())
    q_sqrt_unconstrained = Parameter(tril)

    diff_before_gradient_step = (q_sqrt_constrained - q_sqrt_unconstrained).numpy()
    assert_allclose(diff_before_gradient_step, 0)

    kls = []
    for q_sqrt in [q_sqrt_constrained, q_sqrt_unconstrained]:

        with tf.GradientTape() as tape:
            kl = prior_kl(inducing_points, kernel, mu, q_sqrt, whiten=white)

        grad = tape.gradient(kl, q_sqrt.unconstrained_variable)
        q_sqrt.unconstrained_variable.assign_sub(grad)
        kls.append(kl)

    diff_kls_before_gradient_step = kls[0] - kls[1]

    assert_allclose(diff_kls_before_gradient_step, 0)

    diff_after_gradient_step = (q_sqrt_constrained - q_sqrt_unconstrained).numpy()
    assert_allclose(diff_after_gradient_step, 0)


# tests/gpflow/test_logdensities.py

# Copyright 2018 the GPflow authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
import pytest
import scipy.stats
from numpy.random import randn
from numpy.testing import assert_allclose
from scipy.stats import multivariate_normal as mvn

from gpflow import logdensities
from gpflow.base import AnyNDArray
from gpflow.utilities import to_default_float

rng = np.random.RandomState(1)


@pytest.mark.parametrize("x, mu, var", [(0.9, 0.5, 1.3)])
def test_gaussian(x: float, mu: float, var: float) -> None:
    gpf = logdensities.gaussian(x, mu, var).numpy()
    sps = scipy.stats.norm(loc=mu, scale=np.sqrt(var)).logpdf(x)
    np.testing.assert_allclose(gpf, sps)


@pytest.mark.parametrize("x, mu, var", [(0.9, 0.5, 1.3)])
def test_lognormal(x: float, mu: float, var: float) -> None:
    gpf = logdensities.lognormal(x, mu, var).numpy()
    sps = scipy.stats.lognorm(s=np.sqrt(var), scale=np.exp(mu)).logpdf(x)
    np.testing.assert_allclose(gpf, sps)


@pytest.mark.parametrize(
    "x, p",
    [
        [1, 0.6],
        [0, 0.6],
    ],
)
def test_bernoulli(x: float, p: float) -> None:
    gpf = logdensities.bernoulli(x, p).numpy()
    sps = scipy.stats.bernoulli.logpmf(k=x, p=p)
    np.testing.assert_allclose(gpf, sps)


@pytest.mark.parametrize(
    "x, lam",
    [
        [0, 1.3],
        [1, 1.3],
        [2, 1.3],
    ],
)
def test_poisson(x: float, lam: float) -> None:
    gpf = logdensities.poisson(x, lam).numpy()
    sps = scipy.stats.poisson.logpmf(k=x, mu=lam)
    np.testing.assert_allclose(gpf, sps)


@pytest.mark.parametrize("x, scale", [(0.9, 1.3)])
def test_exponential(x: float, scale: float) -> None:
    gpf = logdensities.exponential(x, scale).numpy()
    sps = scipy.stats.expon(loc=0.0, scale=scale).logpdf(x)
    np.testing.assert_allclose(gpf, sps)


@pytest.mark.parametrize("x, shape, scale", [(0.9, 0.5, 1.3)])
def test_gamma(x: float, shape: float, scale: float) -> None:
    gpf = logdensities.gamma(x, shape, scale).numpy()
    sps = scipy.stats.gamma(a=shape, loc=0.0, scale=scale).logpdf(x)
    np.testing.assert_allclose(gpf, sps)


@pytest.mark.parametrize(
    "x, mean, scale, df",
    [
        (0.9, 0.5, 1.3, 1),
        (0.9, 0.5, 1.3, 2),
        (0.9, 0.5, 1.3, 3),
    ],
)
def test_student_t(x: float, mean: float, scale: float, df: int) -> None:
    cast = to_default_float
    gpf = logdensities.student_t(cast(x), cast(mean), cast(scale), df).numpy()
    sps = scipy.stats.t(df=df, loc=mean, scale=scale).logpdf(x)
    np.testing.assert_allclose(gpf, sps)


@pytest.mark.parametrize("x, alpha, beta", [(0.9, 0.5, 1.3)])
def test_beta(x: float, alpha: float, beta: float) -> None:
    gpf = logdensities.beta(x, alpha, beta).numpy()
    sps = scipy.stats.beta(a=alpha, b=beta).logpdf(x)
    np.testing.assert_allclose(gpf, sps)


@pytest.mark.parametrize("x, mu, sigma", [(0.9, 0.5, 1.3)])
def test_laplace(x: float, mu: float, sigma: float) -> None:
    gpf = logdensities.laplace(x, mu, sigma).numpy()
    sps = scipy.stats.laplace(loc=mu, scale=sigma).logpdf(x)
    np.testing.assert_allclose(gpf, sps)


@pytest.mark.parametrize("x", [randn(4, 10), randn(4, 1)])
@pytest.mark.parametrize("mu", [randn(4, 10), randn(4, 1)])
@pytest.mark.parametrize("cov_sqrt", [randn(4, 4), np.eye(4)])
def test_multivariate_normal(x: AnyNDArray, mu: AnyNDArray, cov_sqrt: AnyNDArray) -> None:
    cov = np.dot(cov_sqrt, cov_sqrt.T)
    L = np.linalg.cholesky(cov)

    gp_result = logdensities.multivariate_normal(x, mu, L)

    if mu.shape[1] > 1:
        if x.shape[1] > 1:
            sp_result = [mvn.logpdf(x[:, i], mu[:, i], cov) for i in range(mu.shape[1])]
        else:
            sp_result = [mvn.logpdf(x.ravel(), mu[:, i], cov) for i in range(mu.shape[1])]
    else:
        sp_result = mvn.logpdf(x.T, mu.ravel(), cov)
    assert_allclose(gp_result, sp_result)


# tests/gpflow/test_monitor.py

from pathlib import Path
from typing import Any, List, Optional, Sequence, Union

import numpy as np
import pytest
import tensorflow as tf
from _pytest.capture import CaptureFixture
from matplotlib.axes import Axes
from matplotlib.figure import Figure

import gpflow
from gpflow.experimental.check_shapes import check_shape as cs
from gpflow.experimental.check_shapes import check_shapes
from gpflow.models import GPR, GPModel
from gpflow.monitor import (
    ExecuteCallback,
    ImageToTensorBoard,
    ModelToTensorBoard,
    Monitor,
    MonitorTask,
    MonitorTaskGroup,
    ScalarToTensorBoard,
)


class Data:
    num_data = 20
    num_steps = 2


class DummyTask(MonitorTask):
    def run(self, **kwargs: Any) -> None:
        pass


class DummyStepCallback:
    current_step = 0

    @check_shapes(
        "variables[all]: [...]",
        "values[all]: [...]",
    )
    def callback(
        self, step: int, variables: Sequence[tf.Variable], values: Sequence[tf.Tensor]
    ) -> None:
        self.current_step = step


@pytest.fixture
@check_shapes()
def model() -> GPModel:
    data = (
        cs(np.random.randn(Data.num_data, 2), "[N, 2]"),
        cs(np.random.randn(Data.num_data, 2), "[N, 2]"),
    )
    kernel = gpflow.kernels.SquaredExponential(lengthscales=[1.0, 2.0])
    return GPR(data, kernel, noise_variance=0.01)


@pytest.fixture
def monitor(model: GPModel, tmp_path: Path) -> Monitor:
    tmp_path_str = str(tmp_path)

    @check_shapes(
        "return: []",
    )
    def lml_callback() -> tf.Tensor:
        return model.log_marginal_likelihood()

    def print_callback() -> None:
        print("foo")

    return Monitor(
        MonitorTaskGroup(
            [
                ModelToTensorBoard(tmp_path_str, model),
                ScalarToTensorBoard(tmp_path_str, lml_callback, "lml"),
            ],
            period=2,
        ),
        MonitorTaskGroup(ExecuteCallback(print_callback), period=1),
    )


def _get_size_directory(d: Path) -> int:
    """Calculating the size of a directory (in Bytes)."""
    return sum(f.stat().st_size for f in d.glob("**/*"))


# Smoke tests for the individual tasks
# #####################################


def test_ExecuteCallback() -> None:
    def callback() -> None:
        print("ExecuteCallback test")

    task = ExecuteCallback(callback)
    task(0)
    compiled_task = tf.function(task)
    compiled_task(0)


def test_ImageToTensorBoard(tmp_path: Path) -> None:
    """Smoke test `ImageToTensorBoard` in Eager and Compiled mode."""
    tmp_path_str = str(tmp_path)

    def plotting_cb(fig: Figure, axes: Axes) -> None:
        axes[0, 0].plot(np.random.randn(2), np.random.randn(2))
        axes[1, 0].plot(np.random.randn(2), np.random.randn(2))
        axes[0, 1].plot(np.random.randn(2), np.random.randn(2))
        axes[1, 1].plot(np.random.randn(2), np.random.randn(2))

    fig_kwargs = dict(figsize=(10, 10))
    subplots_kwargs = dict(sharex=True, nrows=2, ncols=2)
    task = ImageToTensorBoard(
        tmp_path_str, plotting_cb, "image", fig_kw=fig_kwargs, subplots_kw=subplots_kwargs
    )

    task(0)
    compiled_task = tf.function(task)
    compiled_task(0)


def test_ScalarToTensorBoard(tmp_path: Path) -> None:
    """Smoke test `ScalarToTensorBoard` in Eager and Compiled mode."""
    tmp_path_str = str(tmp_path)

    def scalar_cb() -> float:
        return 0.0

    task = ScalarToTensorBoard(tmp_path_str, scalar_cb, "scalar")
    task(0)
    compiled_task = tf.function(task)
    compiled_task(0)


def test_ScalarToTensorBoard_with_argument(tmp_path: Path) -> None:
    """Smoke test `ScalarToTensorBoard` in Eager and Compiled mode."""
    tmp_path_str = str(tmp_path)

    def scalar_cb(x: Optional[float] = None) -> float:
        assert x is not None
        return 2 * x

    task = ScalarToTensorBoard(tmp_path_str, scalar_cb, "scalar")
    compiled_task = tf.function(task)
    task(0, x=1.0)
    compiled_task(0, x=1.0)


def test_ScalarToTensorBoard_with_wrong_keyword_argument(tmp_path: Path) -> None:
    tmp_path_str = str(tmp_path)

    def scalar_cb(x: Optional[float] = None) -> float:
        assert x is not None
        return 2 * x

    task = ScalarToTensorBoard(tmp_path_str, scalar_cb, "scalar")
    compiled_task = tf.function(task)

    with pytest.raises(TypeError, match=r"got an unexpected keyword argument 'y'"):
        task(0, y=1.0)

    with pytest.raises(TypeError, match=r"got an unexpected keyword argument 'y'"):
        compiled_task(0, y=1.0)


def test_ModelToTensorboard(model: GPModel, tmp_path: Path) -> None:
    """Smoke test `ModelToTensorBoard` in Eager and Compiled mode."""
    tmp_path_str = str(tmp_path)
    task = ModelToTensorBoard(tmp_path_str, model)
    task(0)
    compiled_task = tf.function(task)
    compiled_task(0)


def test_ExecuteCallback_arguments(capsys: CaptureFixture[str]) -> None:
    def cb1(x: Optional[int] = None, **_: Any) -> None:
        assert x is not None
        print(x)

    def cb2(**_: Any) -> None:
        print(2)

    def cb3(y: Optional[int] = None, **_: Any) -> None:
        assert y is not None
        print(y)

    group1 = MonitorTaskGroup([ExecuteCallback(cb1), ExecuteCallback(cb2)])
    group2 = MonitorTaskGroup(ExecuteCallback(cb3))
    monitor = Monitor(group1, group2)
    monitor(0, x=1, y=3)
    out, _ = capsys.readouterr()
    assert out == "1\n2\n3\n"


# Smoke test Monitor and MonitorTaskGroup
# ########################################


def none() -> None:
    return None


@pytest.mark.parametrize(
    "task_or_tasks",
    [
        ExecuteCallback(none),
        [ExecuteCallback(none)],
        [ExecuteCallback(none), ExecuteCallback(none)],
    ],
)
def test_MonitorTaskGroup_and_Monitor(task_or_tasks: Union[MonitorTask, List[MonitorTask]]) -> None:
    group = MonitorTaskGroup(task_or_tasks, period=2)

    # check that the tasks is actually a list (custom setter)
    isinstance(group.tasks, list)

    # Smoke test the __call__
    group(0)
    compiled_group = tf.function(group)
    compiled_group(0)

    # Smoke test the Monitor wrapper
    monitor = Monitor(group)
    monitor(0)
    compiled_monitor = tf.function(monitor)
    compiled_monitor(0)


def test_Monitor(monitor: Monitor) -> None:
    monitor(0)
    compiled_monitor = tf.function(monitor)
    compiled_monitor(0)


# Functionality tests
# ###################


def test_compiled_execute_callable(capsys: CaptureFixture[str]) -> None:
    """
    Test that the `ExecuteCallback` when compiled behaves as expected.
    We test that python prints are not executed anymore.
    """
    string_to_print = "Eager mode"

    def callback() -> None:
        print(string_to_print)

    task = ExecuteCallback(callback)

    # Eager mode
    for i in range(Data.num_steps):
        task(i)
    out, _ = capsys.readouterr()

    # We expect a print for each step
    assert out == (f"{string_to_print}\n" * Data.num_steps)

    # Autograph mode
    compiled_task = tf.function(task)
    for i in tf.range(Data.num_steps):
        compiled_task(i)
    out, _ = capsys.readouterr()
    assert out == f"{string_to_print}\n"


def test_periodicity_group(capsys: CaptureFixture[str]) -> None:
    """Test that groups are called at different periods."""

    task_a = ExecuteCallback(lambda: print("a", end=" "))
    task_b = ExecuteCallback(lambda: print("b", end=" "))
    task_X = ExecuteCallback(lambda: print("X", end=" "))

    group_often = MonitorTaskGroup([task_a, task_b], period=1)
    group_seldom = MonitorTaskGroup([task_X], period=3)
    monitor = Monitor(group_often, group_seldom)
    for i in range(7):
        monitor(i)

    out, _ = capsys.readouterr()
    expected = "a b X a b a b a b X a b a b a b X "
    assert out == expected

    # AutoGraph mode
    compiled_monitor = tf.function(monitor)
    for i in tf.range(7):
        compiled_monitor(i)

    # When using TF's range and compiling the monitoring we only expected the python prints once.
    out, _ = capsys.readouterr()
    assert "a b X"


def test_logdir_created(monitor: Monitor, model: GPModel, tmp_path: Path) -> None:
    """
    Check that TensorFlow summaries are written.
    """
    # check existence
    assert tmp_path.is_dir()
    size_before = _get_size_directory(tmp_path)
    assert size_before > 0

    opt = tf.optimizers.Adam()
    for step in range(Data.num_steps):
        opt.minimize(model.training_loss, model.trainable_variables)
        monitor(step)

    size_after = _get_size_directory(tmp_path)
    assert size_after > size_before


def test_compile_monitor(monitor: Monitor, model: GPModel) -> None:
    opt = tf.optimizers.Adam()

    @tf.function
    def tf_func(step: tf.Tensor) -> None:
        opt.minimize(model.training_loss, model.trainable_variables)
        monitor(step)

    for step in tf.range(100):
        tf_func(step)


def test_scipy_monitor(monitor: Monitor, model: GPModel) -> None:
    opt = gpflow.optimizers.Scipy()

    opt.minimize(model.training_loss, model.trainable_variables, step_callback=monitor)


def test_scipy_monitor_called(model: GPModel) -> None:
    task = DummyTask()
    monitor = Monitor(MonitorTaskGroup(task, period=1))
    opt = gpflow.optimizers.Scipy()
    opt.minimize(model.training_loss, model.trainable_variables, step_callback=monitor)
    assert task.current_step > 1


def test_scipy_step_callback_called(model: GPModel) -> None:
    dsc = DummyStepCallback()
    opt = gpflow.optimizers.Scipy()
    opt.minimize(model.training_loss, model.trainable_variables, step_callback=dsc.callback)
    assert dsc.current_step > 1


# tests/integration/test_benchmark.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import json
from pathlib import Path

import pandas as pd
from _pytest.capture import CaptureFixture
from PIL import Image

from benchmark.benchmark_api import BENCHMARK_SUITES
from benchmark.run import run
from benchmark.sharding import ShardingStrategy


def test_benchmark(tmp_path: Path, capsys: CaptureFixture[str]) -> None:
    suite = BENCHMARK_SUITES.get("integration_test")
    dest = tmp_path / "dest"
    yes_subdir = False
    no_subdir = True
    cache_dir = tmp_path / "cache_dir"
    plot = True

    run(ShardingStrategy("start"), suite, dest, yes_subdir, cache_dir, plot)
    out_dir_str = capsys.readouterr().out.strip()
    out_dir = Path(out_dir_str)

    with open(out_dir / "metadata.json", "rt") as f:
        assert isinstance(json.load(f), dict)

    run(ShardingStrategy("0/3"), suite, out_dir, no_subdir, cache_dir, plot)
    with open(out_dir / "metadata_0of3.json", "rt") as f:
        assert isinstance(json.load(f), dict)
    assert isinstance(pd.read_csv(out_dir / "metrics_0of3.csv"), pd.DataFrame)

    run(ShardingStrategy("1/3"), suite, out_dir, no_subdir, cache_dir, plot)
    with open(out_dir / "metadata_1of3.json", "rt") as f:
        assert isinstance(json.load(f), dict)
    assert isinstance(pd.read_csv(out_dir / "metrics_1of3.csv"), pd.DataFrame)

    run(ShardingStrategy("2/3"), suite, out_dir, no_subdir, cache_dir, plot)
    with open(out_dir / "metadata_2of3.json", "rt") as f:
        assert isinstance(json.load(f), dict)
    assert isinstance(pd.read_csv(out_dir / "metrics_2of3.csv"), pd.DataFrame)

    run(ShardingStrategy("collect"), suite, out_dir, no_subdir, cache_dir, plot)
    assert isinstance(pd.read_csv(out_dir / "metrics.csv"), pd.DataFrame)
    assert isinstance(Image.open(out_dir / "metrics_tiny_linear.png"), Image.Image)
    assert isinstance(Image.open(out_dir / "metrics_tiny_sine.png"), Image.Image)


# tests/integration/test_docs.py

# Copyright 2022 the GPflow authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import subprocess
from pathlib import Path

import pytest

BUILD_DOCS_PATH = Path(__file__).parent.parent.parent / "doc" / "build_docs.py"


@pytest.mark.docs
def test_docs(tmp_path: Path) -> None:
    assert BUILD_DOCS_PATH.is_file()
    build_path = tmp_path / "build"

    try:
        subprocess.run(
            [
                "python",
                str(BUILD_DOCS_PATH),
                "develop",
                str(build_path),
                "--limit-notebooks",
                "--fail-on-warning",
            ]
        ).check_returncode()
    except subprocess.CalledProcessError as e:
        raise AssertionError(
            "Documentation build had errors / warnings."
            " Please fix."
            " Check both any .rst files you may have modified directly,"
            " and docstrings of your Python code."
        ) from e


# tests/integration/test_dynamic_shapes.py

# Copyright 2019 the GPflow authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
import pytest
import tensorflow as tf

import gpflow
from gpflow.base import AnyNDArray
from gpflow.config import default_float
from gpflow.experimental.check_shapes import ShapeChecker

rng = np.random.RandomState(0)


class Datum:
    cs = ShapeChecker().check_shape

    n_inputs = 1
    n_outputs = 2
    n_outputs_c = 1

    X: AnyNDArray = cs(rng.rand(20, n_inputs) * 10, "[N, n_inputs]")
    Y = cs(np.sin(X) + 0.9 * np.cos(X * 1.6) + rng.randn(*X.shape) * 0.8, "[N, n_outputs_c]")
    Y = cs(np.tile(Y, n_outputs), "[N, n_outputs]")  # identical columns
    Xtest: AnyNDArray = cs(rng.rand(10, n_inputs) * 10, "[N_new, n_inputs]")
    data = (X, Y)

    # for classification:
    Yc = cs(Y[:, :n_outputs_c], "[N, n_outputs_c]")
    cdata = (X, Yc)


def test_vgp() -> None:
    X = tf.Variable(
        tf.zeros((1, Datum.n_inputs), dtype=default_float()), shape=(None, None), trainable=False
    )
    Y = tf.Variable(
        tf.zeros((1, Datum.n_outputs), dtype=default_float()), shape=(None, None), trainable=False
    )

    model = gpflow.models.VGP(
        (X, Y),
        gpflow.kernels.SquaredExponential(),
        gpflow.likelihoods.Gaussian(),
        num_latent_gps=Datum.n_outputs,
    )

    @tf.function
    def model_closure() -> tf.Tensor:
        return -model.elbo()

    model_closure()  # Trigger compilation.

    gpflow.models.vgp.update_vgp_data(model, (Datum.X, Datum.Y))
    opt = gpflow.optimizers.Scipy()

    # simply test whether it runs without erroring...:
    opt.minimize(
        model_closure,
        variables=model.trainable_variables,
        options=dict(maxiter=3),
        compile=True,
    )


@pytest.mark.parametrize("whiten", [True, False])
@pytest.mark.parametrize("q_diag", [True, False])
def test_svgp(whiten: bool, q_diag: bool) -> None:
    model = gpflow.models.SVGP(
        gpflow.kernels.SquaredExponential(),
        gpflow.likelihoods.Gaussian(),
        inducing_variable=Datum.X.copy(),
        q_diag=q_diag,
        whiten=whiten,
        mean_function=gpflow.mean_functions.Constant(),
        num_latent_gps=Datum.n_outputs,
    )
    gpflow.set_trainable(model.inducing_variable, False)

    # test with explicitly unknown shapes:
    tensor_spec = tf.TensorSpec(shape=None, dtype=default_float())

    # lambda because: https://github.com/GPflow/GPflow/issues/1929
    elbo_lambda = lambda data: model.elbo(data)

    elbo = tf.function(
        elbo_lambda,
        input_signature=[(tensor_spec, tensor_spec)],
    )

    @tf.function
    def model_closure() -> tf.Tensor:
        return -elbo(Datum.data)

    model_closure()  # Trigger compilation.

    opt = gpflow.optimizers.Scipy()

    # simply test whether it runs without erroring...:
    opt.minimize(
        model_closure,
        variables=model.trainable_variables,
        options=dict(maxiter=3),
        compile=True,
    )


def test_vgp_multiclass() -> None:
    X = tf.Variable(
        tf.zeros((1, Datum.n_inputs), dtype=default_float()), shape=(None, None), trainable=False
    )
    Yc = tf.Variable(
        tf.zeros((1, Datum.n_outputs_c), dtype=default_float()), shape=(None, None), trainable=False
    )

    num_classes = 3
    model = gpflow.models.VGP(
        (X, Yc),
        gpflow.kernels.SquaredExponential(),
        gpflow.likelihoods.MultiClass(num_classes=num_classes),
        num_latent_gps=num_classes,
    )

    @tf.function
    def model_closure() -> tf.Tensor:
        return -model.elbo()

    model_closure()  # Trigger compilation.

    gpflow.models.vgp.update_vgp_data(model, (Datum.X, Datum.Yc))
    opt = gpflow.optimizers.Scipy()

    # simply test whether it runs without erroring...:
    opt.minimize(
        model_closure,
        variables=model.trainable_variables,
        options=dict(maxiter=3),
        compile=True,
    )


def test_svgp_multiclass() -> None:
    num_classes = 3
    model = gpflow.models.SVGP(
        gpflow.kernels.SquaredExponential(),
        gpflow.likelihoods.MultiClass(num_classes=num_classes),
        inducing_variable=Datum.X.copy(),
        num_latent_gps=num_classes,
    )
    gpflow.set_trainable(model.inducing_variable, False)

    # lambda because: https://github.com/GPflow/GPflow/issues/1929
    elbo_lambda = lambda data: model.elbo(data)

    # test with explicitly unknown shapes:
    tensor_spec = tf.TensorSpec(shape=None, dtype=default_float())
    elbo = tf.function(
        elbo_lambda,
        input_signature=[(tensor_spec, tensor_spec)],
    )

    @tf.function
    def model_closure() -> tf.Tensor:
        return -elbo(Datum.cdata)

    model_closure()  # Trigger compilation.

    opt = gpflow.optimizers.Scipy()

    # simply test whether it runs without erroring...:
    opt.minimize(
        model_closure,
        variables=model.trainable_variables,
        options=dict(maxiter=3),
        compile=True,
    )


# tests/integration/test_linear_noise.py

# Copyright 2019 the GPflow authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Callable

import numpy as np
import pytest

import gpflow
from gpflow import default_float
from gpflow.base import AnyNDArray, RegressionData
from gpflow.config import default_float
from gpflow.experimental.check_shapes import ShapeChecker
from gpflow.functions import Linear
from gpflow.kernels import Kernel
from gpflow.likelihoods import Gaussian
from gpflow.models import GPR, GPRFITC, SGPR, SVGP, VGP, GPModel, training_loss_closure
from gpflow.models.util import InducingPointsLike


class Datum:
    cs = ShapeChecker().check_shape

    rng = np.random.default_rng(20220630)
    n = 100
    X: AnyNDArray = cs(rng.random((n, 1), dtype=default_float()), "[N, 1]")
    noise_slope = -0.7
    noise_offset = 0.7
    noise = cs(
        (noise_slope * X + noise_offset) * rng.standard_normal((n, 1), dtype=default_float()),
        "[N, 1]",
    )
    Y = cs(np.sin(5 * X) + noise, "[N, 1]")
    data = X, Y


def create_kernel() -> Kernel:
    return gpflow.kernels.RBF(lengthscales=0.2)


def create_inducing() -> InducingPointsLike:
    Z = np.linspace(0.0, 1.0, 10)[:, None]
    iv = gpflow.inducing_variables.InducingPoints(Z)
    gpflow.set_trainable(iv.Z, False)
    return iv


def create_linear_noise() -> Gaussian:
    return Gaussian(scale=Linear())


def gpr(data: RegressionData) -> GPModel:
    return GPR(
        data,
        kernel=create_kernel(),
        likelihood=create_linear_noise(),
    )


def vgp(data: RegressionData) -> GPModel:
    return VGP(
        data,
        kernel=create_kernel(),
        likelihood=create_linear_noise(),
    )


def sgpr(data: RegressionData) -> GPModel:
    return SGPR(
        data,
        kernel=create_kernel(),
        inducing_variable=create_inducing(),
        likelihood=create_linear_noise(),
    )


def gprfitc(data: RegressionData) -> GPModel:
    return GPRFITC(
        data,
        kernel=create_kernel(),
        inducing_variable=create_inducing(),
        likelihood=create_linear_noise(),
    )


def svgp(data: RegressionData) -> GPModel:
    return SVGP(
        kernel=create_kernel(),
        likelihood=create_linear_noise(),
        inducing_variable=create_inducing(),
    )


CREATE_MODELS = (
    gpr,
    vgp,
    sgpr,
    gprfitc,
    svgp,
)


@pytest.mark.parametrize("create_model", CREATE_MODELS)
def test_infer_noise(create_model: Callable[[RegressionData], GPModel]) -> None:
    model = create_model(Datum.data)
    gpflow.optimizers.Scipy().minimize(
        training_loss_closure(model, Datum.data),
        variables=model.trainable_variables,
    )

    noise_scale = model.likelihood.scale
    np.testing.assert_allclose(Datum.noise_slope, noise_scale.A, atol=0.1)
    np.testing.assert_allclose(Datum.noise_offset, noise_scale.b, atol=0.1)


# tests/integration/test_method_equivalence.py

# Copyright 2019 the GPflow authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Optional, Tuple

import numpy as np
import pytest
from numpy.testing import assert_allclose

import gpflow
from gpflow.base import AnyNDArray, TensorType
from gpflow.config import default_jitter
from gpflow.experimental.check_shapes import ShapeChecker
from gpflow.kernels import Kernel
from gpflow.likelihoods import Likelihood
from gpflow.mean_functions import Constant
from gpflow.models import GPModel, maximum_log_likelihood_objective, training_loss_closure

rng = np.random.RandomState(0)


class Datum:
    cs = ShapeChecker().check_shape

    X: AnyNDArray = cs(rng.rand(20, 1) * 10, "[N, D]")
    Y = cs(np.sin(X) + 0.9 * np.cos(X * 1.6) + rng.randn(*X.shape) * 0.8, "[N, 1]")
    Y = cs(np.tile(Y, 2), "[N, P]")  # two identical columns
    Xtest: AnyNDArray = cs(rng.rand(10, 1) * 10, "[N_new, D]")
    data = (X, Y)


class DatumVGP:
    N, Ns, DX, DY = 100, 10, 2, 2
    np.random.seed(1)
    X = np.random.randn(N, DX)
    Xs = np.random.randn(Ns, DX)
    Y = np.random.randn(N, DY)
    q_mu = np.random.randn(N, DY)
    q_sqrt = np.random.randn(DY, N, N)
    q_alpha: AnyNDArray = np.random.randn(N, DX)
    q_lambda: AnyNDArray = np.random.randn(N, DX) ** 2
    data = (X, Y)


def _create_full_gp_model() -> gpflow.models.GPR:
    """
    GP Regression
    """
    full_gp_model = gpflow.models.GPR(
        (Datum.X, Datum.Y),
        kernel=gpflow.kernels.SquaredExponential(),
        mean_function=gpflow.mean_functions.Constant(),
    )
    return full_gp_model


def _create_approximate_models() -> Tuple[
    gpflow.models.VGP,
    gpflow.models.SVGP,
    gpflow.models.SVGP,
    gpflow.models.SGPR,
    gpflow.models.GPRFITC,
]:
    """
    1) Variational GP (with the likelihood set to Gaussian)
    2) Sparse variational GP (likelihood is Gaussian, inducing points
       at the data)
    3) Sparse variational GP (as above, but with the whitening rotation
       of the inducing variables)
    4) Sparse variational GP Regression (as above, but there the inducing
       variables are 'collapsed' out, as in Titsias 2009)
    5) FITC Sparse GP Regression
    """
    model_1 = gpflow.models.VGP(
        (Datum.X, Datum.Y),
        kernel=gpflow.kernels.SquaredExponential(),
        likelihood=gpflow.likelihoods.Gaussian(),
        mean_function=gpflow.mean_functions.Constant(),
    )

    model_2 = gpflow.models.SVGP(
        kernel=gpflow.kernels.SquaredExponential(),
        likelihood=gpflow.likelihoods.Gaussian(),
        inducing_variable=Datum.X.copy(),
        q_diag=False,
        whiten=False,
        mean_function=gpflow.mean_functions.Constant(),
        num_latent_gps=Datum.Y.shape[1],
    )
    gpflow.set_trainable(model_2.inducing_variable, False)

    model_3 = gpflow.models.SVGP(
        kernel=gpflow.kernels.SquaredExponential(),
        likelihood=gpflow.likelihoods.Gaussian(),
        inducing_variable=Datum.X.copy(),
        q_diag=False,
        whiten=True,
        mean_function=gpflow.mean_functions.Constant(),
        num_latent_gps=Datum.Y.shape[1],
    )
    gpflow.set_trainable(model_3.inducing_variable, False)

    model_4 = gpflow.models.SGPR(
        (Datum.X, Datum.Y),
        kernel=gpflow.kernels.SquaredExponential(),
        inducing_variable=Datum.X.copy(),
        mean_function=Constant(),
    )
    gpflow.set_trainable(model_4.inducing_variable, False)

    model_5 = gpflow.models.GPRFITC(
        (Datum.X, Datum.Y),
        kernel=gpflow.kernels.SquaredExponential(),
        inducing_variable=Datum.X.copy(),
        mean_function=Constant(),
    )
    gpflow.set_trainable(model_5.inducing_variable, False)

    return model_1, model_2, model_3, model_4, model_5


def _create_vgp_model(
    kernel: Kernel,
    likelihood: Likelihood,
    q_mu: Optional[TensorType] = None,
    q_sqrt: Optional[TensorType] = None,
) -> gpflow.models.VGP:
    model_vgp = gpflow.models.VGP((DatumVGP.X, DatumVGP.Y), kernel, likelihood)
    if q_mu is not None and q_sqrt is not None:
        model_vgp.q_mu.assign(q_mu)
        model_vgp.q_sqrt.assign(q_sqrt)
    return model_vgp


def _create_vgpao_model(
    kernel: Kernel,
    likelihood: Likelihood,
    q_alpha: Optional[TensorType],
    q_lambda: Optional[TensorType],
) -> gpflow.models.VGPOpperArchambeau:
    model_vgpoa = gpflow.models.VGPOpperArchambeau(
        (DatumVGP.X, DatumVGP.Y), kernel, likelihood, num_latent_gps=DatumVGP.DY
    )
    model_vgpoa.q_alpha.assign(q_alpha)
    model_vgpoa.q_lambda.assign(q_lambda)

    return model_vgpoa


def _create_svgp_model(
    kernel: Kernel,
    likelihood: Likelihood,
    q_mu: Optional[TensorType],
    q_sqrt: Optional[TensorType],
    whiten: bool,
) -> gpflow.models.SVGP:
    model_svgp = gpflow.models.SVGP(
        kernel,
        likelihood,
        DatumVGP.X.copy(),
        whiten=whiten,
        q_diag=False,
        num_latent_gps=DatumVGP.DY,
    )
    model_svgp.q_mu.assign(q_mu)
    model_svgp.q_sqrt.assign(q_sqrt)
    return model_svgp


@pytest.mark.parametrize("approximate_model", _create_approximate_models())
def test_equivalence(approximate_model: GPModel) -> None:
    """
    With a Gaussian likelihood, and inducing points (where appropriate)
    positioned at the data, many of the gpflow methods are equivalent (perhaps
    subject to some optimization).
    """

    def optimize(model: GPModel) -> None:
        opt = gpflow.optimizers.Scipy()
        loss = training_loss_closure(model, Datum.data)
        opt.minimize(loss, model.trainable_variables, options=dict(maxiter=3000))
        if isinstance(model, gpflow.models.SVGP) and not model.whiten:
            # The (S)VGP model in non-whitened representation has significantly
            # worse optimization behaviour. To get the tests to pass, we need
            # to optimize much harder: we set ftol=gtol=0.0 to enforce
            # continued optimization.
            opt.minimize(
                loss, model.trainable_variables, options=dict(maxiter=7000, ftol=0.0, gtol=0.0)
            )

    gpr_model = _create_full_gp_model()
    optimize(gpr_model)
    optimize(approximate_model)

    gpr_likelihood = gpr_model.log_marginal_likelihood()
    approximate_likelihood = maximum_log_likelihood_objective(approximate_model, Datum.data)
    assert_allclose(approximate_likelihood, gpr_likelihood, rtol=1e-6)

    gpr_kernel_ls = gpr_model.kernel.lengthscales.numpy()
    gpr_kernel_var = gpr_model.kernel.variance.numpy()

    approximate_kernel_ls = approximate_model.kernel.lengthscales.numpy()
    approximate_kernel_var = approximate_model.kernel.variance.numpy()

    assert_allclose(gpr_kernel_ls, approximate_kernel_ls, 1e-4)
    assert_allclose(gpr_kernel_var, approximate_kernel_var, 1e-3)

    gpr_mu, gpr_var = gpr_model.predict_y(Datum.Xtest)
    approximate_mu, approximate_var = approximate_model.predict_y(Datum.Xtest)

    assert_allclose(gpr_mu, approximate_mu, 1e-3)
    assert_allclose(gpr_var, approximate_var, 1e-4)


def test_equivalence_vgp_and_svgp() -> None:
    kernel = gpflow.kernels.Matern52()
    likelihood = gpflow.likelihoods.StudentT()

    svgp_model = _create_svgp_model(kernel, likelihood, DatumVGP.q_mu, DatumVGP.q_sqrt, whiten=True)
    vgp_model = _create_vgp_model(kernel, likelihood, DatumVGP.q_mu, DatumVGP.q_sqrt)

    likelihood_svgp = svgp_model.elbo(DatumVGP.data)
    likelihood_vgp = vgp_model.elbo()
    assert_allclose(likelihood_svgp, likelihood_vgp, rtol=1e-2)

    svgp_mu, svgp_var = svgp_model.predict_f(DatumVGP.Xs)
    vgp_mu, vgp_var = vgp_model.predict_f(DatumVGP.Xs)

    assert_allclose(svgp_mu, vgp_mu)
    assert_allclose(svgp_var, vgp_var)


def test_equivalence_vgp_and_opper_archambeau() -> None:
    kernel = gpflow.kernels.Matern52()
    likelihood = gpflow.likelihoods.StudentT()

    vgp_oa_model = _create_vgpao_model(kernel, likelihood, DatumVGP.q_alpha, DatumVGP.q_lambda)

    K = kernel(DatumVGP.X) + np.eye(DatumVGP.N) * default_jitter()
    L = np.linalg.cholesky(K)
    L_inv = np.linalg.inv(L)
    K_inv = np.linalg.inv(K)

    mean = K @ DatumVGP.q_alpha

    prec_dnn = K_inv[None, :, :] + np.array([np.diag(l ** 2) for l in DatumVGP.q_lambda.T])
    var_dnn = np.linalg.inv(prec_dnn)

    svgp_model_unwhitened = _create_svgp_model(
        kernel, likelihood, mean, np.linalg.cholesky(var_dnn), whiten=False
    )

    mean_white_nd = L_inv.dot(mean)
    var_white_dnn = np.einsum("nN,dNM,mM->dnm", L_inv, var_dnn, L_inv)
    q_sqrt_nnd = np.linalg.cholesky(var_white_dnn)

    vgp_model = _create_vgp_model(kernel, likelihood, mean_white_nd, q_sqrt_nnd)

    likelihood_vgp = vgp_model.elbo()
    likelihood_vgp_oa = vgp_oa_model.elbo()
    likelihood_svgp_unwhitened = svgp_model_unwhitened.elbo(DatumVGP.data)

    assert_allclose(likelihood_vgp, likelihood_vgp_oa, rtol=1e-2)
    assert_allclose(likelihood_vgp, likelihood_svgp_unwhitened, rtol=1e-2)

    vgp_oa_mu, vgp_oa_var = vgp_oa_model.predict_f(DatumVGP.Xs)
    svgp_unwhitened_mu, svgp_unwhitened_var = svgp_model_unwhitened.predict_f(DatumVGP.Xs)
    vgp_mu, vgp_var = vgp_model.predict_f(DatumVGP.Xs)

    assert_allclose(vgp_oa_mu, vgp_mu)
    assert_allclose(vgp_oa_var, vgp_var, rtol=1e-4)  # jitter?
    assert_allclose(svgp_unwhitened_mu, vgp_mu)
    assert_allclose(svgp_unwhitened_var, vgp_var, rtol=1e-4)


class DatumUpper:
    cs = ShapeChecker().check_shape

    rng = np.random.default_rng(123)
    X = cs(rng.random((100, 1)), "[N, D]")
    Y = cs(np.sin(1.5 * 2 * np.pi * X) + rng.standard_normal(X.shape) * 0.1 + 5.3, "[N, P]")
    assert Y.mean() > 5.0, "offset ensures a regression test against the bug fixed by PR #1560"
    data = (X, Y)


def test_upper_bound_few_inducing_points() -> None:
    """
    Test for upper bound for regression marginal likelihood
    """
    model_vfe = gpflow.models.SGPR(
        (DatumUpper.X, DatumUpper.Y),
        gpflow.kernels.SquaredExponential(),
        inducing_variable=DatumUpper.X[:10, :].copy(),
        mean_function=Constant(),
        likelihood=gpflow.likelihoods.Gaussian(scale=gpflow.functions.Linear()),
    )
    opt = gpflow.optimizers.Scipy()

    opt.minimize(
        model_vfe.training_loss,
        variables=model_vfe.trainable_variables,
        options=dict(maxiter=500),
    )

    full_gp = gpflow.models.GPR(
        (DatumUpper.X, DatumUpper.Y),
        kernel=model_vfe.kernel,
        mean_function=model_vfe.mean_function,
        likelihood=model_vfe.likelihood,
    )

    lml_upper = model_vfe.upper_bound()
    lml_vfe = model_vfe.elbo()
    lml_full_gp = full_gp.log_marginal_likelihood()

    assert lml_vfe < lml_full_gp
    assert lml_full_gp < lml_upper


# tests/integration/test_model_serialization.py

# Copyright 2022 the GPflow authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pathlib import Path

import numpy as np
import tensorflow as tf

import gpflow


def test_model_serialization(tmp_path: Path) -> None:
    rng = np.random.default_rng(1234)
    X = rng.random((5, 1))
    Y = rng.standard_normal(X.shape)
    kernel = gpflow.kernels.SquaredExponential()
    model = gpflow.models.GPR((X, Y), kernel=kernel)

    frozen_model = gpflow.utilities.freeze(model)
    module_out = tf.Module()
    module_out.predict = tf.function(
        frozen_model.predict_y,
        input_signature=[tf.TensorSpec(shape=[None, 1], dtype=tf.float64)],
    )

    tf.saved_model.save(module_out, str(tmp_path))

    module_in = tf.saved_model.load(str(tmp_path))

    Xnew = rng.random((3, 1))
    np.testing.assert_allclose(
        model.predict_y(Xnew),
        module_in.predict(Xnew),
    )


# tests/integration/test_notebooks.py

# Copyright 2017 the GPflow authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import itertools
import sys
import traceback
from pathlib import Path
from typing import Sequence, Set

import jupytext
import nbformat
import pytest
from nbconvert.preprocessors import ExecutePreprocessor
from nbconvert.preprocessors.execute import CellExecutionError

_NOTEBOOK_DIR = (Path(__file__) / "../../../doc/sphinx/notebooks").resolve()


def test_notebook_dir_exists() -> None:
    assert _NOTEBOOK_DIR.is_dir()


# To ignore a notebook, add its full base name (including .pct.py or .md
# extension, but without any directory component). NOTE that if there are
# several notebooks in different directories with the same base name, they will
# all get ignored (change the ignoring check to something else in that
# case, if need be!)
IGNORED_NOTEBOOKS: Set[str] = set()


def get_notebooks() -> Sequence[Path]:
    """
    Returns all notebooks in `_nbpath` that are not ignored.
    """

    def notebook_ignored(nb: Path) -> bool:
        return nb.name in IGNORED_NOTEBOOKS

    # recursively traverse the notebook directory in search for ipython notebooks
    all_py_notebooks = _NOTEBOOK_DIR.glob("**/*.pct.py")
    all_md_notebooks = _NOTEBOOK_DIR.glob("**/*.md")
    all_notebooks = itertools.chain(all_md_notebooks, all_py_notebooks)
    notebooks_to_test = [nb for nb in all_notebooks if not notebook_ignored(nb)]
    return notebooks_to_test


def _preproc() -> ExecutePreprocessor:
    pythonkernel = "python" + str(sys.version_info[0])
    return ExecutePreprocessor(timeout=300, kernel_name=pythonkernel, interrupt_on_timeout=True)


def _exec_notebook(notebook_path: Path) -> None:
    with notebook_path.open() as notebook_file:
        nb = jupytext.read(notebook_file, as_version=nbformat.current_nbformat)
        try:
            meta_data = {"path": str(notebook_path.parent)}
            _preproc().preprocess(nb, {"metadata": meta_data})
        except CellExecutionError as cell_error:
            traceback.print_exc(file=sys.stdout)
            msg = "Error executing the notebook {0}. See above for error.\nCell error: {1}"
            pytest.fail(msg.format(str(notebook_path), str(cell_error)))


@pytest.mark.notebooks
@pytest.mark.parametrize("notebook_path", get_notebooks(), ids=str)
def test_notebook(notebook_path: Path) -> None:
    _exec_notebook(notebook_path)


def test_has_notebooks() -> None:
    assert len(get_notebooks()) >= 30, "there are probably some notebooks that were not discovered"


# benchmark/benchmark_api.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Classes and other infrastructure for defining which benchmarks to run and what plots to create.

See ``benchmarks.py`` for concrete instances of these.
"""
from dataclasses import dataclass
from typing import Any, Collection, Dict, List, Optional, Tuple

import pandas as pd

from benchmark.dataset_api import DatasetFactory
from benchmark.grouping import GroupingKey, GroupingSpec
from benchmark.model_api import ModelFactory
from benchmark.plotter_api import Plotter
from benchmark.registry import Named, Registry


@dataclass(unsafe_hash=True)
class BenchmarkTask:
    """
    Definition of a single benchmark to run.

    This class is an implementation detail and needs to be a simple datastructure that is easy to
    (de-)serialise.
    """

    dataset_name: str
    """ Name of datasets to benchmark against. """

    model_name: str
    """ Name of models to benchmark. """

    do_compile: bool
    """ Whether to use ``tf.function``. """

    do_optimise: bool
    """ Whether to train the models. """

    do_predict: bool
    """ Whether to benchmark ``model.predict_f``. """

    do_posterior: bool
    """ Whether to benchmark ``model.posterior()``. """

    repetitions: int
    """ Number of times to repeat benchmarks, to estimate noise. """

    @property
    def name(self) -> str:
        tokens = []

        def print_bool(name: str) -> None:
            value_str = "T" if getattr(self, name) else "F"
            tokens.append(f"{name}={value_str}")

        tokens.append(self.dataset_name)
        tokens.append(self.model_name)
        print_bool("do_compile")
        print_bool("do_optimise")
        print_bool("do_predict")
        print_bool("do_posterior")

        return "/".join(tokens)


@dataclass(frozen=True)
class BenchmarkSet:
    """
    Defines a set of plots to produce.

    This will create plots for the cartesian product of the parameters in this class.
    See :class:`BenchmarkSuite` if you want a union of parameters.
    """

    name: str
    """ Name of this benchmark set. """

    datasets: Collection[DatasetFactory]
    """ Datasets to benchmark against. """

    models: Collection[ModelFactory]
    """ Models to benchmark. """

    plots: Collection[Plotter]
    """ Plots to generate. """

    do_compile: Collection[bool]
    """ Whether to use ``tf.function``. """

    do_optimise: Collection[bool]
    """ Whether to train the models. """

    do_predict: bool
    """ Whether to benchmark ``model.predict_f``. """

    do_posterior: bool
    """ Whether to benchmark ``model.posterior()``. """

    file_by: GroupingSpec
    """ How to split plots into different ``.png`` files. """

    column_by: GroupingSpec
    """ How to split plots into different columns within their file. """

    row_by: GroupingSpec
    """ How to split plots into different rows within their file. """

    line_by: Optional[GroupingSpec]
    """
    How to split data into different lines within a plot.

    If ``None`` data will be split by all columns not used by any of the other ``GroupingSpec``\s.
    """

    repetitions: int = 1
    """ Number of times to repeat benchmarks, to estimate noise. """

    def __post_init__(self) -> None:
        def has_unique_values(values: Collection[Any]) -> bool:
            return len(values) == len(set(values))

        def assert_unique_names(attr: str) -> None:
            names = [v.name for v in getattr(self, attr)]
            assert has_unique_values(names), f"'{attr}' must have unique names. Found: {names}."

        assert_unique_names("datasets")
        assert_unique_names("models")
        assert_unique_names("plots")

        def assert_unique_values(attr: str) -> None:
            values = getattr(self, attr)
            assert has_unique_values(values), f"'{attr}' must have unique values. Found: {values}."

        assert_unique_values("do_compile")
        assert_unique_values("do_optimise")

        def assert_disjoint_by(attr1: str, attr2: str) -> None:
            values1 = getattr(self, attr1).by
            values2 = getattr(self, attr2).by
            assert not (set(values1) & set(values2)), (
                f"'{attr1}.by' and '{attr2}.by' must be disjoint."
                f" Found: {attr1}.by={values1} and {attr2}.by={values2}."
            )

        assert_disjoint_by("file_by", "column_by")
        assert_disjoint_by("file_by", "row_by")
        assert_disjoint_by("column_by", "row_by")
        if self.line_by:
            assert_disjoint_by("file_by", "line_by")
            assert_disjoint_by("column_by", "line_by")
            assert_disjoint_by("row_by", "line_by")

        def assert_grouping_by(by: GroupingKey) -> None:
            assert (by in self.file_by.by) or (by in self.column_by.by) or (by in self.row_by.by), (
                f"Must group by '{by}' above the 'line' level. Found:"
                f" file_by={self.file_by.by},"
                f" column_by={self.column_by.by},"
                f" row_by={self.row_by.by}."
            )

        assert_grouping_by(GroupingKey.METRIC)
        assert_grouping_by(GroupingKey.PLOTTER)

    @property
    def safe_line_by(self) -> GroupingSpec:
        """
        Get ``line_by``, or a default value if ``line_by`` is ``None``.
        """
        if self.line_by is None:
            used_by = set(self.file_by.by) | set(self.column_by.by) | set(self.row_by.by)
            line_by = set(GroupingKey) - used_by
            sorted_line_by = sorted(line_by, key=lambda k: k.key_cost)
            return GroupingSpec(sorted_line_by, minimise=True)

        return self.line_by

    def get_tasks(self) -> Collection[BenchmarkTask]:
        """
        Compute ``BenchmarkTask`` objects for the cartesian product of the parameters of this
        object.
        """
        result: List[BenchmarkTask] = []
        for dataset in self.datasets:
            dataset_name = dataset.name
            for model in self.models:
                if not model.dataset_req.satisfied(dataset.tags):
                    continue

                model_name = model.name
                for do_compile in self.do_compile:
                    for do_optimise in self.do_optimise:
                        result.append(
                            BenchmarkTask(
                                dataset_name,
                                model_name,
                                do_compile,
                                do_optimise,
                                self.do_predict,
                                self.do_posterior,
                                self.repetitions,
                            )
                        )
        return result

    def filter_metrics(self, metrics_df: pd.DataFrame) -> pd.DataFrame:
        """
        Filter a dataframe for any metrics that are not relevant to this :class:`BenchmarkSet`.
        """
        dataset_names = set(d.name for d in self.datasets)
        model_names = set(m.name for m in self.models)
        return metrics_df[
            metrics_df.dataset.isin(dataset_names)
            & metrics_df.model.isin(model_names)
            & metrics_df.do_compile.isin(self.do_compile)
            & metrics_df.do_optimise.isin(self.do_optimise)
            & (metrics_df.repetition < self.repetitions)
        ]


@dataclass(frozen=True)
class BenchmarkSuite(Named):
    """
    A union of :class:`BenchmarkSet`\s.

    This is the main definition of work in this framework.

    This will intelligently merge tasks, so if a task is present in multiple benchmark sets, it is
    executed once, and the result reused.
    """

    name: str
    description: str
    sets: Collection[BenchmarkSet]

    def __post_init__(self) -> None:
        set_names = [d.name for d in self.sets]
        assert len(set_names) == len(
            set(set_names)
        ), f"Benchmark sets must have unique names. Got: {set_names}"

    def get_tasks(self) -> Collection[BenchmarkTask]:
        """
        Computes the minimal number of benchmarks to run for this suite.
        """
        result: Dict[Tuple[str, str, bool, bool], BenchmarkTask] = {}
        for benchmark_set in self.sets:
            for task in benchmark_set.get_tasks():
                key = (task.dataset_name, task.model_name, task.do_compile, task.do_optimise)
                if key in result:
                    old_task = result[key]
                    old_task.do_predict |= task.do_predict
                    old_task.do_posterior |= task.do_posterior
                    old_task.repetitions = max(old_task.repetitions, task.repetitions)
                else:
                    result[key] = task
        return list(result.values())


BENCHMARK_SUITES: Registry[BenchmarkSuite] = Registry()


def make_benchmark_suite(**kwargs: Any) -> BenchmarkSuite:
    """
    Create a new :class:`BenchmarkSuite` and register it.

    :param kwargs: Passed on to ``BenchmarkSuite.__init__``.
    """
    result = BenchmarkSuite(**kwargs)
    BENCHMARK_SUITES.add(result)
    return result


# benchmark/benchmarks.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Definitions of benchmark suites.

A benchmark suite defines which datasets and models to use; and how to plot the results.
"""
import benchmark.datasets as ds
import benchmark.models as md
import benchmark.plotters as pl
from benchmark.benchmark_api import BenchmarkSet, make_benchmark_suite
from benchmark.dataset_api import DATASET_FACTORIES
from benchmark.grouping import GroupingKey as GK
from benchmark.grouping import GroupingSpec
from benchmark.model_api import MODEL_FACTORIES
from benchmark.plotter_api import PLOTTERS

make_benchmark_suite(
    name="integration_test",
    description="Suite used in tests/integration/test_benchmark.py.",
    sets=[
        BenchmarkSet(
            name="metrics",
            datasets=[
                ds.tiny_linear,
                ds.tiny_sine,
            ],
            models=[
                md.gpr,
                md.svgp,
            ],
            plots=[
                pl.metrics_box_plot,
            ],
            do_compile=[True],
            do_optimise=[True],
            do_predict=True,
            do_posterior=True,
            file_by=GroupingSpec((GK.DATASET,), minimise=False),
            row_by=GroupingSpec((GK.METRIC,), minimise=False),
            column_by=GroupingSpec((GK.PLOTTER,), minimise=False),
            line_by=None,
            repetitions=2,
        )
    ],
)


make_benchmark_suite(
    name="ci",
    description="Suite that is run in our CI pipeline, to monitor long-term performance.",
    sets=[
        BenchmarkSet(
            name="timeline",
            datasets=[
                ds.boston,
            ],
            models=[
                md.gpr,
                md.sgpr,
                md.vgp,
                md.svgp,
            ],
            plots=[
                pl.time_line,
            ],
            do_compile=[True],
            do_optimise=[True],
            do_predict=True,
            do_posterior=True,
            file_by=GroupingSpec((GK.DATASET, GK.PLOTTER), minimise=True),
            row_by=GroupingSpec((GK.METRIC,), minimise=False),
            column_by=GroupingSpec((), minimise=False),
            line_by=GroupingSpec((GK.MODEL,), minimise=False),
            repetitions=5,
        )
    ],
)


make_benchmark_suite(
    name="full",
    description="Suite that runs everything.",
    sets=[
        BenchmarkSet(
            name="metrics",
            datasets=DATASET_FACTORIES.all(),
            models=MODEL_FACTORIES.all(),
            plots=PLOTTERS.all(),
            do_compile=[False, True],
            do_optimise=[False, True],
            do_predict=True,
            do_posterior=True,
            file_by=GroupingSpec((GK.DATASET,), minimise=False),
            row_by=GroupingSpec((GK.METRIC,), minimise=False),
            column_by=GroupingSpec((GK.PLOTTER,), minimise=False),
            line_by=None,
            repetitions=5,
        )
    ],
)


# benchmark/dataset_api.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Classes and other infrastructure for creating datasets.

For concrete instances of datasets see ``datasets.py``.

Code adapted from:
https://github.com/hughsalimbeni/bayesian_benchmarks/blob/master/bayesian_benchmarks/data.py
"""
import zipfile
from abc import ABC, abstractmethod
from dataclasses import dataclass
from functools import wraps
from pathlib import Path
from typing import AbstractSet, Callable, Optional, Tuple
from urllib.parse import urlparse
from urllib.request import urlopen

import numpy as np
import pandas as pd

from benchmark.registry import TaggedRegistry
from benchmark.tag import Tag, TagReq
from gpflow.base import AnyNDArray
from gpflow.experimental.check_shapes import check_shapes


class DatasetTag(Tag["DatasetTag"]):
    """
    A tag that can be applied to a dataset.
    """


DatasetReq = TagReq[DatasetTag]

REGRESSION = DatasetTag("REGRESSION")
CLASSIFICATION = DatasetTag("CLASSIFICATION")
TINY = DatasetTag("TINY")
MEDIUM = DatasetTag("MEDIUM")
LARGE = DatasetTag("LARGE")
SYNTHETIC = DatasetTag("SYNTHETIC")
REAL_DATA = DatasetTag("REAL_DATA")


@dataclass
class XYData:
    """ Data with corresponding X and Y points. """

    tags: AbstractSet[DatasetTag]
    """ Tags signalling the properties of this data. """

    X: AnyNDArray
    """ Inputs. """

    Y: AnyNDArray
    """ Outputs. """

    @check_shapes(
        "self.X: [N, D]",
        "self.Y: [N, P]",
    )
    def __post_init__(self) -> None:
        pass

    @property
    def N(self) -> int:
        """ Number of rows of data. """
        return self.X.shape[0]  # type: ignore[no-any-return]

    @property
    def D(self) -> int:
        """ Number of inputs. """
        return self.X.shape[1]  # type: ignore[no-any-return]

    @property
    def P(self) -> int:
        """ Number of outputs. """
        return self.Y.shape[1]  # type: ignore[no-any-return]

    @property
    def XY(self) -> Tuple[AnyNDArray, AnyNDArray]:
        """ This data as a (X, Y) tuple. """
        return (self.X, self.Y)


@dataclass
class Dataset:
    """ A set of both training and testing data. """

    name: str
    """ Unique name of this dataset. """

    tags: AbstractSet[DatasetTag]
    """ Tags signalling the properties of this data. """

    train: XYData
    """ Training data. """

    test: XYData
    """ Test data. """

    @check_shapes(
        "self.train.X: [N_train, D]",
        "self.train.Y: [N_train, P]",
        "self.test.X: [N_test, D]",
        "self.test.Y: [N_test, P]",
    )
    def __post_init__(self) -> None:
        assert self.name.isidentifier()
        assert self.tags == self.train.tags
        assert self.tags == self.test.tags

    @property
    def D(self) -> int:
        """ Number of inputs. """
        return self.train.D

    @property
    def P(self) -> int:
        """ Number of outputs. """
        return self.train.P

    @property
    def stats(self) -> pd.DataFrame:
        """
        Get some basic statistics about this dataset; for debugging, logging, etc.
        """
        return pd.DataFrame(
            [
                ("N_train", self.train.N),
                ("N_test", self.test.N),
                ("D", self.D),
                ("P", self.P),
            ],
            columns=["statistic", "value"],
        )


class DatasetFactory(ABC):
    """
    A way to generate a dataset.

    The factory itself should be cheap to create, though it may take some time to actually generate
    the dataset.
    """

    name: str
    """
    Unique name of this dataset / dataset factory.
    """

    tags: AbstractSet[DatasetTag]
    """
    Tags representing properties / capabilities of this dataset.
    """

    @abstractmethod
    def create_dataset(self, cache_dir: Path) -> Dataset:
        """
        Create the dataset.

        This method should be deterministic - if called multiple times it should return the same
        data. If you need randomness, please use a fixed seed.

        :param cache_dir: A directory you can use to read/write any cache data.
        """


DATASET_FACTORIES: TaggedRegistry[DatasetFactory, DatasetTag] = TaggedRegistry()


DatasetFactoryFn = Callable[[Path], Dataset]
"""
A function that can be used as a :class:`ModelFactory`.

Takes a cache directory, and produces a :class:`Dataset`.
"""


class FnDatasetFactory(DatasetFactory):
    """
    Adapter from a function to a :class:`DatasetFactory`.
    """

    def __init__(self, name: str, tags: AbstractSet[DatasetTag], fn: DatasetFactoryFn) -> None:
        self.name = name
        self.tags = tags
        self._fn = fn

    def create_dataset(self, cache_dir: Path) -> Dataset:
        return self._fn(cache_dir)


def make_dataset_factory(
    tags: AbstractSet[DatasetTag],
) -> Callable[[DatasetFactoryFn], FnDatasetFactory]:
    """
    Decorator for turning a function into a :class:`DatasetFactory`.
    """

    def wrap(fn: DatasetFactoryFn) -> FnDatasetFactory:
        name = fn.__name__
        factory = FnDatasetFactory(name, tags, fn)
        DATASET_FACTORIES.add(factory)
        return factory

    return wrap


@check_shapes(
    "X: [N, D]",
    "return: [N, D]",
)
def _normalise(X: AnyNDArray) -> AnyNDArray:
    X_mean = np.average(X, axis=0)[None, :]
    X_std = 1e-6 + np.std(X, axis=0)[None, :]
    return (X - X_mean) / X_std  # type: ignore[no-any-return]


@check_shapes(
    "X: [N, D]",
    "Y: [N, P]",
)
def _create_dataset(
    name: str,
    tags: AbstractSet[DatasetTag],
    test_fraction: float,
    normalise: bool,
    X: AnyNDArray,
    Y: AnyNDArray,
) -> Dataset:
    """
    Randomly splits raw data into a training and testing and create a new :class:`Dataset` from it.

    :param name: Name of the resulting `Dataset`.
    :param tags: Tags of the resulting `Dataset`.
    :param test_fraction: Fraction of data to use for the test set. The rest is used for training.
    :param normalise: Whether to normalise the data to mean 0 and variance 1.
    :param X: Input values.
    :param Y: Output values.
    """
    if normalise:
        X = _normalise(X)
        Y = _normalise(Y)

    N = X.shape[0]
    rng = np.random.default_rng(20220722)
    indices = rng.permutation(N)
    N_test = round(test_fraction * N)
    train_indices = indices[N_test:]
    test_indices = indices[:N_test]
    train = XYData(tags, X=X[train_indices], Y=Y[train_indices])
    test = XYData(tags, X=X[test_indices], Y=Y[test_indices])
    return Dataset(name=name, tags=tags, train=train, test=test)


def _get_data_file_from_url(cache_dir: Path, url: str, archive_member: Optional[Path]) -> Path:
    """
    Downloads a URL to a file, and returns the path to it.

    This will cache the downloaded file, and return the cached result if present.

    :param cache_dir: Where to read / write downloaded data.
    :param url: URL to download.
    :param archive_member: If not `None` the URL is assumed to point to a zip-file, which will be
        un-zipped, and the member with this relative path will be returned.
        If `None` the downloaded file will be returned without unzipping.
    """
    file_name = Path(urlparse(url).path).name
    file_path = cache_dir / file_name
    if not file_path.is_file():
        cache_dir.mkdir(parents=True, exist_ok=True)
        with urlopen(url) as response:
            data = response.read()
        file_path.write_bytes(data)
        if archive_member:
            with zipfile.ZipFile(file_path, "r") as archive:
                archive.extract(str(archive_member), cache_dir)
    if archive_member:
        return cache_dir / archive_member
    return file_path


UnsplitArrayFactory = Callable[[np.random.Generator], Tuple[AnyNDArray, AnyNDArray]]
"""
Generate synthetic data and return it as a `(X, Y)` tuple.

The returned data is assumed to not yet have been split into training/test data.
"""


def make_unsplit_array_dataset_factory(
    tags: AbstractSet[DatasetTag],
    test_fraction: float,
    normalise: bool,
) -> Callable[[UnsplitArrayFactory], FnDatasetFactory]:
    """
    Decorator for creating a :class:`DatasetFactory` that generates synthetic data.

    The generated data is assumed to not yet have been split into training/test data.
    """

    def wrap(array_factory: UnsplitArrayFactory) -> FnDatasetFactory:
        name = array_factory.__name__

        @wraps(array_factory)
        def wrapper(cache_dir: Path) -> Dataset:
            rng = np.random.default_rng(20220808)
            X, Y = array_factory(rng)
            return _create_dataset(name, tags, test_fraction, normalise, X, Y)

        return make_dataset_factory(tags)(wrapper)

    return wrap


DataReader = Callable[[Path], Tuple[AnyNDArray, AnyNDArray]]
"""
A function that reads data from a file.

Read data from the given path and return it as a `(X, Y)` tuple.
"""


def make_url_dataset_factory(
    tags: AbstractSet[DatasetTag],
    url: str,
    archive_member: Optional[Path] = None,
    test_fraction: float = 0.1,
    normalise: bool = True,
) -> Callable[[DataReader], FnDatasetFactory]:
    """
    Decorator for creating a :class:`DatasetFactory` that downloads data from a URL.

    :param tags: Tags that applies to the data.
    :param url: URL to download.
    :param archive_member: If not `None` the URL is assumed to point to a zip-file, which will be
        un-zipped, and the member with this relative path will be used.
        If `None` the downloaded file will be used without unzipping.
    :param test_fraction: Fraction of data to use for the test set. The rest is used for training.
    :param normalise: Whether to normalise the data to mean 0 and variance 1.
    """

    def wrap(data_reader: DataReader) -> FnDatasetFactory:
        name = data_reader.__name__

        @wraps(data_reader)
        def wrapper(cache_dir: Path) -> Dataset:
            dataset_dir = cache_dir / name
            data_path = _get_data_file_from_url(dataset_dir, url, archive_member)
            X, Y = data_reader(data_path)
            return _create_dataset(name, tags, test_fraction, normalise, X, Y)

        return make_dataset_factory(tags)(wrapper)

    return wrap


# benchmark/datasets.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Definitions of our datasets.

Code adapted from:
https://github.com/hughsalimbeni/bayesian_benchmarks/blob/master/bayesian_benchmarks/data.py
"""
from pathlib import Path
from typing import Tuple

import numpy as np
import pandas as pd

from benchmark.dataset_api import (
    LARGE,
    MEDIUM,
    REAL_DATA,
    REGRESSION,
    SYNTHETIC,
    TINY,
    make_unsplit_array_dataset_factory,
    make_url_dataset_factory,
)
from gpflow.base import AnyNDArray


@make_unsplit_array_dataset_factory(
    tags={REGRESSION, TINY, SYNTHETIC},
    test_fraction=0.5,
    normalise=False,
)
def tiny_linear(rng: np.random.Generator) -> Tuple[AnyNDArray, AnyNDArray]:
    X = rng.random((20, 1))
    F = 0.3 * X - 0.1
    noise = rng.normal(0.0, 0.1, F.shape)
    Y = F + noise
    return X, Y


@make_unsplit_array_dataset_factory(
    tags={REGRESSION, TINY, SYNTHETIC},
    test_fraction=0.5,
    normalise=False,
)
def tiny_sine(rng: np.random.Generator) -> Tuple[AnyNDArray, AnyNDArray]:
    X = rng.random((20, 1))
    F = np.sin(10 * X)
    noise = rng.normal(0.0, 0.1, F.shape)
    Y = F + noise
    return X, Y


UCI_BASE_URL = "https://archive.ics.uci.edu/ml/machine-learning-databases/"


@make_url_dataset_factory(
    tags={REGRESSION, MEDIUM, REAL_DATA},
    url=UCI_BASE_URL + "housing/housing.data",
)
def boston(path: Path) -> Tuple[AnyNDArray, AnyNDArray]:
    data = pd.read_fwf(path, header=None).values
    return data[:, :-1], data[:, -1].reshape(-1, 1)


@make_url_dataset_factory(
    tags={REGRESSION, MEDIUM, REAL_DATA},
    url=UCI_BASE_URL + "concrete/compressive/Concrete_Data.xls",
)
def concrete(path: Path) -> Tuple[AnyNDArray, AnyNDArray]:
    data = pd.read_excel(path).values
    return data[:, :-1], data[:, -1].reshape(-1, 1)


@make_url_dataset_factory(
    tags={REGRESSION, MEDIUM, REAL_DATA},
    url=UCI_BASE_URL + "00242/ENB2012_data.xlsx",
)
def energy(path: Path) -> Tuple[AnyNDArray, AnyNDArray]:
    data = pd.read_excel(path, engine="openpyxl", usecols=np.arange(9)).dropna().values
    return data[:, :-1], data[:, -1].reshape(-1, 1)


@make_url_dataset_factory(
    tags={REGRESSION, LARGE, REAL_DATA},
    url=UCI_BASE_URL + "00316/UCI%20CBM%20Dataset.zip",
    archive_member=Path("UCI CBM Dataset/data.txt"),
)
def naval(path: Path) -> Tuple[AnyNDArray, AnyNDArray]:
    data = pd.read_fwf(path, header=None).values
    # NB this is the first output
    X = data[:, :-2]
    Y = data[:, -2].reshape(-1, 1)

    # dims 8 and 11 have std=0:
    X = np.delete(X, [8, 11], axis=1)
    return X, Y


@make_url_dataset_factory(
    tags={REGRESSION, LARGE, REAL_DATA},
    url=UCI_BASE_URL + "00294/CCPP.zip",
    archive_member=Path("CCPP/Folds5x2_pp.xlsx"),
)
def power(path: Path) -> Tuple[AnyNDArray, AnyNDArray]:
    data = pd.read_excel(path, engine="openpyxl").values
    return data[:, :-1], data[:, -1].reshape(-1, 1)


@make_url_dataset_factory(
    tags={REGRESSION, LARGE, REAL_DATA},
    url=UCI_BASE_URL + "00265/CASP.csv",
)
def protein(path: Path) -> Tuple[AnyNDArray, AnyNDArray]:
    data = pd.read_csv(path).values
    return data[:, 1:], data[:, 0].reshape(-1, 1)


@make_url_dataset_factory(
    tags={REGRESSION, MEDIUM, REAL_DATA},
    url=UCI_BASE_URL + "wine-quality/winequality-red.csv",
)
def red_wine(path: Path) -> Tuple[AnyNDArray, AnyNDArray]:
    data = pd.read_csv(path, delimiter=";").values
    return data[:, :-1], data[:, -1].reshape(-1, 1)


@make_url_dataset_factory(
    tags={REGRESSION, LARGE, REAL_DATA},
    url=UCI_BASE_URL + "wine-quality/winequality-white.csv",
)
def white_wine(path: Path) -> Tuple[AnyNDArray, AnyNDArray]:
    data = pd.read_csv(path, delimiter=";").values
    return data[:, :-1], data[:, -1].reshape(-1, 1)


@make_url_dataset_factory(
    tags={REGRESSION, MEDIUM, REAL_DATA},
    url=UCI_BASE_URL + "/00243/yacht_hydrodynamics.data",
)
def yacht(path: Path) -> Tuple[AnyNDArray, AnyNDArray]:
    data = np.loadtxt(path)
    return data[:, :-1], data[:, -1].reshape(-1, 1)


# benchmark/grouping.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Code for grouping data for plotting.
"""
import heapq
from dataclasses import dataclass
from enum import Enum
from typing import Any, Collection, Iterator, List, Sequence, Tuple, Type

import pandas as pd

from benchmark.metadata import BenchmarkMetadata
from benchmark.plotter_api import Plotter


class GroupingKeySource(str, Enum):
    """ Sources of data to group by. """

    METADATA = "metadata"
    METRICS_DF = "metrics_df"
    PLOTTER = "plotter"


class GroupingKey(Tuple[GroupingKeySource, type, bool, float], Enum):
    """ Keys we can group by. """

    USER = (GroupingKeySource.METADATA, str, False, 1.13)
    HOSTNAME = (GroupingKeySource.METADATA, str, False, 2.01)
    TIMESTAMP = (GroupingKeySource.METADATA, str, False, 3.02)
    PY_VER = (GroupingKeySource.METADATA, str, True, 1.12)
    TF_VER = (GroupingKeySource.METADATA, str, True, 1.10)
    NP_VER = (GroupingKeySource.METADATA, str, True, 1.11)
    RAM = (GroupingKeySource.METADATA, int, False, 1.15)
    CPU_NAME = (GroupingKeySource.METADATA, str, False, 1.09)
    CPU_COUNT = (GroupingKeySource.METADATA, int, True, 1.14)
    CPU_FREQUENCY = (GroupingKeySource.METADATA, int, True, 3.01)
    GPU_NAME = (GroupingKeySource.METADATA, str, False, 1.07)
    GIT_BRANCH_NAME = (GroupingKeySource.METADATA, str, False, 1.08)

    DATASET = (GroupingKeySource.METRICS_DF, str, False, 1.03)
    MODEL = (GroupingKeySource.METRICS_DF, str, False, 1.04)
    DO_COMPILE = (GroupingKeySource.METRICS_DF, bool, True, 1.06)
    DO_OPTIMISE = (GroupingKeySource.METRICS_DF, bool, True, 1.05)
    METRIC = (GroupingKeySource.METRICS_DF, str, False, 1.02)

    PLOTTER = (GroupingKeySource.PLOTTER, Plotter, False, 1.01)

    @property
    def key_name(self) -> str:
        """ Name of this key. """
        return self.name.lower()  # pylint: disable=no-member

    @property
    def key_source(self) -> GroupingKeySource:
        """ Source of data for this key. """
        return self.value[0]  # type: ignore[no-any-return]

    @property
    def key_type(self) -> Type[Any]:
        """ Type/class of the data of this key. """
        return self.value[1]  # type: ignore[no-any-return]

    @property
    def prefix_key_to_value(self) -> bool:
        """
        Whether we should append the name of the key to the value of the key before printing.

        Some key values, such as dataset and model names, are easy to understand out of context.
        Other keys, often bools or ints, need some context to make sense. This property indicates
        whether values of this key need context.
        """
        return self.value[2]  # type: ignore[no-any-return]

    @property
    def key_cost(self) -> float:
        """
        How much we don't want to show this key.

        Should be positive.

        When trying to find a "minimal" group we will pick the group with smallest sum of costs.

        Keys the user is more likely to care about should have a small cost, while keys the user is
        unlikely to care about should have a larger cost.
        """
        return self.value[3]  # type: ignore[no-any-return]

    def __repr__(self) -> str:
        return self.name

    def __lt__(self, other: "GroupingKey") -> bool:  # type: ignore[override]
        return self.key_cost < other.key_cost


@dataclass(frozen=True)
class GroupingSpec:
    by: Sequence[GroupingKey]
    """ Keys to group by. """

    minimise: bool
    """ Whether to ignore redundant keys. """

    def __post_init__(self) -> None:
        assert len(self.by) == len(set(self.by)), f"'by' must have unique values. Found {self.by}."


def group(
    metrics_df: pd.DataFrame,
    plotters: Collection[Plotter],
    metadata: Collection[BenchmarkMetadata],
    spec: GroupingSpec,
) -> Sequence[Tuple[Tuple[str, ...], pd.DataFrame, Collection[Plotter]]]:
    """
    Group the given data according to the given specification.
    """
    result = _group(metrics_df, plotters, metadata, spec.by)

    if spec.minimise:
        # Terribly inefficient, but it's probably fast enough.
        for candidate_by in _iter_by(spec.by):
            candidate_result = _group(metrics_df, plotters, metadata, candidate_by)
            if len(result) == len(candidate_result):
                return candidate_result

    return result


def _group(
    metrics_df: pd.DataFrame,
    plotters: Collection[Plotter],
    metadata: Collection[BenchmarkMetadata],
    by: Sequence[GroupingKey],
) -> Sequence[Tuple[Tuple[str, ...], pd.DataFrame, Collection[Plotter]]]:
    """
    Group the given data by the given columns.
    """
    if not by:
        return [((), metrics_df, plotters)]

    joined_df = metrics_df

    metadata_by_columns = [k.key_name for k in by if k.key_source == GroupingKeySource.METADATA]
    if metadata_by_columns:
        metadata_columns = ["run_id"] + metadata_by_columns
        metadata_df = pd.DataFrame(
            [[getattr(md, c) for c in metadata_columns] for md in metadata],
            columns=metadata_columns,
        )
        joined_df = pd.merge(joined_df, metadata_df, on="run_id")

    plotter_by_columns = [k.key_name for k in by if k.key_source == GroupingKeySource.PLOTTER]
    if plotter_by_columns:
        plotter_key_name = GroupingKey.PLOTTER.key_name
        assert plotter_by_columns == [plotter_key_name]
        plotter_df = pd.DataFrame({plotter_key_name: [p.name for p in plotters]})
        joined_df = joined_df.merge(plotter_df, how="cross")

    result = []
    for keys, df in joined_df.groupby([k.key_name for k in by]):
        keys, key_plotters = _sanitise_keys(by, keys, plotters)
        if metadata_by_columns:
            df = df.drop(columns=metadata_by_columns)
        if plotter_by_columns:
            df = df.drop(columns=plotter_by_columns)
        result.append((keys, df, key_plotters))
    return result


def _sanitise_keys(
    by: Sequence[GroupingKey], keys: Any, plotters: Collection[Plotter]
) -> Tuple[Tuple[str, ...], Collection[Plotter]]:
    """
    Sanitises keys from Pandas group_by.

    Takes the keys from Pandas group_by, which can have all sorts of types, and:
    * Separates the plotters from the "regular" keys.
    * Converts the "regular" keys to strings.
    """
    result_keys: List[str] = []
    result_plotters = plotters
    if not isinstance(keys, tuple):
        keys = (keys,)

    for b, k in zip(by, keys):
        if b.key_type == Plotter:
            assert result_plotters == plotters
            result_plotters = tuple(p for p in plotters if p.name == k)
        else:
            key_repr = k if b.key_type == str else repr(k)
            if b.prefix_key_to_value:
                key_repr = f"{b.key_name}={key_repr}"
            result_keys.append(key_repr)

    return tuple(result_keys), result_plotters


def _iter_by(all_by: Sequence[GroupingKey]) -> Iterator[Sequence[GroupingKey]]:
    """
    Iterate over all possible subsets of grouping keys, sorted by cost.
    """
    todo: List[Tuple[float, int, List[GroupingKey]]] = [(0.0, 0, [])]
    while todo:
        _, prev_begin, prev_by = heapq.heappop(todo)
        yield prev_by
        for i, next_key in enumerate(all_by[prev_begin:]):
            next_begin = prev_begin + i + 1
            next_by = prev_by + [next_key]
            next_cost = sum(k.key_cost for k in next_by)
            heapq.heappush(todo, (next_cost, next_begin, next_by))


# benchmark/metadata.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Code for determining metadata for a run.
"""
import getpass
import sys
from dataclasses import asdict, dataclass, replace
from datetime import datetime
from pathlib import Path
from platform import python_version
from socket import gethostname
from typing import Any, Mapping, Sequence, Tuple

import numpy as np
import psutil
import tensorflow as tf
from cpuinfo import get_cpu_info
from git.exc import InvalidGitRepositoryError
from git.repo import Repo
from tensorflow.python.client import device_lib  # pylint: disable=no-name-in-module

TIMESTAMP_FORMAT = "%Y%m%d_%H%M%S.%f"


@dataclass(frozen=True)
class BenchmarkMetadata:
    """
    Struct with metadata about a script execution.
    """

    suite_name: str
    argv: Sequence[str]
    user: str
    hostname: str
    timestamp: str
    py_ver: str
    tf_ver: str
    np_ver: str
    ram: int
    cpu_name: str
    cpu_count: int
    cpu_frequency: int
    gpu_names: Sequence[str]
    git_branch_name: str
    git_commit: str
    run_id: str

    @staticmethod
    def create(suite_name: str) -> "BenchmarkMetadata":
        """
        Collects information about the current environment and creates a new `BenchmarkMetadata`
        with it.
        """
        suite_name = suite_name
        argv = sys.argv
        script_name = Path(argv[0]).stem
        user = getpass.getuser()
        hostname = gethostname()
        timestamp = _get_timestamp()
        py_ver = python_version()
        tf_ver = tf.__version__
        np_ver = np.__version__
        ram = psutil.virtual_memory().total
        cpu_name, cpu_count, cpu_frequency = _get_cpu_name_count_frequency()
        gpu_names = _get_gpu_names()
        git_branch_name, git_commit = _get_git_branch_and_commit()
        run_id = f"{script_name}_{suite_name}_{git_branch_name}_{timestamp}".replace("/", "_")
        return BenchmarkMetadata(
            suite_name=suite_name,
            argv=argv,
            user=user,
            hostname=hostname,
            timestamp=timestamp,
            py_ver=py_ver,
            tf_ver=tf_ver,
            np_ver=np_ver,
            ram=ram,
            cpu_name=cpu_name,
            cpu_count=cpu_count,
            cpu_frequency=cpu_frequency,
            gpu_names=gpu_names,
            git_branch_name=git_branch_name,
            git_commit=git_commit,
            run_id=run_id,
        )

    def for_shard(self) -> "BenchmarkMetadata":
        """
        Update this metadata with current machine information.

        Used in sharding, to create a machine-specific metadata, while retaining the process
        timestamp and id.
        """
        new_metadata = BenchmarkMetadata.create(self.suite_name)
        kwargs = asdict(new_metadata)
        keep_new_fields = [
            "argv",
            "user",
            "hostname",
            "ram",
            "cpu_name",
            "cpu_count",
            "cpu_frequency",
            "gpu_names",
        ]
        keep_old_fields = ["timestamp", "run_id"]
        must_match_fields = [
            "suite_name",
            "py_ver",
            "tf_ver",
            "np_ver",
            "git_branch_name",
            "git_commit",
        ]
        for field in keep_old_fields:
            del kwargs[field]
        for field in must_match_fields:
            assert getattr(self, field) == kwargs[field], (
                f"Field {field} must match between new and old metadata."
                f" Found {getattr(self, field)} and {kwargs[field]}"
            )
            del kwargs[field]
        assert keep_new_fields == list(kwargs)

        return replace(self, **kwargs)

    @property
    def gpu_name(self) -> str:
        if not self.gpu_names:
            return "No GPU"
        return " + ".join(self.gpu_names)

    def to_json(self) -> Mapping[str, Any]:
        return asdict(self)

    @staticmethod
    def from_json(data: Mapping[str, Any]) -> "BenchmarkMetadata":
        # This only works as long as our field types are sufficiently simple. We may need to update
        # this in the future:
        return BenchmarkMetadata(**data)


def _get_timestamp() -> str:
    """ Get the current date/time. """
    return datetime.utcnow().strftime(TIMESTAMP_FORMAT)


def parse_timestamp(timestamp: str) -> datetime:
    """ Parse a timestamp, as formatted in :class:`BenchmarkMetadata` into a ``datetime``. """
    return datetime.strptime(timestamp, TIMESTAMP_FORMAT)


def _get_cpu_name_count_frequency() -> Tuple[str, int, int]:
    """ Get the name, number and clock frequency of the available CPU. """
    cpu_info = get_cpu_info()
    frequency, what = cpu_info["hz_actual"]
    assert what == 0, "I have no idea what this is, but it seem to always be 0..."
    return cpu_info["brand_raw"], cpu_info["count"], frequency


def _get_gpu_names() -> Sequence[str]:
    """ Get a list of GPUs that are available to TensorFlow. """
    result = []
    for device in device_lib.list_local_devices():
        if device.device_type != "GPU":
            continue
        desc = device.physical_device_desc

        fields = desc.split(",")
        for field in fields:
            name, value = field.split(":", maxsplit=1)
            name = name.strip()
            value = value.strip()
            if name == "name":
                result.append(value)
    return result


def _get_git_branch_and_commit() -> Tuple[str, str]:
    """ Get a `(branch_name, commit_id)` tuple from the current direcory. """
    branch_name = "NO_BRANCH"
    commit = "NO_COMMIT"
    try:
        repo = Repo(__file__, search_parent_directories=True)
        try:
            branch_name = str(repo.active_branch)
        except TypeError:
            pass  # Keep current/default branch_name
        commit = str(repo.commit())
        if repo.is_dirty():
            commit += " + uncomitted changes"
    except InvalidGitRepositoryError:
        pass  # Keep current/default branch_name and commit
    return branch_name, commit


# benchmark/metric_api.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Classes and other infrastructure for defining metrics.

Concrete metrics are defined in `metrics.py`.
"""
from dataclasses import dataclass
from enum import Enum
from typing import Any, Optional

from benchmark.registry import Named, Registry


class MetricOrientation(Enum):
    """
    Whether we like small or large values.
    """

    LOWER_IS_BETTER = "lower_is_better"
    GREATER_IS_BETTER = "greater_is_better"


@dataclass(frozen=True)
class Metric(Named):
    """
    Metadata about a metric.

    Super useful when plotting it.
    """

    name: str
    pretty_name: str
    lower_bound: Optional[float]
    upper_bound: Optional[float]
    orientation: MetricOrientation
    unit: Optional[str]

    def __repr__(self) -> str:
        return self.name


METRICS: Registry[Metric] = Registry()


def make_metric(**kwargs: Any) -> Metric:
    """
    Create a new :class:`Metric` and register it.

    :param kwargs: Passed on to ``Metric.__init__``.
    """
    result = Metric(**kwargs)
    METRICS.add(result)
    return result


# benchmark/metrics.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Concrete definitions of our metrics.
"""
from benchmark.metric_api import MetricOrientation, make_metric

n_training_iterations = make_metric(
    name="n_training_iterations",
    pretty_name="Training iterations",
    lower_bound=0,
    upper_bound=None,
    orientation=MetricOrientation.LOWER_IS_BETTER,
    unit=None,
)


training_time = make_metric(
    name="training_time",
    pretty_name="Time to train",
    lower_bound=0,
    upper_bound=None,
    orientation=MetricOrientation.LOWER_IS_BETTER,
    unit="s",
)


training_iteration_time = make_metric(
    name="training_iteration_time",
    pretty_name="Time run one iteration",
    lower_bound=0,
    upper_bound=None,
    orientation=MetricOrientation.LOWER_IS_BETTER,
    unit="s",
)


prediction_time = make_metric(
    name="prediction_time",
    pretty_name="Time to predict Y",
    lower_bound=0,
    upper_bound=None,
    orientation=MetricOrientation.LOWER_IS_BETTER,
    unit="s",
)


nlpd = make_metric(
    name="nlpd",
    pretty_name="Negative Log Predictive Density",
    lower_bound=None,
    upper_bound=None,
    orientation=MetricOrientation.GREATER_IS_BETTER,
    unit=None,
)


mae = make_metric(
    name="mae",
    pretty_name="Mean Absolute Error",
    lower_bound=0,
    upper_bound=None,
    orientation=MetricOrientation.LOWER_IS_BETTER,
    unit=None,
)


rmse = make_metric(
    name="rmse",
    pretty_name="Root Mean Squared Error",
    lower_bound=0,
    upper_bound=None,
    orientation=MetricOrientation.LOWER_IS_BETTER,
    unit=None,
)


posterior_build_time = make_metric(
    name="posterior_build_time",
    pretty_name="Time to build posterior",
    lower_bound=0,
    upper_bound=None,
    orientation=MetricOrientation.LOWER_IS_BETTER,
    unit="s",
)


posterior_prediction_time = make_metric(
    name="posterior_prediction_time",
    pretty_name="Time to predict test Y (posterior)",
    lower_bound=0,
    upper_bound=None,
    orientation=MetricOrientation.LOWER_IS_BETTER,
    unit="s",
)


posterior_nlpd = make_metric(
    name="posterior_nlpd",
    pretty_name="Negative Log Predictive Density (posterior)",
    lower_bound=None,
    upper_bound=None,
    orientation=MetricOrientation.GREATER_IS_BETTER,
    unit=None,
)


posterior_mae = make_metric(
    name="posterior_mae",
    pretty_name="Mean Absolute Error (posterior)",
    lower_bound=0,
    upper_bound=None,
    orientation=MetricOrientation.LOWER_IS_BETTER,
    unit=None,
)


posterior_rmse = make_metric(
    name="posterior_rmse",
    pretty_name="Root Mean Squared Error (posterior)",
    lower_bound=0,
    upper_bound=None,
    orientation=MetricOrientation.LOWER_IS_BETTER,
    unit=None,
)


# benchmark/model_api.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Classes and other infrastructure for defining and creating models.

Concrete models factories are found in ``models.py``.
"""
from abc import ABC, abstractmethod
from typing import AbstractSet, Callable

import numpy as np

from benchmark.dataset_api import DatasetReq, XYData
from benchmark.registry import TaggedRegistry
from benchmark.tag import Tag, TagReq
from gpflow.models import GPModel


class ModelTag(Tag["ModelTag"]):
    pass


ModelReq = TagReq[ModelTag]


REGRESSION = ModelTag("REGRESSION")
SPARSE = ModelTag("SPARSE")
VARIATIONAL = ModelTag("VARIATIONAL")


class ModelFactory(ABC):
    """
    A way to instantiate a model.

    The factory itself should be cheap to create, though it may take some time to actually create
    the model.
    """

    name: str
    """
    Name of this model / model factory.
    """

    tags: AbstractSet[ModelTag]
    """
    Tags representing properties / capabilities of this model.
    """

    dataset_req: DatasetReq
    """
    This model is compatible with datasets that has these tags.
    """

    @abstractmethod
    def create_model(self, data: XYData, rng: np.random.Generator) -> GPModel:
        """
        Create the model.

        Any model parameters should be randomly initialised, using the `rng`.
        """


MODEL_FACTORIES: TaggedRegistry[ModelFactory, ModelTag] = TaggedRegistry()


ModelFactoryFn = Callable[[XYData, np.random.Generator], GPModel]
"""
A function that can be used as a :class:`ModelFactory`.

Any model parameters should be randomly initialised, using the ``Generator``.
"""


class FnModelFactory(ModelFactory):
    """
    Adapter from a function to a :class:`ModelFactory`.
    """

    def __init__(
        self, name: str, tags: AbstractSet[ModelTag], dataset_req: DatasetReq, fn: ModelFactoryFn
    ) -> None:
        self.name = name
        self.tags = tags
        self.dataset_req = dataset_req
        self._fn = fn

    def create_model(self, data: XYData, rng: np.random.Generator) -> GPModel:
        assert self.dataset_req.satisfied(data.tags)
        return self._fn(data, rng)


def make_model_factory(
    tags: AbstractSet[ModelTag], dataset_req: DatasetReq
) -> Callable[[ModelFactoryFn], FnModelFactory]:
    """
    Decorator for turning a function into a :class:`ModelFactory`.
    """

    def wrap(fn: ModelFactoryFn) -> FnModelFactory:
        name = fn.__name__
        factory = FnModelFactory(name, tags, dataset_req, fn)
        MODEL_FACTORIES.add(factory)
        return factory

    return wrap


# benchmark/models.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Concrete code for creating models.
"""
import numpy as np
from scipy.stats.qmc import Halton  # pylint: disable=no-name-in-module

import benchmark.dataset_api as ds
from benchmark.dataset_api import XYData
from benchmark.model_api import REGRESSION, SPARSE, VARIATIONAL, make_model_factory
from gpflow.inducing_variables import InducingPoints
from gpflow.kernels import RBF, Kernel
from gpflow.likelihoods import Gaussian
from gpflow.models import GPR, SGPR, SVGP, VGP, GPModel
from gpflow.models.util import InducingPointsLike


def create_rbf(data: XYData, rng: np.random.Generator) -> Kernel:
    return RBF(
        variance=rng.gamma(5.0, 0.2, []),
        lengthscales=rng.gamma(5.0, 0.2, [data.D]),
    )


def create_inducing(data: XYData, rng: np.random.Generator) -> InducingPointsLike:
    n = min(data.N // 2, 200)
    Z = Halton(data.D, scramble=False).random(n)
    lower = np.min(data.X, axis=0)
    upper = np.max(data.X, axis=0)
    Z = Z * (upper - lower) + lower
    return InducingPoints(Z)


def create_gaussian(data: XYData, rng: np.random.Generator) -> Gaussian:
    return Gaussian(variance=rng.gamma(5.0, 0.2, []))


@make_model_factory(tags={REGRESSION}, dataset_req=ds.REGRESSION & ~ds.LARGE)
def gpr(data: XYData, rng: np.random.Generator) -> GPModel:
    return GPR(
        data.XY,
        kernel=create_rbf(data, rng),
        noise_variance=rng.gamma(5.0, 0.2, []),
    )


@make_model_factory(tags={REGRESSION, VARIATIONAL}, dataset_req=ds.REGRESSION & ~ds.LARGE)
def vgp(data: XYData, rng: np.random.Generator) -> GPModel:
    return VGP(
        data.XY,
        kernel=create_rbf(data, rng),
        likelihood=create_gaussian(data, rng),
    )


@make_model_factory(tags={REGRESSION, SPARSE}, dataset_req=ds.REGRESSION)
def sgpr(data: XYData, rng: np.random.Generator) -> GPModel:
    return SGPR(
        data.XY,
        kernel=create_rbf(data, rng),
        inducing_variable=create_inducing(data, rng),
        noise_variance=rng.gamma(5.0, 0.2, []),
    )


@make_model_factory(tags={REGRESSION, SPARSE, VARIATIONAL}, dataset_req=ds.REGRESSION)
def svgp(data: XYData, rng: np.random.Generator) -> GPModel:
    return SVGP(
        kernel=create_rbf(data, rng),
        likelihood=create_gaussian(data, rng),
        inducing_variable=create_inducing(data, rng),
    )


# benchmark/paths.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Utilities for dealing with paths.
"""
from pathlib import Path

from benchmark.metadata import BenchmarkMetadata


def setup_dest(metadata: BenchmarkMetadata, dest: Path, no_subdir: bool) -> Path:
    """
    Set up destination/output directory.

    :param metadata: Data about this run.
    :param dest: Root destination directory.
    :param no_subdir: Whether to NOT create a subdirectory of ``dest``.
    :return: Path to new destination directory.
    """
    if not no_subdir:
        root = dest
        root.mkdir(parents=True, exist_ok=True)

        relative_dest = Path(metadata.run_id)
        dest = root / relative_dest
        dest.mkdir()

        latest_dir = root / "latest"
        if latest_dir.is_symlink():
            latest_dir.unlink()
        latest_dir.symlink_to(relative_dest)
    else:
        dest.mkdir(parents=True)
    return dest


# benchmark/plot.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Script for plotting data created by ``run.py``.

The ``run.py`` script generally generates its own plots, but this scripet is useful if you want to
plot results from multiple runs, or if you have updated some code and want to regenerate the plots.

Usage::

    python -m benchmark.plot <benchmark_suite> <input_directories...> <output_directory>

Multiple input directories can be defined, and all of them will be recursively crawled to find all
data. This data is then plotted according to the given benchmark suite.
"""
import argparse
import json
from pathlib import Path
from typing import Collection, Sequence, Tuple

import matplotlib.pyplot as plt
import pandas as pd

import benchmark.benchmarks  # pylint: disable=unused-import  # Make sure registry is loaded.
import benchmark.metrics  # pylint: disable=unused-import  # Make sure registry is loaded.
from benchmark.benchmark_api import BENCHMARK_SUITES, BenchmarkSuite
from benchmark.grouping import group
from benchmark.metadata import BenchmarkMetadata
from benchmark.paths import setup_dest


def _crawl_sources(sources: Sequence[Path]) -> Tuple[pd.DataFrame, Collection[BenchmarkMetadata]]:
    results_df_list = []
    results_metadata = []
    visited = set()
    queue = list(sources)
    while queue:
        source = queue.pop()

        source = source.resolve()
        assert source.is_dir()
        # Sym-links can cause cycles - don't visit the same directory multiple times...
        if source in visited:
            continue
        visited.add(source)

        metrics_df_path = source / "metrics.csv"
        metadata_path = source / "metadata.json"
        if metrics_df_path.exists() or metadata_path.exists():
            assert metrics_df_path.exists()
            assert metadata_path.exists()
            results_df_list.append(pd.read_csv(metrics_df_path))
            with open(metadata_path, "rt") as f:
                results_metadata.append(BenchmarkMetadata.from_json(json.load(f)))

        for child in source.iterdir():
            if child.is_dir():
                queue.append(child)

    results_df = pd.concat(results_df_list, axis="index", ignore_index=True)
    return results_df, results_metadata


def plot(
    metadata: BenchmarkMetadata,
    suite: BenchmarkSuite,
    results_df: pd.DataFrame,
    results_metadata: Collection[BenchmarkMetadata],
    dest: Path,
) -> None:
    """ Plot previously collected metrics. """
    for benchmark_set in suite.sets:
        benchmark_set_df = benchmark_set.filter_metrics(results_df)
        for file_key, file_df, file_plotters in group(
            benchmark_set_df, benchmark_set.plots, results_metadata, benchmark_set.file_by
        ):
            row_groups = group(file_df, file_plotters, results_metadata, benchmark_set.row_by)
            n_rows = len(row_groups)
            n_columns = len(
                group(file_df, file_plotters, results_metadata, benchmark_set.column_by)
            )

            width = 6 * n_columns
            height = 4 * n_rows
            fig, axes = plt.subplots(
                ncols=n_columns,
                nrows=n_rows,
                figsize=(width, height),
                squeeze=False,
                dpi=100,
            )

            for row_axes, (row_key, row_df, row_plotters) in zip(axes, row_groups):
                column_groups = group(
                    row_df, row_plotters, results_metadata, benchmark_set.column_by
                )
                for column_ax, (column_key, column_df, column_plotters) in zip(
                    row_axes, column_groups
                ):
                    (plotter,) = column_plotters
                    plotter.plot(
                        column_ax,
                        file_key,
                        column_key,
                        row_key,
                        benchmark_set.safe_line_by,
                        column_df,
                        results_metadata,
                    )

            fig.tight_layout()
            file_tokens = (benchmark_set.name,) + file_key
            fig.savefig(dest / f"{'_'.join(file_tokens)}.png")
            plt.close(fig)


def plot_from_argv() -> None:
    """ Plot previously collected metrics, based on command line arguments. """
    parser = argparse.ArgumentParser(description="Plot a benchmark suite.")
    parser.add_argument(
        "--no_subdir",
        "--no-subdir",
        action="store_true",
        help="By default this script will create a subdirectory inside `dest` and write output to"
        " that. Setting this flag will suppress the subdirectory and use `dest` directly.",
    )
    parser.add_argument(
        "suite",
        type=str,
        choices=BENCHMARK_SUITES.names(),
        help="Which benchmark suite to run.",
    )
    parser.add_argument(
        "sources",
        type=Path,
        nargs="+",
        help="Directory(s) to read data from."
        " These will be recursively crawled to find all `metrics.csv`s.",
    )
    parser.add_argument(
        "dest",
        type=Path,
        help="Directory to write results to.",
    )
    args = parser.parse_args()

    metadata = BenchmarkMetadata.create(args.suite)
    suite = BENCHMARK_SUITES.get(args.suite)
    results_df, results_metadata = _crawl_sources(args.sources)
    dest = setup_dest(metadata, args.dest, args.no_subdir)
    plot(metadata, suite, results_df, results_metadata, dest)


if __name__ == "__main__":
    plot_from_argv()


# benchmark/plotter_api.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Classes and other infrastructre for plotting results.

Concrete plotters are found in ``plotters.py``.
"""
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any, Callable, Collection, Tuple

import pandas as pd
from matplotlib.axes import Axes

from benchmark.metadata import BenchmarkMetadata
from benchmark.registry import Registry

# Avoid cyclic imports:
if TYPE_CHECKING:
    from benchmark.grouping import GroupingSpec
else:
    GroupingSpec = Any


class Plotter(ABC):
    """
    Strategy for how to create a plot.
    """

    name: str
    """
    Name of this plotter.
    """

    @abstractmethod
    def plot(
        self,
        ax: Axes,
        file_key: Tuple[str, ...],
        column_key: Tuple[str, ...],
        row_key: Tuple[str, ...],
        line_by: GroupingSpec,
        metrics_df: pd.DataFrame,
        metadata: Collection[BenchmarkMetadata],
    ) -> None:
        """ Plot the given data to the given `ax`. """


PLOTTERS: Registry[Plotter] = Registry()


PlotterFn = Callable[
    [
        Axes,
        Tuple[str, ...],  # file_key
        Tuple[str, ...],  # column_key
        Tuple[str, ...],  # row_key
        GroupingSpec,  # line_by
        pd.DataFrame,  # metrics_df
        Collection[BenchmarkMetadata],
    ],
    None,
]
"""
A function that can be used as a :class:`Plotter`.

Plot the given data to the given `Axes`.
"""


class FnPlotter(Plotter):
    """
    Adapter from a function to a :class:`Plotter`.
    """

    def __init__(self, name: str, fn: PlotterFn) -> None:
        self.name = name
        self._fn = fn

    def plot(
        self,
        ax: Axes,
        file_key: Tuple[str, ...],
        column_key: Tuple[str, ...],
        row_key: Tuple[str, ...],
        line_by: GroupingSpec,
        metrics_df: pd.DataFrame,
        metadata: Collection[BenchmarkMetadata],
    ) -> None:
        self._fn(ax, file_key, column_key, row_key, line_by, metrics_df, metadata)


def make_plotter(fn: PlotterFn) -> FnPlotter:
    """
    Decorator for turning a function into a :class:`Plotter`.
    """
    name = fn.__name__
    result = FnPlotter(name, fn)
    PLOTTERS.add(result)
    return result


# benchmark/plotters.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Concrete code for plotting results.
"""
from dataclasses import replace
from math import isnan
from typing import Collection, Tuple

import pandas as pd
from matplotlib.axes import Axes

from benchmark.grouping import GroupingKey, GroupingSpec, group
from benchmark.metadata import BenchmarkMetadata, parse_timestamp
from benchmark.metric_api import METRICS, Metric
from benchmark.plotter_api import make_plotter


def _join_key(key: Tuple[str, ...]) -> str:
    return "; ".join(key)


def _get_metric(metrics_df: pd.DataFrame) -> Metric:
    (metric_name,) = metrics_df.metric.unique()
    return METRICS.get(metric_name)


def _shared_ax_config(
    ax: Axes,
    file_key: Tuple[str, ...],
    column_key: Tuple[str, ...],
    row_key: Tuple[str, ...],
    metric: Metric,
) -> None:

    title_key = column_key + row_key
    if metric.name in title_key:
        title_key = tuple(k for k in title_key if k != metric.name)
    if title_key:
        ax.set_title(_join_key(title_key))

    ax.tick_params(axis="x", labelrotation=30)

    if metric.unit is None:
        y_label = metric.pretty_name
    else:
        y_label = f"{metric.pretty_name} ({metric.unit})"
    ax.set_ylabel(y_label)

    if metric.lower_bound is not None:
        ax.set_ylim(bottom=metric.lower_bound)
    if metric.upper_bound is not None:
        ax.set_ylim(top=metric.upper_bound)


@make_plotter
def metrics_box_plot(
    ax: Axes,
    file_key: Tuple[str, ...],
    column_key: Tuple[str, ...],
    row_key: Tuple[str, ...],
    line_by: GroupingSpec,
    metrics_df: pd.DataFrame,
    metadata: Collection[BenchmarkMetadata],
) -> None:
    """
    Creates box-plots of metrics.
    """
    assert GroupingKey.PLOTTER not in line_by.by
    metric = _get_metric(metrics_df)
    labels = []
    values = []
    for key, df, _ in group(metrics_df, [], metadata, line_by):
        labels.append(_join_key(key))
        values.append(df.value)
    ax.boxplot(values, labels=labels)

    _shared_ax_config(ax, file_key, column_key, row_key, metric)


@make_plotter
def time_line(
    ax: Axes,
    file_key: Tuple[str, ...],
    column_key: Tuple[str, ...],
    row_key: Tuple[str, ...],
    line_by: GroupingSpec,
    metrics_df: pd.DataFrame,
    metadata: Collection[BenchmarkMetadata],
) -> None:
    """
    Creates a plot of performance over time.
    """
    assert GroupingKey.PLOTTER not in line_by.by
    metric = _get_metric(metrics_df)
    line_by = replace(line_by, by=[k for k in line_by.by if k != GroupingKey.TIMESTAMP])
    line_groups = group(metrics_df, [], metadata, line_by)
    for key, df, _ in line_groups:
        line_xs = []
        line_y_means = []
        line_y_uppers = []
        line_y_lowers = []

        scatter_xs = []
        scatter_ys = []

        timestamp_by = GroupingSpec([GroupingKey.TIMESTAMP], minimise=False)
        timestamp_groups = group(df, [], metadata, timestamp_by)
        for (timestamp,), timestamp_df, _ in timestamp_groups:
            parsed_timestamp = parse_timestamp(timestamp)
            line_xs.append(parsed_timestamp)
            y_mean = timestamp_df.value.mean()
            y_std = timestamp_df.value.std()
            if isnan(y_std):
                y_std = 0.0
            line_y_means.append(y_mean)
            line_y_uppers.append(y_mean + 1.96 * y_std)
            line_y_lowers.append(y_mean - 1.96 * y_std)

            scatter_xs.extend(len(timestamp_df) * [parsed_timestamp])
            scatter_ys.extend(timestamp_df.value)

        (mean_line,) = ax.plot(line_xs, line_y_means, label=_join_key(key))
        color = mean_line.get_color()
        ax.fill_between(line_xs, line_y_lowers, line_y_uppers, color=color, alpha=0.3)
        ax.scatter(scatter_xs, scatter_ys, color=color)

    if len(line_groups) > 1:
        ax.legend()

    _shared_ax_config(ax, file_key, column_key, row_key, metric)


# benchmark/registry.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Code for registering stuff by name.

We have *a lot* of stuff with names in this project...
"""
from typing import AbstractSet, Any, Dict, Generic, Sequence, TypeVar

from typing_extensions import Protocol

from benchmark.tag import Tag, TagReq


class Named(Protocol):
    """
    Something with a name.
    """

    name: str


N = TypeVar("N", bound=Named)


class Registry(Generic[N]):
    """
    A registry of stuff with names.
    """

    def __init__(self) -> None:
        self._members: Dict[str, N] = {}

    def add(self, member: N) -> N:
        """
        Add the given member to this registry.

        Returns the input, for easy inlining.
        """
        assert (
            member.name not in self._members
        ), f"{member} already registred under name {member.name}."
        self._members[member.name] = member
        return member

    def get(self, name: str) -> N:
        """
        Get a member by name.
        """
        return self._members[name]

    def names(self) -> Sequence[str]:
        """
        Get a list of all registered names.
        """
        return tuple(self._members)

    def all(self) -> Sequence[N]:
        """
        Get a list of all registered values.
        """
        return tuple(self._members.values())


TTag = TypeVar("TTag", bound=Tag[Any])


class Tagged(Named, Protocol[TTag]):
    """
    Something with a name and some tags.
    """

    tags: AbstractSet[TTag]


T = TypeVar("T", bound=Tagged[Any])


class TaggedRegistry(Registry[T], Generic[T, TTag]):
    """
    A registry of stuff with names and tags.
    """

    def where(self, req: TagReq[TTag]) -> Sequence[T]:
        """
        Get all memebers that fulfill the given tag requirements.
        """
        return tuple(m for m in self._members.values() if req.satisfied(m.tags))


# benchmark/run.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Main script for running benchmarks.

Usage::

    python -m benchmark.run <benchmark_suite> <output_directory>

Runs all benchmarks defined in the given suite, and plots the results. If you want to plot results
from across multiple runs, use the ``plot.py`` script.
"""
import argparse
import json
import multiprocessing as mp
import traceback
from datetime import timedelta
from pathlib import Path
from time import perf_counter
from typing import Mapping, Optional

import numpy as np
import pandas as pd
from tabulate import tabulate

import benchmark.benchmarks  # pylint: disable=unused-import  # Make sure registry is loaded.
import benchmark.metrics as mt
import gpflow
from benchmark.benchmark_api import BENCHMARK_SUITES, BenchmarkSuite, BenchmarkTask
from benchmark.dataset_api import DATASET_FACTORIES
from benchmark.metadata import BenchmarkMetadata
from benchmark.metric_api import Metric
from benchmark.model_api import MODEL_FACTORIES
from benchmark.paths import setup_dest
from benchmark.plot import plot as plot_
from benchmark.sharding import ShardingStrategy

BASE_SEED = 20220721
ONE_HOUR_S = 60 * 60


def _collect_metrics(
    task: BenchmarkTask,
    data_cache_dir: Path,
    random_seed: int,
) -> Optional[Mapping[Metric, float]]:
    print("Loading dataset:", task.dataset_name)
    dataset_fac = DATASET_FACTORIES.get(task.dataset_name)
    dataset = dataset_fac.create_dataset(data_cache_dir)
    print(tabulate(dataset.stats, showindex="never"))
    train_data = dataset.train
    test_data = dataset.test

    rng = np.random.default_rng(random_seed)
    model_fac = MODEL_FACTORIES.get(task.model_name)
    model = model_fac.create_model(train_data, rng)

    metrics = {}

    model.predict_y(test_data.X)  # Warm-up TF.

    print("Model before training:")
    gpflow.utilities.print_summary(model)

    if task.do_optimise:
        t_before = perf_counter()
        loss_fn = gpflow.models.training_loss_closure(model, train_data.XY, compile=task.do_compile)
        opt_log = gpflow.optimizers.Scipy().minimize(
            loss_fn,
            variables=model.trainable_variables,
            compile=task.do_compile,
            options={"disp": 10, "maxiter": 1_000},
        )
        t_after = perf_counter()
        n_iter = opt_log.nit
        t_train = t_after - t_before
        metrics[mt.n_training_iterations] = n_iter
        metrics[mt.training_time] = t_train
        metrics[mt.training_iteration_time] = t_train / n_iter
        print(f"Training took {t_after - t_before}s for {n_iter} iterations.")
        print("Model after training:")
        gpflow.utilities.print_summary(model)

    likelihood = model.likelihood

    if task.do_predict:
        t_before = perf_counter()
        f_m, f_v = model.predict_f(test_data.X)
        t_after = perf_counter()
        metrics[mt.prediction_time] = t_after - t_before

        metrics[mt.nlpd] = -np.sum(
            likelihood.predict_log_density(test_data.X, f_m, f_v, test_data.Y)
        )

        y_m, _y_v = likelihood.predict_mean_and_var(test_data.X, f_m, f_v)
        error = test_data.Y - y_m
        metrics[mt.mae] = np.average(np.abs(error))
        metrics[mt.rmse] = np.average(error ** 2) ** 0.5

    if task.do_posterior:
        t_before = perf_counter()
        posterior = model.posterior()
        posterior.predict_f(np.zeros((0, test_data.D)))
        t_after = perf_counter()
        metrics[mt.posterior_build_time] = t_after - t_before

        t_before = perf_counter()
        f_m, f_v = posterior.predict_f(test_data.X)
        t_after = perf_counter()
        metrics[mt.posterior_prediction_time] = t_after - t_before

        metrics[mt.posterior_nlpd] = -np.sum(
            likelihood.predict_log_density(test_data.X, f_m, f_v, test_data.Y)
        )

        y_m, _y_v = likelihood.predict_mean_and_var(test_data.X, f_m, f_v)
        error = test_data.Y - y_m
        metrics[mt.posterior_mae] = np.average(np.abs(error))
        metrics[mt.posterior_rmse] = np.average(error ** 2) ** 0.5

    metrics_values = list(metrics.values())

    # pylint: disable=no-member
    assert np.isfinite(metrics_values).all(), f"{metrics} not all finite."

    return metrics


def _collect_metrics_process(
    task: BenchmarkTask,
    data_cache_dir: Path,
    random_seed: int,
    metrics_queue: "mp.Queue[Optional[Mapping[Metric, float]]]",
) -> None:
    metrics: Optional[Mapping[Metric, float]] = None
    try:
        metrics = _collect_metrics(task, data_cache_dir, random_seed)
    except Exception:  # pylint: disable=broad-except
        traceback.print_exc()
    metrics_queue.put(metrics)


def _run_benchmarks(
    metadata: BenchmarkMetadata,
    shard: ShardingStrategy,
    suite: BenchmarkSuite,
    dest: Path,
    cache_dir: Path,
) -> Optional[pd.DataFrame]:
    data_cache_dir = cache_dir / "data"

    # Run all tasks in separate processes, to make sure all memory etc. is released.
    mp_ctx = mp.get_context("spawn")

    tasks = suite.get_tasks()
    global_repetition = 0
    total_repetitions = sum(t.repetitions for t in tasks)
    max_task_repetitions = max(t.repetitions for t in tasks)

    metrics_list = []
    metrics_df: Optional[pd.DataFrame] = None

    t_before = perf_counter()
    for task_repetition in range(max_task_repetitions):
        for task in tasks:
            if task_repetition >= task.repetitions:
                continue

            if global_repetition % shard.shard_n != shard.shard_i:
                global_repetition += 1
                continue

            for retry in range(3):
                t_now = perf_counter()
                fraction_done = global_repetition / total_repetitions
                if global_repetition == 0:
                    eta_str = "???"
                else:
                    velocity = (t_now - t_before) / global_repetition
                    t_end = velocity * total_repetitions + t_before
                    t_remaining = t_end - t_now
                    eta_str = f"{timedelta(seconds=t_remaining)}"
                print("XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX")
                print(f"Running: {task.name}/rep={task_repetition} (retry {retry})")
                print(f"{fraction_done: .2%} done. ETA: {eta_str}.")
                print("XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX")

                random_seed = BASE_SEED + task_repetition + retry * task.repetitions
                metrics_queue: "mp.Queue[Optional[Mapping[Metric, float]]]" = mp_ctx.Queue()
                metrics_process = mp_ctx.Process(
                    target=_collect_metrics_process,
                    args=(task, data_cache_dir, random_seed, metrics_queue),
                )
                metrics_process.start()
                metrics = metrics_queue.get(timeout=ONE_HOUR_S)
                metrics_process.join()
                assert metrics_process.exitcode == 0
                if metrics is None:
                    print("Model failed. Retrying with another random seed.")
                    continue

                for metric, metric_value in metrics.items():
                    metrics_list.append(
                        (
                            metadata.run_id,
                            task.dataset_name,
                            task.model_name,
                            task.do_compile,
                            task.do_optimise,
                            task_repetition,
                            metric.name,
                            metric_value,
                        )
                    )
                metrics_df = pd.DataFrame(
                    metrics_list,
                    columns=[
                        "run_id",
                        "dataset",
                        "model",
                        "do_compile",
                        "do_optimise",
                        "repetition",
                        "metric",
                        "value",
                    ],
                )
                print(tabulate(metrics_df, headers="keys", showindex="never"))
                metrics_path = dest / f"metrics{shard.file_suffix}.csv"
                metrics_df.to_csv(metrics_path, index=False)
                break
            global_repetition += 1

    return metrics_df


def run(
    shard: ShardingStrategy,
    suite: BenchmarkSuite,
    dest: Path,
    no_subdir: bool,
    cache_dir: Path,
    plot: bool,
) -> None:
    """ Collect data for a benchmark suite. """
    if shard.setup_dest:
        metadata = BenchmarkMetadata.create(suite.name)
        dest = setup_dest(metadata, dest, no_subdir)
    else:
        assert no_subdir, "Must use 'no-subdir=True' with 'shard={shard}'."
        with open(dest / "metadata.json", "rt") as f:
            metadata = BenchmarkMetadata.from_json(json.load(f))
        assert metadata.suite_name == suite.name
        metadata = metadata.for_shard()

    if shard.print_dest:
        print(dest.resolve())

    if shard.write_metadata:
        metadata_path = dest / f"metadata{shard.file_suffix}.json"
        with open(metadata_path, "wt") as f:
            json.dump(metadata.to_json(), f, indent=4)

    metrics_df: Optional[pd.DataFrame] = None

    if shard.run_benchmarks:
        metrics_df = _run_benchmarks(
            metadata,
            shard,
            suite,
            dest,
            cache_dir,
        )

    if shard.collect:
        assert metrics_df is None
        metrics_path = dest / "metrics.csv"
        metrics_df_list = [pd.read_csv(p) for p in shard.find_shards(metrics_path)]
        metrics_df = pd.concat(metrics_df_list, axis="index", ignore_index=True)
        metrics_df.to_csv(metrics_path, index=False)

    if shard.plot and plot and (metrics_df is not None):
        plot_(
            metadata=metadata,
            suite=suite,
            results_df=metrics_df,
            results_metadata=[metadata],
            dest=dest,
        )


def run_from_argv() -> None:
    """ Collect data for a benchmark suite, based on command line arguments. """
    parser = argparse.ArgumentParser(description="Run a benchmark suite.")
    parser.add_argument(
        "--no_subdir",
        "--no-subdir",
        action="store_true",
        help="By default this script will create a subdirectory inside `dest` and write output to"
        " that. Setting this flag will suppress the subdirectory and use `dest` directly.",
    )
    parser.add_argument(
        "--no_plot",
        "--no-plot",
        action="store_true",
        help="By default this script will create plots of data."
        " Setting this flag will suppress plotting.",
    )
    parser.add_argument(
        "--cache_dir",
        "--cache-dir",
        type=Path,
        default=Path("/tmp/gpflow_benchmark_data"),
        help="Where to cache stuff, like datasets.",
    )
    parser.add_argument(
        "--shard",
        default=ShardingStrategy("no"),
        type=ShardingStrategy,
        help=(
            "Sharding strategy:"
            " If set to 'no' this script performs all necessary work."
            " If set to 'start' this script creates an output directory for the shards, and"
            " prints its path."
            " If set to the format '<i>/<n>', where 0 <= i < n then this script only runs"
            " benchmarks for shard 'i' out of 'n', and stores partial results in 'dest'."
            " This requires '--no_subdir'."
            " If set to 'collect' then this script assumes all benchmarks already have been"
            " computed, using the '<i>/<n>' commands, and merges their results into one file."
            " This requires '--no_subdir'."
        ),
    )
    parser.add_argument(
        "suite",
        type=str,
        choices=BENCHMARK_SUITES.names(),
        help="Which benchmark suite to run.",
    )
    parser.add_argument(
        "dest",
        type=Path,
        help="Directory to write results to.",
    )

    args = parser.parse_args()
    suite = BENCHMARK_SUITES.get(args.suite)
    run(
        args.shard,
        suite,
        args.dest,
        args.no_subdir,
        args.cache_dir,
        not args.no_plot,
    )


if __name__ == "__main__":
    run_from_argv()


# benchmark/sharding.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Code for managing sharding of benchmarks.

This will only create one metadata for the entire sharded run, and assume that all shards are run on
similar machines.
"""
from pathlib import Path
from typing import Collection


class ShardingStrategy:
    """
    Strategy for how to shard (split) the work.
    """

    def __init__(self, spec: str) -> None:
        """
        Valid ``spec``\s are:

        - ``no``: this script performs all necessary work.
        - ``start``: this script creates an output directory for the shards, and
          prints its path."
        - ``<i>/<n>``:, where 0 <= i < n this script only runs benchmarks for shard ``i`` out of
          ``n``, and stores partial results.
        - ``collect``: this script assumes all benchmarks already have been computed, using the
          ``<i>/<n>`` commands, and merges their results into one file.
        """
        self.spec = spec
        if spec == "no":
            self.setup_dest = True
            self.print_dest = False
            self.write_metadata = True
            self.run_benchmarks = True
            self.shard_i = 0
            self.shard_n = 1
            self.add_file_suffix = False
            self.collect = False
            self.plot = True
        elif spec == "start":
            self.setup_dest = True
            self.print_dest = True
            self.write_metadata = True
            self.run_benchmarks = False
            self.shard_i = 0
            self.shard_n = 1
            self.add_file_suffix = False
            self.collect = False
            self.plot = False
        elif spec == "collect":
            self.setup_dest = False
            self.print_dest = False
            self.write_metadata = False
            self.run_benchmarks = False
            self.shard_i = 0
            self.shard_n = 1
            self.add_file_suffix = False
            self.collect = True
            self.plot = True
        else:
            i_str, n_str = spec.split("/")
            self.setup_dest = False
            self.print_dest = False
            self.write_metadata = True
            self.run_benchmarks = True
            self.shard_i = int(i_str)
            self.shard_n = int(n_str)
            self.add_file_suffix = True
            self.collect = False
            self.plot = False
        assert not (self.run_benchmarks and self.collect)
        assert 0 <= self.shard_i < self.shard_n, (self.shard_i, self.shard_n)

    @property
    def file_suffix(self) -> str:
        return f"_{self.shard_i}of{self.shard_n}" if self.add_file_suffix else ""

    def find_shards(self, base_path: Path) -> Collection[Path]:
        parent = base_path.parent
        dot_suffixes = "".join(base_path.suffixes)
        glob = f"{base_path.stem}_*of*{dot_suffixes}"
        return tuple(parent.glob(glob))

    def __repr__(self) -> str:
        return self.spec


# benchmark/tag.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Infrastructure for tagging stuff, and checking whether a set of tags fulfill some criteria.
"""
from abc import ABC, abstractmethod
from typing import Any, Callable, Collection, Generic, Iterable, List, TypeVar

T = TypeVar("T")


class TagReq(Generic[T], ABC):
    """
    Requirements on tags.

    A ``TagReq`` may or may not be "satisfied" by a set of tags.

    A ``Tag`` is itself a ``TagReq`` that is satisfied by itself::

        t1.satisfied({t1})  # True
        t1.satisfied({t2})  # False

    A ``TagReq`` can be inverted by the ``~`` operator::

        (~t1).satisfied({t2})  # True
        (~t1).satisfied({t1})  # False

    You can take the disjunction of two ``TagReq`` with the ``&`` operator::

        (t1 & t2).satisfied({t1, t2})  # True
        (t1 & t2).satisfied({t1})  # False

    You can take the union of two ``TagReq`` with the ``|`` operator::

        (t1 | t2).satisfied({t1})  # True
        (t1 | t2).satisfied({t3})  # False

    Finally there's a special ``TagReq`` called ``NO_REQ`` that always is satisfied:

        NO_REQ.satisfied({})  # True

    """

    @abstractmethod
    def satisfied(self, tags: Collection["Tag[T]"]) -> bool:
        """ Return whether this requirement is satisfied by the given ``tags``. """

    def child_repr(self) -> str:
        """
        Do a pretty ``repr`` of a wrapped child requirement.

        Basically wraps the child in ``()`` if it is "complicated".
        """
        if isinstance(self, (Tag, NoTagReq)):
            return repr(self)
        return "(" + repr(self) + ")"

    def __and__(self, other: "TagReq[T]") -> "TagReq[T]":
        return ReducingTagReq.create_and(self, other)

    def __or__(self, other: "TagReq[T]") -> "TagReq[T]":
        return ReducingTagReq.create_or(self, other)

    def __invert__(self) -> "TagReq[T]":
        return NotTagReq(self)


class Tag(TagReq[T], Generic[T]):
    """
    A tag that can be included in requirements.

    Usually you want to subclass this for type safety::

        class MyTag(Tag["MyTag"]):
            pass

        my_tag_1 = MyTag("my_tag_1")
        my_tag_2 = MyTag("my_tag_2")
        ...

    See :class:``TagReq`` for how to define requirements on tags.
    """

    def __init__(self, name: str) -> None:
        self.name = name

    def satisfied(self, tags: Collection["Tag[T]"]) -> bool:
        return self in tags

    def __repr__(self) -> str:
        return self.name


class ReducingTagReq(TagReq[T], Generic[T]):
    """
    ``TagReq`` that combines several children.
    """

    def __init__(
        self, name: str, reducer: Callable[[Iterable[bool]], bool], children: Collection[TagReq[T]]
    ) -> None:
        self.name = name
        self.reducer = reducer
        self.children: List[TagReq[T]] = []
        for child in children:
            if isinstance(child, ReducingTagReq) and self.reducer == child.reducer:
                self.children.extend(child.children)
            else:
                self.children.append(child)

    @staticmethod
    def create_and(*children: TagReq[T]) -> TagReq[T]:
        return ReducingTagReq("&", all, children)

    @staticmethod
    def create_or(*children: TagReq[T]) -> TagReq[T]:
        return ReducingTagReq("|", any, children)

    def satisfied(self, tags: Collection[Tag[T]]) -> bool:
        return self.reducer(c.satisfied(tags) for c in self.children)

    def __repr__(self) -> str:
        return f" {self.name} ".join(f"{c.child_repr()}" for c in self.children)


class NotTagReq(TagReq[T], Generic[T]):
    """
    ``TagReq`` that inverts a child.
    """

    def __init__(self, child: TagReq[T]) -> None:
        self._child = child

    def satisfied(self, tags: Collection[Tag[T]]) -> bool:
        return not self._child.satisfied(tags)

    def __repr__(self) -> str:
        return f"~{self._child.child_repr()}"


class NoTagReq(TagReq[Any]):
    """
    ``TagReq`` that always is satisfied.
    """

    def satisfied(self, tags: Collection[Tag[T]]) -> bool:
        return True

    def __repr__(self) -> str:
        return "NO_REQ"


NO_REQ = NoTagReq()
""" Singleton no-op ``TagReq`` that always is satisfied. """


# gpflow/__init__.py

# Copyright 2016-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# flake8: noqa

from . import (
    conditionals,
    config,
    covariances,
    expectations,
    experimental,
    functions,
    inducing_variables,
    kernels,
    kullback_leiblers,
    likelihoods,
    logdensities,
    mean_functions,
    models,
    monitor,
    optimizers,
    posteriors,
    probability_distributions,
    quadrature,
    utilities,
)
from .base import Module, Parameter
from .config import default_float, default_int, default_jitter
from .utilities import set_trainable
from .versions import __version__

__all__ = [
    "Module",
    "Parameter",
    "__version__",
    "base",
    "ci_utils",
    "conditionals",
    "config",
    "covariances",
    "default_float",
    "default_int",
    "default_jitter",
    "expectations",
    "experimental",
    "functions",
    "inducing_variables",
    "kernels",
    "kullback_leiblers",
    "likelihoods",
    "logdensities",
    "mean_functions",
    "models",
    "monitor",
    "mypy_flags",
    "optimizers",
    "posteriors",
    "probability_distributions",
    "quadrature",
    "set_trainable",
    "type_flags",
    "utilities",
    "versions",
]


# gpflow/base.py

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from enum import Enum
from typing import TYPE_CHECKING, Any, List, Optional, Sequence, Tuple, Union

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
from typing_extensions import Final

from .config import default_float, default_summary_fmt
from .experimental.check_shapes import check_shapes
from .type_flags import GENERIC_NP_ARRAYS, NP_TYPE_CHECKING

if TYPE_CHECKING:  # pragma: no cover
    from IPython.lib import pretty

DType = Union[np.dtype, tf.DType]


if TYPE_CHECKING and (not NP_TYPE_CHECKING):  # pragma: no cover
    AnyNDArray = Any
else:
    if GENERIC_NP_ARRAYS:
        # It would be nice to use something more interesting than `Any` here, but it looks like
        # the infrastructure in the rest of the ecosystem isn't really set up for this
        # yet. Maybe when we get Python 3.11?
        AnyNDArray = np.ndarray[Any, Any]  # type: ignore[misc]
    else:
        AnyNDArray = Union[np.ndarray]  # type: ignore[misc]

VariableData = Union[List[Any], Tuple[Any], AnyNDArray, int, float]  # deprecated
Transform = Union[tfp.bijectors.Bijector]
Prior = Union[tfp.distributions.Distribution]


# We've left this as object until we've tested the performance consequences of using the full set
# (np.ndarray, tf.Tensor, tf.Variable, Parameter), see https://github.com/GPflow/GPflow/issues/1434
TensorLike: Final[Tuple[type, ...]] = (object,)
"""
:var TensorLike: Collection of tensor-like types for registering implementations with
    `multipledispatch` dispatchers.
"""


_NativeScalar = Union[int, float]
_Array = Sequence[Any]  # a nested array of int, float, bool etc. kept simple for readability

MeanAndVariance = Tuple[tf.Tensor, tf.Tensor]
SamplesMeanAndVariance = Tuple[tf.Tensor, tf.Tensor, tf.Tensor]


def _IS_PARAMETER(o: object) -> bool:
    return isinstance(o, Parameter)


def _IS_TRAINABLE_PARAMETER(o: object) -> bool:
    return isinstance(o, Parameter) and o.trainable


class Module(tf.Module):
    @property
    def parameters(self) -> Tuple["Parameter", ...]:
        return tuple(self._flatten(predicate=_IS_PARAMETER))

    @property
    def trainable_parameters(self) -> Tuple["Parameter", ...]:
        return tuple(self._flatten(predicate=_IS_TRAINABLE_PARAMETER))

    def _representation_table(self, object_name: str, tablefmt: Optional[str]) -> str:
        from .utilities import leaf_components, tabulate_module_summary

        repr_components = [object_name]
        if leaf_components(self):
            repr_components.append(tabulate_module_summary(self, tablefmt=tablefmt))
        return "\n".join(repr_components)

    def _repr_html_(self) -> str:
        """ Nice representation of GPflow objects in IPython/Jupyter notebooks """
        from html import escape

        return self._representation_table(escape(repr(self)), "html")

    def _repr_pretty_(self, p: "pretty.RepresentationPrinter", cycle: bool) -> None:
        """ Nice representation of GPflow objects in the IPython shell """
        repr_str = self._representation_table(repr(self), default_summary_fmt())
        p.text(repr_str)


class PriorOn(Enum):
    CONSTRAINED = "constrained"
    UNCONSTRAINED = "unconstrained"


class Parameter(tfp.util.TransformedVariable):
    def __init__(
        self,
        value: "TensorData",
        *,
        transform: Optional[Transform] = None,
        prior: Optional[Prior] = None,
        prior_on: Optional[Union[str, PriorOn]] = None,
        trainable: Optional[bool] = None,
        dtype: Optional[DType] = None,
        name: Optional[str] = None,
        unconstrained_shape: Optional[Sequence[Optional[int]]] = None,
        constrained_shape: Optional[Sequence[Optional[int]]] = None,
        shape: Optional[Sequence[Optional[int]]] = None,
    ):
        """A parameter retains both constrained and unconstrained representations. If no transform
        is provided, these two values will be the same.  It is often challenging to operate with
        unconstrained parameters. For example, a variance cannot be negative, therefore we need a
        positive constraint and it is natural to use constrained values.  A prior can be imposed
        either on the constrained version (default) or on the unconstrained version of the
        parameter.

        :param unconstrained_shape: Declare the shape of the unconstrained / pre-transformed values.
            Useful for setting dynamic shapes.
        :param constrained_shape: Declare the shape of the constrained / transformed values. Useful
            for setting dynamic shapes.
        :param shape: Convenience shortcut for setting both `unconstrained_shape` and
            `constrained_shape` to the same value.
        """
        if transform:
            name = name or transform.name

        if isinstance(value, Parameter):
            transform = transform or value.transform
            prior = prior or value.prior
            prior_on = prior_on or value.prior_on
            name = name or value.bijector.name
            trainable = value.trainable if trainable is None else trainable

            if dtype:
                tensor_value: TensorType = _cast_to_dtype(value, dtype)
            else:
                tensor_value = value
        else:
            if transform is None:
                transform = tfp.bijectors.Identity()

            prior_on = prior_on if prior_on else PriorOn.CONSTRAINED
            trainable = trainable if trainable is not None else True

            tensor_value = _cast_to_dtype(value, dtype)

        _validate_unconstrained_value(tensor_value, transform, dtype)

        if shape is not None:
            assert unconstrained_shape is None, "Cannot set both `shape` and `unconstrained_shape`."
            assert constrained_shape is None, "Cannot set both `shape` and `constrained_shape`."
            unconstrained_shape = shape
            constrained_shape = shape

        super().__init__(
            tensor_value,
            transform,
            dtype=tensor_value.dtype,
            trainable=trainable,
            name=name,
            shape=unconstrained_shape,
        )

        # TransformedVariable.__init__ doesn't allow us to pass an unconstrained / pre-transformed
        # shape, so we manually override it.
        if constrained_shape is not None:
            self._shape = tf.TensorShape(constrained_shape)

        self.prior: Optional[Prior] = prior
        self.prior_on = prior_on  # type: ignore[assignment]  # see https://github.com/python/mypy/issues/3004

    @check_shapes("return: []")
    def log_prior_density(self) -> tf.Tensor:
        """ Log of the prior probability density of the constrained variable. """

        if self.prior is None:
            return tf.convert_to_tensor(0.0, dtype=self.dtype)

        y = self

        if self.prior_on == PriorOn.CONSTRAINED:
            # evaluation is in same space as prior
            return tf.reduce_sum(self.prior.log_prob(y))

        else:
            # prior on unconstrained, but evaluating log-prior in constrained space
            x = self.unconstrained_variable
            log_p = tf.reduce_sum(self.prior.log_prob(x))

            if self.transform is not None:
                # need to include log|Jacobian| to account for coordinate transform
                log_det_jacobian = self.transform.inverse_log_det_jacobian(y, y.shape.ndims)
                log_p += tf.reduce_sum(log_det_jacobian)

            return log_p

    @property
    def prior_on(self) -> PriorOn:
        return self._prior_on

    @prior_on.setter
    def prior_on(self, value: Union[str, PriorOn]) -> None:
        self._prior_on = PriorOn(value)

    @property
    def unconstrained_variable(self) -> tf.Variable:
        return self._pretransformed_input

    @property
    def transform(self) -> Optional[Transform]:
        return self.bijector

    @property
    def trainable(self) -> bool:
        """
        `True` if this instance is trainable, else `False`.

        This attribute cannot be set directly. Use :func:`gpflow.set_trainable`.
        """
        return self.unconstrained_variable.trainable  # type: ignore[no-any-return]

    def assign(
        self,
        value: "TensorData",
        use_locking: bool = False,
        name: Optional[str] = None,
        read_value: bool = True,
    ) -> tf.Tensor:
        """
        Assigns constrained `value` to the unconstrained parameter's variable.
        It passes constrained value through parameter's transform first.

        Example::

            a = Parameter(2.0, transform=tfp.bijectors.Softplus())
            b = Parameter(3.0)

            a.assign(4.0)               # `a` parameter to `2.0` value.
            a.assign(tf.constant(5.0))  # `a` parameter to `5.0` value.
            a.assign(b)                 # `a` parameter to constrained value of `b`.


        :param value: Constrained tensor-like value.
        :param use_locking: If `True`, use locking during the assignment.
        :param name: The name of the operation to be created.
        :param read_value: if True, will return something which evaluates to the new
            value of the variable; if False will return the assign op.
        """
        unconstrained_value = _validate_unconstrained_value(value, self.transform, self.dtype)
        return self.unconstrained_variable.assign(
            unconstrained_value, use_locking=use_locking, name=name, read_value=read_value
        )


# These types are defined after "Parameter" to avoid forward references that breaks our
# documentation build:
TensorType = Union[AnyNDArray, tf.Tensor, tf.Variable, Parameter]
"""
Type alias for tensor-like types that are supported by most TensorFlow and GPflow operations.

NOTE: Union types like this do not work with the `register` method of `multipledispatch`'s
`Dispatcher` class. Instead use `TensorLike`.
"""

TensorData = Union[_NativeScalar, _Array, TensorType]
InputData = Union[TensorType]
OutputData = Union[TensorType]
RegressionData = Tuple[InputData, OutputData]


def _cast_to_dtype(
    value: TensorData, dtype: Optional[DType] = None
) -> Union[tf.Tensor, tf.Variable]:
    if dtype is None:
        dtype = default_float()

    if tf.is_tensor(value):
        # NOTE(awav) TF2.2 resolves issue with cast.
        # From TF2.2, `tf.cast` can be used alone instead of this auxiliary function.
        # workaround for https://github.com/tensorflow/tensorflow/issues/35938
        return tf.cast(value, dtype)
    else:
        return tf.convert_to_tensor(value, dtype=dtype)


def _validate_unconstrained_value(
    value: TensorData, transform: tfp.bijectors.Bijector, dtype: DType
) -> tf.Tensor:
    value = _cast_to_dtype(value, dtype)
    unconstrained_value = _to_unconstrained(value, transform)
    if unconstrained_value.dtype.is_integer:
        return unconstrained_value
    message = (
        "gpflow.Parameter: the value to be assigned is incompatible with this parameter's "
        "transform (the corresponding unconstrained value has NaN or Inf) and hence cannot be "
        "assigned."
    )
    return tf.debugging.assert_all_finite(unconstrained_value, message=message)


def _to_constrained(value: TensorType, transform: Optional[Transform]) -> TensorType:
    if transform is not None:
        return transform.forward(value)
    return value


def _to_unconstrained(value: TensorType, transform: Optional[Transform]) -> TensorType:
    if transform is not None:
        return transform.inverse(value)
    return value


# gpflow/ci_utils.py

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# pylint: skip-file

import os
from typing import Any, Iterable, List, Sequence, Type, TypeVar


def is_continuous_integration() -> bool:
    """
    Determines whether we are running on the Continuous Integration system for
    notebook integration tests. This is used to speed up notebook integration
    tests (built on every pull request commit) by capping all expensive loops
    at a small number, rather than running until convergence. When building the
    docs (indicated by the presence of the `DOCS` environment variable), we
    need to run notebooks to completion, and this function returns `False`.
    Whether we are running on CI is determined by the presence of the `CI`
    environment variable.
    """
    if "DOCS" in os.environ:
        return False

    return "CI" in os.environ


def reduce_in_tests(n: int, test_n: int = 2) -> int:
    return test_n if is_continuous_integration() else n


def subclasses(cls: Type[Any]) -> Iterable[Type[Any]]:
    """
    Generator that returns all (not just direct) subclasses of `cls`
    """
    for subclass in cls.__subclasses__():
        yield from subclasses(subclass)
        yield subclass


# gpflow/functions.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Throughout GPflow, by default, latent functions being modelled with Gaussian
processes are assumed to have zero mean, f ~ GP(0, k(x,x')).

In some cases we may wish to model only the deviation from a fixed function
with a Gaussian process.  For flexibility this fixed function could be both
input dependent and parameterised function, μ(x; θ),
with some unknown parameters θ, resulting in f ~ GP(μ(x;θ), k(x,x')).

The GPflow :class:`MeanFunction <gpflow.mean_functions.MeanFunction>` class
allows this to be done whilst additionally learning parameters of the
parametric function.
"""
from typing import Collection, Optional, Sequence, Tuple

import numpy as np
import tensorflow as tf

from .base import Module, Parameter, TensorType
from .config import default_float, default_int
from .experimental.check_shapes import check_shape as cs
from .experimental.check_shapes import check_shapes, inherit_check_shapes


class Function(Module):
    """
    The base function class.
    To implement a function, write the ``__call__`` method. This takes a
    tensor X and returns a tensor f(X). In accordance with the GPflow
    standard, each row of X represents one datum, and each row of Y is computed
    independently for each row of X.

    :class:`Function` classes can have parameters, see the :class:`Linear` class for an
    example.
    """

    @check_shapes(
        "X: [batch..., D]",
        "return: [batch..., Q]",
    )
    def __call__(self, X: TensorType) -> tf.Tensor:
        raise NotImplementedError("Implement the __call__ method for this mean function")

    def __add__(self, other: "Function") -> "Function":
        return Additive(self, other)

    def __mul__(self, other: "Function") -> "Function":
        return Product(self, other)


class MeanFunction(Function):
    """
    Mixin for :class:`Function`\ s that are appropriate for use as mean functions.
    """


class Additive(MeanFunction, Function):
    def __init__(self, first_part: Function, second_part: Function) -> None:
        MeanFunction.__init__(self)
        self.add_1 = first_part
        self.add_2 = second_part

    @inherit_check_shapes
    def __call__(self, X: TensorType) -> tf.Tensor:
        return tf.add(self.add_1(X), self.add_2(X))


class Product(MeanFunction, Function):
    def __init__(self, first_part: Function, second_part: Function):
        MeanFunction.__init__(self)

        self.prod_1 = first_part
        self.prod_2 = second_part

    @inherit_check_shapes
    def __call__(self, X: TensorType) -> tf.Tensor:
        return tf.multiply(self.prod_1(X), self.prod_2(X))


class Linear(MeanFunction, Function):
    """
    y_i = A x_i + b
    """

    @check_shapes(
        "A: [broadcast D, broadcast Q]",
        "b: [broadcast Q]",
    )
    def __init__(self, A: TensorType = None, b: TensorType = None) -> None:
        """
        A is a matrix which maps each element of X to Y, b is an additive
        constant.
        """
        MeanFunction.__init__(self)
        A = np.ones((1, 1), dtype=default_float()) if A is None else A
        b = np.zeros(1, dtype=default_float()) if b is None else b
        self.A = Parameter(np.atleast_2d(A))
        self.b = Parameter(b)

    @inherit_check_shapes
    def __call__(self, X: TensorType) -> tf.Tensor:
        return tf.tensordot(X, self.A, [[-1], [0]]) + self.b


class Identity(Linear, Function):
    """
    y_i = x_i
    """

    # The many type-ignores in this class is because we replace a field in the super class with a
    # property, which mypy doesn't like.

    def __init__(self, input_dim: Optional[int] = None) -> None:
        Linear.__init__(self)
        self.input_dim = input_dim

    @inherit_check_shapes
    def __call__(self, X: TensorType) -> tf.Tensor:
        return X

    @property
    def A(self) -> tf.Tensor:  # type: ignore[override]
        if self.input_dim is None:
            raise ValueError(
                "An input_dim needs to be specified when using the "
                "`Identity` mean function in combination with expectations."
            )
        return tf.eye(self.input_dim, dtype=default_float())

    @property
    def b(self) -> tf.Tensor:  # type: ignore[override]
        if self.input_dim is None:
            raise ValueError(
                "An input_dim needs to be specified when using the "
                "`Identity` mean function in combination with expectations."
            )

        return tf.zeros(self.input_dim, dtype=default_float())

    @A.setter  # type: ignore[attr-defined, no-redef]
    def A(self, A: tf.Tensor) -> None:
        pass

    @b.setter  # type: ignore[attr-defined, no-redef]
    def b(self, b: tf.Tensor) -> None:
        pass


class Constant(MeanFunction, Function):
    @check_shapes(
        "c: [broadcast Q]",
    )
    def __init__(self, c: TensorType = None) -> None:
        super().__init__()
        c = np.zeros(1) if c is None else c
        self.c = Parameter(c)

    @inherit_check_shapes
    def __call__(self, X: TensorType) -> tf.Tensor:
        tile_shape = tf.concat(
            [tf.shape(X)[:-1], [1]],
            axis=0,
        )
        reshape_shape = tf.concat(
            [tf.ones(shape=(tf.rank(X) - 1), dtype=default_int()), [-1]],
            axis=0,
        )
        return tf.tile(tf.reshape(self.c, reshape_shape), tile_shape)


class Zero(Constant, Function):
    def __init__(self, output_dim: int = 1) -> None:
        Constant.__init__(self)
        self.output_dim = output_dim
        del self.c

    @inherit_check_shapes
    def __call__(self, X: TensorType) -> tf.Tensor:
        output_shape = tf.concat([tf.shape(X)[:-1], [self.output_dim]], axis=0)
        return tf.zeros(output_shape, dtype=X.dtype)


class Polynomial(MeanFunction, Function):
    """
    A generic polynomial mean function.
    """

    @check_shapes("w: [broadcast output_dim, broadcast n_terms]")
    def __init__(
        self, degree: int, input_dim: int = 1, output_dim: int = 1, w: Optional[TensorType] = None
    ) -> None:
        """
        :param degree: The degree of the polynomial.
        :param input_dim: Number of inputs / variables this polynomial is defined over.
        :param output_dim: Number of outputs / polynomials.
        :param w: Initial weights of the terms of the polynomial. The inner dimension (``n_terms``)
            should correspond to the powers returned by ``compute_powers``.
        """
        powers = cs(tuple(self.compute_powers(degree, input_dim)), "[n_terms, input_dim]")
        if w is None:
            w = cs([1.0] + (len(powers) - 1) * [0.0], "[n_terms]")
        w_shape = (output_dim, len(powers))
        self.powers = tf.constant(powers, dtype=default_float())
        self.w = Parameter(tf.broadcast_to(w, w_shape))

    @staticmethod
    def compute_powers(degree: int, input_dim: int) -> Sequence[Tuple[int, ...]]:
        """
        Computes integer tuples corresponding to the powers to raise inputs to.

        Specifically this returns, in lexicographical order, all tuples where:

        * The tuple has length `input_dim`.
        * The values are non-negative integers.
        * The sum of the tuple is no greater than `degree`.

        For example::

            compute_powers(degree=2, input_dim=3)

        returns::

            (0, 0, 0)
            (0, 0, 1)
            (0, 0, 2)
            (0, 1, 0)
            (0, 1, 1)
            (0, 2, 0)
            (1, 0, 0)
            (1, 0, 1)
            (1, 1, 0)
            (2, 0, 0)

        where a tuple::

            (1, 0, 2)

        will translate to a the term::

            w[i] * (x[0]**1) * (x[1]**0) * (x[2]**2)
        """
        if not input_dim:
            return [()]
        result = []
        for i in range(degree + 1):
            for inner in Polynomial.compute_powers(degree - i, input_dim - 1):
                result.append((i,) + inner)
        return result

    @inherit_check_shapes
    def __call__(self, X: TensorType) -> tf.Tensor:
        raised = cs(tf.pow(X[..., None, :], self.powers), "[batch..., n_terms, input_dim]")
        prod = cs(tf.math.reduce_prod(raised, axis=-1), "[batch..., n_terms]")
        return tf.einsum("...i,ji->...j", prod, self.w)


class SwitchedFunction(MeanFunction, Function):
    """
    This class enables to use different (independent) functions respective
    to the data 'label'.
    We assume the 'label' is stored in the extra column of X.
    """

    def __init__(self, function_list: Collection[Function]) -> None:
        super().__init__()
        self.functions = function_list

    @inherit_check_shapes
    def __call__(self, X: TensorType) -> tf.Tensor:
        ind = tf.gather(tf.transpose(X), tf.shape(X)[1] - 1)  # ind = X[:,-1]
        ind = tf.cast(ind, tf.int32)
        X = tf.transpose(
            tf.gather(tf.transpose(X), tf.range(0, tf.shape(X)[1] - 1))
        )  # X = X[:,:-1]

        # split up X into chunks corresponding to the relevant likelihoods
        x_list = tf.dynamic_partition(X, ind, len(self.functions))
        # apply the likelihood-function to each section of the data
        results = [m(x) for x, m in zip(x_list, self.functions)]
        # stitch the results back together
        partitions = tf.dynamic_partition(tf.range(0, tf.size(ind)), ind, len(self.functions))
        return tf.dynamic_stitch(partitions, results)


class SwitchedMeanFunction(SwitchedFunction):
    """
    Renamed :class:`SwitchedFunction` for backwards compatibility.
    """

    def __init__(self, meanfunction_list: Collection[MeanFunction]) -> None:
        super().__init__(function_list=meanfunction_list)

    @property
    def meanfunctions(self) -> Collection[MeanFunction]:
        return self.functions

    @meanfunctions.setter
    def meanfunctions(self, value: Collection[MeanFunction]) -> None:
        self.function_list = value


# gpflow/kullback_leiblers.py

# Copyright 2016-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# -*- coding: utf-8 -*-

import tensorflow as tf
from packaging.version import Version

from .base import TensorType
from .config import default_float, default_jitter
from .covariances import Kuu
from .experimental.check_shapes import check_shapes
from .inducing_variables import InducingVariables
from .kernels import Kernel
from .utilities import Dispatcher, to_default_float

prior_kl = Dispatcher("prior_kl")


@prior_kl.register(InducingVariables, Kernel, object, object)
@check_shapes(
    "inducing_variable: [N, D, broadcast L]",
    "q_mu: [M, L]",
    "q_sqrt: [M_L_or_L_M_M...]",
    "return: []",
)
def _(
    inducing_variable: InducingVariables,
    kernel: Kernel,
    q_mu: TensorType,
    q_sqrt: TensorType,
    whiten: bool = False,
) -> tf.Tensor:
    if whiten:
        return gauss_kl(q_mu, q_sqrt, None)
    else:
        K = Kuu(inducing_variable, kernel, jitter=default_jitter())  # [P, M, M] or [M, M]
        return gauss_kl(q_mu, q_sqrt, K)


@check_shapes(
    "q_mu: [M, L]",
    "q_sqrt: [M_L_or_L_M_M...]",
    "K: [broadcast L, M, M]",
    "K_cholesky: [broadcast L, M, M]",
    "return: []",
)
def gauss_kl(
    q_mu: TensorType, q_sqrt: TensorType, K: TensorType = None, *, K_cholesky: TensorType = None
) -> tf.Tensor:
    """
    Compute the KL divergence KL[q || p] between::

          q(x) = N(q_mu, q_sqrt^2)

    and::

          p(x) = N(0, K)    if K is not None
          p(x) = N(0, I)    if K is None

    We assume L multiple independent distributions, given by the columns of
    q_mu and the first or last dimension of q_sqrt. Returns the *sum* of the
    divergences.

    q_mu is a matrix ([M, L]), each column contains a mean.

    - q_sqrt can be a 3D tensor ([L, M, M]), each matrix within is a lower
      triangular square-root matrix of the covariance of q.
    - q_sqrt can be a matrix ([M, L]), each column represents the diagonal of a
      square-root matrix of the covariance of q.

    K is the covariance of p (positive-definite matrix).  The K matrix can be
    passed either directly as `K`, or as its Cholesky factor, `K_cholesky`.  In
    either case, it can be a single matrix [M, M], in which case the sum of the
    L KL divergences is computed by broadcasting, or L different covariances
    [L, M, M].

    Note: if no K matrix is given (both `K` and `K_cholesky` are None),
    `gauss_kl` computes the KL divergence from p(x) = N(0, I) instead.
    """

    if (K is not None) and (K_cholesky is not None):
        raise ValueError(
            "Ambiguous arguments: gauss_kl() must only be passed one of `K` or `K_cholesky`."
        )

    is_white = (K is None) and (K_cholesky is None)
    is_diag = len(q_sqrt.shape) == 2

    M, L = tf.shape(q_mu)[0], tf.shape(q_mu)[1]

    if is_white:
        alpha = q_mu  # [M, L]
    else:
        if K is not None:
            Lp = tf.linalg.cholesky(K)  # [L, M, M] or [M, M]
        elif K_cholesky is not None:
            Lp = K_cholesky  # [L, M, M] or [M, M]

        is_batched = len(Lp.shape) == 3

        q_mu = tf.transpose(q_mu)[:, :, None] if is_batched else q_mu  # [L, M, 1] or [M, L]
        alpha = tf.linalg.triangular_solve(Lp, q_mu, lower=True)  # [L, M, 1] or [M, L]

    if is_diag:
        Lq = Lq_diag = q_sqrt
        Lq_full = tf.linalg.diag(tf.transpose(q_sqrt))  # [L, M, M]
    else:
        Lq = Lq_full = tf.linalg.band_part(q_sqrt, -1, 0)  # force lower triangle # [L, M, M]
        Lq_diag = tf.linalg.diag_part(Lq)  # [M, L]

    # Mahalanobis term: μqᵀ Σp⁻¹ μq
    mahalanobis = tf.reduce_sum(tf.square(alpha))

    # Constant term: - L * M
    constant = -to_default_float(tf.size(q_mu, out_type=tf.int64))

    # Log-determinant of the covariance of q(x):
    logdet_qcov = tf.reduce_sum(tf.math.log(tf.square(Lq_diag)))

    # Trace term: tr(Σp⁻¹ Σq)
    if is_white:
        trace = tf.reduce_sum(tf.square(Lq))
    else:
        if is_diag and not is_batched:
            # K is [M, M] and q_sqrt is [M, L]: fast specialisation
            LpT = tf.transpose(Lp)  # [M, M]
            Lp_inv = tf.linalg.triangular_solve(
                Lp, tf.eye(M, dtype=default_float()), lower=True
            )  # [M, M]
            K_inv = tf.linalg.diag_part(tf.linalg.triangular_solve(LpT, Lp_inv, lower=False))[
                :, None
            ]  # [M, M] -> [M, 1]
            trace = tf.reduce_sum(K_inv * tf.square(q_sqrt))
        else:
            if is_batched or Version(tf.__version__) >= Version("2.2"):
                Lp_full = Lp
            else:
                # workaround for segfaults when broadcasting in TensorFlow<2.2
                Lp_full = tf.tile(tf.expand_dims(Lp, 0), [L, 1, 1])
            LpiLq = tf.linalg.triangular_solve(Lp_full, Lq_full, lower=True)
            trace = tf.reduce_sum(tf.square(LpiLq))

    twoKL = mahalanobis + constant - logdet_qcov + trace

    # Log-determinant of the covariance of p(x):
    if not is_white:
        log_sqdiag_Lp = tf.math.log(tf.square(tf.linalg.diag_part(Lp)))
        sum_log_sqdiag_Lp = tf.reduce_sum(log_sqdiag_Lp)
        # If K is [L, M, M], num_latent_gps is no longer implicit, no need to multiply the single kernel logdet
        scale = 1.0 if is_batched else to_default_float(L)
        twoKL += scale * sum_log_sqdiag_Lp

    return 0.5 * twoKL


# gpflow/logdensities.py

# Copyright 2016-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
import tensorflow as tf

from .base import TensorType
from .experimental.check_shapes import check_shapes
from .utilities import to_default_float


@check_shapes(
    "x: [broadcast shape...]",
    "mu: [broadcast shape...]",
    "var: [broadcast shape...]",
    "return: [shape...]",
)
def gaussian(x: TensorType, mu: TensorType, var: TensorType) -> tf.Tensor:
    return -0.5 * (np.log(2 * np.pi) + tf.math.log(var) + tf.square(mu - x) / var)


@check_shapes(
    "x: [broadcast shape...]",
    "mu: [broadcast shape...]",
    "var: [broadcast shape...]",
    "return: [shape...]",
)
def lognormal(x: TensorType, mu: TensorType, var: TensorType) -> tf.Tensor:
    lnx = tf.math.log(x)
    return gaussian(lnx, mu, var) - lnx


@check_shapes(
    "x: [broadcast shape...]",
    "p: [broadcast shape...]",
    "return: [shape...]",
)
def bernoulli(x: TensorType, p: TensorType) -> tf.Tensor:
    return tf.math.log(tf.where(tf.equal(x, 1), p, 1 - p))


@check_shapes(
    "x: [broadcast shape...]",
    "lam: [broadcast shape...]",
    "return: [shape...]",
)
def poisson(x: TensorType, lam: TensorType) -> tf.Tensor:
    return x * tf.math.log(lam) - lam - tf.math.lgamma(x + 1.0)


@check_shapes(
    "x: [broadcast shape...]",
    "scale: [broadcast shape...]",
    "return: [shape...]",
)
def exponential(x: TensorType, scale: TensorType) -> tf.Tensor:
    return -x / scale - tf.math.log(scale)


@check_shapes(
    "x: [broadcast shape...]",
    "shape: [broadcast shape...]",
    "scale: [broadcast shape...]",
    "return: [shape...]",
)
def gamma(x: TensorType, shape: TensorType, scale: TensorType) -> tf.Tensor:
    return (
        -shape * tf.math.log(scale)
        - tf.math.lgamma(shape)
        + (shape - 1.0) * tf.math.log(x)
        - x / scale
    )


@check_shapes(
    "x: [broadcast shape...]",
    "mean: [broadcast shape...]",
    "scale: [broadcast shape...]",
    "df: [broadcast shape...]",
    "return: [shape...]",
)
def student_t(x: TensorType, mean: TensorType, scale: TensorType, df: TensorType) -> tf.Tensor:
    df = to_default_float(df)
    const = (
        tf.math.lgamma((df + 1.0) * 0.5)
        - tf.math.lgamma(df * 0.5)
        - 0.5 * (tf.math.log(tf.square(scale)) + tf.math.log(df) + np.log(np.pi))
    )
    return const - 0.5 * (df + 1.0) * tf.math.log(
        1.0 + (1.0 / df) * (tf.square((x - mean) / scale))
    )


@check_shapes(
    "x: [broadcast shape...]",
    "alpha: [broadcast shape...]",
    "beta: [broadcast shape...]",
    "return: [shape...]",
)
def beta(x: TensorType, alpha: TensorType, beta: TensorType) -> tf.Tensor:
    # need to clip x, since log of 0 is nan...
    x = tf.clip_by_value(x, 1e-6, 1 - 1e-6)
    return (
        (alpha - 1.0) * tf.math.log(x)
        + (beta - 1.0) * tf.math.log(1.0 - x)
        + tf.math.lgamma(alpha + beta)
        - tf.math.lgamma(alpha)
        - tf.math.lgamma(beta)
    )


@check_shapes(
    "x: [broadcast shape...]",
    "mu: [broadcast shape...]",
    "sigma: [broadcast shape...]",
    "return: [shape...]",
)
def laplace(x: TensorType, mu: TensorType, sigma: TensorType) -> tf.Tensor:
    return -tf.abs(mu - x) / sigma - tf.math.log(2.0 * sigma)


@check_shapes(
    "x: [D, broadcast N]",
    "mu: [D, broadcast N]",
    "L: [D, D]",
    "return: [N]",
)
def multivariate_normal(x: TensorType, mu: TensorType, L: TensorType) -> tf.Tensor:
    """
    Computes the log-density of a multivariate normal.

    :param x: sample(s) for which we want the density
    :param mu: mean(s) of the normal distribution
    :param L: Cholesky decomposition of the covariance matrix
    :return: log densities
    """

    d = x - mu
    alpha = tf.linalg.triangular_solve(L, d, lower=True)
    num_dims = tf.cast(tf.shape(d)[0], L.dtype)
    p = -0.5 * tf.reduce_sum(tf.square(alpha), 0)
    p -= 0.5 * num_dims * np.log(2 * np.pi)
    p -= tf.reduce_sum(tf.math.log(tf.linalg.diag_part(L)))

    return p


# gpflow/mean_functions.py

# Copyright 2016-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Alias for (deprecated) mean_functions.py.

Use functions.py instead.
"""

from .functions import (
    Additive,
    Constant,
    Identity,
    Linear,
    MeanFunction,
    Polynomial,
    Product,
    SwitchedMeanFunction,
    Zero,
)

__all__ = [
    "Additive",
    "Constant",
    "Identity",
    "Linear",
    "MeanFunction",
    "Polynomial",
    "Product",
    "SwitchedMeanFunction",
    "Zero",
]


# gpflow/mypy_flags.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Code for printing flags for mypy, depending on library versions.
"""
from .type_flags import compute_mypy_flags


def print_mypy_flags() -> None:
    print(compute_mypy_flags())


if __name__ == "__main__":
    print_mypy_flags()


# gpflow/posteriors.py

# Copyright 2016-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import enum
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Optional, Tuple, Type, Union, cast

import tensorflow as tf

from . import covariances, kernels, mean_functions
from .base import MeanAndVariance, Module, RegressionData, TensorType
from .conditionals.util import (
    base_conditional,
    base_conditional_with_lm,
    expand_independent_outputs,
    fully_correlated_conditional,
    independent_interdomain_conditional,
    mix_latent_gp,
    separate_independent_conditional_implementation,
)
from .config import default_float, default_jitter
from .covariances import Kuf, Kuu
from .experimental.check_shapes import (
    ErrorContext,
    Shape,
    check_shapes,
    get_shape,
    inherit_check_shapes,
    register_get_shape,
)
from .inducing_variables import (
    FallbackSeparateIndependentInducingVariables,
    FallbackSharedIndependentInducingVariables,
    InducingPoints,
    InducingVariables,
    SeparateIndependentInducingVariables,
    SharedIndependentInducingVariables,
)
from .kernels import Kernel
from .likelihoods import Gaussian
from .mean_functions import MeanFunction
from .utilities import Dispatcher, add_likelihood_noise_cov, assert_params_false
from .utilities.ops import eye


class _QDistribution(Module):
    """
    Base class for our parametrization of q(u) in the `AbstractPosterior`.
    Internal - do not rely on this outside of GPflow.
    """


class _DeltaDist(_QDistribution):
    @check_shapes(
        "q_mu: [M, L]",
    )
    def __init__(self, q_mu: TensorType) -> None:
        self.q_mu = q_mu

    @property
    def q_sqrt(self) -> Optional[tf.Tensor]:
        return None


class _DiagNormal(_QDistribution):
    @check_shapes(
        "q_mu: [M, L]",
        "q_sqrt: [M, L]",
    )
    def __init__(self, q_mu: TensorType, q_sqrt: TensorType) -> None:
        self.q_mu = q_mu
        self.q_sqrt = q_sqrt


class _MvNormal(_QDistribution):
    @check_shapes(
        "q_mu: [M, L]",
        "q_sqrt: [L, M, M]  # lower-triangular",
    )
    def __init__(self, q_mu: TensorType, q_sqrt: TensorType) -> None:
        self.q_mu = q_mu
        self.q_sqrt = q_sqrt


class PrecomputeCacheType(enum.Enum):
    """
    - `PrecomputeCacheType.TENSOR` (or `"tensor"`): Precomputes the cached
      quantities and stores them as tensors (which allows differentiating
      through the prediction). This is the default.
    - `PrecomputeCacheType.VARIABLE` (or `"variable"`): Precomputes the cached
      quantities and stores them as variables, which allows for updating
      their values without changing the compute graph (relevant for AOT
      compilation).
    - `PrecomputeCacheType.NOCACHE` (or `"nocache"` or `None`): Avoids
      immediate cache computation. This is useful for avoiding extraneous
      computations when you only want to call the posterior's
      `fused_predict_f` method.
    """

    TENSOR = "tensor"
    VARIABLE = "variable"
    NOCACHE = "nocache"


@dataclass
class PrecomputedValue:
    value: tf.Tensor
    """
    The precomputed value itself.
    """

    axis_dynamic: Tuple[bool, ...]
    """
    A tuple with one element per dimension of `value`. That element is `True` if that dimension
    of `value` might change size.
    """

    def __post_init__(self) -> None:
        tf.debugging.assert_rank(
            self.value,
            len(self.axis_dynamic),
            "axis_dynamic must have one element per dimension of value.",
        )

    @staticmethod
    @check_shapes(
        "alpha: [M_L_or_L_M_M...]",
        "Qinv: [M_M_or_L_M_M...]",
    )
    def wrap_alpha_Qinv(alpha: TensorType, Qinv: TensorType) -> Tuple["PrecomputedValue", ...]:
        """
        Wraps `alpha` and `Qinv` in `PrecomputedValue`\ s.
        """
        one_dynamic = False
        L_dynamic = False
        M_dynamic = False  # TODO(jesper): Support variable number of inducing points?

        alpha_rank = tf.rank(alpha)
        if alpha_rank == 2:
            alpha_dynamic: Tuple[bool, ...] = (M_dynamic, L_dynamic)
        elif alpha_rank == 3:
            alpha_dynamic = (L_dynamic, M_dynamic, one_dynamic)
        else:
            raise AssertionError(f"Unknown rank of alpha {alpha_rank}.")

        Qinv_rank = tf.rank(Qinv)
        if Qinv_rank == 2:
            Qinv_dynamic: Tuple[bool, ...] = (M_dynamic, M_dynamic)
        elif Qinv_rank == 3:
            Qinv_dynamic = (L_dynamic, M_dynamic, M_dynamic)
        else:
            raise AssertionError(f"Unknown rank of Qinv {Qinv_rank}.")

        return (
            PrecomputedValue(alpha, alpha_dynamic),
            PrecomputedValue(Qinv, Qinv_dynamic),
        )


@register_get_shape(PrecomputedValue)
def get_precomputed_value_shape(shaped: PrecomputedValue, context: ErrorContext) -> Shape:
    return get_shape(shaped.value, context)


def _validate_precompute_cache_type(
    value: Union[None, PrecomputeCacheType, str]
) -> PrecomputeCacheType:
    if value is None:
        return PrecomputeCacheType.NOCACHE
    elif isinstance(value, PrecomputeCacheType):
        return value
    elif isinstance(value, str):
        return PrecomputeCacheType(value.lower())
    else:
        raise ValueError(
            f"{value} is not a valid PrecomputeCacheType."
            " Valid options: 'tensor', 'variable', 'nocache' (or None)."
        )


class AbstractPosterior(Module, ABC):
    @check_shapes(
        "X_data: [N_D_or_M_D_P...]",
    )
    def __init__(
        self,
        kernel: Kernel,
        X_data: Union[tf.Tensor, InducingVariables],
        cache: Optional[Tuple[tf.Tensor, ...]] = None,
        mean_function: Optional[mean_functions.MeanFunction] = None,
    ) -> None:
        """
        Users should use `create_posterior` to create instances of concrete
        subclasses of this AbstractPosterior class instead of calling this
        constructor directly. For `create_posterior` to be able to correctly
        instantiate subclasses, developers need to ensure their subclasses
        don't change the constructor signature.
        """
        super().__init__()

        self.kernel = kernel
        self.X_data = X_data
        self.cache = cache
        self.mean_function = mean_function

        self._precompute_cache: Optional[PrecomputeCacheType] = None

    @check_shapes(
        "Xnew: [batch..., D]",
        "mean: [batch..., Q]",
        "return: [batch..., Q]",
    )
    def _add_mean_function(self, Xnew: TensorType, mean: TensorType) -> tf.Tensor:
        if self.mean_function is None:
            return mean
        else:
            return mean + self.mean_function(Xnew)

    @abstractmethod
    def _precompute(self) -> Tuple[PrecomputedValue, ...]:
        """
        Precompute a cache.

        The result of this method will later be passed to `_conditional_with_precompute` as the
        `cache` argument.
        """

    @check_shapes(
        "Xnew: [batch..., N, D]",
        "return[0]: [batch..., N, P]",
        "return[1]: [batch..., N, P, N, P] if full_cov and full_output_cov",
        "return[1]: [batch..., P, N, N] if full_cov and (not full_output_cov)",
        "return[1]: [batch..., N, P, P] if (not full_cov) and full_output_cov",
        "return[1]: [batch..., N, P] if (not full_cov) and (not full_output_cov)",
    )
    def fused_predict_f(
        self, Xnew: TensorType, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        Computes predictive mean and (co)variance at Xnew, including mean_function
        Does not make use of caching
        """
        mean, cov = self._conditional_fused(
            Xnew, full_cov=full_cov, full_output_cov=full_output_cov
        )
        return self._add_mean_function(Xnew, mean), cov

    @abstractmethod
    @check_shapes(
        "Xnew: [batch..., N, D]",
        "return[0]: [batch..., N, P]",
        "return[1]: [batch..., N, P, N, P] if full_cov and full_output_cov",
        "return[1]: [batch..., P, N, N] if full_cov and (not full_output_cov)",
        "return[1]: [batch..., N, P, P] if (not full_cov) and full_output_cov",
        "return[1]: [batch..., N, P] if (not full_cov) and (not full_output_cov)",
    )
    def _conditional_fused(
        self, Xnew: TensorType, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        Computes predictive mean and (co)variance at Xnew, *excluding* mean_function
        Does not make use of caching
        """

    @check_shapes(
        "Xnew: [batch..., N, D]",
        "return[0]: [batch..., N, P]",
        "return[1]: [batch..., N, P, N, P] if full_cov and full_output_cov",
        "return[1]: [batch..., P, N, N] if full_cov and (not full_output_cov)",
        "return[1]: [batch..., N, P, P] if (not full_cov) and full_output_cov",
        "return[1]: [batch..., N, P] if (not full_cov) and (not full_output_cov)",
    )
    def predict_f(
        self, Xnew: TensorType, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        Computes predictive mean and (co)variance at Xnew, including mean_function.
        Relies on precomputed alpha and Qinv (see _precompute method)
        """
        if self.cache is None:
            raise ValueError(
                "Cache has not been precomputed yet. Call update_cache first or use fused_predict_f"
            )
        mean, cov = self._conditional_with_precompute(
            self.cache, Xnew, full_cov=full_cov, full_output_cov=full_output_cov
        )
        return self._add_mean_function(Xnew, mean), cov

    @abstractmethod
    @check_shapes(
        "Xnew: [batch..., N, D]",
        "return[0]: [batch..., N, P]",
        "return[1]: [batch..., N, P, N, P] if full_cov and full_output_cov",
        "return[1]: [batch..., P, N, N] if full_cov and (not full_output_cov)",
        "return[1]: [batch..., N, P, P] if (not full_cov) and full_output_cov",
        "return[1]: [batch..., N, P] if (not full_cov) and (not full_output_cov)",
    )
    def _conditional_with_precompute(
        self,
        cache: Tuple[tf.Tensor, ...],
        Xnew: TensorType,
        full_cov: bool = False,
        full_output_cov: bool = False,
    ) -> MeanAndVariance:
        """
        Computes predictive mean and (co)variance at Xnew, *excluding* mean_function.
        Relies on cached alpha and Qinv.
        """

    def update_cache(self, precompute_cache: Optional[PrecomputeCacheType] = None) -> None:
        """
        Sets the cache depending on the value of `precompute_cache` to a
        `tf.Tensor`, `tf.Variable`, or clears the cache. If `precompute_cache`
        is not given, the setting defaults to the most-recently-used one.
        """
        if precompute_cache is None:
            if self._precompute_cache is None:
                raise ValueError(
                    "You must pass precompute_cache explicitly"
                    " (the cache had not been updated before)."
                )
            precompute_cache = self._precompute_cache
        else:
            self._precompute_cache = precompute_cache

        if precompute_cache is PrecomputeCacheType.NOCACHE:
            self.cache = None

        elif precompute_cache is PrecomputeCacheType.TENSOR:
            self.cache = tuple(c.value for c in self._precompute())

        elif precompute_cache is PrecomputeCacheType.VARIABLE:
            cache = self._precompute()

            if self.cache is not None and all(isinstance(c, tf.Variable) for c in self.cache):
                # re-use existing variables
                for cache_var, c in zip(self.cache, cache):
                    cache_var.assign(c.value)
            else:  # create variables
                shapes = [
                    [None if d else s for d, s in zip(c.axis_dynamic, tf.shape(c.value))]
                    for c in cache
                ]
                self.cache = tuple(
                    tf.Variable(c.value, trainable=False, shape=s) for c, s in zip(cache, shapes)
                )


class GPRPosterior(AbstractPosterior):
    @check_shapes(
        "data[0]: [N, D]",
        "data[1]: [N, Q]",
    )
    def __init__(
        self,
        kernel: Kernel,
        data: RegressionData,
        likelihood: Gaussian,
        mean_function: MeanFunction,
        *,
        precompute_cache: Optional[PrecomputeCacheType],
    ) -> None:
        X, Y = data
        super().__init__(kernel, X, mean_function=mean_function)
        self.Y_data = Y
        self.likelihood = likelihood

        if precompute_cache is not None:
            self.update_cache(precompute_cache)

    @inherit_check_shapes
    def _conditional_with_precompute(
        self,
        cache: Tuple[tf.Tensor, ...],
        Xnew: TensorType,
        full_cov: bool = False,
        full_output_cov: bool = False,
    ) -> MeanAndVariance:
        """
        Computes predictive mean and (co)variance at Xnew, *excluding* mean_function.
        Relies on cached alpha and Qinv.
        """
        assert_params_false(self._conditional_with_precompute, full_output_cov=full_output_cov)
        err, Lm = cache

        Knn = self.kernel(Xnew, full_cov=full_cov)
        Kmn = self.kernel(self.X_data, Xnew)

        return base_conditional_with_lm(
            Kmn=Kmn,
            Lm=Lm,
            Knn=Knn,
            f=err,
            full_cov=full_cov,
            q_sqrt=None,
            white=False,
        )

    @check_shapes(
        "return[0]: [M, D]",
        "return[1]: [M, M]",
    )
    def _precompute(self) -> Tuple[PrecomputedValue, ...]:
        assert self.mean_function is not None
        X_data = cast(tf.Tensor, self.X_data)
        err = self.Y_data - self.mean_function(X_data)

        Kmm = self.kernel(X_data)
        Kmm_plus_s = add_likelihood_noise_cov(Kmm, self.likelihood, X_data)
        Lm = tf.linalg.cholesky(Kmm_plus_s)

        D = err.shape[1]
        M = X_data.shape[0]
        D_dynamic = D is None
        M_dynamic = M is None

        return (
            PrecomputedValue(err, (M_dynamic, D_dynamic)),
            PrecomputedValue(Lm, (M_dynamic, M_dynamic)),
        )

    @inherit_check_shapes
    def _conditional_fused(
        self, Xnew: TensorType, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        Computes predictive mean and (co)variance at Xnew, *excluding* mean_function
        Does not make use of caching
        """
        temp_cache = tuple(c.value for c in self._precompute())
        return self._conditional_with_precompute(temp_cache, Xnew, full_cov, full_output_cov)


class SGPRPosterior(AbstractPosterior):
    """
    This class represents posteriors which can be derived from SGPR
    models to compute faster predictions on unseen points.
    """

    @check_shapes(
        "data[0]: [N, D]",
        "data[1]: [N, Q]",
        "inducing_variable: [M, D, 1]",
    )
    def __init__(
        self,
        kernel: Kernel,
        data: RegressionData,
        inducing_variable: InducingPoints,
        likelihood: Gaussian,
        num_latent_gps: int,
        mean_function: MeanFunction,
        *,
        precompute_cache: Optional[PrecomputeCacheType],
    ) -> None:
        X, Y = data
        super().__init__(kernel, X, mean_function=mean_function)
        self.Y_data = Y
        self.likelihood = likelihood
        self.inducing_variable = inducing_variable
        self.num_latent_gps = num_latent_gps

        if precompute_cache is not None:
            self.update_cache(precompute_cache)

    @inherit_check_shapes
    def _conditional_with_precompute(
        self,
        cache: Tuple[tf.Tensor, ...],
        Xnew: TensorType,
        full_cov: bool = False,
        full_output_cov: bool = False,
    ) -> MeanAndVariance:
        """
        Computes predictive mean and (co)variance at Xnew, *excluding* mean_function.
        Relies on cached alpha and Qinv.
        """
        assert_params_false(self._conditional_with_precompute, full_output_cov=full_output_cov)

        L, LB, c = cache

        Kus = Kuf(self.inducing_variable, self.kernel, Xnew)
        tmp1 = tf.linalg.triangular_solve(L, Kus, lower=True)
        tmp2 = tf.linalg.triangular_solve(LB, tmp1, lower=True)
        mean = tf.linalg.matmul(tmp2, c, transpose_a=True)
        if full_cov:
            var = (
                self.kernel(Xnew)
                + tf.linalg.matmul(tmp2, tmp2, transpose_a=True)
                - tf.linalg.matmul(tmp1, tmp1, transpose_a=True)
            )
            var = tf.tile(var[None, ...], [self.num_latent_gps, 1, 1])  # [P, N, N]
        else:
            var = (
                self.kernel(Xnew, full_cov=False)
                + tf.reduce_sum(tf.square(tmp2), 0)
                - tf.reduce_sum(tf.square(tmp1), 0)
            )
            var = tf.tile(var[:, None], [1, self.num_latent_gps])

        return mean, var

    @check_shapes(
        "return[0]: [M, M]",
        "return[1]: [M, M]",
        "return[2]: [M, D]",
    )
    def _precompute(self) -> Tuple[PrecomputedValue, ...]:
        assert self.mean_function is not None

        X_data = cast(tf.Tensor, self.X_data)
        num_inducing = self.inducing_variable.num_inducing
        err = self.Y_data - self.mean_function(X_data)

        kuf = Kuf(self.inducing_variable, self.kernel, X_data)
        kuu = Kuu(self.inducing_variable, self.kernel, jitter=default_jitter())

        sigma_sq = tf.squeeze(self.likelihood.variance_at(X_data), axis=-1)
        sigma = tf.sqrt(sigma_sq)

        L = tf.linalg.cholesky(kuu)  # cache alpha, qinv
        A = tf.linalg.triangular_solve(L, kuf / sigma, lower=True)
        B = tf.linalg.matmul(A, A, transpose_b=True) + tf.eye(
            num_inducing, dtype=default_float()
        )  # cache qinv
        LB = tf.linalg.cholesky(B)  # cache alpha
        Aerr = tf.linalg.matmul(A, err / sigma[..., None])
        c = tf.linalg.triangular_solve(LB, Aerr, lower=True)

        D = err.shape[1]
        M = X_data.shape[0]
        D_dynamic = D is None
        M_dynamic = M is None

        return (
            PrecomputedValue(L, (M_dynamic, M_dynamic)),
            PrecomputedValue(LB, (M_dynamic, M_dynamic)),
            PrecomputedValue(c, (M_dynamic, D_dynamic)),
        )

    @inherit_check_shapes
    def _conditional_fused(
        self, Xnew: TensorType, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        Compute the mean and variance of the latent function at some new points
        Xnew. Does not make use of caching
        """
        temp_cache = tuple(c.value for c in self._precompute())
        return self._conditional_with_precompute(temp_cache, Xnew, full_cov, full_output_cov)


class VGPPosterior(AbstractPosterior):
    @check_shapes(
        "X: [N, D]",
        "q_mu: [N, P]",
        "q_sqrt: [N_P_or_P_N_N...]",
    )
    def __init__(
        self,
        kernel: Kernel,
        X: tf.Tensor,
        q_mu: tf.Tensor,
        q_sqrt: tf.Tensor,
        mean_function: Optional[mean_functions.MeanFunction] = None,
        white: bool = True,
        *,
        precompute_cache: Optional[PrecomputeCacheType],
    ) -> None:
        super().__init__(kernel, X, mean_function=mean_function)
        self.q_mu = q_mu
        self.q_sqrt = q_sqrt
        self.white = white

        if precompute_cache is not None:
            self.update_cache(precompute_cache)

    @inherit_check_shapes
    def _conditional_with_precompute(
        self,
        cache: Tuple[tf.Tensor, ...],
        Xnew: TensorType,
        full_cov: bool = False,
        full_output_cov: bool = False,
    ) -> MeanAndVariance:
        assert_params_false(self._conditional_with_precompute, full_output_cov=full_output_cov)

        (Lm,) = cache
        Kmn = self.kernel(self.X_data, Xnew)  # [M, ..., N]
        Knn = self.kernel(
            Xnew, full_cov=full_cov
        )  # [..., N] (full_cov = False) or [..., N, N] (True)

        return base_conditional_with_lm(
            Kmn=Kmn,
            Lm=Lm,
            Knn=Knn,
            f=self.q_mu,
            full_cov=full_cov,
            q_sqrt=self.q_sqrt,
            white=self.white,
        )

    @check_shapes(
        "return[0]: [M, M]",
    )
    def _precompute(self) -> Tuple[PrecomputedValue, ...]:
        X_data = cast(tf.Tensor, self.X_data)
        Kmm = self.kernel(X_data) + eye(
            tf.shape(X_data)[-2], value=default_jitter(), dtype=X_data.dtype
        )  # [..., M, M]
        Lm = tf.linalg.cholesky(Kmm)

        M = X_data.shape[0]
        M_dynamic = M is None

        return (PrecomputedValue(Lm, (M_dynamic, M_dynamic)),)

    @inherit_check_shapes
    def _conditional_fused(
        self, Xnew: TensorType, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        temp_cache = tuple(c.value for c in self._precompute())
        return self._conditional_with_precompute(temp_cache, Xnew, full_cov, full_output_cov)


class BasePosterior(AbstractPosterior):
    @check_shapes(
        "inducing_variable: [M, D, broadcast P]",
        "q_mu: [N, P]",
        "q_sqrt: [N_P_or_P_N_N...]",
    )
    def __init__(
        self,
        kernel: Kernel,
        inducing_variable: InducingVariables,
        q_mu: tf.Tensor,
        q_sqrt: tf.Tensor,
        whiten: bool = True,
        mean_function: Optional[mean_functions.MeanFunction] = None,
        *,
        precompute_cache: Optional[PrecomputeCacheType],
    ):

        super().__init__(kernel, inducing_variable, mean_function=mean_function)
        self.whiten = whiten
        self._set_qdist(q_mu, q_sqrt)

        if precompute_cache is not None:
            self.update_cache(precompute_cache)

    @property  # type: ignore[misc]
    @check_shapes(
        "return: [N, P]",
    )
    def q_mu(self) -> tf.Tensor:
        return self._q_dist.q_mu

    @property  # type: ignore[misc]
    @check_shapes(
        "return: [N_P_or_P_N_N...]",
    )
    def q_sqrt(self) -> tf.Tensor:
        return self._q_dist.q_sqrt

    @check_shapes(
        "q_mu: [N, P]",
        "q_sqrt: [N_P_or_P_N_N...]",
    )
    def _set_qdist(self, q_mu: TensorType, q_sqrt: TensorType) -> None:
        if q_sqrt is None:
            self._q_dist = _DeltaDist(q_mu)
        elif len(q_sqrt.shape) == 2:  # q_diag
            self._q_dist = _DiagNormal(q_mu, q_sqrt)
        else:
            self._q_dist = _MvNormal(q_mu, q_sqrt)

    @check_shapes(
        "return[0]: [M_L_or_L_M_M...]",
        "return[1]: [L, M, M]",
    )
    def _precompute(self) -> Tuple[PrecomputedValue, ...]:
        Kuu = covariances.Kuu(self.X_data, self.kernel, jitter=default_jitter())  # [(R), M, M]
        q_mu = self._q_dist.q_mu

        if Kuu.shape.ndims == 4:
            ML = tf.reduce_prod(tf.shape(Kuu)[:2])
            Kuu = tf.reshape(Kuu, [ML, ML])
        if Kuu.shape.ndims == 3:
            q_mu = tf.linalg.adjoint(self._q_dist.q_mu)[..., None]  # [..., R, M, 1]
        L = tf.linalg.cholesky(Kuu)

        if not self.whiten:
            # alpha = Kuu⁻¹ q_mu
            alpha = tf.linalg.cholesky_solve(L, q_mu)
        else:
            # alpha = L⁻ᵀ q_mu
            alpha = tf.linalg.triangular_solve(L, q_mu, adjoint=True)
        # predictive mean = Kfu alpha
        # predictive variance = Kff - Kfu Qinv Kuf
        # S = q_sqrt q_sqrtᵀ
        I = tf.eye(tf.shape(L)[-1], dtype=L.dtype)
        if isinstance(self._q_dist, _DeltaDist):
            B = I
        else:
            if not self.whiten:
                # Qinv = Kuu⁻¹ - Kuu⁻¹ S Kuu⁻¹
                #      = Kuu⁻¹ - L⁻ᵀ L⁻¹ S L⁻ᵀ L⁻¹
                #      = L⁻ᵀ (I - L⁻¹ S L⁻ᵀ) L⁻¹
                #      = L⁻ᵀ B L⁻¹
                if isinstance(self._q_dist, _DiagNormal):
                    q_sqrt = tf.linalg.diag(tf.linalg.adjoint(self._q_dist.q_sqrt))
                elif isinstance(self._q_dist, _MvNormal):
                    q_sqrt = self._q_dist.q_sqrt
                Linv_qsqrt = tf.linalg.triangular_solve(L, q_sqrt)
                Linv_cov_u_LinvT = tf.matmul(Linv_qsqrt, Linv_qsqrt, transpose_b=True)
            else:
                if isinstance(self._q_dist, _DiagNormal):
                    Linv_cov_u_LinvT = tf.linalg.diag(tf.linalg.adjoint(self._q_dist.q_sqrt ** 2))
                elif isinstance(self._q_dist, _MvNormal):
                    q_sqrt = self._q_dist.q_sqrt
                    Linv_cov_u_LinvT = tf.matmul(q_sqrt, q_sqrt, transpose_b=True)
                # Qinv = Kuu⁻¹ - L⁻ᵀ S L⁻¹
                # Linv = (L⁻¹ I) = solve(L, I)
                # Kinv = Linvᵀ @ Linv
            B = I - Linv_cov_u_LinvT
        LinvT_B = tf.linalg.triangular_solve(L, B, adjoint=True)
        B_Linv = tf.linalg.adjoint(LinvT_B)
        Qinv = tf.linalg.triangular_solve(L, B_Linv, adjoint=True)

        M, L = tf.unstack(tf.shape(self._q_dist.q_mu), num=2)
        Qinv = tf.broadcast_to(Qinv, [L, M, M])

        return PrecomputedValue.wrap_alpha_Qinv(alpha, Qinv)


class IndependentPosterior(BasePosterior):
    @check_shapes(
        "mean: [batch..., N, P]",
        "cov: [batch..., P, N, N] if full_cov",
        "cov: [batch..., N, P] if not full_cov",
        "return[0]: [batch..., N, P]",
        "return[1]: [batch..., N, P, N, P] if full_cov and full_output_cov",
        "return[1]: [batch..., N, P, P] if (not full_cov) and full_output_cov",
        "return[1]: [batch..., P, N, N] if full_cov and (not full_output_cov)",
        "return[1]: [batch..., N, P] if (not full_cov) and (not full_output_cov)",
    )
    def _post_process_mean_and_cov(
        self, mean: TensorType, cov: TensorType, full_cov: bool, full_output_cov: bool
    ) -> MeanAndVariance:
        return mean, expand_independent_outputs(cov, full_cov, full_output_cov)

    @check_shapes(
        "Xnew: [N, D]",
        "return: [broadcast P, N, N] if full_cov",
        "return: [broadcast P, N] if (not full_cov)",
    )
    def _get_Kff(self, Xnew: TensorType, full_cov: bool) -> tf.Tensor:

        # TODO: this assumes that Xnew has shape [N, D] and no leading dims

        if isinstance(self.kernel, (kernels.SeparateIndependent, kernels.IndependentLatent)):
            # NOTE calling kernel(Xnew, full_cov=full_cov, full_output_cov=False) directly would
            # return
            # if full_cov: [P, N, N] -- this is what we want
            # else: [N, P] instead of [P, N] as we get from the explicit stack below
            Kff = tf.stack([k(Xnew, full_cov=full_cov) for k in self.kernel.kernels], axis=0)
        elif isinstance(self.kernel, kernels.MultioutputKernel):
            # effectively, SharedIndependent path
            Kff = self.kernel.kernel(Xnew, full_cov=full_cov)
            # NOTE calling kernel(Xnew, full_cov=full_cov, full_output_cov=False) directly would
            # return
            # if full_cov: [P, N, N] instead of [N, N]
            # else: [N, P] instead of [N]
        else:
            # standard ("single-output") kernels
            Kff = self.kernel(Xnew, full_cov=full_cov)  # [N, N] if full_cov else [N]

        return Kff

    @inherit_check_shapes
    def _conditional_with_precompute(
        self,
        cache: Tuple[tf.Tensor, ...],
        Xnew: TensorType,
        full_cov: bool = False,
        full_output_cov: bool = False,
    ) -> MeanAndVariance:
        # Qinv: [L, M, M]
        # alpha: [M, L]
        alpha, Qinv = cache

        Kuf = covariances.Kuf(self.X_data, self.kernel, Xnew)  # [(R), M, N]
        Kff = self._get_Kff(Xnew, full_cov)

        mean = tf.matmul(Kuf, alpha, transpose_a=True)
        if Kuf.shape.ndims == 3:
            mean = tf.linalg.adjoint(tf.squeeze(mean, axis=-1))

        if full_cov:
            Kfu_Qinv_Kuf = tf.matmul(Kuf, Qinv @ Kuf, transpose_a=True)
            cov = Kff - Kfu_Qinv_Kuf
        else:
            # [Aᵀ B]_ij = Aᵀ_ik B_kj = A_ki B_kj
            # TODO check whether einsum is faster now?
            Kfu_Qinv_Kuf = tf.reduce_sum(Kuf * tf.matmul(Qinv, Kuf), axis=-2)
            cov = Kff - Kfu_Qinv_Kuf
            cov = tf.linalg.adjoint(cov)

        return self._post_process_mean_and_cov(mean, cov, full_cov, full_output_cov)


class IndependentPosteriorSingleOutput(IndependentPosterior):
    # could almost be the same as IndependentPosteriorMultiOutput ...
    @inherit_check_shapes
    def _conditional_fused(
        self, Xnew: TensorType, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        # same as IndependentPosteriorMultiOutput, Shared~/Shared~ branch, except for following
        # line:
        Knn = self.kernel(Xnew, full_cov=full_cov)

        Kmm = covariances.Kuu(self.X_data, self.kernel, jitter=default_jitter())  # [M, M]
        Kmn = covariances.Kuf(self.X_data, self.kernel, Xnew)  # [M, N]

        fmean, fvar = base_conditional(
            Kmn, Kmm, Knn, self.q_mu, full_cov=full_cov, q_sqrt=self.q_sqrt, white=self.whiten
        )  # [N, P],  [P, N, N] or [N, P]
        return self._post_process_mean_and_cov(fmean, fvar, full_cov, full_output_cov)


class IndependentPosteriorMultiOutput(IndependentPosterior):
    @inherit_check_shapes
    def _conditional_fused(
        self, Xnew: TensorType, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        if isinstance(self.X_data, SharedIndependentInducingVariables) and isinstance(
            self.kernel, kernels.SharedIndependent
        ):
            # same as IndependentPosteriorSingleOutput except for following line
            Knn = self.kernel.kernel(Xnew, full_cov=full_cov)
            # we don't call self.kernel() directly as that would do unnecessary tiling

            Kmm = covariances.Kuu(self.X_data, self.kernel, jitter=default_jitter())  # [M, M]
            Kmn = covariances.Kuf(self.X_data, self.kernel, Xnew)  # [M, N]

            fmean, fvar = base_conditional(
                Kmn, Kmm, Knn, self.q_mu, full_cov=full_cov, q_sqrt=self.q_sqrt, white=self.whiten
            )  # [N, P],  [P, N, N] or [N, P]
        else:
            # this is the messy thing with tf.map_fn, cleaned up by the
            # st/clean_up_broadcasting_conditionals branch

            # Following are: [P, M, M]  -  [P, M, N]  -  [P, N](x N)
            Kmms = covariances.Kuu(self.X_data, self.kernel, jitter=default_jitter())  # [P, M, M]
            Kmns = covariances.Kuf(self.X_data, self.kernel, Xnew)  # [P, M, N]
            if isinstance(self.kernel, kernels.Combination):
                kernel_list = self.kernel.kernels
            else:
                kernel_list = [self.kernel.kernel] * len(self.X_data.inducing_variable_list)
            Knns = tf.stack(
                [k.K(Xnew) if full_cov else k.K_diag(Xnew) for k in kernel_list], axis=0
            )

            fmean, fvar = separate_independent_conditional_implementation(
                Kmns,
                Kmms,
                Knns,
                self.q_mu,
                q_sqrt=self.q_sqrt,
                full_cov=full_cov,
                white=self.whiten,
            )

        return self._post_process_mean_and_cov(fmean, fvar, full_cov, full_output_cov)


class LinearCoregionalizationPosterior(IndependentPosteriorMultiOutput):
    @check_shapes(
        "mean: [batch..., N, L]",
        "cov: [batch..., L, N, N] if full_cov",
        "cov: [batch..., N, L] if not full_cov",
        "return[0]: [batch..., N, P]",
        "return[1]: [batch..., N, P, N, P] if full_cov and full_output_cov",
        "return[1]: [batch..., N, P, P] if (not full_cov) and full_output_cov",
        "return[1]: [batch..., P, N, N] if full_cov and (not full_output_cov)",
        "return[1]: [batch..., N, P] if (not full_cov) and (not full_output_cov)",
    )
    def _post_process_mean_and_cov(
        self, mean: TensorType, cov: TensorType, full_cov: bool, full_output_cov: bool
    ) -> MeanAndVariance:
        cov = expand_independent_outputs(cov, full_cov, full_output_cov=False)
        mean, cov = mix_latent_gp(self.kernel.W, mean, cov, full_cov, full_output_cov)
        return mean, cov


class FullyCorrelatedPosterior(BasePosterior):
    @inherit_check_shapes
    def _conditional_with_precompute(
        self,
        cache: Tuple[tf.Tensor, ...],
        Xnew: TensorType,
        full_cov: bool = False,
        full_output_cov: bool = False,
    ) -> MeanAndVariance:
        # TODO: this assumes that Xnew has shape [N, D] and no leading dims

        # Qinv: [L, M, M]
        # alpha: [M, L]
        alpha, Qinv = cache

        Kuf = covariances.Kuf(self.X_data, self.kernel, Xnew)
        assert Kuf.shape.ndims == 4
        M, L, N, K = tf.unstack(tf.shape(Kuf), num=Kuf.shape.ndims, axis=0)
        Kuf = tf.reshape(Kuf, (M * L, N * K))

        kernel: kernels.MultioutputKernel = self.kernel
        Kff = kernel(Xnew, full_cov=full_cov, full_output_cov=full_output_cov)
        # full_cov=True and full_output_cov=True: [N, P, N, P]
        # full_cov=True and full_output_cov=False: [P, N, N]
        # full_cov=False and full_output_cov=True: [N, P, P]
        # full_cov=False and full_output_cov=False: [N, P]
        if full_cov == full_output_cov:
            new_shape = (N * K, N * K) if full_cov else (N * K,)
            Kff = tf.reshape(Kff, new_shape)

        N = tf.shape(Xnew)[0]
        K = tf.shape(Kuf)[-1] // N

        mean = tf.matmul(Kuf, alpha, transpose_a=True)
        if Kuf.shape.ndims == 3:
            mean = tf.linalg.adjoint(tf.squeeze(mean, axis=-1))

        if not full_cov and not full_output_cov:
            # fully diagonal case in both inputs and outputs
            # [Aᵀ B]_ij = Aᵀ_ik B_kj = A_ki B_kj
            # TODO check whether einsum is faster now?
            Kfu_Qinv_Kuf = tf.reduce_sum(Kuf * tf.matmul(Qinv, Kuf), axis=-2)
        else:
            Kfu_Qinv_Kuf = tf.matmul(Kuf, Qinv @ Kuf, transpose_a=True)
            if not (full_cov and full_output_cov):
                # diagonal in either inputs or outputs
                new_shape = tf.concat([tf.shape(Kfu_Qinv_Kuf)[:-2], (N, K, N, K)], axis=0)
                Kfu_Qinv_Kuf = tf.reshape(Kfu_Qinv_Kuf, new_shape)
                if full_cov:
                    # diagonal in outputs: move outputs to end
                    tmp = tf.linalg.diag_part(tf.einsum("...ijkl->...ikjl", Kfu_Qinv_Kuf))
                elif full_output_cov:
                    # diagonal in inputs: move inputs to end
                    tmp = tf.linalg.diag_part(tf.einsum("...ijkl->...jlik", Kfu_Qinv_Kuf))
                Kfu_Qinv_Kuf = tf.einsum("...ijk->...kij", tmp)  # move diagonal dim to [-3]
        cov = Kff - Kfu_Qinv_Kuf

        if not full_cov and not full_output_cov:
            cov = tf.linalg.adjoint(cov)

        mean = tf.reshape(mean, (N, K))
        if full_cov == full_output_cov:
            cov_shape = (N, K, N, K) if full_cov else (N, K)
        else:
            cov_shape = (K, N, N) if full_cov else (N, K, K)
        cov = tf.reshape(cov, cov_shape)

        return mean, cov

    @inherit_check_shapes
    def _conditional_fused(
        self, Xnew: TensorType, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        Kmm = covariances.Kuu(self.X_data, self.kernel, jitter=default_jitter())  # [M, L, M, L]
        Kmn = covariances.Kuf(self.X_data, self.kernel, Xnew)  # [M, L, N, P]
        kernel: kernels.MultioutputKernel = self.kernel
        Knn = kernel(
            Xnew, full_cov=full_cov, full_output_cov=full_output_cov
        )  # [N, P](x N)x P  or  [N, P](x P)

        M, L, N, K = tf.unstack(tf.shape(Kmn), num=Kmn.shape.ndims, axis=0)
        Kmm = tf.reshape(Kmm, (M * L, M * L))

        if full_cov == full_output_cov:
            Kmn = tf.reshape(Kmn, (M * L, N * K))
            Knn = tf.reshape(Knn, (N * K, N * K)) if full_cov else tf.reshape(Knn, (N * K,))
            mean, cov = base_conditional(
                Kmn, Kmm, Knn, self.q_mu, full_cov=full_cov, q_sqrt=self.q_sqrt, white=self.whiten
            )  # [K, 1], [1, K](x NK)
            mean = tf.reshape(mean, (N, K))
            cov = tf.reshape(cov, (N, K, N, K) if full_cov else (N, K))
        else:
            Kmn = tf.reshape(Kmn, (M * L, N, K))
            mean, cov = fully_correlated_conditional(
                Kmn,
                Kmm,
                Knn,
                self.q_mu,
                full_cov=full_cov,
                full_output_cov=full_output_cov,
                q_sqrt=self.q_sqrt,
                white=self.whiten,
            )
        return mean, cov


class FallbackIndependentLatentPosterior(FullyCorrelatedPosterior):  # XXX
    @inherit_check_shapes
    def _conditional_fused(
        self, Xnew: TensorType, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        Kmm = covariances.Kuu(self.X_data, self.kernel, jitter=default_jitter())  # [L, M, M]
        Kmn = covariances.Kuf(self.X_data, self.kernel, Xnew)  # [M, L, N, P]
        kernel: kernels.IndependentLatent = self.kernel
        Knn = kernel(
            Xnew, full_cov=full_cov, full_output_cov=full_output_cov
        )  # [N, P](x N)x P  or  [N, P](x P)

        return independent_interdomain_conditional(
            Kmn,
            Kmm,
            Knn,
            self.q_mu,
            full_cov=full_cov,
            full_output_cov=full_output_cov,
            q_sqrt=self.q_sqrt,
            white=self.whiten,
        )


get_posterior_class = Dispatcher("get_posterior_class")


@get_posterior_class.register(kernels.Kernel, InducingVariables)
def _get_posterior_base_case(
    kernel: Kernel, inducing_variable: InducingVariables
) -> Type[BasePosterior]:
    # independent single output
    return IndependentPosteriorSingleOutput


@get_posterior_class.register(kernels.MultioutputKernel, InducingPoints)
def _get_posterior_fully_correlated_mo(
    kernel: Kernel, inducing_variable: InducingVariables
) -> Type[BasePosterior]:
    return FullyCorrelatedPosterior


@get_posterior_class.register(
    (kernels.SharedIndependent, kernels.SeparateIndependent),
    (SeparateIndependentInducingVariables, SharedIndependentInducingVariables),
)
def _get_posterior_independent_mo(
    kernel: Kernel, inducing_variable: InducingVariables
) -> Type[BasePosterior]:
    # independent multi-output
    return IndependentPosteriorMultiOutput


@get_posterior_class.register(
    kernels.IndependentLatent,
    (FallbackSeparateIndependentInducingVariables, FallbackSharedIndependentInducingVariables),
)
def _get_posterior_independentlatent_mo_fallback(
    kernel: Kernel, inducing_variable: InducingVariables
) -> Type[BasePosterior]:
    return FallbackIndependentLatentPosterior


@get_posterior_class.register(
    kernels.LinearCoregionalization,
    (SeparateIndependentInducingVariables, SharedIndependentInducingVariables),
)
def _get_posterior_linearcoregionalization_mo_efficient(
    kernel: Kernel, inducing_variable: InducingVariables
) -> Type[BasePosterior]:
    # Linear mixing---efficient multi-output
    return LinearCoregionalizationPosterior


def create_posterior(
    kernel: Kernel,
    inducing_variable: InducingVariables,
    q_mu: TensorType,
    q_sqrt: TensorType,
    whiten: bool,
    mean_function: Optional[MeanFunction] = None,
    precompute_cache: Union[PrecomputeCacheType, str, None] = PrecomputeCacheType.TENSOR,
) -> BasePosterior:
    posterior_class = get_posterior_class(kernel, inducing_variable)
    precompute_cache = _validate_precompute_cache_type(precompute_cache)
    return posterior_class(  # type: ignore[no-any-return]
        kernel,
        inducing_variable,
        q_mu,
        q_sqrt,
        whiten,
        mean_function,
        precompute_cache=precompute_cache,
    )


# gpflow/probability_distributions.py

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


# Eventually, it would be nice to not have to have our own classes for
# probability distributions. The TensorFlow "distributions" framework would
# be a good replacement.
from abc import ABC, abstractmethod

from .base import TensorType
from .experimental.check_shapes import ErrorContext, Shape, check_shapes, register_get_shape


class ProbabilityDistribution(ABC):
    """
    This is the base class for a probability distributions,
    over which we take the expectations in the expectations framework.
    """

    @property
    @abstractmethod
    def shape(self) -> Shape:
        """
        Return the shape of this distribution.

        Shape should be some variation of ``[N, D]``, where:

        * ``N`` is the number of data points.
        * ``D`` is the number of input dimensions.
        """


@register_get_shape(ProbabilityDistribution)
def get_probability_distribution_shape(
    shaped: ProbabilityDistribution, context: ErrorContext
) -> Shape:
    return shaped.shape


class Gaussian(ProbabilityDistribution):
    @check_shapes(
        "mu: [N, D]",
        "cov: [N, D, D]",
    )
    def __init__(self, mu: TensorType, cov: TensorType):
        self.mu = mu
        self.cov = cov

    @property
    def shape(self) -> Shape:
        return self.mu.shape  # type: ignore[no-any-return]


class DiagonalGaussian(ProbabilityDistribution):
    @check_shapes(
        "mu: [N, D]",
        "cov: [N, D]",
    )
    def __init__(self, mu: TensorType, cov: TensorType):
        self.mu = mu
        self.cov = cov

    @property
    def shape(self) -> Shape:
        return self.mu.shape  # type: ignore[no-any-return]


class MarkovGaussian(ProbabilityDistribution):
    """
    Gaussian distribution with Markov structure.
    Only covariances and covariances between t and t+1 need to be
    parameterised. We use the solution proposed by Carl Rasmussen, i.e. to
    represent
    Var[x_t] = cov[x_t, :, :] * cov[x_t, :, :].T
    Cov[x_t, x_{t+1}] = cov[t, :, :] * cov[t+1, :, :]
    """

    @check_shapes(
        "mu: [N_plus_1, D]",
        "cov: [2, N_plus_1, D, D]",
    )
    def __init__(self, mu: TensorType, cov: TensorType):
        self.mu = mu
        self.cov = cov

    @property
    def shape(self) -> Shape:
        shape = self.mu.shape
        if shape is None:
            return shape
        N_plus_1, D = shape
        N = N_plus_1 - 1
        return N, D


# gpflow/type_flags.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Code for setting type hints, depending on library versions.
"""
import sys

from numpy import __version__ as np_version
from packaging.version import Version

NP_VERSION = Version(np_version)

NP_TYPE_CHECKING = False
"""
Whether to type-check numpy arrays at all. Defaults to False, because we don't know which
versions a client might use.
"""

GENERIC_NP_ARRAYS = (sys.version_info >= (3, 9)) and (NP_VERSION >= Version("1.22.0"))
"""
Whether to use generic numpy arrays. This is not applied at all, if type checking and not
`NP_TYPE_CHECKING`.
"""


def compute_mypy_flags() -> str:  # pragma: no cover
    from mypy.version import __version__ as mypy_version

    MYPY_VERSION = Version(mypy_version)

    flags = []

    def set_always(variable: str, value: bool) -> None:
        if value:
            flags.append("--always-true")
        else:
            flags.append("--always-false")
        flags.append(variable)

    set_always("NP_TYPE_CHECKING", True)
    set_always("GENERIC_NP_ARRAYS", NP_VERSION >= Version("1.21.0"))
    if MYPY_VERSION >= Version("0.940"):
        flags.extend(["--enable-error-code", "ignore-without-code"])

    return " ".join(flags)


# gpflow/versions.py

import pkg_resources

try:
    __version__ = str(pkg_resources.get_distribution("gpflow").parsed_version)
except pkg_resources.DistributionNotFound:
    __version__ = "develop"


# tests/__init__.py

import warnings

import numpy as np
import tensorflow as tf

import gpflow

warnings.filterwarnings("ignore")
gpflow.config.set_default_float(np.float64)


# tests/conftest.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Iterable

import pytest
from _pytest.logging import LogCaptureFixture

from gpflow.experimental.check_shapes.config import (
    DocstringFormat,
    ShapeCheckingState,
    get_enable_check_shapes,
    get_enable_function_call_precompute,
    get_rewrite_docstrings,
    set_enable_check_shapes,
    set_enable_function_call_precompute,
    set_rewrite_docstrings,
)


@pytest.fixture(autouse=True)
def test_auto_graph_compile(caplog: LogCaptureFixture) -> Iterable[None]:
    yield

    for when in ["setup", "call", "teardown"]:
        for record in caplog.get_records(when):
            assert not record.msg.startswith("AutoGraph could not transform"), record.getMessage()


@pytest.fixture(autouse=True)
def enable_shape_checks() -> Iterable[None]:
    # Ensure that:
    # 1: `check_shapes` is enabled when running these tests.
    # 2: If a test manipulates `check_shapes` settings, they are reset after the test.
    # See also: tests/conftest.py
    old_enable = get_enable_check_shapes()
    old_rewrite_docstrings = get_rewrite_docstrings()
    old_function_call_precompute = get_enable_function_call_precompute()
    set_enable_check_shapes(ShapeCheckingState.ENABLED)
    set_rewrite_docstrings(DocstringFormat.SPHINX)
    set_enable_function_call_precompute(True)
    yield
    set_enable_function_call_precompute(old_function_call_precompute)
    set_rewrite_docstrings(old_rewrite_docstrings)
    set_enable_check_shapes(old_enable)


# setup.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

# pylint: skip-file

from setuptools import find_packages, setup

##### Dependencies of GPflow

requirements = [
    "deprecated",
    "lark>=1.1.0",
    "multipledispatch>=0.6",
    "numpy",
    "packaging",
    "scipy",
    "setuptools>=41.0.0",  # to satisfy dependency constraints
    "tabulate",
    "tensorflow-probability>=0.12.0",
    "tensorflow>=2.4.0; platform_system!='Darwin' or platform_machine!='arm64'",
    # NOTE: Support of Apple Silicon MacOS platforms is in an experimental mode
    "tensorflow-macos>=2.4.0; platform_system=='Darwin' and platform_machine=='arm64'",
    # NOTE: once we require tensorflow-probability>=0.12, we can remove our custom deepcopy handling
    "typing_extensions",
]


def read_file(filename: str) -> str:
    with open(filename, encoding="utf-8") as f:
        return f.read().strip()


version = read_file("VERSION")
readme_text = read_file("README.md")

packages = find_packages(".", exclude=["tests"])

setup(
    name="gpflow",
    version=version,
    author="James Hensman, Alex Matthews",
    author_email="james.hensman@gmail.com",
    description="Gaussian process methods in TensorFlow",
    long_description=readme_text,
    long_description_content_type="text/markdown",
    license="Apache License 2.0",
    keywords="machine-learning gaussian-processes kernels tensorflow",
    url="https://www.gpflow.org",
    project_urls={
        "Source on GitHub": "https://github.com/GPflow/GPflow",
        "Documentation": "https://gpflow.github.io/GPflow/",
    },
    packages=packages,
    package_data={"": ["*.lark"]},
    include_package_data=True,
    install_requires=requirements,
    extras_require={"ImageToTensorBoard": ["matplotlib"]},
    python_requires=">=3.7",
    classifiers=[
        "Development Status :: 5 - Production/Stable",
        "Environment :: GPU :: NVIDIA CUDA",
        "Intended Audience :: Developers",
        "Intended Audience :: Education",
        "Intended Audience :: Science/Research",
        "License :: OSI Approved :: Apache Software License",
        "Natural Language :: English",
        "Operating System :: MacOS :: MacOS X",
        "Operating System :: Microsoft :: Windows",
        "Operating System :: POSIX :: Linux",
        "Programming Language :: Python :: 3.7",
        "Programming Language :: Python :: 3.8",
        "Programming Language :: Python :: 3.9",
        "Programming Language :: Python :: 3.10",
        "Topic :: Scientific/Engineering :: Artificial Intelligence",
        "Typing :: Typed",
    ],
)


# doc/build_docs.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Code for building GPflow's documentation for a specified branch.
"""
import argparse
import shutil
import subprocess
from itertools import chain
from pathlib import Path
from time import perf_counter
from typing import Collection, Optional

from generate_module_rst import generate_module_rst
from tabulate import tabulate
from versions import Branch

import gpflow

_SRC = Path(__file__).parent
_SPHINX_SRC = _SRC / "sphinx"
_NOTEBOOKS_SRC = _SPHINX_SRC / "notebooks"

_TMP = Path("/tmp/gpflow_build_docs")
_BUILD_TMP = _TMP / "build"
_NOTEBOOKS_TMP = _BUILD_TMP / "notebooks"
_DOCTREE_TMP = _TMP / "doctree"


class ShardingStrategy:
    """
    Strategy for how to shard (split) the work.
    """

    def __init__(self, spec: str) -> None:
        """
        Valid ``spec``\s are:

        - ``no``: No sharding will happen, and a single run of this script does all necessary work.
        - ``<i>/<n>``, where 0 <= i < n: Build a subset of notebooks, corresponding to job ``i`` out
          of ``n``.
        - ``collect``: Collect data generated by previous shards, and finish the work.
        """
        self.spec = spec
        if spec == "no":
            self.setup_tmp = True
            self.build_notebooks = True
            self.build_other = True
            self.shard_i = 0
            self.shard_n = 1
        elif spec == "collect":
            self.setup_tmp = False
            self.build_notebooks = False
            self.build_other = True
            self.shard_i = 0
            self.shard_n = 1
        else:
            i_str, n_str = spec.split("/")
            self.setup_tmp = False
            self.build_notebooks = True
            self.build_other = False
            self.shard_i = int(i_str)
            self.shard_n = int(n_str)
        assert 0 <= self.shard_i < self.shard_n, (self.shard_i, self.shard_n)

    def __repr__(self) -> str:
        return self.spec


def _create_fake_notebook(
    destination_relative_path: Path, limit_notebooks: Collection[str]
) -> None:
    limiting_command = f"--limit_notebooks {' '.join(limit_notebooks)}"
    print(f"Generating fake, due to {limiting_command}")

    destination = _NOTEBOOKS_TMP / destination_relative_path
    title = f"Fake {destination.name}"
    title_line = "#" * len(title)

    destination.write_text(
        f"""{title}
{title_line}

Fake {destination.name} due to::

   {limiting_command}
"""
    )


def _build_notebooks(
    limit_notebooks: Optional[Collection[str]], sharding: ShardingStrategy
) -> None:
    # Building the notebooks is really slow. Let's time it so we know which notebooks we can /
    # should optimise.
    timings = []
    all_notebooks = sorted(
        chain(_NOTEBOOKS_TMP.glob("**/*.pct.py"), _NOTEBOOKS_TMP.glob("**/*.md"))
    )
    for i, source_path in enumerate(all_notebooks):
        before = perf_counter()
        print()
        print("XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX")
        print("Building:", source_path)
        print("XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX")

        source_relative_path = source_path.relative_to(_NOTEBOOKS_TMP)
        destination_relative_path = source_relative_path
        while destination_relative_path.suffix:  # .pct.py has several suffixes. Remove all of them.
            destination_relative_path = destination_relative_path.with_suffix("")
        destination_relative_path = destination_relative_path.with_suffix(".ipynb")

        if i % sharding.shard_n != sharding.shard_i:
            print("Skipping due to sharding...")
        elif limit_notebooks is None or destination_relative_path.stem in limit_notebooks:
            subprocess.run(
                [
                    "jupytext",
                    "--execute",
                    "--to",
                    "notebook",
                    "-o",
                    str(destination_relative_path),
                    str(source_relative_path),
                ],
                cwd=_NOTEBOOKS_TMP,
            ).check_returncode()
        else:
            _create_fake_notebook(destination_relative_path, limit_notebooks)

        after = perf_counter()
        timings.append((after - before, source_relative_path))

    timings.sort(reverse=True)
    print()
    print("Notebooks by build-time:")
    print(tabulate(timings, headers=["Time", "Notebook"]))
    print()


def main() -> None:
    parser = argparse.ArgumentParser(description="Build the GPflow documentation.")
    parser.add_argument(
        "branch",
        nargs="?",
        default=None,
        type=str,
        choices=[b.value for b in Branch],
        help="Git branch that is currently being built.",
    )
    parser.add_argument(
        "destination",
        nargs="?",
        default=None,
        type=Path,
        help="Directory to write docs to.",
    )
    parser.add_argument(
        "--limit_notebooks",
        "--limit-notebooks",
        type=str,
        nargs="*",
        help="Only process the notebooks with this base/stem name. Useful when debugging.",
    )
    parser.add_argument(
        "--fail_on_warning",
        "--fail-on-warning",
        default=False,
        action="store_true",
        help="If set, crash if there were any warnings while generating documentation.",
    )
    parser.add_argument(
        "--shard",
        default=ShardingStrategy("no"),
        type=ShardingStrategy,
        help=(
            "Sharding strategy:"
            " If set to 'no' this script performs all necessary work."
            " If set to the format <i>/<n>, where 0 <= i < n then this script only computes"
            f" notebooks for shard <i> out of <n> shards. This requires that {_TMP} has manually"
            " been created, and is empty."
            " If set to 'collect' then this script assumes all notebooks already have been"
            " computed, using the <i>/<n> commands, and finishes the work."
        ),
    )
    args = parser.parse_args()
    sharding = args.shard

    if sharding.setup_tmp:
        shutil.rmtree(_TMP, ignore_errors=True)
        _TMP.mkdir(parents=True)
    else:
        assert _TMP.is_dir()

    # Type-ignore below is because the `dirs_exist_ok` parameter was added in Python 3.8, and we
    # still support Python 3.7. However, we only build our documentation using Python3.10+, so
    # actually this is ok.
    shutil.copytree(_SPHINX_SRC, _BUILD_TMP, dirs_exist_ok=True)  # type: ignore[call-arg]

    if sharding.build_notebooks:
        _build_notebooks(args.limit_notebooks, sharding)

    if sharding.build_other:
        branch = Branch(args.branch)
        assert branch, "'branch' command line argument missing."
        dest = args.destination
        assert dest, "'destination' command line argument missing."
        version_dest = dest / branch.version
        shutil.rmtree(version_dest, ignore_errors=True)

        (_BUILD_TMP / "build_version.txt").write_text(branch.version)
        generate_module_rst(gpflow, _BUILD_TMP / "api")

        sphinx_commands = [
            "sphinx-build",
            "-b",
            "html",
            "-d",
            str(_DOCTREE_TMP),
            str(_BUILD_TMP),
            str(version_dest),
        ]
        if args.fail_on_warning:
            sphinx_commands.extend(
                [
                    "-W",
                    "--keep-going",
                ]
            )

        subprocess.run(sphinx_commands).check_returncode()


if __name__ == "__main__":
    main()


# doc/generate_module_rst.py

# Copyright 2019 GPflow Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Script to autogenerate .rst files for autodocumentation of classes and modules in GPflow.
To be run by the CI system to update docs.
"""
import inspect
from dataclasses import dataclass
from pathlib import Path
from types import ModuleType
from typing import Any, Callable, Deque, Dict, List, Mapping, Set, TextIO, Type, Union

from gpflow.utilities import Dispatcher

RST_LEVEL_SYMBOLS = ["=", "-", "~", '"', "'", "^"]

IGNORE_MODULES = {
    "gpflow.covariances.dispatch",
    "gpflow.conditionals.dispatch",
    "gpflow.expectations.dispatch",
    "gpflow.kullback_leiblers.dispatch",
    "gpflow.versions",
}


def _header(header: str, level: int) -> str:
    return f"{header}\n{RST_LEVEL_SYMBOLS[level] * len(header)}"


@dataclass
class DocumentableDispatcher:

    name: str
    obj: Dispatcher

    def implementations(self) -> Mapping[Callable[..., Any], List[Type[Any]]]:
        implementations: Dict[Callable[..., Any], List[Type[Any]]] = {}
        for args, impl in self.obj.funcs.items():
            implementations.setdefault(impl, []).append(args)
        return implementations

    def write(self, out: TextIO) -> None:
        out.write(
            f"""
{_header(self.name, 2)}

This function uses multiple dispatch, which will depend on the type of argument passed in:
"""
        )
        for impl, argss in self.implementations().items():
            impl_name = f"{impl.__module__}.{impl.__name__}"

            out.write(
                """
.. code-block:: python

"""
            )
            for args in argss:
                arg_names = ", ".join([a.__name__ for a in args])
                out.write(f"    {self.name}( {arg_names} )\n")
            out.write(f"    # dispatch to -> {impl_name}(...)\n")
            out.write(
                f"""
.. autofunction:: {impl_name}
"""
            )


@dataclass
class DocumentableClass:

    name: str
    obj: Type[Any]

    def write(self, out: TextIO) -> None:
        out.write(
            f"""
{_header(self.name, 2)}

.. autoclass:: {self.name}
   :show-inheritance:
   :members:
"""
        )


@dataclass
class DocumentableFunction:

    name: str
    obj: Callable[..., Any]

    def write(self, out: TextIO) -> None:
        out.write(
            f"""
{_header(self.name, 2)}

.. autofunction:: {self.name}
"""
        )


@dataclass
class DocumentableModule:

    name: str
    obj: ModuleType
    modules: List["DocumentableModule"]
    classes: List[DocumentableClass]
    functions: List[Union[DocumentableDispatcher, DocumentableFunction]]

    @staticmethod
    def collect(
        root: ModuleType,
    ) -> "DocumentableModule":
        root_name = root.__name__
        exported_names = set(getattr(root, "__all__", []))

        modules: List["DocumentableModule"] = []
        classes: List[DocumentableClass] = []
        functions: List[Union[DocumentableDispatcher, DocumentableFunction]] = []

        for key in dir(root):
            if key.startswith("_"):
                continue

            child = getattr(root, key)
            child_name = root_name + "." + key
            if child_name in IGNORE_MODULES:
                continue

            def _should_ignore(child: Union[Callable[..., Any], Type[Any]]) -> bool:
                declared_in_root = child.__module__ == root_name
                explicitly_exported = key in exported_names
                return not (declared_in_root or explicitly_exported)

            if isinstance(child, Dispatcher):
                functions.append(DocumentableDispatcher(child_name, child))
            elif inspect.ismodule(child):
                if child.__name__ != child_name:  # Ignore imports of modules.
                    continue
                modules.append(DocumentableModule.collect(child))
            elif inspect.isclass(child):
                if _should_ignore(child):
                    continue
                classes.append(DocumentableClass(child_name, child))
            elif inspect.isfunction(child):
                if _should_ignore(child):
                    continue
                functions.append(DocumentableFunction(child_name, child))

        return DocumentableModule(root_name, root, modules, classes, functions)

    def seen_in_dispatchers(self, seen: Set[int]) -> None:
        for module in self.modules:
            module.seen_in_dispatchers(seen)
        for function in self.functions:
            if isinstance(function, DocumentableDispatcher):
                impls = function.obj.funcs.values()
                for impl in impls:
                    seen.add(id(impl))

    def prune_duplicates(self) -> None:
        seen: Set[int] = set()
        self.seen_in_dispatchers(seen)

        # Breadth-first search so that we prefer objects with shorter names.
        todo = Deque([self])
        while todo:
            module = todo.popleft()

            new_classes = []
            for c in module.classes:
                if id(c.obj) not in seen:
                    seen.add(id(c.obj))
                    new_classes.append(c)
            module.classes = new_classes

            new_functions = []
            for f in module.functions:
                if id(f.obj) not in seen:
                    seen.add(id(f.obj))
                    new_functions.append(f)
            module.functions = new_functions

            todo.extend(module.modules)

    def prune_empty_modules(self) -> None:
        new_modules = []
        for m in self.modules:
            m.prune_empty_modules()

            if m.modules or m.classes or m.functions:
                new_modules.append(m)
        self.modules = new_modules

    def prune(self) -> None:
        self.prune_duplicates()
        self.prune_empty_modules()

    def write_modules(self, out: TextIO) -> None:
        if not self.modules:
            return

        out.write(
            f"""
{_header('Modules', 1)}

.. toctree::
   :maxdepth: 1

"""
        )
        for module in self.modules:
            out.write(f"   {module.name} <{module.name.split('.')[-1]}/index>\n")

    def write_classes(self, out: TextIO) -> None:
        if not self.classes:
            return

        out.write(
            f"""
{_header('Classes', 1)}
"""
        )
        for cls in self.classes:
            cls.write(out)

    def write_functions(self, out: TextIO) -> None:
        if not self.functions:
            return

        out.write(
            f"""
{_header('Functions', 1)}
"""
        )
        for function in self.functions:
            function.write(out)

    def write(self, path: Path) -> None:
        dir_path = path / f"{self.name.replace('.', '/')}"
        dir_path.mkdir(parents=True, exist_ok=True)
        index_path = dir_path / "index.rst"
        with index_path.open("wt") as out:
            print("Writing", index_path)
            out.write(
                f"""{_header(self.name, 0)}

.. THIS IS AN AUTOGENERATED RST FILE

.. automodule:: {self.name}
"""
            )
            self.write_modules(out)
            self.write_classes(out)
            self.write_functions(out)

        for module in self.modules:
            module.write(path)

    def str_into(self, indent: int, lines: List[str]) -> None:
        lines.append(2 * indent * " " + "Module: " + self.name)
        for module in self.modules:
            module.str_into(indent + 1, lines)
        for cls in self.classes:
            lines.append(2 * (indent + 1) * " " + "Class: " + cls.name)
        for function in self.functions:
            lines.append(2 * (indent + 1) * " " + "Function: " + function.name)

    def __str__(self) -> str:
        lines: List[str] = []
        self.str_into(0, lines)
        return "\n".join(lines)


def generate_module_rst(module: ModuleType, dest: Path) -> None:
    """
    Traverses the given `module` and generates `.rst` files for Sphinx.
    """
    docs = DocumentableModule.collect(module)
    docs.prune()
    docs.write(dest)


# doc/update_versions.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Tool for updating the `versions.json` and `index.html` in our documentation.
"""
import argparse
import json
from pathlib import Path
from typing import Any, List, Mapping, Optional

from packaging.version import Version
from versions import Branch


class _Versions:
    def __init__(self, versions: List[Mapping[str, Any]]) -> None:
        self._versions = versions
        self._sorted = False

    @staticmethod
    def _versions_json(dest: Path) -> Path:
        return dest / "versions.json"

    @staticmethod
    def read(docs_dir: Path) -> "_Versions":
        if _Versions._versions_json(docs_dir).exists():
            with _Versions._versions_json(docs_dir).open("rt") as versions_file:
                versions = json.load(versions_file)
        else:
            versions = {}
        return _Versions(versions)

    def write(self, docs_dir: Path) -> None:
        with _Versions._versions_json(docs_dir).open("wt") as versions_file:
            json.dump(self._versions, versions_file, indent=2)

    def add(self, version: str) -> None:
        self._versions = [v for v in self._versions if v["version"] != version]
        self._versions.append(
            {
                "version": version,
                "url": f"https://gpflow.github.io/GPflow/{version}/index.html",
            }
        )
        self._sorted = False

    def sort(self) -> None:
        develop: Optional[Mapping[str, Any]] = None
        for v in self._versions:
            if v["version"] == "develop":
                develop = v
        versions = [v for v in self._versions if v["version"] != "develop"]
        versions.sort(key=lambda v: Version(v["version"]), reverse=True)
        if develop is not None:
            # Insert `develop` in the second spot, after the latest release, but before older
            # releases:
            versions.insert(1, develop)
        self._versions = versions
        self._sorted = True

    @property
    def latest(self) -> str:
        assert self._sorted, "Must sort versions to find latest."
        latest = self._versions[0]["version"]
        assert isinstance(latest, str)  # Hint for mypy.
        return latest


def _create_root_redirect(versions: _Versions, dest: Path) -> None:

    (dest / "index.html").write_text(
        f"""<html xmlns="http://www.w3.org/1999/xhtml">
  <head>
    <title>GPflow documentation</title>
    <meta http-equiv="refresh" content="0;URL='{versions.latest}/index.html'" />
  </head>
  <body>
    <p>Redirecting to the most recent release.</p>
  </body>
</html>
"""
    )


def main() -> None:
    parser = argparse.ArgumentParser(description="Build the GPflow documentation.")
    parser.add_argument(
        "branch",
        type=str,
        choices=[b.value for b in Branch],
        help="Git branch that is currently being built.",
    )
    parser.add_argument("docs_dir", type=Path, help="To read / write docs versions from.")
    args = parser.parse_args()
    branch = Branch(args.branch)
    docs_dir = args.docs_dir

    versions = _Versions.read(docs_dir)
    versions.add(branch.version)
    versions.sort()
    versions.write(docs_dir)

    _create_root_redirect(versions, docs_dir)


if __name__ == "__main__":
    main()


# doc/versions.py

# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Code for keeping track of different versions of documentation.
"""
from enum import Enum
from pathlib import Path

# gpflow.__version__ does not seem reliable when used from a branch, so we read the VERSION file
# directly. (That's also a lot faster than importing gpflow...)
_GPFLOW_VERSION_FILE = Path(__file__).parent.parent / "VERSION"


class Branch(Enum):
    DEVELOP = "develop"
    """
    Building the `develop` branch means we're creating temporary documentation.

    Use this if in doubt.
    """

    MASTER = "master"
    """
    Building the `master` branch means we're creating a new release.
    """

    @property
    def version(self) -> str:
        """
        Returns the version of code / documentation this branch refers to.
        """
        if self == Branch.DEVELOP:
            return "develop"
        else:
            assert self == Branch.MASTER
            return _GPFLOW_VERSION_FILE.read_text().strip()


# doc/sphinx/notebooks/advanced/changepoints.pct.py

# ---
# jupyter:
#   jupytext:
#     formats: ipynb,.pct.py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.3.3
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Change points
#
# *Joseph Hall (October 2019)*
#
# This notebook demonstrates the use of the `ChangePoints` kernel, which can be used to describe one-dimensional functions that contain a number of change-points, or regime changes. The kernel makes use of sigmoids ($\sigma$) to blend smoothly between different kernels. For example, a single change-point kernel is defined by:
#
# \begin{equation}
# \textrm{cov}(f(x), f(y)) = k_1(x, y)\cdot\bar{\sigma}(x, y) + k_2(x, y)\cdot\sigma(x, y)
# \end{equation}
#
# where $\sigma(x, y) = \sigma(x)\cdot\sigma(y)$ and $\bar{\sigma}(x, y) = (1 - \sigma(x))\cdot(1 - \sigma(y))$. The sigmoid ($\sigma$) is parameterized by a location ($l$) and a width ($w$).

# %%
import matplotlib.pyplot as plt
import numpy as np

import gpflow

np.random.seed(123)  # for reproducibility of this notebook

plt.style.use("ggplot")
# %matplotlib inline


def plotkernelsample(k, ax, xmin=-3, xmax=3, title=None):
    xx = np.linspace(xmin, xmax, 100)[:, None]
    ax.plot(xx, np.random.multivariate_normal(np.zeros(100), k(xx), 3).T)
    ax.set_title(title)


# %% [markdown]
# We demonstrate the use of the kernel by drawing a number of samples from different parameterizations. Firstly, a simple single change-point between two kernels of differing lengthscales.

# %%
np.random.seed(1)

base_k1 = gpflow.kernels.Matern32(lengthscales=0.2)
base_k2 = gpflow.kernels.Matern32(lengthscales=2.0)
k = gpflow.kernels.ChangePoints([base_k1, base_k2], [0.0], steepness=5.0)

f, ax = plt.subplots(1, 1, figsize=(10, 3))
plotkernelsample(k, ax)

# %% [markdown]
# Secondly, an implementation of a "change window" in which we change from one kernel to another, then back to the original.

# %%
np.random.seed(3)

base_k1 = gpflow.kernels.Matern32(lengthscales=0.3)
base_k2 = gpflow.kernels.Constant()
k = gpflow.kernels.ChangePoints(
    [base_k1, base_k2, base_k1], locations=[-1, 1], steepness=10.0
)

f, ax = plt.subplots(1, 1, figsize=(10, 3))
plotkernelsample(k, ax)

# %% [markdown]
# And finally, allowing different change-points to occur more or less abruptly by defining different steepness parameters.

# %%
np.random.seed(2)

base_k1 = gpflow.kernels.Matern32(lengthscales=0.3)
base_k2 = gpflow.kernels.Constant()
k = gpflow.kernels.ChangePoints(
    [base_k1, base_k2, base_k1], locations=[-1, 1], steepness=[5.0, 50.0]
)

f, ax = plt.subplots(1, 1, figsize=(10, 3))
plotkernelsample(k, ax)


# doc/sphinx/notebooks/advanced/convolutional.pct.py

# ---
# jupyter:
#   jupytext:
#     formats: ipynb,.pct.py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.3.3
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Convolutional Gaussian Processes
# Mark van der Wilk (July 2019)
#
# Here we show a simple example of the rectangles experiment, where we compare a normal squared exponential GP, and a convolutional GP. This is similar to the experiment in [1].
#
# [1] Van der Wilk, Rasmussen, Hensman (2017). Convolutional Gaussian Processes. *Advances in Neural Information Processing Systems 30*.

# %% [markdown]
# ## Generate dataset
# Generate a simple dataset of rectangles. We want to classify whether they are tall or wide. **NOTE:** Here we take care to make sure that the rectangles don't touch the edge, which is different to the original paper. We do this to avoid needing to use patch weights, which are needed to correctly account for edge effects.

# %%
import time

import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp

import gpflow
from gpflow import set_trainable
from gpflow.ci_utils import is_continuous_integration

gpflow.config.set_default_float(np.float64)
gpflow.config.set_default_jitter(1e-4)
gpflow.config.set_default_summary_fmt("notebook")

# for reproducibility of this notebook:
np.random.seed(123)
tf.random.set_seed(42)

MAXITER = 2 if is_continuous_integration() else 100
NUM_TRAIN_DATA = (
    5 if is_continuous_integration() else 100
)  # This is less than in the original rectangles dataset
NUM_TEST_DATA = 7 if is_continuous_integration() else 300
H = W = 14  # width and height. In the original paper this is 28
IMAGE_SHAPE = [H, W]


# %%
def affine_scalar_bijector(shift=None, scale=None):
    scale_bijector = (
        tfp.bijectors.Scale(scale) if scale else tfp.bijectors.Identity()
    )
    shift_bijector = (
        tfp.bijectors.Shift(shift) if shift else tfp.bijectors.Identity()
    )
    return shift_bijector(scale_bijector)


def make_rectangle(arr, x0, y0, x1, y1):
    arr[y0:y1, x0] = 1
    arr[y0:y1, x1] = 1
    arr[y0, x0:x1] = 1
    arr[y1, x0 : x1 + 1] = 1


def make_random_rectangle(arr):
    x0 = np.random.randint(1, arr.shape[1] - 3)
    y0 = np.random.randint(1, arr.shape[0] - 3)
    x1 = np.random.randint(x0 + 2, arr.shape[1] - 1)
    y1 = np.random.randint(y0 + 2, arr.shape[0] - 1)
    make_rectangle(arr, x0, y0, x1, y1)
    return x0, y0, x1, y1


def make_rectangles_dataset(num, w, h):
    d, Y = np.zeros((num, h, w)), np.zeros((num, 1))
    for i, img in enumerate(d):
        for j in range(1000):  # Finite number of tries
            x0, y0, x1, y1 = make_random_rectangle(img)
            rw, rh = y1 - y0, x1 - x0
            if rw == rh:
                img[:, :] = 0
                continue
            Y[i, 0] = rw > rh
            break
    return (
        d.reshape(num, w * h).astype(gpflow.config.default_float()),
        Y.astype(gpflow.config.default_float()),
    )


# %%
X, Y = data = make_rectangles_dataset(NUM_TRAIN_DATA, *IMAGE_SHAPE)
Xt, Yt = test_data = make_rectangles_dataset(NUM_TEST_DATA, *IMAGE_SHAPE)

# %%
plt.figure(figsize=(8, 3))
for i in range(4):
    plt.subplot(1, 4, i + 1)
    plt.imshow(X[i, :].reshape(*IMAGE_SHAPE))
    plt.title(Y[i, 0])

# %% [markdown]
# ## Squared Exponential kernel

# %%
rbf_m = gpflow.models.SVGP(
    gpflow.kernels.SquaredExponential(),
    gpflow.likelihoods.Bernoulli(),
    gpflow.inducing_variables.InducingPoints(X.copy()),
)

# %%
rbf_training_loss_closure = rbf_m.training_loss_closure(data, compile=True)
rbf_elbo = lambda: -rbf_training_loss_closure().numpy()
print("RBF elbo before training: %.4e" % rbf_elbo())

# %%
set_trainable(rbf_m.inducing_variable, False)
start_time = time.time()
res = gpflow.optimizers.Scipy().minimize(
    rbf_training_loss_closure,
    variables=rbf_m.trainable_variables,
    method="l-bfgs-b",
    options={"disp": True, "maxiter": MAXITER},
)
print(f"{res.nfev / (time.time() - start_time):.3f} iter/s")

# %%
train_acc = np.mean((rbf_m.predict_y(X)[0] > 0.5).numpy().astype("float") == Y)
test_acc = np.mean((rbf_m.predict_y(Xt)[0] > 0.5).numpy().astype("float") == Yt)
print(f"Train acc: {train_acc * 100}%\nTest acc : {test_acc*100}%")
print("RBF elbo after training: %.4e" % rbf_elbo())

# %% [markdown]
# ## Convolutional kernel

# %%
f64 = lambda x: np.array(x, dtype=np.float64)
positive_with_min = lambda: affine_scalar_bijector(shift=f64(1e-4))(
    tfp.bijectors.Softplus()
)
constrained = lambda: affine_scalar_bijector(shift=f64(1e-4), scale=f64(100.0))(
    tfp.bijectors.Sigmoid()
)
max_abs_1 = lambda: affine_scalar_bijector(shift=f64(-2.0), scale=f64(4.0))(
    tfp.bijectors.Sigmoid()
)

patch_shape = [3, 3]
conv_k = gpflow.kernels.Convolutional(
    gpflow.kernels.SquaredExponential(), IMAGE_SHAPE, patch_shape
)
conv_k.base_kernel.lengthscales = gpflow.Parameter(
    1.0, transform=positive_with_min()
)
# Weight scale and variance are non-identifiable. We also need to prevent variance from shooting off crazily.
conv_k.base_kernel.variance = gpflow.Parameter(1.0, transform=constrained())
conv_k.weights = gpflow.Parameter(conv_k.weights.numpy(), transform=max_abs_1())
conv_f = gpflow.inducing_variables.InducingPatches(
    np.unique(conv_k.get_patches(X).numpy().reshape(-1, 9), axis=0)
)

# %%
conv_m = gpflow.models.SVGP(conv_k, gpflow.likelihoods.Bernoulli(), conv_f)

# %%
set_trainable(conv_m.inducing_variable, False)
set_trainable(conv_m.kernel.base_kernel.variance, False)
set_trainable(conv_m.kernel.weights, False)

# %%
conv_training_loss_closure = conv_m.training_loss_closure(data, compile=True)
conv_elbo = lambda: -conv_training_loss_closure().numpy()
print("conv elbo before training: %.4e" % conv_elbo())

# %%
start_time = time.time()
res = gpflow.optimizers.Scipy().minimize(
    conv_training_loss_closure,
    variables=conv_m.trainable_variables,
    method="l-bfgs-b",
    options={"disp": True, "maxiter": MAXITER / 10},
)
print(f"{res.nfev / (time.time() - start_time):.3f} iter/s")

# %%
set_trainable(conv_m.kernel.base_kernel.variance, True)
res = gpflow.optimizers.Scipy().minimize(
    conv_training_loss_closure,
    variables=conv_m.trainable_variables,
    method="l-bfgs-b",
    options={"disp": True, "maxiter": MAXITER},
)
train_acc = np.mean((conv_m.predict_y(X)[0] > 0.5).numpy().astype("float") == Y)
test_acc = np.mean(
    (conv_m.predict_y(Xt)[0] > 0.5).numpy().astype("float") == Yt
)
print(f"Train acc: {train_acc * 100}%\nTest acc : {test_acc*100}%")
print("conv elbo after training: %.4e" % conv_elbo())

# %%
res = gpflow.optimizers.Scipy().minimize(
    conv_training_loss_closure,
    variables=conv_m.trainable_variables,
    method="l-bfgs-b",
    options={"disp": True, "maxiter": MAXITER},
)
train_acc = np.mean((conv_m.predict_y(X)[0] > 0.5).numpy().astype("float") == Y)
test_acc = np.mean(
    (conv_m.predict_y(Xt)[0] > 0.5).numpy().astype("float") == Yt
)
print(f"Train acc: {train_acc * 100}%\nTest acc : {test_acc*100}%")
print("conv elbo after training: %.4e" % conv_elbo())

# %%
set_trainable(conv_m.kernel.weights, True)
res = gpflow.optimizers.Scipy().minimize(
    conv_training_loss_closure,
    variables=conv_m.trainable_variables,
    method="l-bfgs-b",
    options={"disp": True, "maxiter": MAXITER},
)
train_acc = np.mean((conv_m.predict_y(X)[0] > 0.5).numpy().astype("float") == Y)
test_acc = np.mean(
    (conv_m.predict_y(Xt)[0] > 0.5).numpy().astype("float") == Yt
)
print(f"Train acc: {train_acc * 100}%\nTest acc : {test_acc*100}%")
print("conv elbo after training: %.4e" % conv_elbo())

# %%
gpflow.utilities.print_summary(rbf_m)

# %%
gpflow.utilities.print_summary(conv_m)

# %% [markdown]
# ## Conclusion
# The convolutional kernel performs much better in this simple task. It demonstrates non-local generalization of the strong assumptions in the kernel.


# doc/sphinx/notebooks/advanced/coregionalisation.pct.py

# ---
# jupyter:
#   jupytext:
#     formats: ipynb,.pct.py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.3.3
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %% [markdown]
# # A simple demonstration of coregionalization
#
# This notebook shows how to construct a multi-output GP model using GPflow. We will consider a regression problem for functions $f: \mathbb{R}^D \rightarrow \mathbb{R}^P$. We assume that the dataset is of the form $(X_1, f_1), \dots, (X_P, f_P)$, that is, we do not necessarily observe all the outputs for a particular input location (for cases where there are fully observed outputs for each input, see [Multi-output Gaussian processes in GPflow](./multioutput.ipynb) for a more efficient implementation). We allow each $f_i$ to have a different noise distribution by assigning a different likelihood to each.
#
# For this problem, we model $f$ as a *coregionalized* Gaussian process, which assumes a kernel of the form:
#
# \begin{equation}
# \textrm{cov}(f_i(X), f_j(X^\prime)) = k(X, X^\prime) \cdot B[i, j].
# \end{equation}
#
# The covariance of the $i$th function at $X$ and the $j$th function at $X^\prime$ is a kernel applied at $X$ and $X^\prime$, times the $(i, j)$th entry of a positive definite $P \times P$ matrix $B$. This is known as the **intrinsic model of coregionalization (ICM)** (Bonilla and Williams, 2008).
#
# To make sure that B is positive-definite, we parameterize it as:
#
# \begin{equation}
# B = W W^\top + \textrm{diag}(\kappa).
# \end{equation}
#
# To build such a model in GPflow, we need to perform the following two steps:
#
#  * Create the kernel function defined previously, using the `Coregion` kernel class.
#  * Augment the training data X with an extra column that contains an integer index to indicate which output an observation is associated with. This is essential to make the data work with the `Coregion` kernel.
#  * Create a likelihood for each output using the `SwitchedLikelihood` class, which is a container for other likelihoods.
#  * Augment the training data Y with an extra column that contains an integer index to indicate which likelihood an observation is associated with.

# %%
import matplotlib.pyplot as plt
import numpy as np

import gpflow
from gpflow.ci_utils import reduce_in_tests

# %matplotlib inline

plt.rcParams["figure.figsize"] = (12, 6)
np.random.seed(123)

# %% [markdown]
# ## Data preparation
# We start by generating some training data to fit the model with. For this example, we choose the following two correlated functions for our outputs:
#
# \begin{align}
# y_1 &= \sin(6x) + \epsilon_1, \qquad \epsilon_1 \sim \mathcal{N}(0, 0.009) \\
# y_2 &= \sin(6x + 0.7) + \epsilon_2, \qquad \epsilon_2 \sim \mathcal{N}(0, 0.01)
# \end{align}
#

# %%
# make a dataset with two outputs, correlated, heavy-tail noise. One has more noise than the other.
X1 = np.random.rand(100, 1)  # Observed locations for first output
X2 = np.random.rand(50, 1) * 0.5  # Observed locations for second output

Y1 = np.sin(6 * X1) + np.random.randn(*X1.shape) * 0.03
Y2 = np.sin(6 * X2 + 0.7) + np.random.randn(*X2.shape) * 0.1

plt.figure(figsize=(8, 4))
plt.plot(X1, Y1, "x", mew=2)
_ = plt.plot(X2, Y2, "x", mew=2)

# %% [markdown]
# ## Data formatting for the coregionalized model
# We add an extra column to our training dataset that contains an index that specifies which output is observed.

# %%
# Augment the input with ones or zeros to indicate the required output dimension
X_augmented = np.vstack(
    (np.hstack((X1, np.zeros_like(X1))), np.hstack((X2, np.ones_like(X2))))
)

# Augment the Y data with ones or zeros that specify a likelihood from the list of likelihoods
Y_augmented = np.vstack(
    (np.hstack((Y1, np.zeros_like(Y1))), np.hstack((Y2, np.ones_like(Y2))))
)

# %% [markdown]
# ## Building the coregionalization kernel
# We build a coregionalization kernel with the Matern 3/2 kernel as the base kernel. This acts on the leading ([0]) data dimension of the augmented X values. The `Coregion` kernel indexes the outputs, and acts on the last ([1]) data dimension (indices) of the augmented X values. To specify these dimensions, we use the built-in `active_dims` argument in the kernel constructor. To construct the full multi-output kernel, we take the product of the two kernels (for a more in-depth tutorial on kernel combination, see [Manipulating kernels](./kernels.ipynb)).

# %%
output_dim = 2  # Number of outputs
rank = 1  # Rank of W

# Base kernel
k = gpflow.kernels.Matern32(active_dims=[0])

# Coregion kernel
coreg = gpflow.kernels.Coregion(
    output_dim=output_dim, rank=rank, active_dims=[1]
)

kern = k * coreg

# %% [markdown]
# **Note:** W = 0 is a saddle point in the objective, which would result in the value of `W` not being optimized to fit the data.
# Hence, by default, the `W` matrix is initialized with 0.1. Alternatively, you could re-initialize the matrix to random entries.

# %% [markdown]
# ## Constructing the model
# The final element in building the model is to specify the likelihood for each output dimension. To do this, use the `SwitchedLikelihood` object in GPflow.

# %%
# This likelihood switches between Gaussian noise with different variances for each f_i:
lik = gpflow.likelihoods.SwitchedLikelihood(
    [gpflow.likelihoods.Gaussian(), gpflow.likelihoods.Gaussian()]
)

# now build the GP model as normal
m = gpflow.models.VGP((X_augmented, Y_augmented), kernel=kern, likelihood=lik)

# fit the covariance function parameters
maxiter = reduce_in_tests(10000)
gpflow.optimizers.Scipy().minimize(
    m.training_loss,
    m.trainable_variables,
    options=dict(maxiter=maxiter),
    method="L-BFGS-B",
)


# %% [markdown]
# That's it: the model is trained. Let's plot the model fit to see what's happened.

# %%
def plot_gp(x, mu, var, color, label):
    plt.plot(x, mu, color=color, lw=2, label=label)
    plt.fill_between(
        x[:, 0],
        (mu - 2 * np.sqrt(var))[:, 0],
        (mu + 2 * np.sqrt(var))[:, 0],
        color=color,
        alpha=0.4,
    )


def plot(m):
    plt.figure(figsize=(8, 4))
    Xtest = np.linspace(0, 1, 100)[:, None]
    (line,) = plt.plot(X1, Y1, "x", mew=2)
    mu, var = m.predict_f(np.hstack((Xtest, np.zeros_like(Xtest))))
    plot_gp(Xtest, mu, var, line.get_color(), "Y1")

    (line,) = plt.plot(X2, Y2, "x", mew=2)
    mu, var = m.predict_f(np.hstack((Xtest, np.ones_like(Xtest))))
    plot_gp(Xtest, mu, var, line.get_color(), "Y2")

    plt.legend()


plot(m)

# %% [markdown]
# From the plots, we can see:
#
#  - The first function (blue) has low posterior variance everywhere because there are so many observations, and the noise variance is small.
#  - The second function (orange) has higher posterior variance near the data, because the data are more noisy, and very high posterior variance where there are no observations (x > 0.5).
#  - The model has done a reasonable job of estimating the noise variance and lengthscale.
#  - The model recognises the correlation between the two functions and is able to suggest (with uncertainty) that because x > 0.5 the orange curve should follow the blue curve (which we know to be the truth from the data-generating procedure).
#
# The covariance matrix between outputs is as follows:

# %%
B = coreg.output_covariance().numpy()
print("B =", B)
_ = plt.imshow(B)

# %% [markdown]
# ## References
#
# Bonilla, Edwin V., Kian M. Chai, and Christopher Williams. "Multi-task Gaussian process prediction." _Advances in neural information processing systems_. 2008.


# doc/sphinx/notebooks/advanced/fast_predictions.pct.py

# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.9.1
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# + [markdown] id="eYrSpUncKGSk"
# # Faster predictions by caching

# + [markdown] id="PLuPjfS7KLQ-"
# The default behaviour of `predict_f` in GPflow models is to compute the predictions from scratch on each call. This is convenient when predicting and training are interleaved, and simplifies the use of these models. There are some use cases, such as Bayesian optimisation, where prediction (at different test points) happens much more frequently than training. In these cases it is convenient to cache parts of the calculation which do not depend upon the test points, and reuse those parts between predictions.
#
# There are three models to which we want to add this caching capability: GPR, (S)VGP and SGPR. The VGP and SVGP can be considered together; the difference between the models is whether to condition on the full training data set (VGP) or on the inducing variables (SVGP).

# + [markdown] id="EACkO-iRKM5T"
# ## Posterior predictive distribution
#
# The posterior predictive distribution evaluated at a set of test points $\mathbf{x}_*$ for a Gaussian process model is given by:
# \begin{equation*}
# p(\mathbf{f}_*|X, Y) = \mathcal{N}(\mu, \Sigma)
# \end{equation*}
#
# In the case of the GPR model, the parameters $\mu$ and $\Sigma$ are given by:
# \begin{equation*}
# \mu = K_{nm}[K_{mm} + \sigma^2I]^{-1}\mathbf{y}
# \end{equation*}
# and
# \begin{equation*}
# \Sigma = K_{nn} - K_{nm}[K_{mm} + \sigma^2I]^{-1}K_{mn}
# \end{equation*}
#
# The posterior predictive distribution for the VGP and SVGP model is parameterised as follows:
# \begin{equation*}
# \mu = K_{nu}K_{uu}^{-1}\mathbf{u}
# \end{equation*}
# and
# \begin{equation*}
# \Sigma = K_{nn} - K_{nu}K_{uu}^{-1}K_{un}
# \end{equation*}
#
# Finally, the parameters for the SGPR model are:
# \begin{equation*}
# \mu = K_{nu}L^{-T}L_B^{-T}\mathbf{c}
# \end{equation*}
# and
# \begin{equation*}
# \Sigma = K_{nn} - K_{nu}L^{-T}(I - B^{-1})L^{-1}K_{un}
# \end{equation*}
#
# Where the mean function is not the zero function, the predictive mean should have the mean function evaluated at the test points added to it.

# + [markdown] id="GX1U-fYPKPrt"
# ## What can be cached?
#
# We cache two separate values: $\alpha$ and $Q^{-1}$. These correspond to the parts of the mean and covariance functions respectively which do not depend upon the test points. In the case of the GPR these are the same value:
# \begin{equation*}
# \alpha = Q^{-1} = [K_{mm} + \sigma^2I]^{-1}
# \end{equation*}
# in the case of the VGP and SVGP model these are:
# \begin{equation*}
# \alpha = K_{uu}^{-1}\mathbf{u}\\ Q^{-1} = K_{uu}^{-1}
# \end{equation*}
# and in the case of the SGPR model these are:
# \begin{equation*}
# \alpha = L^{-T}L_B^{-T}\mathbf{c}\\ Q^{-1} = L^{-T}(I - B^{-1})L^{-1}
# \end{equation*}
#
#
# Note that in the (S)VGP case, $\alpha$ is the parameter as proposed by Opper and Archambeau for the mean of the predictive distribution.

# +
import numpy as np

import gpflow
from gpflow.ci_utils import reduce_in_tests

# Create some data
n_data = reduce_in_tests(1000)
X = np.linspace(-1.1, 1.1, n_data)[:, None]
Y = np.sin(X)
Xnew = np.linspace(-1.1, 1.1, n_data)[:, None]
inducing_points = Xnew

# + [markdown] id="FzCgor4nKUcW"
#
# ## GPR Example
#
# We will construct a GPR model to demonstrate the faster predictions from using the cached data in the GPFlow posterior classes (subclasses of `gpflow.posteriors.AbstractPosterior`).

# + id="BMnIdXNiKU6t"
model = gpflow.models.GPR(
    (X, Y),
    gpflow.kernels.SquaredExponential(),
)
# -

# The `predict_f` method on the `GPModel` class performs no caching.

# %%timeit
model.predict_f(Xnew)

# To make use of the caching, first retrieve the posterior class from the model. The posterior class has methods to predict the parameters of marginal distributions at test points, in the same way as the `predict_f` method of the `GPModel`.
posterior = model.posterior()

# %%timeit
posterior.predict_f(Xnew)

# ## SVGP Example
#
# Likewise, we will construct an SVGP model to demonstrate the faster predictions from using the cached data in the GPFlow posterior classes.

# + id="BMnIdXNiKU6t"
model = gpflow.models.SVGP(
    gpflow.kernels.SquaredExponential(),
    gpflow.likelihoods.Gaussian(),
    inducing_points,
)
# -

# The `predict_f` method on the `GPModel` class performs no caching.

# %%timeit
model.predict_f(Xnew)

# And again using the posterior object and caching

posterior = model.posterior()

# %%timeit
posterior.predict_f(Xnew)

# ## SGPR Example
#
# And finally, we follow the same approach this time for the SGPR case.

model = gpflow.models.SGPR(
    (X, Y), gpflow.kernels.SquaredExponential(), inducing_points
)

# The predict_f method on the instance performs no caching.

# %%timeit
model.predict_f(Xnew)

# Using the posterior object instead:

posterior = model.posterior()

# %%timeit
posterior.predict_f(Xnew)


# doc/sphinx/notebooks/advanced/gps_for_big_data.pct.py

# ---
# jupyter:
#   jupytext:
#     formats: ipynb,.pct.py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.3.3
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Stochastic Variational Inference for scalability with SVGP

# %% [markdown]
# One of the main criticisms of Gaussian processes is their scalability to large datasets. In this notebook, we illustrate how to use the state-of-the-art Stochastic Variational Gaussian Process (SVGP) (*Hensman, et. al. 2013*) to overcome this problem.

# %%
# %matplotlib inline
import itertools
import time

import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf

import gpflow
from gpflow.ci_utils import reduce_in_tests

plt.style.use("ggplot")

# for reproducibility of this notebook:
rng = np.random.RandomState(123)
tf.random.set_seed(42)

# %% [markdown]
# ## Generating data
# For this notebook example, we generate 10,000 noisy observations from a test function:
# \begin{equation}
# f(x) = \sin(3\pi x) + 0.3\cos(9\pi x) + \frac{\sin(7 \pi x)}{2}
# \end{equation}

# %%
def func(x):
    return (
        np.sin(x * 3 * 3.14)
        + 0.3 * np.cos(x * 9 * 3.14)
        + 0.5 * np.sin(x * 7 * 3.14)
    )


N = 10000  # Number of training observations

X = rng.rand(N, 1) * 2 - 1  # X values
Y = func(X) + 0.2 * rng.randn(N, 1)  # Noisy Y values
data = (X, Y)

# %% [markdown]
# We plot the data along with the noiseless generating function:

# %%
plt.plot(X, Y, "x", alpha=0.2)
Xt = np.linspace(-1.1, 1.1, 1000)[:, None]
Yt = func(Xt)
_ = plt.plot(Xt, Yt, c="k")

# %% [markdown]
# ## Building the model
# The main idea behind SVGP is to approximate the true GP posterior with a GP conditioned on a small set of "inducing" values. This smaller set can be thought of as summarizing the larger dataset. For this example, we will select a set of 50 inducing locations that are initialized from the training dataset:

# %%
M = 50  # Number of inducing locations

kernel = gpflow.kernels.SquaredExponential()
Z = X[
    :M, :
].copy()  # Initialize inducing locations to the first M inputs in the dataset

m = gpflow.models.SVGP(kernel, gpflow.likelihoods.Gaussian(), Z, num_data=N)

# %% [markdown]
# ## Likelihood computation: batch vs. minibatch
# First we showcase the model's performance using the whole dataset to compute the ELBO.

# %%
elbo = tf.function(m.elbo)

# %%
# TensorFlow re-traces & compiles a `tf.function`-wrapped method at *every* call if the arguments are numpy arrays instead of tf.Tensors. Hence:
tensor_data = tuple(map(tf.convert_to_tensor, data))
elbo(tensor_data)  # run it once to trace & compile

# %%
# %%timeit
elbo(tensor_data)

# %% [markdown]
# We can speed up this calculation by using minibatches of the data. For this example, we use minibatches of size 100.

# %%
minibatch_size = 100

train_dataset = tf.data.Dataset.from_tensor_slices((X, Y)).repeat().shuffle(N)

train_iter = iter(train_dataset.batch(minibatch_size))

ground_truth = elbo(tensor_data).numpy()

# %%
# %%timeit
elbo(next(train_iter))

# %% [markdown]
# ### Stochastical estimation of ELBO
# The minibatch estimate should be an unbiased estimator of the `ground_truth`. Here we show a histogram of the value from different evaluations, together with its mean and the ground truth. The small difference between the mean of the minibatch estimations and the ground truth shows that the minibatch estimator is working as expected.

# %%
evals = [
    elbo(minibatch).numpy() for minibatch in itertools.islice(train_iter, 100)
]

# %%
plt.hist(evals, label="Minibatch estimations")
plt.axvline(ground_truth, c="k", label="Ground truth")
plt.axvline(np.mean(evals), c="g", ls="--", label="Minibatch mean")
plt.legend()
plt.title("Histogram of ELBO evaluations using minibatches")
print(
    "Discrepancy between ground truth and minibatch estimate:",
    ground_truth - np.mean(evals),
)

# %% [markdown]
# ### Minibatches speed up computation
# The reason for using minibatches is that it decreases the time needed to make an optimization step, because estimating the objective is computationally cheaper with fewer data points. Here we plot the change in time required with the size of the minibatch. We see that smaller minibatches result in a cheaper estimate of the objective.

# %%
# Evaluate objective for different minibatch sizes
minibatch_proportions = np.logspace(-2, 0, 10)
times = []
objs = []
for mbp in minibatch_proportions:
    batchsize = int(N * mbp)
    train_iter = iter(train_dataset.batch(batchsize))
    start_time = time.time()
    objs.append(
        [elbo(minibatch) for minibatch in itertools.islice(train_iter, 20)]
    )
    times.append(time.time() - start_time)

# %%
f, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))
ax1.plot(minibatch_proportions, times, "x-")
ax1.set_xlabel("Minibatch proportion")
ax1.set_ylabel("Time taken")

ax2.plot(minibatch_proportions, np.array(objs), "kx")
ax2.set_xlabel("Minibatch proportion")
ax2.set_ylabel("ELBO estimates")


# %% [markdown]
# ### Running stochastic optimization

# %% [markdown]
# First we create a utility function that plots the model's predictions:

# %%
def plot(title=""):
    plt.figure(figsize=(12, 4))
    plt.title(title)
    pX = np.linspace(-1, 1, 100)[:, None]  # Test locations
    pY, pYv = m.predict_y(pX)  # Predict Y values at test locations
    plt.plot(X, Y, "x", label="Training points", alpha=0.2)
    (line,) = plt.plot(pX, pY, lw=1.5, label="Mean of predictive posterior")
    col = line.get_color()
    plt.fill_between(
        pX[:, 0],
        (pY - 2 * pYv ** 0.5)[:, 0],
        (pY + 2 * pYv ** 0.5)[:, 0],
        color=col,
        alpha=0.6,
        lw=1.5,
    )
    Z = m.inducing_variable.Z.numpy()
    plt.plot(Z, np.zeros_like(Z), "k|", mew=2, label="Inducing locations")
    plt.legend(loc="lower right")


plot(title="Predictions before training")

# %% [markdown]
# Now we can train our model. For optimizing the ELBO, we use the Adam Optimizer *(Kingma and Ba 2015)* which is designed for stochastic objective functions. We create a `run_adam` utility function  to perform the optimization.

# %%
minibatch_size = 100

# We turn off training for inducing point locations
gpflow.set_trainable(m.inducing_variable, False)


def run_adam(model, iterations):
    """
    Utility function running the Adam optimizer

    :param model: GPflow model
    :param interations: number of iterations
    """
    # Create an Adam Optimizer action
    logf = []
    train_iter = iter(train_dataset.batch(minibatch_size))
    training_loss = model.training_loss_closure(train_iter, compile=True)
    optimizer = tf.optimizers.Adam()

    @tf.function
    def optimization_step():
        optimizer.minimize(training_loss, model.trainable_variables)

    for step in range(iterations):
        optimization_step()
        if step % 10 == 0:
            elbo = -training_loss().numpy()
            logf.append(elbo)
    return logf


# %% [markdown]
# Now we run the optimization loop for 20,000 iterations.

# %%
maxiter = reduce_in_tests(20000)

logf = run_adam(m, maxiter)
plt.plot(np.arange(maxiter)[::10], logf)
plt.xlabel("iteration")
_ = plt.ylabel("ELBO")

# %% [markdown]
# Finally, we plot the model's predictions.

# %%
plot("Predictions after training")

# %% [markdown]
# ## Further reading
#
# Several notebooks expand on this one:
#
# - [Advanced Sparse GP regression](../advanced/advanced_many_points.ipynb), which goes into deeper detail on sparse Gaussian process methods.
# - [Optimization](../advanced/optimisation.ipynb) discussing optimizing GP models.
# - [Natural gradients](../advanced/natural_gradients.ipynb) for optimizing SVGP models efficiently.

# %% [markdown]
# ## References:
# Hensman, James, Nicolo Fusi, and Neil D. Lawrence. "Gaussian processes for big data." Uncertainty in Artificial Intelligence (2013).
#
# Kingma, Diederik P., and Jimmy Ba. "Adam: A method for stochastic optimization." arXiv preprint arXiv:1412.6980 (2014).


# doc/sphinx/notebooks/advanced/heteroskedastic.pct.py

# ---
# jupyter:
#   jupytext:
#     formats: ipynb,.pct.py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.6.0
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Heteroskedastic Likelihood and Multi-Latent GP

# %% [markdown]
# ## Standard (Homoskedastic) Regression
# In standard GP regression, the GP latent function is used to learn the location parameter of a likelihood distribution (usually a Gaussian) as a function of the input $x$, whereas the scale parameter is considered constant. This is a homoskedastic model, which is unable to capture variations of the noise distribution with the input $x$.
#
#
# ## Heteroskedastic Regression
# This notebooks shows how to construct a model which uses multiple (here two) GP latent functions to learn both the location and the scale of the Gaussian likelihood distribution. It does so by connecting a **Multi-Output Kernel**, which generates multiple GP latent functions, to a **Heteroskedastic Likelihood**, which maps the latent GPs into a single likelihood.
#
# The generative model is described as:
#
# $$ f_1(x) \sim \mathcal{GP}(0, k_1(\cdot, \cdot)) $$
# $$ f_2(x) \sim \mathcal{GP}(0, k_2(\cdot, \cdot)) $$
# $$ \text{loc}(x) = f_1(x) $$
# $$ \text{scale}(x) = \text{transform}(f_2(x)) $$
# $$ y_i|f_1, f_2, x_i \sim \mathcal{N}(\text{loc}(x_i),\;\text{scale}(x_i)^2)$$
#
# The function $\text{transform}$ is used to map from the unconstrained GP $f_2$ to **positive-only values**, which is required as it represents the $\text{scale}$ of a Gaussian likelihood. In this notebook, the $\exp$ function will be used as the $\text{transform}$. Other positive transforms such as the $\text{softplus}$ function can also be used.

# %%
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp

import gpflow as gpf

# %% [markdown]
# ## Data Generation
# We generate heteroskedastic data by substituting the random latent functions $f_1$ and $f_2$ of the generative model by deterministic $\sin$ and $\cos$ functions. The input $X$ is built with $N=1001$ uniformly spaced values in the interval $[0, 4\pi]$. The outputs $Y$ are still sampled from a Gaussian likelihood.
#
# $$ x_i \in [0, 4\pi], \quad i = 1,\dots,N $$
# $$ f_1(x) = \sin(x) $$
# $$ f_2(x) = \cos(x) $$
# $$ \text{loc}(x) = f_1(x) $$
# $$ \text{scale}(x) = \exp(f_2(x)) $$
# $$ y_i|x_i \sim \mathcal{N}(\text{loc}(x_i),\;\text{scale}(x_i)^2)$$

# %%
N = 1001

np.random.seed(0)
tf.random.set_seed(0)

# Build inputs X
X = np.linspace(0, 4 * np.pi, N)[:, None]  # X must be of shape [N, 1]

# Deterministic functions in place of latent ones
f1 = np.sin
f2 = np.cos

# Use transform = exp to ensure positive-only scale values
transform = np.exp

# Compute loc and scale as functions of input X
loc = f1(X)
scale = transform(f2(X))

# Sample outputs Y from Gaussian Likelihood
Y = np.random.normal(loc, scale)

# %% [markdown]
# ### Plot Data
# Note how the distribution density (shaded area) and the outputs $Y$ both change depending on the input $X$.

# %%
def plot_distribution(X, Y, loc, scale):
    plt.figure(figsize=(15, 5))
    x = X.squeeze()
    for k in (1, 2):
        lb = (loc - k * scale).squeeze()
        ub = (loc + k * scale).squeeze()
        plt.fill_between(x, lb, ub, color="silver", alpha=1 - 0.05 * k ** 3)
    plt.plot(x, lb, color="silver")
    plt.plot(x, ub, color="silver")
    plt.plot(X, loc, color="black")
    plt.scatter(X, Y, color="gray", alpha=0.8)
    plt.show()
    plt.close()


plot_distribution(X, Y, loc, scale)


# %% [markdown]
# ## Build Model

# %% [markdown]
# ### Likelihood
# This implements the following part of the generative model:
# $$ \text{loc}(x) = f_1(x) $$
# $$ \text{scale}(x) = \text{transform}(f_2(x)) $$
# $$ y_i|f_1, f_2, x_i \sim \mathcal{N}(\text{loc}(x_i),\;\text{scale}(x_i)^2)$$

# %%
likelihood = gpf.likelihoods.HeteroskedasticTFPConditional(
    distribution_class=tfp.distributions.Normal,  # Gaussian Likelihood
    scale_transform=tfp.bijectors.Exp(),  # Exponential Transform
)

print(f"Likelihood's expected latent_dim: {likelihood.latent_dim}")

# %% [markdown]
# ### Kernel
# This implements the following part of the generative model:
# $$ f_1(x) \sim \mathcal{GP}(0, k_1(\cdot, \cdot)) $$
# $$ f_2(x) \sim \mathcal{GP}(0, k_2(\cdot, \cdot)) $$
#
# with both kernels being modeled as separate and independent $\text{SquaredExponential}$ kernels.

# %%
kernel = gpf.kernels.SeparateIndependent(
    [
        gpf.kernels.SquaredExponential(),  # This is k1, the kernel of f1
        gpf.kernels.SquaredExponential(),  # this is k2, the kernel of f2
    ]
)
# The number of kernels contained in gpf.kernels.SeparateIndependent must be the same as likelihood.latent_dim

# %% [markdown]
# ### Inducing Points
# Since we will use the **SVGP** model to perform inference, we need to implement the inducing variables $U_1$ and $U_2$, both with size $M=20$, which are used to approximate $f_1$ and $f_2$ respectively, and initialize the inducing points positions $Z_1$ and $Z_2$. This gives a total of $2M=40$ inducing variables and inducing points.
#
# The inducing variables and their corresponding inputs will be Separate and Independent, but both $Z_1$ and $Z_2$ will be initialized as $Z$, which are placed as $M=20$ equally spaced points in $[\min(X), \max(X)]$.
#

# %%
M = 20  # Number of inducing variables for each f_i

# Initial inducing points position Z
Z = np.linspace(X.min(), X.max(), M)[:, None]  # Z must be of shape [M, 1]

inducing_variable = gpf.inducing_variables.SeparateIndependentInducingVariables(
    [
        gpf.inducing_variables.InducingPoints(Z),  # This is U1 = f1(Z1)
        gpf.inducing_variables.InducingPoints(Z),  # This is U2 = f2(Z2)
    ]
)

# %% [markdown]
# ### SVGP Model
# Build the **SVGP** model by composing the **Kernel**, the **Likelihood** and the **Inducing Variables**.
#
# Note that the model needs to be instructed about the number of latent GPs by passing `num_latent_gps=likelihood.latent_dim`.

# %%
model = gpf.models.SVGP(
    kernel=kernel,
    likelihood=likelihood,
    inducing_variable=inducing_variable,
    num_latent_gps=likelihood.latent_dim,
)

model

# %% [markdown]
# ## Model Optimization
#
# ### Build Optimizers (NatGrad + Adam)

# %%
data = (X, Y)
loss_fn = model.training_loss_closure(data)

gpf.utilities.set_trainable(model.q_mu, False)
gpf.utilities.set_trainable(model.q_sqrt, False)

variational_vars = [(model.q_mu, model.q_sqrt)]
natgrad_opt = gpf.optimizers.NaturalGradient(gamma=0.1)

adam_vars = model.trainable_variables
adam_opt = tf.optimizers.Adam(0.01)


@tf.function
def optimisation_step():
    natgrad_opt.minimize(loss_fn, variational_vars)
    adam_opt.minimize(loss_fn, adam_vars)


# %% [markdown]
# ### Run Optimization Loop

# %%
epochs = 100
log_freq = 20

for epoch in range(1, epochs + 1):
    optimisation_step()

    # For every 'log_freq' epochs, print the epoch and plot the predictions against the data
    if epoch % log_freq == 0 and epoch > 0:
        print(f"Epoch {epoch} - Loss: {loss_fn().numpy() : .4f}")
        Ymean, Yvar = model.predict_y(X)
        Ymean = Ymean.numpy().squeeze()
        Ystd = tf.sqrt(Yvar).numpy().squeeze()
        plot_distribution(X, Y, Ymean, Ystd)

model

# %% [markdown]
# ## Further reading
#
# See [Chained Gaussian Processes](http://proceedings.mlr.press/v51/saul16.html) by Saul et al. (AISTATS 2016).


# doc/sphinx/notebooks/advanced/mcmc.pct.py

# ---
# jupyter:
#   jupytext:
#     formats: ipynb,.pct.py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.13.8
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %% [markdown]
# # MCMC (Markov Chain Monte Carlo)

# %% [markdown]
# GPflow allows you to approximate the posterior over the latent functions of its models (and over the hyperparameters after setting a prior for those) using Hamiltonian Monte Carlo (HMC)

# %%
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
from multiclass_classification import colors, plot_from_samples
from tensorflow_probability import distributions as tfd

import gpflow
from gpflow import set_trainable
from gpflow.ci_utils import reduce_in_tests

gpflow.config.set_default_float(np.float64)
gpflow.config.set_default_jitter(1e-4)
gpflow.config.set_default_summary_fmt("notebook")
# convert to float64 for tfp to play nicely with gpflow in 64
f64 = gpflow.utilities.to_default_float

tf.random.set_seed(123)


# %matplotlib inline

# %% [markdown]
#
# In this notebook, we provide three examples:
#
# * [Example 1](#Example-1:-GP-regression): Sampling hyperparameters in Gaussian process regression
# * [Example 2](#Example-2:-Sparse-MC-for-multiclass-classification): Sparse Variational MC applied to the multiclass classification problem
# * [Example 3](#Example-3:-Fully-Bayesian-inference-for-generalized-GP-models-with-HMC): Full Bayesian inference for Gaussian process models

# %% [markdown]
# ## Example 1: GP regression

# %% [markdown]
# We first consider the GP regression (with Gaussian noise) for which the marginal likelihood $p(\mathbf y\,|\,\theta)$ can be computed exactly.
#
# The GPR model parameterized by $\theta = [\tau]$ is given by
# \begin{equation}
# Y_i = f(X_i) + \varepsilon_i
# \end{equation}
# where $f \sim \mathcal{GP}(\mu(.), k(., .))$, and $\varepsilon \sim \mathcal{N}(0, \tau^2 I)$.
#
# See the [Basic (Gaussian likelihood) GP regression model](../basics/regression.ipynb) for more details on GPR and for a treatment of the direct likelihood maximization.
#
#

# %% [markdown]
# ### Data for a one-dimensional regression problem

# %%
rng = np.random.RandomState(42)

N = 15


def synthetic_data(num: int, rng: np.random.RandomState):
    X = rng.rand(num, 1)
    Y = np.sin(12 * X) + 0.66 * np.cos(25 * X) + rng.randn(num, 1) * 0.1 + 3
    return X, Y


data = (X, Y) = synthetic_data(N, rng)

plt.figure(figsize=(12, 6))
plt.plot(X, Y, "kx", mew=2)
plt.xlabel("$X$")
plt.ylabel("$Y$")
plt.title("toy data")
plt.show()

# %% [markdown]
# ### MCMC for hyperparameters $\theta$
#
# We now want to sample from the posterior over $\theta$:
# \begin{equation}
# p(\theta|\mathbf{y}) \propto p(\mathbf{y}|\theta)p(\theta)
# \end{equation}
#
# Firstly, we build the GPR model.

# %%
kernel = gpflow.kernels.Matern52(lengthscales=0.3)
mean_function = gpflow.mean_functions.Linear(1.0, 0.0)
model = gpflow.models.GPR(data, kernel, mean_function, noise_variance=0.01)

# %% [markdown]
# Secondly, we initialize the model to the maximum likelihood solution.

# %%
optimizer = gpflow.optimizers.Scipy()
optimizer.minimize(model.training_loss, model.trainable_variables)

print(f"log posterior density at optimum: {model.log_posterior_density()}")

# %% [markdown]
# Thirdly, we add priors to the hyperparameters.

# %%
# tfp.distributions dtype is inferred from parameters - so convert to 64-bit
model.kernel.lengthscales.prior = tfd.Gamma(f64(1.0), f64(1.0))
model.kernel.variance.prior = tfd.Gamma(f64(1.0), f64(1.0))
model.likelihood.variance.prior = tfd.Gamma(f64(1.0), f64(1.0))
model.mean_function.A.prior = tfd.Normal(f64(0.0), f64(10.0))
model.mean_function.b.prior = tfd.Normal(f64(0.0), f64(10.0))

gpflow.utilities.print_summary(model)

# %% [markdown]
# We now sample from the posterior using HMC.

# %%
num_burnin_steps = reduce_in_tests(300)
num_samples = reduce_in_tests(1000)

# Note that here we need model.trainable_parameters, not trainable_variables - only parameters can have priors!
hmc_helper = gpflow.optimizers.SamplingHelper(
    model.log_posterior_density, model.trainable_parameters
)

hmc = tfp.mcmc.HamiltonianMonteCarlo(
    target_log_prob_fn=hmc_helper.target_log_prob_fn,
    num_leapfrog_steps=10,
    step_size=0.01,
)
adaptive_hmc = tfp.mcmc.SimpleStepSizeAdaptation(
    hmc,
    num_adaptation_steps=10,
    target_accept_prob=f64(0.75),
    adaptation_rate=0.1,
)


@tf.function
def run_chain_fn():
    return tfp.mcmc.sample_chain(
        num_results=num_samples,
        num_burnin_steps=num_burnin_steps,
        current_state=hmc_helper.current_state,
        kernel=adaptive_hmc,
        trace_fn=lambda _, pkr: pkr.inner_results.is_accepted,
    )


samples, traces = run_chain_fn()
parameter_samples = hmc_helper.convert_to_constrained_values(samples)

param_to_name = {
    param: name
    for name, param in gpflow.utilities.parameter_dict(model).items()
}


# %% [markdown]
# **NOTE:** All the Hamiltonian MCMC sampling takes place in an unconstrained space (where constrained parameters have been mapped via a bijector to an unconstrained space). This makes the optimization, as required in the gradient step, much easier.
#
# However, we often wish to sample the constrained parameter values, not the unconstrained one. The `SamplingHelper` helps us convert our unconstrained values to constrained parameter ones.
#

# %%
def plot_samples(samples, parameters, y_axis_label):
    plt.figure(figsize=(8, 4))
    for val, param in zip(samples, parameters):
        plt.plot(tf.squeeze(val), label=param_to_name[param])
    plt.legend(bbox_to_anchor=(1.0, 1.0))
    plt.xlabel("HMC iteration")
    plt.ylabel(y_axis_label)


plot_samples(samples, model.trainable_parameters, "unconstrained values")
plot_samples(
    parameter_samples,
    model.trainable_parameters,
    "constrained parameter values",
)


# %% [markdown]
# You can also inspect the marginal distribution of samples.

# %%
def marginal_samples(samples, parameters, y_axis_label):
    fig, axes = plt.subplots(
        1, len(param_to_name), figsize=(15, 3), constrained_layout=True
    )
    for ax, val, param in zip(axes, samples, parameters):
        ax.hist(np.stack(val).flatten(), bins=20)
        ax.set_title(param_to_name[param])
    fig.suptitle(y_axis_label)
    plt.show()


marginal_samples(
    samples, model.trainable_parameters, "unconstrained variable samples"
)
marginal_samples(
    parameter_samples,
    model.trainable_parameters,
    "constrained parameter samples",
)


# %% [markdown]
#
#
# **NOTE:** The sampler runs in unconstrained space (so that positive parameters remain positive, and parameters that are not trainable are ignored).
#
# For serious analysis you most certainly want to run the sampler longer, with multiple chains and convergence checks. This will do for illustration though!
#

# %%
def plot_joint_marginals(samples, parameters, y_axis_label):
    name_to_index = {
        param_to_name[param]: i for i, param in enumerate(parameters)
    }
    f, axs = plt.subplots(1, 3, figsize=(12, 4), constrained_layout=True)

    axs[0].plot(
        samples[name_to_index[".likelihood.variance"]],
        samples[name_to_index[".kernel.variance"]],
        "k.",
        alpha=0.15,
    )
    axs[0].set_xlabel("noise_variance")
    axs[0].set_ylabel("signal_variance")

    axs[1].plot(
        samples[name_to_index[".likelihood.variance"]],
        samples[name_to_index[".kernel.lengthscales"]],
        "k.",
        alpha=0.15,
    )
    axs[1].set_xlabel("noise_variance")
    axs[1].set_ylabel("lengthscale")

    axs[2].plot(
        samples[name_to_index[".kernel.lengthscales"]],
        samples[name_to_index[".kernel.variance"]],
        "k.",
        alpha=0.1,
    )
    axs[2].set_xlabel("lengthscale")
    axs[2].set_ylabel("signal_variance")
    f.suptitle(y_axis_label)
    plt.show()


plot_joint_marginals(
    samples, model.trainable_parameters, "unconstrained variable samples"
)
plot_joint_marginals(
    parameter_samples, model.trainable_parameters, "parameter samples"
)


# %% [markdown]
# To plot the posterior of predictions, we'll iterate through the samples and set the model state with each sample. Then, for that state (set of hyperparameters) we'll draw some samples from the prediction function.

# %%
# plot the function posterior
xx = np.linspace(-0.1, 1.1, 100)[:, None]
plt.figure(figsize=(12, 6))

for i in range(0, num_samples, 20):
    for var, var_samples in zip(hmc_helper.current_state, samples):
        var.assign(var_samples[i])
    f = model.predict_f_samples(xx, 1)
    plt.plot(xx, f[0, :, :], "C0", lw=2, alpha=0.3)

plt.plot(X, Y, "kx", mew=2)
_ = plt.xlim(xx.min(), xx.max())
_ = plt.ylim(0, 6)
plt.xlabel("$x$")
plt.ylabel("$f|X,Y$")
plt.title("Posterior GP samples")

plt.show()

# %% [markdown]
# This example demonstrates the multi-modality of the true posterior over hyperparameters: the data could be explained both by long lengthscale, small signal variance, and high noise variance, as well as by shorter lengthscale, higher signal variance, and low noise variance.

# %% [markdown]
# ## Example 2: Sparse MC for multiclass classification

# %% [markdown]
# We now consider the multiclass classification problem (see the [Multiclass classification](../advanced/multiclass_classification.ipynb) notebook). Here the marginal likelihood is not available in closed form. Instead we use a sparse variational approximation where we approximate the posterior for each GP as $q(f_c) \propto p(f_c|\mathbf{u}_c)q(\mathbf{u}_c)$
#
# In the standard Sparse Variational GP (SVGP) formulation, $q(\mathbf{u_c})$ is parameterized as a multivariate Gaussian.
#
# An alternative is to directly sample from the optimal $q(\mathbf{u}_c)$; this is what Sparse Variational GP using MCMC (SGPMC) does.

# %% [markdown]
# We first build a multiclass classification dataset.

# %%
# Generate data by sampling from SquaredExponential kernel, and classifying with the argmax
rng = np.random.RandomState(42)

C, N = 3, 100
X = rng.rand(N, 1)
kernel = gpflow.kernels.SquaredExponential(lengthscales=0.1)
K = kernel.K(X) + np.eye(N) * 1e-6

f = rng.multivariate_normal(mean=np.zeros(N), cov=K, size=(C)).T
Y = np.argmax(f, 1)[:, None].astype(int)
# One-hot encoding
Y_hot = np.zeros((N, C), dtype=bool)
Y_hot[np.arange(N), Y.squeeze(-1)] = 1

data = (X, Y)

# %%
plt.figure(figsize=(12, 6))
order = np.argsort(X.flatten())

for c in range(C):
    plt.plot(X[order], f[order, c], ".", color=colors[c], label=str(c))
    plt.plot(X[order], Y_hot[order, c], "-", color=colors[c])

plt.legend()
plt.xlabel("$X$")
plt.ylabel("Latent (dots) and one-hot labels (lines)")
plt.title("Sample from the joint $p(Y, \mathbf{f})$")
plt.grid()
plt.show()

# %% [markdown]
# We then build the SGPMC model.

# %%
kernel = gpflow.kernels.Matern32(lengthscales=0.1) + gpflow.kernels.White(
    variance=0.01
)

model = gpflow.models.SGPMC(
    data,
    kernel=kernel,
    likelihood=gpflow.likelihoods.MultiClass(3),
    inducing_variable=X[::5].copy(),
    num_latent_gps=3,
)
model.kernel.kernels[0].variance.prior = tfd.Gamma(f64(1.0), f64(1.0))
model.kernel.kernels[0].lengthscales.prior = tfd.Gamma(f64(2.0), f64(2.0))
set_trainable(model.kernel.kernels[1].variance, False)

gpflow.utilities.print_summary(model)

# %%
# The inducing point locations Z should not be included in the MCMC (see [Hensman et al. (2015)](https://papers.nips.cc/paper/5875-mcmc-for-variationally-sparse-gaussian-processes), hence we set them to non-trainable.
set_trainable(model.inducing_variable, False)

# %% [markdown]
# The chain of samples for $\mathbf{u}_c, \theta$ is initialized at the value maximizing  $p(Y|\mathbf{u}_c, \theta)$.

# %%
optimizer = gpflow.optimizers.Scipy()
optimizer.minimize(
    model.training_loss, model.trainable_variables, options={"maxiter": 20}
)
print(f"log posterior density at optimum: {model.log_posterior_density()}")

# %% [markdown]
# Sampling starts with a 'burn in' period.

# %%
num_burnin_steps = reduce_in_tests(100)
num_samples = reduce_in_tests(500)

# Note that here we need model.trainable_parameters, not trainable_variables - only parameters can have priors!
hmc_helper = gpflow.optimizers.SamplingHelper(
    model.log_posterior_density, model.trainable_parameters
)

hmc = tfp.mcmc.HamiltonianMonteCarlo(
    target_log_prob_fn=hmc_helper.target_log_prob_fn,
    num_leapfrog_steps=10,
    step_size=0.01,
)

adaptive_hmc = tfp.mcmc.SimpleStepSizeAdaptation(
    hmc,
    num_adaptation_steps=10,
    target_accept_prob=f64(0.75),
    adaptation_rate=0.1,
)


@tf.function
def run_chain_fn():
    return tfp.mcmc.sample_chain(
        num_results=num_samples,
        num_burnin_steps=num_burnin_steps,
        current_state=hmc_helper.current_state,
        kernel=adaptive_hmc,
        trace_fn=lambda _, pkr: pkr.inner_results.is_accepted,
    )


samples, _ = run_chain_fn()
constrained_samples = hmc_helper.convert_to_constrained_values(samples)

# %% [markdown]
# Statistics of the posterior samples can now be reported.

# %%
plot_from_samples(
    model, X, Y, model.trainable_parameters, constrained_samples, thin=10
)

# %% [markdown]
# You can also display the sequence of sampled hyperparameters.

# %%
param_to_name = {
    param: name
    for name, param in gpflow.utilities.parameter_dict(model).items()
}
name_to_index = {
    param_to_name[param]: i
    for i, param in enumerate(model.trainable_parameters)
}
hyperparameters = [
    ".kernel.kernels[0].lengthscales",
    ".kernel.kernels[0].variance",
]

plt.figure(figsize=(8, 4))
for param_name in hyperparameters:
    plt.plot(constrained_samples[name_to_index[param_name]], label=param_name)
plt.legend(bbox_to_anchor=(1.0, 1.0))
plt.xlabel("HMC iteration")
_ = plt.ylabel("hyperparameter value")


# %% [markdown]
# ## Example 3: Fully Bayesian inference for generalized GP models with HMC

# %% [markdown]
# You can construct very flexible models with Gaussian processes by combining them with different likelihoods (sometimes called 'families' in the GLM literature). This makes inference of the GP intractable because the likelihoods are not generally conjugate to the Gaussian process. The general form of the model is
# \begin{align}
# \theta &\sim p(\theta) \\
# f &\sim \mathcal {GP}(m(x; \theta),\, k(x, x'; \theta)) \\
# y_i &\sim p(y | g(f(x_i))\,.
# \end{align}
#
#
# To perform inference in this model, we'll run MCMC using Hamiltonian Monte Carlo (HMC) over the function values and the parameters $\theta$ jointly. The key to an effective scheme is rotation of the field using the Cholesky decomposition. We write:
#
# \begin{align}
# \theta &\sim p(\theta) \\
# v &\sim \mathcal {N}(0,\, I) \\
# LL^\top &= K \\
# f &= m + Lv \\
# y_i &\sim p(y | g(f(x_i))\,.
# \end{align}
#
# Joint HMC over $v$ and the function values is not widely adopted in the literature because of the difficulty in differentiating $LL^\top=K$. We've made this derivative available in TensorFlow, and so application of HMC is relatively straightforward.

# %% [markdown]
# ### Exponential Regression
# We consider an exponential regression model:
# \begin{align}
# \theta &\sim p(\theta) \\
# f &\sim \mathcal {GP}(0, k(x, x'; \theta)) \\
# f_i &= f(x_i) \\
# y_i &\sim \mathcal {Exp} (e^{f_i})
# \end{align}
#
# We'll use MCMC to deal with both the kernel parameters $\theta$ and the latent function values $f$. Firstly, generate a data set.

# %%
rng = np.random.RandomState(14)

X = np.linspace(-3, 3, 20)
Y = rng.exponential(np.sin(X) ** 2)

plt.figure()
plt.plot(X, Y, "x")
plt.xlabel("input $X$")
plt.ylabel("output $Y$")
plt.title("toy dataset")
plt.show()

data = (X[:, None], Y[:, None])

# %% [markdown]
# GPflow's model for fully-Bayesian MCMC is called GPMC. It's constructed like any other model, but contains a parameter `V` which represents the centered values of the function.

# %%
kernel = gpflow.kernels.Matern32() + gpflow.kernels.Constant()
likelihood = gpflow.likelihoods.Exponential()
model = gpflow.models.GPMC(data, kernel, likelihood)

# %% [markdown]
# The `V` parameter already has a prior applied. We'll add priors to the parameters also (these are rather arbitrary, for illustration).

# %%
model.kernel.kernels[0].lengthscales.prior = tfd.Gamma(f64(1.0), f64(1.0))
model.kernel.kernels[0].variance.prior = tfd.Gamma(f64(1.0), f64(1.0))
model.kernel.kernels[1].variance.prior = tfd.Gamma(f64(1.0), f64(1.0))

gpflow.utilities.print_summary(model)


# %% [markdown]
# Running HMC is pretty similar to optimizing a model. GPflow builds on top of [tensorflow_probability's mcmc module](https://www.tensorflow.org/probability/api_docs/python/tfp/mcmc) and provides a SamplingHelper class to make interfacing easier.

# %% [markdown]
# We initialize HMC at the maximum a posteriori parameter values of the model.

# %%
optimizer = gpflow.optimizers.Scipy()
maxiter = reduce_in_tests(3000)
_ = optimizer.minimize(
    model.training_loss,
    model.trainable_variables,
    options=dict(maxiter=maxiter),
)
# We can now start HMC near maximum a posteriori (MAP)

# %% [markdown]
# We then run the sampler,

# %%
num_burnin_steps = reduce_in_tests(600)
num_samples = reduce_in_tests(1000)

# Note that here we need model.trainable_parameters, not trainable_variables - only parameters can have priors!
hmc_helper = gpflow.optimizers.SamplingHelper(
    model.log_posterior_density, model.trainable_parameters
)

hmc = tfp.mcmc.HamiltonianMonteCarlo(
    target_log_prob_fn=hmc_helper.target_log_prob_fn,
    num_leapfrog_steps=10,
    step_size=0.01,
)

adaptive_hmc = tfp.mcmc.SimpleStepSizeAdaptation(
    hmc,
    num_adaptation_steps=10,
    target_accept_prob=f64(0.75),
    adaptation_rate=0.1,
)


@tf.function
def run_chain_fn():
    return tfp.mcmc.sample_chain(
        num_results=num_samples,
        num_burnin_steps=num_burnin_steps,
        current_state=hmc_helper.current_state,
        kernel=adaptive_hmc,
        trace_fn=lambda _, pkr: pkr.inner_results.is_accepted,
    )


samples, _ = run_chain_fn()

# %% [markdown]
# And compute the posterior prediction on a grid for plotting purposes.

# %%
Xtest = np.linspace(-4, 4, 100)[:, None]
f_samples = []

for i in range(num_samples):
    # Note that hmc_helper.current_state contains the unconstrained variables
    for var, var_samples in zip(hmc_helper.current_state, samples):
        var.assign(var_samples[i])
    f = model.predict_f_samples(Xtest, 5)
    f_samples.append(f)
f_samples = np.vstack(f_samples)

# %%
rate_samples = np.exp(f_samples[:, :, 0])

(line,) = plt.plot(Xtest, np.mean(rate_samples, 0), lw=2)
plt.fill_between(
    Xtest[:, 0],
    np.percentile(rate_samples, 5, axis=0),
    np.percentile(rate_samples, 95, axis=0),
    color=line.get_color(),
    alpha=0.2,
)

plt.plot(X, Y, "kx", mew=2)
_ = plt.ylim(-0.1, np.max(np.percentile(rate_samples, 95, axis=0)))

# %% [markdown]
# You can also display the sequence of sampled hyperparameters.

# %%
parameter_samples = hmc_helper.convert_to_constrained_values(samples)
param_to_name = {
    param: name
    for name, param in gpflow.utilities.parameter_dict(model).items()
}
name_to_index = {
    param_to_name[param]: i
    for i, param in enumerate(model.trainable_parameters)
}
hyperparameters = [
    ".kernel.kernels[0].lengthscales",
    ".kernel.kernels[0].variance",
    ".kernel.kernels[1].variance",
]


plt.figure(figsize=(8, 4))
for param_name in hyperparameters:
    plt.plot(parameter_samples[name_to_index[param_name]], label=param_name)
plt.legend(bbox_to_anchor=(1.0, 1.0))
plt.xlabel("HMC iteration")
_ = plt.ylabel("hyperparameter value")


# %% [markdown]
# You can also inspect the marginal of the posterior samples.

# %%
fig, axes = plt.subplots(1, len(hyperparameters), sharex=True, figsize=(12, 4))
for ax, param_name in zip(axes, hyperparameters):
    ax.hist(parameter_samples[name_to_index[param_name]], bins=20)
    ax.set_title(param_name)
plt.tight_layout()

# %% [markdown]
# ## Prior on constrained and unconstrained parameters

# %% [markdown]
# GPflow's `Parameter` class provides options for setting a prior. `Parameter` wraps a constrained tensor and
# provides computation of the gradient with respect to unconstrained transformation of that tensor.
# The user can set a prior either in **constrained** space or **unconstrained** space.

# %% [markdown]
# By default, the prior for the `Parameter` is set on the _constrained_ space.
# To explicitly set the space on which the prior is defined, use the `prior_on` keyword argument:

# %%
prior_distribution = tfd.Normal(f64(0.0), f64(1.0))
_ = gpflow.Parameter(1.0, prior_on="unconstrained", prior=prior_distribution)
_ = gpflow.Parameter(1.0, prior_on="constrained", prior=prior_distribution)

# %% [markdown]
# `gpflow.optimizers.SamplingHelper` makes sure that the prior density correctly reflects the space in which the prior is defined.

# %% [markdown]
# Below we repeat the same experiment as before, but with some priors defined in the `unconstrained` space.
# We are using the exponential transform to ensure positivity of the kernel parameters (`set_default_positive_bijector("exp")`),
# so a log-normal prior on a constrained parameter corresponds to a normal prior on the unconstrained space:

# %%
gpflow.config.set_default_positive_bijector("exp")
gpflow.config.set_default_positive_minimum(1e-6)

rng = np.random.RandomState(42)
data = synthetic_data(30, rng)

kernel = gpflow.kernels.Matern52(lengthscales=0.3)
meanf = gpflow.mean_functions.Linear(1.0, 0.0)
model = gpflow.models.GPR(data, kernel, meanf)
model.likelihood.variance.assign(0.01)

mu = f64(0.0)
std = f64(4.0)
one = f64(1.0)

model.kernel.lengthscales.prior_on = "unconstrained"
model.kernel.lengthscales.prior = tfd.Normal(mu, std)
model.kernel.variance.prior_on = "unconstrained"
model.kernel.variance.prior = tfd.Normal(mu, std)
model.likelihood.variance.prior_on = "unconstrained"
model.likelihood.variance.prior = tfd.Normal(mu, std)

model.mean_function.A.prior_on = "constrained"
model.mean_function.A.prior = tfd.Normal(mu, std)
model.mean_function.b.prior_on = "constrained"
model.mean_function.b.prior = tfd.Normal(mu, std)

model.kernel.lengthscales.prior_on

# %% [markdown]
# Let's run HMC and plot chain traces:

# %%
num_burnin_steps = reduce_in_tests(300)
num_samples = reduce_in_tests(500)

hmc_helper = gpflow.optimizers.SamplingHelper(
    model.log_posterior_density, model.trainable_parameters
)

hmc = tfp.mcmc.HamiltonianMonteCarlo(
    target_log_prob_fn=hmc_helper.target_log_prob_fn,
    num_leapfrog_steps=10,
    step_size=0.01,
)
adaptive_hmc = tfp.mcmc.SimpleStepSizeAdaptation(
    hmc,
    num_adaptation_steps=10,
    target_accept_prob=f64(0.75),
    adaptation_rate=0.1,
)


@tf.function
def run_chain_fn_unconstrained():
    return tfp.mcmc.sample_chain(
        num_results=num_samples,
        num_burnin_steps=num_burnin_steps,
        current_state=hmc_helper.current_state,
        kernel=adaptive_hmc,
        trace_fn=lambda _, pkr: pkr.inner_results.is_accepted,
    )


samples, traces = run_chain_fn_unconstrained()
parameter_samples = hmc_helper.convert_to_constrained_values(samples)

param_to_name = {
    param: name
    for name, param in gpflow.utilities.parameter_dict(model).items()
}
marginal_samples(
    samples, model.trainable_parameters, "unconstrained variable samples"
)
marginal_samples(
    parameter_samples,
    model.trainable_parameters,
    "constrained parameter samples",
)


# doc/sphinx/notebooks/advanced/multiclass_classification.pct.py

# ---
# jupyter:
#   jupytext:
#     formats: ipynb,.pct.py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.3.3
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Multiclass classification

# %% [markdown]
# The multiclass classification problem is a regression problem from an input $x \in {\cal X}$ to discrete labels $y\in {\cal Y}$, where ${\cal Y}$ is a discrete set of size $C$ bigger than two (for $C=2$ it is the more usual binary classification).
#
# Labels are encoded in a one-hot fashion, that is if $C=4$ and $y=2$, we note $\bar{y} = [0,1,0,0]$.
#
# The generative model for this problem consists of:
#
# * $C$ latent functions $\mathbf{f} = [f_1,...,f_C]$ with an independent Gaussian Process prior
# * a deterministic function that builds a discrete distribution $\pi(\mathbf{f}) = [\pi_1(f_1),...,\pi_C(f_C)]$ from the latents such that $\sum_c \pi_c(f_c) = 1$
# * a discrete likelihood $p(y|\mathbf{f}) = Discrete(y;\pi(\mathbf{f})) = \prod_c \pi_c(f_c)^{\bar{y}_c}$
#
# A typical example of $\pi$ is the softmax function:
#
# \begin{equation}
# \pi_c (f_c) \propto \exp( f_c)
# \end{equation}
#
# Another convenient one is the robust max:
# \begin{equation}
# \pi_c(\mathbf{f}) = \begin{cases} 1 - \epsilon, & \mbox{if } c = \arg \max_c f_c \\
# \epsilon /(C-1), & \mbox{ otherwise} \end{cases}
# \end{equation}
#
#
#
#

# %%
import warnings

warnings.filterwarnings("ignore")  # ignore DeprecationWarnings from tensorflow

import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
from multiclass_classification import colors, plot_posterior_predictions

import gpflow
from gpflow.ci_utils import reduce_in_tests
from gpflow.utilities import print_summary, set_trainable

# %matplotlib inline


# reproducibility:
np.random.seed(0)
tf.random.set_seed(123)

# %% [markdown]
# ## Sampling from the GP multiclass generative model

# %% [markdown]
# ### Declaring model parameters and input

# %%
# Number of functions and number of data points
C = 3
N = 100

# Lengthscale of the SquaredExponential kernel (isotropic -- change to `[0.1] * C` for ARD)
lengthscales = 0.1

# Jitter
jitter_eye = np.eye(N) * 1e-6

# Input
X = np.random.rand(N, 1)

# %% [markdown]
# ### Sampling

# %%
# SquaredExponential kernel matrix
kernel_se = gpflow.kernels.SquaredExponential(lengthscales=lengthscales)
K = kernel_se(X) + jitter_eye

# Latents prior sample
f = np.random.multivariate_normal(mean=np.zeros(N), cov=K, size=(C)).T

# Hard max observation
Y = np.argmax(f, 1).flatten().astype(int)

# One-hot encoding
Y_hot = np.zeros((N, C), dtype=bool)
Y_hot[np.arange(N), Y] = 1

data = (X, Y)

# %% [markdown]
# ### Plotting

# %%
plt.figure(figsize=(12, 6))
order = np.argsort(X.flatten())

for c in range(C):
    plt.plot(X[order], f[order, c], ".", color=colors[c], label=str(c))
    plt.plot(X[order], Y_hot[order, c], "-", color=colors[c])


plt.legend()
plt.xlabel("$X$")
plt.ylabel("Latent (dots) and one-hot labels (lines)")
plt.title("Sample from the joint $p(Y, \mathbf{f})$")
plt.grid()
plt.show()

# %% [markdown]
# ## Inference
#

# %% [markdown]
# Inference here consists of computing the posterior distribution over the latent functions given the data $p(\mathbf{f}|Y, X)$.
#
# You can use different inference methods. Here we perform variational inference.
# For a treatment of the multiclass classification problem using MCMC sampling, see [Markov Chain Monte Carlo (MCMC)](../advanced/mcmc.ipynb).
#
#

# %% [markdown]
# ### Approximate inference: Sparse Variational Gaussian Process

# %% [markdown]
# #### Declaring the SVGP model (see [GPs for big data](../advanced/gps_for_big_data.ipynb))

# %%
# sum kernel: Matern32 + White
kernel = gpflow.kernels.Matern32() + gpflow.kernels.White(variance=0.01)

# Robustmax Multiclass Likelihood
invlink = gpflow.likelihoods.RobustMax(C)  # Robustmax inverse link function
likelihood = gpflow.likelihoods.MultiClass(
    3, invlink=invlink
)  # Multiclass likelihood
Z = X[::5].copy()  # inducing inputs

m = gpflow.models.SVGP(
    kernel=kernel,
    likelihood=likelihood,
    inducing_variable=Z,
    num_latent_gps=C,
    whiten=True,
    q_diag=True,
)

# Only train the variational parameters
set_trainable(m.kernel.kernels[1].variance, False)
set_trainable(m.inducing_variable, False)
print_summary(m, fmt="notebook")

# %% [markdown]
# #### Running inference

# %%
opt = gpflow.optimizers.Scipy()

opt_logs = opt.minimize(
    m.training_loss_closure(data),
    m.trainable_variables,
    options=dict(maxiter=reduce_in_tests(1000)),
)
print_summary(m, fmt="notebook")

# %%
plot_posterior_predictions(m, X, Y)


# doc/sphinx/notebooks/advanced/multiclass_classification.py

# Function to plot the predictions of the trained
# multi-class classification model in multiclass_classification.ipynb

import matplotlib.pyplot as plt
import numpy as np

colors = ["#1f77b4", "#ff7f0e", "#2ca02c"]


def plot_posterior_predictions(m, X, Y):

    f = plt.figure(figsize=(12, 6))
    a1 = f.add_axes([0.05, 0.05, 0.9, 0.6])
    a2 = f.add_axes([0.05, 0.7, 0.9, 0.1])
    a3 = f.add_axes([0.05, 0.85, 0.9, 0.1])

    xx = np.linspace(X.min(), X.max(), 200).reshape(-1, 1)
    mu, var = m.predict_f(xx)
    p, _ = m.predict_y(xx)

    a3.set_xticks([])
    a3.set_yticks([])

    for c in range(m.likelihood.num_classes):
        x = X[Y.flatten() == c]

        color = colors[c]
        a3.plot(x, x * 0, ".", color=color)
        a1.plot(xx, mu[:, c], color=color, lw=2, label="%d" % c)
        a1.plot(xx, mu[:, c] + 2 * np.sqrt(var[:, c]), "--", color=color)
        a1.plot(xx, mu[:, c] - 2 * np.sqrt(var[:, c]), "--", color=color)
        a2.plot(xx, p[:, c], "-", color=color, lw=2)

    a2.set_ylim(-0.1, 1.1)
    a2.set_yticks([0, 1])
    a2.set_xticks([])

    a3.set_title("inputs X")
    a2.set_title(
        "predicted mean label value \
                 $\mathbb{E}_{q(\mathbf{u})}[y^*|x^*, Z, \mathbf{u}]$"
    )
    a1.set_title(
        "posterior process \
                $\int d\mathbf{u} q(\mathbf{u})p(f^*|\mathbf{u}, Z, x^*)$"
    )

    handles, labels = a1.get_legend_handles_labels()
    a1.legend(handles, labels)
    f.tight_layout()
    plt.show()


def plot_from_samples(m, X, Y, parameters, parameter_samples, thin):

    f = plt.figure(figsize=(12, 6))
    a1 = f.add_axes([0.05, 0.05, 0.9, 0.6])
    a2 = f.add_axes([0.05, 0.7, 0.9, 0.1])
    a3 = f.add_axes([0.05, 0.85, 0.9, 0.1])

    xx = np.linspace(X.min(), X.max(), 200).reshape(-1, 1)

    Fpred, Ypred = [], []
    num_samples = len(parameter_samples[0])
    for i in range(0, num_samples, thin):
        for parameter, var_samples in zip(parameters, parameter_samples):
            parameter.assign(var_samples[i])
        Ypred.append(m.predict_y(xx)[0])
        Fpred.append(np.squeeze(m.predict_f_samples(xx, 1)))

    for i in range(m.likelihood.num_classes):
        x = X[Y.reshape(-1) == i]
        (points,) = a3.plot(x, x * 0, ".")
        color = points.get_color()
        for F in Fpred:
            a1.plot(xx, F[:, i], color=color, lw=0.2, alpha=1.0)
        for Yp in Ypred:
            a2.plot(xx, Yp[:, i], color=color, lw=0.5, alpha=1.0)

    a2.set_ylim(-0.1, 1.1)
    a2.set_yticks([0, 1])
    a2.set_xticks([])

    a3.set_xticks([])
    a3.set_yticks([])

    a3.set_title("inputs X")
    a2.set_title(
        "predicted mean label value \
                 $\mathbb{E}_{q(\mathbf{u})}[y^*|x^*, Z, \mathbf{u}]$"
    )
    a1.set_title(
        "posterior process samples \
                $\int d\mathbf{u} q(\mathbf{u})p(f^*|\mathbf{u}, Z, x^*)$"
    )


# doc/sphinx/notebooks/advanced/multioutput.pct.py

# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,.pct.py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.4.1
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Multi-output Gaussian processes in GPflow

# %% [markdown]
# This notebook shows how to construct a multi-output GP model using GPflow, together with different interdomain inducing variables which lead to different approximation properties. GPflow provides a framework for specifying multioutput GP priors, and interdomain approximations which is
# - modular, by providing a consistent interface for the user of the resulting `SVGP` model,
# - extensible, by allowing new interdomain variables and kernels to be specified while reusing exising code where possible,
# - efficient, by allowing the most efficient custom code path to be specified where desired.
#
# Getting to grips with the maths and code can be a bit daunting, so to accompany the documentation there is an [in-depth review on arXiv](https://arxiv.org/abs/2003.01115), which provides a unified mathematical framework, together with a high-level description of software design choices in GPflow.
#
# This notebook shows the various design choices that can be made, to show the reader the flexibility of the framework. This is done in the hope that an example is provided that can be easily adapted to the special case that the reader wants to implement.
#
# A reader who just wants to use a multioutput kernel should simply choose the most efficient set of inducing variables.
#
# To cite this framework, please reference our [arXiv paper](https://arxiv.org/abs/2003.01115).
# ```
# @article{GPflow2020multioutput,
#   author = {{van der Wilk}, Mark and Dutordoir, Vincent and John, ST and
#             Artemev, Artem and Adam, Vincent and Hensman, James},
#   title = {A Framework for Interdomain and Multioutput {G}aussian Processes},
#   year = {2020},
#   journal = {arXiv:2003.01115},
#   url = {https://arxiv.org/abs/2003.01115}
# }
# ```
#
# \begin{equation}
# \newcommand{\GP}{\mathcal{GP}}
# \newcommand{\NN}{\mathcal{N}}
# \newcommand{\LL}{\mathcal{L}}
# \newcommand{\RR}{\mathbb{R}}
# \newcommand{\EE}{\mathbb{E}}
# \newcommand{\valpha}{\boldsymbol\alpha}
# \newcommand{\vf}{\mathbf{f}}
# \newcommand{\vF}{\mathbf{F}}
# \newcommand{\vg}{\mathbf{g}}
# \newcommand{\vW}{\mathbf{W}}
# \newcommand{\vI}{\mathbf{I}}
# \newcommand{\vZ}{\mathbf{Z}}
# \newcommand{\vu}{\mathbf{u}}
# \newcommand{\vU}{\mathbf{U}}
# \newcommand{\vX}{\mathbf{X}}
# \newcommand{\vY}{\mathbf{Y}}
# \newcommand{\identity}{\mathbb{I}}
# \end{equation}
#
#
#
# ## Task
# We will consider a regression problem for functions $f: \mathbb{R}^D \rightarrow \mathbb{R}^P$. We assume that the dataset is of the form $(X, f_1), \dots, (X, f_P)$, that is, we observe all the outputs for a particular input location (for cases where there are **not** fully observed outputs for each input, see [A simple demonstration of coregionalization](./coregionalisation.ipynb)).
#
# Here we assume a model of the form:
# \begin{equation}
# f(x) = W g(x),
# \end{equation}
# where $g(x) \in \mathbb{R}^L$, $f(x) \in \mathbb{R}^P$ and $W \in \mathbb{R}^{P \times L}$. We assume that the outputs of $g$ are uncorrelated, and that by *mixing* them with $W$ they become correlated. In this notebook, we show how to build this model using Sparse Variational Gaussian Process (SVGP) for $g$, which scales well with the numbers of data points and outputs.
#
# Here we have two options for $g$:
# 1. The output dimensions of $g$ share the same kernel.
# 1. Each output of $g$ has a separate kernel.
#
#
# In addition, we have two further suboptions for the inducing inputs of $g$:
# 1. The instances of $g$ share the same inducing inputs.
# 1. Each output of $g$ has its own set of inducing inputs.
#
# The notation is as follows:
# - $X \in \mathbb{R}^{N \times D}$ denotes the input
# - $Y \in \RR^{N \times P}$ denotes the output
# - $k_{1..L}$, $L$ are kernels on $\RR^{N \times D}$
# - $g_{1..L}$, $L$ are independent $\GP$s  with $g_l \sim \GP(0,k_l)$
# - $f_{1..P}$, $P$ are correlated  $\GP$s  with $\vf = \vW \vg$

# %%
import matplotlib.pyplot as plt
import numpy as np

import gpflow as gpf
from gpflow.ci_utils import reduce_in_tests
from gpflow.utilities import print_summary

gpf.config.set_default_float(np.float64)
gpf.config.set_default_summary_fmt("notebook")
np.random.seed(0)
# %matplotlib inline

MAXITER = reduce_in_tests(2000)

# %% [markdown]
# ## Generate synthetic data
# We create a utility function to generate synthetic data. We assume that:

# %%
N = 100  # number of points
D = 1  # number of input dimensions
M = 15  # number of inducing points
L = 2  # number of latent GPs
P = 3  # number of observations = output dimensions


# %%
def generate_data(N=100):
    X = np.random.rand(N)[:, None] * 10 - 5  # Inputs = N x D
    G = np.hstack((0.5 * np.sin(3 * X) + X, 3.0 * np.cos(X) - X))  # G = N x L
    W = np.array([[0.5, -0.3, 1.5], [-0.4, 0.43, 0.0]])  # L x P
    F = np.matmul(G, W)  # N x P
    Y = F + np.random.randn(*F.shape) * [0.2, 0.2, 0.2]

    return X, Y


# %%
X, Y = data = generate_data(N)
Zinit = np.linspace(-5, 5, M)[:, None]


# %% [markdown]
# We create a utility function for plotting:

# %%
def plot_model(m, lower=-8.0, upper=8.0):
    pX = np.linspace(lower, upper, 100)[:, None]
    pY, pYv = m.predict_y(pX)
    if pY.ndim == 3:
        pY = pY[:, 0, :]
    plt.plot(X, Y, "x")
    plt.gca().set_prop_cycle(None)
    plt.plot(pX, pY)
    for i in range(pY.shape[1]):
        top = pY[:, i] + 2.0 * pYv[:, i] ** 0.5
        bot = pY[:, i] - 2.0 * pYv[:, i] ** 0.5
        plt.fill_between(pX[:, 0], top, bot, alpha=0.3)
    plt.xlabel("X")
    plt.ylabel("f")
    plt.title(f"ELBO: {m.elbo(data):.3}")
    plt.plot(Z, Z * 0.0, "o")


# %% [markdown]
# ## Model the outputs of $f(x)$ directly
# The three following examples show how to model the outputs of the model $f(x)$ directly. Mathematically, this case is equivalent to having:
# \begin{equation}
# f(x) = I g(x),
# \end{equation}
# i.e. $W = I$ and $P = L$.

# %% [markdown]
# ### 1. Shared independent multi-output kernel (MOK) and shared independent inducing variables
# Here the priors on all outputs are constrained to have the same kernel hyperparameters. We also share the inducing inputs between all outputs. The different GPs are independent both in the prior and the approximate posterior.

# %%
# create multi-output kernel
kernel = gpf.kernels.SharedIndependent(
    gpf.kernels.SquaredExponential() + gpf.kernels.Linear(), output_dim=P
)
# initialization of inducing input locations (M random points from the training inputs)
Z = Zinit.copy()
# create multi-output inducing variables from Z
iv = gpf.inducing_variables.SharedIndependentInducingVariables(
    gpf.inducing_variables.InducingPoints(Z)
)

# %%
# create SVGP model as usual and optimize
m = gpf.models.SVGP(
    kernel, gpf.likelihoods.Gaussian(), inducing_variable=iv, num_latent_gps=P
)
print_summary(m)


# %%
def optimize_model_with_scipy(model):
    optimizer = gpf.optimizers.Scipy()
    optimizer.minimize(
        model.training_loss_closure(data),
        variables=model.trainable_variables,
        method="l-bfgs-b",
        options={"disp": 50, "maxiter": MAXITER},
    )


optimize_model_with_scipy(m)

# %%
print_summary(m)

# %%
# Plot predictions and observations
plot_model(m)

# %%
print_summary(m.kernel)
m.kernel.kernel.kernels[0].lengthscales

# %% [markdown]
# ### 2. Separate independent MOK and shared independent inducing variables
# Here we allow different hyperparameters for the priors of each output. We still share the inducing inputs between all outputs.

# %%
# Create list of kernels for each output
kern_list = [
    gpf.kernels.SquaredExponential() + gpf.kernels.Linear() for _ in range(P)
]
# Create multi-output kernel from kernel list
kernel = gpf.kernels.SeparateIndependent(kern_list)
# initialization of inducing input locations (M random points from the training inputs)
Z = Zinit.copy()
# create multi-output inducing variables from Z
iv = gpf.inducing_variables.SharedIndependentInducingVariables(
    gpf.inducing_variables.InducingPoints(Z)
)

# %%
# create SVGP model as usual and optimize
m = gpf.models.SVGP(
    kernel, gpf.likelihoods.Gaussian(), inducing_variable=iv, num_latent_gps=P
)

# %%
optimize_model_with_scipy(m)

# %%
print_summary(m.kernel)

# %%
plot_model(m)

# %%
[k.kernels[0].lengthscales for k in m.kernel.kernels]

# %% [markdown]
# ### 3. Separate independent kernel and separate independent inducing variables
# Here we allow different hyperparameters for the priors of each output. We now allow different inducing inputs for each output.

# %%
# Create list of kernels for each output
kern_list = [
    gpf.kernels.SquaredExponential() + gpf.kernels.Linear() for _ in range(P)
]
# Create multi-output kernel from kernel list
kernel = gpf.kernels.SeparateIndependent(kern_list)
# initialization of inducing input locations, one set of locations per output
Zs = [Zinit.copy() for _ in range(P)]
# initialize as list inducing inducing variables
iv_list = [gpf.inducing_variables.InducingPoints(Z) for Z in Zs]
# create multi-output inducing variables from iv_list
iv = gpf.inducing_variables.SeparateIndependentInducingVariables(iv_list)

# %% [markdown]
# **NOTE:** While the inducing points are independent, there needs to be the same number of inducing points per dimension.

# %%
# create SVGP model as usual and optimize
m = gpf.models.SVGP(
    kernel, gpf.likelihoods.Gaussian(), inducing_variable=iv, num_latent_gps=P
)

# %%
optimize_model_with_scipy(m)

# %%
plot_model(m)

# %% [markdown]
# The following plot shows that we use different inducing *inputs* in each output dimension.

# %%
for i in range(len(m.inducing_variable.inducing_variable_list)):
    q_mu_unwhitened, q_var_unwhitened = m.predict_f(
        m.inducing_variable.inducing_variable_list[i].Z
    )
    plt.plot(
        m.inducing_variable.inducing_variable_list[i].Z.numpy(),
        q_mu_unwhitened[:, i, None].numpy(),
        "o",
    )
plt.gca().set_xticks(np.linspace(-6, 6, 20), minor=True)
plt.gca().set_yticks(np.linspace(-9, 9, 20), minor=True)
plt.grid(which="minor")

# %%
m.inducing_variable.inducing_variable_list

# %% [markdown]
# ## Model $f(x)$ by doing inference in the $g$ space
# ### Mixed kernel and uncorrelated inducing variables
#
# Remember the general case: $f(x) = W g(x)$, where $g(x) \in \mathbb{R}^L$, $f(x) \in \mathbb{R}^P$ and $W \in \mathbb{R}^{P \times L}$, where $L \leq P$.
# We assume that the outputs of $g$ are uncorrelated, and by *mixing* them with $W$ they become correlated.
# With this setup we perform the optimal routine to calculate the conditional. Namely, calculate the conditional of the uncorrelated latent $g$ and afterwards project the mean and variance using the mixing matrix: $\mu_f = W \mu_g$ and $\Sigma_f = W\Sigma_g W^\top$
#
# - $K_{uu} = L \times M \times M$
# - $K_{uf} = L \times M \times N$

# %%
# Create list of kernels for each output
kern_list = [
    gpf.kernels.SquaredExponential() + gpf.kernels.Linear() for _ in range(L)
]
# Create multi-output kernel from kernel list
kernel = gpf.kernels.LinearCoregionalization(
    kern_list, W=np.random.randn(P, L)
)  # Notice that we initialise the mixing matrix W
# initialisation of inducing input locations (M random points from the training inputs)
Z = Zinit.copy()
# create multi-output inducing variables from Z
iv = gpf.inducing_variables.SharedIndependentInducingVariables(
    gpf.inducing_variables.InducingPoints(Z)
)

# %%
# initialize mean of variational posterior to be of shape MxL
q_mu = np.zeros((M, L))
# initialize \sqrt(Σ) of variational posterior to be of shape LxMxM
q_sqrt = np.repeat(np.eye(M)[None, ...], L, axis=0) * 1.0

# create SVGP model as usual and optimize
m = gpf.models.SVGP(
    kernel,
    gpf.likelihoods.Gaussian(),
    inducing_variable=iv,
    q_mu=q_mu,
    q_sqrt=q_sqrt,
)

# %%
optimize_model_with_scipy(m)

# %%
plot_model(m)


# %% [markdown]
# ## Illustration of GPflow's multi-output capabilities
# This section shows the inheritance structure in GPflow's multi-output framework.

# %% [markdown]
# ### Multi-output kernels (MOK) class diagram
# We include three multi-output kernels:
# - `SharedIndependent`: This kernel is included mainly as an illustration of specifying a conditional using the multiple dispatch framework. The same functionality is provided by using a normal kernel and passing in multiple approximate posteriors by stacking `q_mu` and `q_sqrt`.
# - `SeparateIndependent`: This kernel allows you to use different priors for each output GP.
# - `LinearCoregionalization`: This kernel describes the prior of the linear model of coregionalization. As shown previously, this implementation supports various inducing point approximations.
# ![Multi-output kernels](./multioutput_kernels.svg)
#
# We include several base classes. Two are noteworthy:
# - `MultioutputKernel` is included to be the base class for all multi-output kernels.
# - `IndepedentLatent` is the base class for all multi-output kernels which are constructed from independent latent processes. Including this kernel allows the specification of a default approximation method which, while not the most efficient, does take advantage of _some_ structure. It can be applied to _any_ kernel constructed from independent latent processes.
#
# There is a similarity in the meaning of `SeparateIndependent` and `IndependentLatent`. Both kernels indicate that independent processes are used, and that $\mathbf{K}_{\bf uu}$ can therefore be represented as a `[L, M, M]` tensor. It could therefore be suggested that `SeparateIndependent` be the parent class of all "independent latent" kernels, instead of having a separate `IndependentLatent` class. We decided against this because:
# - this would increase the complexity in specifying `conditionals()` for the otherwise simple multi-output kernels `SeparateIndependent` and `SharedIndependent`.
# - we did not want to specify too much of an implementation in `IndependentLatent`, leaving implementation details to child classes. Using `SeparateIndependent` as the base class would force all child classes to be a `Combination` kernel.

# %% [markdown]
# ### Multi-output inducing variables class diagram
# ![Multi-output features](./multioutput_features.svg)
#
# #### Inducing points
# The goal of this class is to provide inducing variables that can be used with _any_ kernel, even if the method ends up being slow.
#
# The multiouput framework extends `InducingPoints` to work with multi-output kernels. Just like for single-output kernels, we want `InducingPoints` to work for all `MultioutputKernel`s. We do this by defining `InducingPoints` to take _all_ outputs for specific inducing inputs as inducing variables.
#
# #### Fallback shared/separate independent inducing variables
# The goal of these classes is to provide a reasonably efficient implementation for kernels that give exploitable independence structure in the prior of inducing variables (that is, subclasses of `IndependentLatent`), while only needing to implement `Kuu()` and `Kuf()` methods.
#
# #### Shared/separate independent inducing variables
# The goal of these classes is to provide the most efficient code path for kernels that allow exploiting independence structure in the prior of inducing variables.
#
# For more specialized multi-output kernels (i.e. `{Shared|Separate}Independent` or `LinearCoregionalization`) we define `{Shared|Separate}IndependentInducingVariables`. These wrap (a list of) single-output inducing variables to define groups of a-priori independent inducing variables, which leads to a $\mathbf{K}_{\bf uu}$ that can be represented as a `[L, M, M]` tensor. We saw the use of these previously.
#
# `{Shared|Separate}IndependentInducingVariables` inherit from `Fallback{Shared|Separate}IndependentInducingVariables`, so the multiple dispatch will fall back on the slower but general implementation.

# %% [markdown]
# ### Implemented combinations
# Multiple dispatch is applied to both `Kuu()`, `Kuf()`, and `conditional()`. The return values of the covariances can therefore be tailored to a specific implementation of `conditional()`. The following table lists combinations which are currently available in GPflow. Thanks to the multiple dispatch code, implementing your own outside of GPflow should require only a small amount of code!
#
# |  Inducing variable class  |  Kernel  |  Kuu  |  Kuf  |  conditional  |  note  |
# |  -----------------------  |  ------  |  ---  |  ---  |  -----------  |  ----  |
# |  `InducingPoints`                                |  `MultioutputKernel`        |  `[M, P, M, P]`  |  `[M, P, N, P]`  |  `inducing_point_conditional()`, which calls `fully_correlated_conditional()`  |  Works for all kernels, but might be very inefficient. In this case `q_mu` and `q_sqrt` should have shapes of `[1, MP]` and `[1, MP, MP]`  |
# |  `SharedIndependentInducingVariables`            |  `SharedIndependent`        |  `[M, M]`        |  `[M, N]`        |  `shared_independent_conditional()`, which calls `base_conditional()`  |  The combination of these two classes is in a sense redundant, because we can achieve the same behavior by using the single output Kernel and InducingVariable classes. They are added for illustrative purposes. Thanks to the conditional dispatch, the most efficient code path is used.  |
# |  `SeparateIndependentInducingVariables`          |  `SharedIndependent`        |  `[P, M, M]`     |  `[P, M, N]`     |  `separate_independent_conditional()`, which calls `base_conditional()` P times  |  We loop P times over the `base_conditional()`  |
# |  `SeparateIndependentInducingVariable`           |  `SeparateIndependent`      |  `[P, M, M]`     |  `[P, M, N]`     |   `separate_independent_conditional()`, which calls `base_conditional()` P times  |  We loop P times over the `base_conditional()`  |
# |  `SharedIndependentInducingVariables`            |  `SeparateIndependent`      |  `[P, M, M]`     |  `[P, M, N]`     |  `separate_independent_conditional()`, which calls `base_conditional()` P times  |  We loop P times over the `base_conditional()`  |
# |  `FallbackSharedIndependentInducingVariables`    |  `IndependentLatent`        |  `[L, M, M]`     |  `[M, L, N, P]`  |  `fallback_independent_latent_conditional()`, which calls `independent_interdomain_conditional()`  |  Implementation which only requires custom `Kuu()` and `Kuf()`  |
# |  `FallbackSeparateIndependentInducingVariable`   |  `IndependentLatent`        |  `[L, M, M]`     |  `[M, L, N, P]`  |  `fallback_independent_latent_conditional()`, which calls `independent_interdomain_conditional()`  |  Implementation which only requires custom `Kuu()` and `Kuf()`  |
# |  `SharedIndependentInducingVariables`            |  `LinearCoregionalization`  |  `[L, M, M]`     |  `[L, M, N]`     |  `coregionalization_conditional()`, which calls `base_conditional()`  |  This is the most efficient implementation for linear coregionalization. The inducing outputs live in g-space. Here we use the output of the base conditional and project the mean and covariance with the mixing matrix W.  |
# |  `SeparateIndependentInducingVariables`          |  `LinearCoregionalization`  |  `[L, M, M]`      |  `[L, M, N]`      |  `base_conditional()`  |  This is the most efficient implementation for linear coregionalization. The inducing outputs live in g-space. Here we use the output of the base conditional and project the mean and covariance with the mixing matrix W.  |

# %% [markdown]
# ## Debugging: introspect
#
# Given all these possibilities it can be hard to determine which conditional will be called for which set of kernel and inducing variable. The following method lets you proactively introspect which implementation will be executed. This can be useful when debugging new code.

# %%
def inspect_conditional(inducing_variable_type, kernel_type):
    """
    Helper function returning the exact implementation called
    by the multiple dispatch `conditional` given the type of
    kernel and inducing variable.

    :param inducing_variable_type:
        Type of the inducing variable
    :param kernel_type:
        Type of the kernel

    :return: String
        Contains the name, the file and the linenumber of the
        implementation.
    """
    import inspect

    from gpflow.conditionals import conditional

    implementation = conditional.dispatch(
        object, inducing_variable_type, kernel_type, object
    )
    info = dict(inspect.getmembers(implementation))
    return info["__code__"]


# Example:
inspect_conditional(
    gpf.inducing_variables.SharedIndependentInducingVariables,
    gpf.kernels.SharedIndependent,
)

# %% [markdown]
# ## Further Reading:
# - [A simple demonstration of coregionalization](./coregionalisation.ipynb), which details other GPflow features for multi-output prediction without fully observed outputs.


# doc/sphinx/notebooks/advanced/natural_gradients.pct.py

# ---
# jupyter:
#   jupytext:
#     formats: ipynb,.pct.py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.3.3
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Natural gradients
#
# This notebook shows some basic usage of the natural gradient optimizer, both on its own and in combination with Adam optimizer.

# %%
import numpy as np
import tensorflow as tf

import gpflow
from gpflow import set_trainable
from gpflow.ci_utils import reduce_in_tests
from gpflow.models import GPR, SGPR, SVGP, VGP
from gpflow.optimizers import NaturalGradient
from gpflow.optimizers.natgrad import XiSqrtMeanVar

# %matplotlib inline
# %precision 4

np.random.seed(0)
tf.random.set_seed(0)

N, D = 100, 2
batch_size = 50

# inducing points
M = 10

x = np.random.uniform(size=(N, D))
y = np.sin(10 * x[:, :1]) + 5 * x[:, 1:] ** 2

data = (x, y)
inducing_variable = tf.random.uniform((M, D))
adam_learning_rate = 0.01
iterations = reduce_in_tests(5)
autotune = tf.data.experimental.AUTOTUNE

# %% [markdown]
# ## VGP is a GPR

# %% [markdown]
# The following section demonstrates how natural gradients can turn VGP into GPR *in a single step, if the likelihood is Gaussian*.

# %% [markdown]
# Let's start by first creating a standard GPR model with Gaussian likelihood:

# %%
gpr = GPR(data, kernel=gpflow.kernels.Matern52())

# %% [markdown]
# The log marginal likelihood of the exact GP model is:

# %%
gpr.log_marginal_likelihood().numpy()

# %% [markdown]
# Now we will create an approximate model which approximates the true posterior via a variational Gaussian distribution.<br>We initialize the distribution to be zero mean and unit variance.

# %%
vgp = VGP(
    data,
    kernel=gpflow.kernels.Matern52(),
    likelihood=gpflow.likelihoods.Gaussian(),
)
# (Note that GPflow's NaturalGradient optimizer does not implement diagonal covariance parametrization, i.e., it does not work for `q_diag=True`.)

# %% [markdown]
# The log marginal likelihood lower bound (evidence lower bound or ELBO) of the approximate GP model is:

# %%
vgp.elbo().numpy()

# %% [markdown]
# Obviously, our initial guess for the variational distribution is not correct, which results in a lower bound to the likelihood of the exact GPR model. We can optimize the variational parameters in order to get a tighter bound.

# %% [markdown]
# In fact, we only need to take **one step** in the natural gradient direction to recover the exact posterior:

# %%
natgrad_opt = NaturalGradient(gamma=1.0)
variational_params = [(vgp.q_mu, vgp.q_sqrt)]
natgrad_opt.minimize(vgp.training_loss, var_list=variational_params)

# %% [markdown]
# The ELBO of the approximate GP model after a single NatGrad step:

# %%
vgp.elbo().numpy()

# %% [markdown]
# ## Optimize both variational parameters and kernel hyperparameters together
#
# In the Gaussian likelihood case we can iterate between an Adam update for the hyperparameters and a NatGrad update for the variational parameters. That way, we achieve optimization of hyperparameters as if the model were a GPR.

# %% [markdown]
# The trick is to forbid Adam from updating the variational parameters by setting them to not trainable.

# %%
# Stop Adam from optimizing the variational parameters
set_trainable(vgp.q_mu, False)
set_trainable(vgp.q_sqrt, False)

adam_opt_for_vgp = tf.optimizers.Adam(adam_learning_rate)
adam_opt_for_gpr = tf.optimizers.Adam(adam_learning_rate)

# %%
for i in range(iterations):
    adam_opt_for_gpr.minimize(
        gpr.training_loss, var_list=gpr.trainable_variables
    )
    likelihood = gpr.log_marginal_likelihood()
    tf.print(f"GPR with Adam: iteration {i + 1} likelihood {likelihood:.04f}")

# %%
for i in range(iterations):
    natgrad_opt.minimize(vgp.training_loss, var_list=variational_params)
    adam_opt_for_vgp.minimize(
        vgp.training_loss, var_list=vgp.trainable_variables
    )
    likelihood = vgp.elbo()
    tf.print(
        f"VGP with NaturalGradient and Adam: iteration {i + 1} likelihood {likelihood:.04f}"
    )

# %% [markdown]
# Compare GPR and VGP lengthscales after optimization:

# %%
print(f"GPR lengthscales = {gpr.kernel.lengthscales.numpy():.04f}")
print(f"VGP lengthscales = {vgp.kernel.lengthscales.numpy():.04f}")

# %% [markdown]
# ## Natural gradients also work for the sparse model
# Similarly, natural gradients turn SVGP into SGPR in the Gaussian likelihood case. <br>
# We can again combine natural gradients with Adam to update both variational parameters and hyperparameters too.<br>
# Here we'll just do a single natural step demonstration.

# %%
svgp = SVGP(
    kernel=gpflow.kernels.Matern52(),
    likelihood=gpflow.likelihoods.Gaussian(),
    inducing_variable=inducing_variable,
)
sgpr = SGPR(
    data, kernel=gpflow.kernels.Matern52(), inducing_variable=inducing_variable
)

for model in svgp, sgpr:
    model.likelihood.variance.assign(0.1)

# %% [markdown]
# Analytically optimal sparse model ELBO:

# %%
sgpr.elbo().numpy()

# %% [markdown]
# SVGP ELBO before natural gradient step:

# %%
svgp.elbo(data).numpy()

# %%
variational_params = [(svgp.q_mu, svgp.q_sqrt)]

natgrad_opt = NaturalGradient(gamma=1.0)
natgrad_opt.minimize(
    svgp.training_loss_closure(data), var_list=variational_params
)

# %% [markdown]
# SVGP ELBO after a single natural gradient step:

# %%
svgp.elbo(data).numpy()

# %% [markdown]
# ## Minibatches
# A crucial property of the natural gradient method is that it still works with minibatches.
# In practice though, we need to use a smaller gamma.

# %%
natgrad_opt = NaturalGradient(gamma=0.1)

data_minibatch = (
    tf.data.Dataset.from_tensor_slices(data)
    .prefetch(autotune)
    .repeat()
    .shuffle(N)
    .batch(batch_size)
)
data_minibatch_it = iter(data_minibatch)


svgp_objective = svgp.training_loss_closure(data_minibatch_it)
for _ in range(reduce_in_tests(100)):
    natgrad_opt.minimize(svgp_objective, var_list=variational_params)

# %% [markdown]
# Minibatch SVGP ELBO after NatGrad optimization:

# %%
np.average(
    [svgp.elbo(next(data_minibatch_it)) for _ in range(reduce_in_tests(100))]
)

# %% [markdown]
# ## Comparison with ordinary gradients in the conjugate case
#
# (Take home message: natural gradients are always better)
#
# Compared to SVGP with ordinary gradients with minibatches, the natural gradient optimizer is much faster in the Gaussian case.
#
# Here we'll do hyperparameter learning together with optimization of the variational parameters, comparing the interleaved natural gradient approach and the one using ordinary gradients for the hyperparameters and variational parameters jointly.
#
# **NOTE:** Again we need to compromise for smaller gamma value, which we'll keep *fixed* during the optimization.

# %%
svgp_ordinary = SVGP(
    kernel=gpflow.kernels.Matern52(),
    likelihood=gpflow.likelihoods.Gaussian(),
    inducing_variable=inducing_variable,
)
svgp_natgrad = SVGP(
    kernel=gpflow.kernels.Matern52(),
    likelihood=gpflow.likelihoods.Gaussian(),
    inducing_variable=inducing_variable,
)

# ordinary gradients with Adam for SVGP
ordinary_adam_opt = tf.optimizers.Adam(adam_learning_rate)

# NatGrads and Adam for SVGP
# Stop Adam from optimizing the variational parameters
set_trainable(svgp_natgrad.q_mu, False)
set_trainable(svgp_natgrad.q_sqrt, False)

# Create the optimize_tensors for SVGP
natgrad_adam_opt = tf.optimizers.Adam(adam_learning_rate)

natgrad_opt = NaturalGradient(gamma=0.1)
variational_params = [(svgp_natgrad.q_mu, svgp_natgrad.q_sqrt)]

# %% [markdown]
# Let's optimize the models:

# %%
data_minibatch = (
    tf.data.Dataset.from_tensor_slices(data)
    .prefetch(autotune)
    .repeat()
    .shuffle(N)
    .batch(batch_size)
)
data_minibatch_it = iter(data_minibatch)


svgp_ordinary_loss = svgp_ordinary.training_loss_closure(data_minibatch_it)
svgp_natgrad_loss = svgp_natgrad.training_loss_closure(data_minibatch_it)


for _ in range(reduce_in_tests(100)):
    ordinary_adam_opt.minimize(
        svgp_ordinary_loss, var_list=svgp_ordinary.trainable_variables
    )


for _ in range(reduce_in_tests(100)):
    natgrad_adam_opt.minimize(
        svgp_natgrad_loss, var_list=svgp_natgrad.trainable_variables
    )
    natgrad_opt.minimize(svgp_natgrad_loss, var_list=variational_params)

# %% [markdown]
# SVGP ELBO after ordinary `Adam` optimization:

# %%
np.average(
    [
        svgp_ordinary.elbo(next(data_minibatch_it))
        for _ in range(reduce_in_tests(100))
    ]
)

# %% [markdown]
# SVGP ELBO after `NaturalGradient` and `Adam` optimization:

# %%
np.average(
    [
        svgp_natgrad.elbo(next(data_minibatch_it))
        for _ in range(reduce_in_tests(100))
    ]
)

# %% [markdown]
# ## Comparison with ordinary gradients in the non-conjugate case
#
# ### Binary classification
#
# (Take home message: natural gradients are usually better)
#
# We can use natural gradients even when the likelihood isn't Gaussian. It isn't guaranteed to be better, but it usually is better in practical situations.

# %%
y_binary = np.random.choice([1.0, -1], size=x.shape)
vgp_data = (x, y_binary)

vgp_bernoulli = VGP(
    vgp_data,
    kernel=gpflow.kernels.Matern52(),
    likelihood=gpflow.likelihoods.Bernoulli(),
)
vgp_bernoulli_natgrad = VGP(
    vgp_data,
    kernel=gpflow.kernels.Matern52(),
    likelihood=gpflow.likelihoods.Bernoulli(),
)

# ordinary gradients with Adam for VGP with Bernoulli likelihood
adam_opt = tf.optimizers.Adam(adam_learning_rate)

# NatGrads and Adam for VGP with Bernoulli likelihood
# Stop Adam from optimizing the variational parameters
set_trainable(vgp_bernoulli_natgrad.q_mu, False)
set_trainable(vgp_bernoulli_natgrad.q_sqrt, False)

# Create the optimize_tensors for VGP with natural gradients
natgrad_adam_opt = tf.optimizers.Adam(adam_learning_rate)
natgrad_opt = NaturalGradient(gamma=0.1)
variational_params = [
    (vgp_bernoulli_natgrad.q_mu, vgp_bernoulli_natgrad.q_sqrt)
]

# %%
# Optimize vgp_bernoulli
for _ in range(reduce_in_tests(100)):
    adam_opt.minimize(
        vgp_bernoulli.training_loss, var_list=vgp_bernoulli.trainable_variables
    )

# Optimize vgp_bernoulli_natgrad
for _ in range(reduce_in_tests(100)):
    adam_opt.minimize(
        vgp_bernoulli_natgrad.training_loss,
        var_list=vgp_bernoulli_natgrad.trainable_variables,
    )
    natgrad_opt.minimize(
        vgp_bernoulli_natgrad.training_loss, var_list=variational_params
    )

# %% [markdown]
# VGP ELBO after ordinary `Adam` optimization:

# %%
vgp_bernoulli.elbo().numpy()

# %% [markdown]
# VGP ELBO after `NaturalGradient` + `Adam` optimization:

# %%
vgp_bernoulli_natgrad.elbo().numpy()

# %% [markdown]
# We can also choose to run natural gradients in another parameterization.<br>
# The sensible choice is the model parameters (q_mu, q_sqrt), which is already in GPflow.

# %%
vgp_bernoulli_natgrads_xi = VGP(
    vgp_data,
    kernel=gpflow.kernels.Matern52(),
    likelihood=gpflow.likelihoods.Bernoulli(),
)

# Stop Adam from optimizing the variational parameters
set_trainable(vgp_bernoulli_natgrads_xi.q_mu, False)
set_trainable(vgp_bernoulli_natgrads_xi.q_sqrt, False)

# Create the optimize_tensors for VGP with Bernoulli likelihood
adam_opt = tf.optimizers.Adam(adam_learning_rate)
natgrad_opt = NaturalGradient(gamma=0.01)

variational_params = [
    (
        vgp_bernoulli_natgrads_xi.q_mu,
        vgp_bernoulli_natgrads_xi.q_sqrt,
        XiSqrtMeanVar(),
    )
]

# %%
# Optimize vgp_bernoulli_natgrads_xi
for _ in range(reduce_in_tests(100)):
    adam_opt.minimize(
        vgp_bernoulli_natgrads_xi.training_loss,
        var_list=vgp_bernoulli_natgrads_xi.trainable_variables,
    )

    natgrad_opt.minimize(
        vgp_bernoulli_natgrads_xi.training_loss, var_list=variational_params
    )

# %% [markdown]
# VGP ELBO after `NaturalGradient` with `XiSqrtMeanVar` + `Adam` optimization:

# %%
vgp_bernoulli_natgrads_xi.elbo().numpy()

# %% [markdown]
# With sufficiently small steps, it shouldn't make a difference which transform is used, but for large
# steps this can make a difference in practice.


# doc/sphinx/notebooks/advanced/ordinal_regression.pct.py

# ---
# jupyter:
#   jupytext:
#     formats: ipynb,.pct.py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.3.3
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %% [markdown]
# Ordinal regression
# --
# Ordinal regression aims to fit a model to some data $(X, Y)$, where $Y$ is an ordinal variable. To do so, we use a `VPG` model with a specific likelihood (`gpflow.likelihoods.Ordinal`).

# %%
import matplotlib.pyplot as plt
import numpy as np

import gpflow

# %matplotlib inline
plt.rcParams["figure.figsize"] = (12, 6)

np.random.seed(123)  # for reproducibility


# %%
# make a one-dimensional ordinal regression problem

# This function generates a set of inputs X,
# quantitative output f (latent) and ordinal values Y


def generate_data(num_data):
    # First generate random inputs
    X = np.random.rand(num_data, 1)

    # Now generate values of a latent GP
    kern = gpflow.kernels.SquaredExponential(lengthscales=0.1)
    K = kern(X)
    f = np.random.multivariate_normal(mean=np.zeros(num_data), cov=K).reshape(
        -1, 1
    )

    # Finally convert f values into ordinal values Y
    Y = np.round((f + f.min()) * 3)
    Y = Y - Y.min()
    Y = np.asarray(Y, np.float64)

    return X, f, Y


np.random.seed(1)
num_data = 20
X, f, Y = generate_data(num_data)

plt.figure(figsize=(11, 6))
plt.plot(X, f, ".")
plt.ylabel("latent function value")

plt.twinx()
plt.plot(X, Y, "kx", mew=1.5)
plt.ylabel("observed data value")

# %%
# construct ordinal likelihood - bin_edges is the same as unique(Y) but centered
bin_edges = np.array(np.arange(np.unique(Y).size + 1), dtype=float)
bin_edges = bin_edges - bin_edges.mean()
likelihood = gpflow.likelihoods.Ordinal(bin_edges)

# build a model with this likelihood
m = gpflow.models.VGP(
    data=(X, Y), kernel=gpflow.kernels.Matern32(), likelihood=likelihood
)

# fit the model
opt = gpflow.optimizers.Scipy()
opt.minimize(m.training_loss, m.trainable_variables, options=dict(maxiter=100))

# %%
# here we'll plot the expected value of Y +- 2 std deviations, as if the distribution were Gaussian
plt.figure(figsize=(11, 6))
X_data, Y_data = (m.data[0].numpy(), m.data[1].numpy())
Xtest = np.linspace(X_data.min(), X_data.max(), 100).reshape(-1, 1)
mu, var = m.predict_y(Xtest)
(line,) = plt.plot(Xtest, mu, lw=2)
col = line.get_color()
plt.plot(Xtest, mu + 2 * np.sqrt(var), "--", lw=2, color=col)
plt.plot(Xtest, mu - 2 * np.sqrt(var), "--", lw=2, color=col)
plt.plot(X_data, Y_data, "kx", mew=2)


# %%
## to see the predictive density, try predicting every possible discrete value for Y.
def pred_log_density(m):
    Xtest = np.linspace(X_data.min(), X_data.max(), 100).reshape(-1, 1)
    ys = np.arange(Y_data.max() + 1)
    densities = []
    for y in ys:
        Ytest = np.full_like(Xtest, y)
        # Predict the log density
        densities.append(m.predict_log_density((Xtest, Ytest)))
    return np.vstack(densities)


# %%
fig = plt.figure(figsize=(14, 6))
plt.imshow(
    np.exp(pred_log_density(m)),
    interpolation="nearest",
    extent=[X_data.min(), X_data.max(), -0.5, Y_data.max() + 0.5],
    origin="lower",
    aspect="auto",
    cmap=plt.cm.viridis,
)
plt.colorbar()
plt.plot(X, Y, "kx", mew=2, scalex=False, scaley=False)

# %%
# Predictive density for a single input x=0.5
x_new = 0.5
Y_new = np.arange(np.max(Y_data + 1)).reshape([-1, 1])
X_new = np.full_like(Y_new, x_new)
# for predict_log_density x and y need to have the same number of rows
dens_new = np.exp(m.predict_log_density((X_new, Y_new)))
fig = plt.figure(figsize=(8, 4))
plt.bar(x=Y_new.flatten(), height=dens_new.flatten())


# doc/sphinx/notebooks/advanced/variational_fourier_features.pct.py

# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,.pct.py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.6.0
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Variational Fourier Features in the GPflow framework
#
# In this notebook we demonstrate how new types of inducing variables can easily be incorporated in the GPflow framework. As an example case, we use the variational Fourier features from [Hensman, Durrande, and Solin (JMLR 2018)](http://jmlr.csail.mit.edu/papers/v18/16-579). All equation and table references are to this paper.
#
# **Note:** This implementation is meant as an example, not as a feature-complete implementation. For more features, such as multi-dimensional inputs, use the [GPflow 2 version of the original VFF code](https://github.com/st--/VFF).
#
# We cannot directly use Fourier features within the multi-output framework without losing the computational advantages, as `Kuu` and `Kuf` for SharedIndependent and SeparateIndependent inducing variables assume that the sub-inducing variable's covariances are simply computed as dense Tensors. However, there is nothing preventing a dedicated implementation of multi-output Fourier features that is computationally more efficient - feel free to discuss this within [the GPflow community](https://github.com/GPflow/GPflow/#the-gpflow-community)!

# %%
import numpy as np
import tensorflow as tf

import gpflow
from gpflow import covariances as cov
from gpflow import kullback_leiblers as kl
from gpflow.base import TensorLike
from gpflow.ci_utils import reduce_in_tests
from gpflow.experimental.check_shapes import Shape
from gpflow.inducing_variables import InducingVariables
from gpflow.utilities import to_default_float

# %%
# VFF give structured covariance matrices that are computationally efficient.
# We take advantage of this using TensorFlow's LinearOperators:
BlockDiag = tf.linalg.LinearOperatorBlockDiag
Diag = tf.linalg.LinearOperatorDiag
LowRank = tf.linalg.LinearOperatorLowRankUpdate

# %%
import matplotlib.pyplot as plt

# %matplotlib inline

# %% [markdown]
# The VFF inducing variables are defined as a projection $u_m = \mathcal{P}_{\phi_m}(f)$ (eq. (59)) of the GP $f(\cdot)$ onto a truncated Fourier basis, $\phi_m = [1, \cos(\omega_1(x-a)),\dots,\cos(\omega_M(x-a)),\sin(\omega_1(x-a)),\dots,\sin(\omega_M(x-a))]$ (eq. (47)). To represent this we define a new inducing variables class that derives from the `InducingVariables` base class.

# %%
class FourierFeatures1D(InducingVariables):
    def __init__(self, a, b, M):
        """
        `a` and `b` define the interval [a, b] of the Fourier representation.
        `M` specifies the number of frequencies to use.
        """
        # [a, b] defining the interval of the Fourier representation:
        self.a = gpflow.Parameter(a, dtype=gpflow.default_float())
        self.b = gpflow.Parameter(b, dtype=gpflow.default_float())
        self.M = M
        # integer array defining the frequencies, ω_m = 2π (b - a)/m:
        self.ms = np.arange(M)

    @property
    def num_inducing(self):
        """ number of inducing variables (defines dimensionality of q(u)) """
        return 2 * self.M - 1  # `M` cosine and `M-1` sine components

    @property
    def shape(self) -> Shape:
        M = 2 * self.M - 1
        D = 1  # Input size.
        P = 1  # Output size.
        return (M, D, P)


# %% [markdown]
# Next, we need to define how to compute $\mathrm{K}_\mathbf{uu} = \operatorname{cov}(u_m, u_{m'})$ (eq. (61)) and $\mathrm{K}_\mathbf{uf} = \operatorname{cov}(u_m, f(x_n))$ (eq. (60)).

# %%
@cov.Kuu.register(FourierFeatures1D, gpflow.kernels.Matern12)
def Kuu_matern12_fourierfeatures1d(inducing_variable, kernel, jitter=None):
    a, b, ms = (lambda u: (u.a, u.b, u.ms))(inducing_variable)
    omegas = 2.0 * np.pi * ms / (b - a)

    # Cosine block:
    lamb = 1.0 / kernel.lengthscales
    two_or_four = to_default_float(tf.where(omegas == 0, 2.0, 4.0))
    d_cos = (
        (b - a)
        * (tf.square(lamb) + tf.square(omegas))
        / lamb
        / kernel.variance
        / two_or_four
    )  # eq. (111)
    v_cos = tf.ones_like(d_cos) / tf.sqrt(kernel.variance)  # eq. (110)
    cosine_block = LowRank(
        Diag(d_cos, is_positive_definite=True), v_cos[:, None]
    )

    # Sine block:
    omegas = omegas[
        tf.not_equal(omegas, 0)
    ]  # the sine block does not include omega=0
    d_sin = (
        (b - a)
        * (tf.square(lamb) + tf.square(omegas))
        / lamb
        / kernel.variance
        / 4.0
    )  # eq. (113)
    sine_block = Diag(d_sin, is_positive_definite=True)

    return BlockDiag([cosine_block, sine_block])


@cov.Kuf.register(FourierFeatures1D, gpflow.kernels.Matern12, TensorLike)
def Kuf_matern12_fourierfeatures1d(inducing_variable, kernel, X):
    X = tf.squeeze(X, axis=1)
    a, b, ms = (lambda u: (u.a, u.b, u.ms))(inducing_variable)

    omegas = 2.0 * np.pi * ms / (b - a)
    Kuf_cos = tf.cos(omegas[:, None] * (X[None, :] - a))
    omegas_sin = omegas[tf.not_equal(omegas, 0)]  # don't compute zero frequency
    Kuf_sin = tf.sin(omegas_sin[:, None] * (X[None, :] - a))

    # correct Kuf outside [a, b] -- see Table 1
    Kuf_sin = tf.where(
        (X < a) | (X > b), tf.zeros_like(Kuf_sin), Kuf_sin
    )  # just zero

    left_tail = tf.exp(-tf.abs(X - a) / kernel.lengthscales)[None, :]
    right_tail = tf.exp(-tf.abs(X - b) / kernel.lengthscales)[None, :]
    Kuf_cos = tf.where(X < a, left_tail, Kuf_cos)  # replace with left tail
    Kuf_cos = tf.where(X > b, right_tail, Kuf_cos)  # replace with right tail

    return tf.concat([Kuf_cos, Kuf_sin], axis=0)


@cov.Kuu.register(FourierFeatures1D, gpflow.kernels.Matern32)
def Kuu_matern32_fourierfeatures1d(inducing_variable, kernel, jitter=None):
    a, b, ms = (lambda u: (u.a, u.b, u.ms))(inducing_variable)
    omegas = 2.0 * np.pi * ms / (b - a)

    # Cosine block: eq. (114)
    lamb = np.sqrt(3.0) / kernel.lengthscales
    four_or_eight = to_default_float(tf.where(omegas == 0, 4.0, 8.0))
    d_cos = (
        (b - a)
        * tf.square(tf.square(lamb) + tf.square(omegas))
        / tf.pow(lamb, 3)
        / kernel.variance
        / four_or_eight
    )
    v_cos = tf.ones_like(d_cos) / tf.sqrt(kernel.variance)
    cosine_block = LowRank(
        Diag(d_cos, is_positive_definite=True), v_cos[:, None]
    )

    # Sine block: eq. (115)
    omegas = omegas[tf.not_equal(omegas, 0)]  # don't compute omega=0
    d_sin = (
        (b - a)
        * tf.square(tf.square(lamb) + tf.square(omegas))
        / tf.pow(lamb, 3)
        / kernel.variance
        / 8.0
    )
    v_sin = omegas / lamb / tf.sqrt(kernel.variance)
    sine_block = LowRank(Diag(d_sin, is_positive_definite=True), v_sin[:, None])

    return BlockDiag([cosine_block, sine_block])  # eq. (116)


@cov.Kuf.register(FourierFeatures1D, gpflow.kernels.Matern32, TensorLike)
def Kuf_matern32_fourierfeatures1d(inducing_variable, kernel, X):
    X = tf.squeeze(X, axis=1)
    a, b, ms = (lambda u: (u.a, u.b, u.ms))(inducing_variable)
    omegas = 2.0 * np.pi * ms / (b - a)

    Kuf_cos = tf.cos(omegas[:, None] * (X[None, :] - a))
    omegas_sin = omegas[tf.not_equal(omegas, 0)]  # don't compute zeros freq.
    Kuf_sin = tf.sin(omegas_sin[:, None] * (X[None, :] - a))

    # correct Kuf outside [a, b] -- see Table 1

    def tail_cos(delta_X):
        arg = np.sqrt(3) * tf.abs(delta_X) / kernel.lengthscales
        return (1 + arg) * tf.exp(-arg)[None, :]

    Kuf_cos = tf.where(X < a, tail_cos(X - a), Kuf_cos)
    Kuf_cos = tf.where(X > b, tail_cos(X - b), Kuf_cos)

    def tail_sin(delta_X):
        arg = np.sqrt(3) * tf.abs(delta_X) / kernel.lengthscales
        return delta_X[None, :] * tf.exp(-arg) * omegas_sin[:, None]

    Kuf_sin = tf.where(X < a, tail_sin(X - a), Kuf_sin)
    Kuf_sin = tf.where(X > b, tail_sin(X - b), Kuf_sin)

    return tf.concat([Kuf_cos, Kuf_sin], axis=0)


# %% [markdown]
# In principle, this is all we need; however, to be able to take advantage of the structure of `Kuu`, we need to also implement new versions of the KL divergence from the prior to the approximate posterior (`prior_kl`) and the computation of the Gaussian process conditional (posterior) equations:

# %%
@kl.prior_kl.register(
    FourierFeatures1D, gpflow.kernels.Kernel, TensorLike, TensorLike
)
def prior_kl_vff(inducing_variable, kernel, q_mu, q_sqrt, whiten=False):
    if whiten:
        raise NotImplementedError
    K = cov.Kuu(inducing_variable, kernel)
    return gauss_kl_vff(q_mu, q_sqrt, K)


def gauss_kl_vff(q_mu, q_sqrt, K):
    """
    Compute the KL divergence from

          q(x) = N(q_mu, q_sqrt^2)
    to
          p(x) = N(0, K)

    q_mu is a vector [N, 1] that contains the mean.
    q_sqrt is a matrix that is the lower triangular square-root matrix of the covariance of q.

    K is a positive definite matrix: the covariance of p.
    NOTE: K is a LinearOperator that provides efficient methjods
        for solve(), log_abs_determinant(), and trace()
    """
    # KL(N₀ || N₁) = ½ [tr(Σ₁⁻¹ Σ₀) + (μ₁ - μ₀)ᵀ Σ₁⁻¹ (μ₁ - μ₀) - k + ln(det(Σ₁)/det(Σ₀))]
    # N₀ = q; μ₀ = q_mu, Σ₀ = q_sqrt q_sqrtᵀ
    # N₁ = p; μ₁ = 0, Σ₁ = K
    # KL(q || p) =
    #     ½ [tr(K⁻¹ q_sqrt q_sqrtᵀA + q_muᵀ K⁻¹ q_mu - k + logdet(K) - logdet(q_sqrt q_sqrtᵀ)]
    # k = number of dimensions, if q_sqrt is m x m this is m²
    Kinv_q_mu = K.solve(q_mu)

    mahalanobis_term = tf.squeeze(tf.matmul(q_mu, Kinv_q_mu, transpose_a=True))

    # GPflow: q_sqrt is num_latent_gps x N x N
    num_latent_gps = to_default_float(tf.shape(q_mu)[1])
    logdet_prior = num_latent_gps * K.log_abs_determinant()

    product_of_dimensions__int = tf.reduce_prod(
        tf.shape(q_sqrt)[:-1]
    )  # dimensions are integers
    constant_term = to_default_float(product_of_dimensions__int)

    Lq = tf.linalg.band_part(q_sqrt, -1, 0)  # force lower triangle
    logdet_q = tf.reduce_sum(tf.math.log(tf.square(tf.linalg.diag_part(Lq))))

    # S = tf.matmul(q_sqrt, q_sqrt, transpose_b=True)
    # trace_term = tf.trace(K.solve(S))
    trace_term = tf.squeeze(
        tf.reduce_sum(Lq * K.solve(Lq), axis=[-1, -2])
    )  # [O(N²) instead of O(N³)

    twoKL = (
        trace_term + mahalanobis_term - constant_term + logdet_prior - logdet_q
    )
    return 0.5 * twoKL


# %%
import gpflow.posteriors


class VFFPosterior(gpflow.posteriors.BasePosterior):
    def _conditional_fused(self, Xnew, full_cov, full_output_cov):
        """
        Xnew is a tensor with the points of the data or minibatch, shape N x D
        """
        if full_output_cov:
            raise NotImplementedError

        f = self._q_dist.q_mu
        q_sqrt = self._q_dist.q_sqrt

        # num_data = tf.shape(Xnew)[0]  # M
        num_func = tf.shape(f)[1]  # K

        Kuu = cov.Kuu(self.X_data, self.kernel)  # this is now a LinearOperator
        Kuf = cov.Kuf(self.X_data, self.kernel, Xnew)  # still a Tensor

        KuuInv_Kuf = Kuu.solve(Kuf)

        # compute the covariance due to the conditioning
        if full_cov:
            fvar = self.kernel(Xnew) - tf.matmul(
                Kuf, KuuInv_Kuf, transpose_a=True
            )
            shape = (num_func, 1, 1)
        else:
            KufT_KuuInv_Kuf_diag = tf.reduce_sum(Kuf * KuuInv_Kuf, axis=-2)
            fvar = self.kernel(Xnew, full_cov=False) - KufT_KuuInv_Kuf_diag
            shape = (num_func, 1)
        fvar = tf.expand_dims(fvar, 0) * tf.ones(
            shape, dtype=gpflow.default_float()
        )  # K x N x N or K x N

        if self.whiten:
            raise NotImplementedError

        A = KuuInv_Kuf

        # construct the conditional mean
        fmean = tf.matmul(A, f, transpose_a=True)

        if q_sqrt is not None:
            if q_sqrt.get_shape().ndims == 2:
                # LTA = A * tf.expand_dims(q_sqrt, 2)  # K x M x N
                # won't work  # make ticket for this?
                raise NotImplementedError
            elif q_sqrt.get_shape().ndims == 3:
                # L = tf.matrix_band_part(tf.transpose(q_sqrt, (2, 0, 1)), -1, 0)  # K x M x M

                # K x M x N
                # A_tiled = tf.expand_dims(A.get(), 0) * tf.ones((num_func, 1, 1), dtype=float_type)

                # LTA = tf.matmul(L, A_tiled, transpose_a=True)  # K x M x N
                # TODO the following won't work for K > 1
                assert q_sqrt.shape[0] == 1
                # LTA = (A.T @ DenseMatrix(q_sqrt[:,:,0])).T.get()[None, :, :]
                ATL = tf.matmul(A, q_sqrt, transpose_a=True)
            else:
                raise ValueError(
                    "Bad dimension for q_sqrt: %s"
                    % str(q_sqrt.get_shape().ndims)
                )
            if full_cov:
                # fvar = fvar + tf.matmul(LTA, LTA, transpose_a=True)  # K x N x N
                fvar = fvar + tf.matmul(ATL, ATL, transpose_b=True)  # K x N x N
            else:
                # fvar = fvar + tf.reduce_sum(tf.square(LTA), 1)  # K x N
                fvar = fvar + tf.reduce_sum(tf.square(ATL), 2)  # K x N

        if not full_cov:
            fvar = tf.transpose(fvar)  # N x K

        return fmean, fvar

    # We can also provide a conditional that precomputes as much as possible,
    # to speed up predictions:

    def _precompute(self):
        Kuu = cov.Kuu(self.X_data, self.kernel)  # this is now a LinearOperator

        q_mu = self._q_dist.q_mu
        q_sqrt = self._q_dist.q_sqrt

        if self.whiten:
            raise NotImplementedError
        else:
            # alpha = Kuu⁻¹ q_mu
            alpha = Kuu.solve(q_mu)  # type: tf.Tensor

        if self.whiten:
            raise NotImplementedError
        else:
            # Qinv = Kuu⁻¹ - Kuu⁻¹ S Kuu⁻¹
            KuuInv_qsqrt = Kuu.solve(q_sqrt)
            KuuInv_covu_KuuInv = tf.matmul(
                KuuInv_qsqrt, KuuInv_qsqrt, transpose_b=True
            )

        Qinv = Kuu.inverse().to_dense() - KuuInv_covu_KuuInv

        return gpflow.posteriors.PrecomputedValue.wrap_alpha_Qinv(alpha, Qinv)

    def _conditional_with_precompute(
        self, cache, Xnew, full_cov, full_output_cov
    ):
        alpha, Qinv = cache

        if full_output_cov:
            raise NotImplementedError

        Kuf = cov.Kuf(self.X_data, self.kernel, Xnew)  # still a Tensor

        # construct the conditional mean
        fmean = tf.matmul(Kuf, alpha, transpose_a=True)

        num_func = tf.shape(alpha)[1]  # K
        Qinv_Kuf = tf.matmul(Qinv, Kuf)

        # compute the covariance due to the conditioning
        if full_cov:
            fvar = self.kernel(Xnew) - tf.matmul(
                Kuf, Qinv_Kuf, transpose_a=True
            )
        else:
            KufT_Qinv_Kuf_diag = tf.reduce_sum(Kuf * Qinv_Kuf, axis=-2)
            fvar = self.kernel(Xnew, full_cov=False) - KufT_Qinv_Kuf_diag
            fvar = tf.transpose(fvar)

        return fmean, fvar


# %% [markdown]
# We now have to register our Posterior object:

# %%
@gpflow.posteriors.get_posterior_class.register(
    gpflow.kernels.Kernel, FourierFeatures1D
)
def _get_posterior_vff(kernel, inducing_variable):
    return VFFPosterior


# %% [markdown]
# `gpflow.conditionals.conditional` is a short-hand for calling the fused prediction code path:

# %%
Mf = 2
M = 2 * Mf - 1
kernel = gpflow.kernels.Matern32()
inducing_variable = FourierFeatures1D(-0.5, 1.5, Mf)

Xnew = np.random.rand(7, 1)
f = np.random.randn(M, 1)
q_sqrt = tf.convert_to_tensor(np.tril(np.random.randn(1, M, M)))

conditional_f_mean, conditional_f_var = gpflow.conditionals.conditional(
    Xnew,
    inducing_variable,
    kernel,
    f,
    q_sqrt=q_sqrt,
    white=False,
    full_cov=True,
)

posterior = VFFPosterior(
    kernel, inducing_variable, f, q_sqrt, whiten=False, precompute_cache=None
)
posterior_f_mean, posterior_f_var = posterior.fused_predict_f(
    Xnew, full_cov=True
)

np.testing.assert_array_equal(conditional_f_mean, posterior_f_mean)
np.testing.assert_array_equal(conditional_f_var, posterior_f_var)

# %% [markdown]
# We can also call the cached path:
#

# %%
posterior.update_cache(gpflow.posteriors.PrecomputeCacheType.TENSOR)
precomputed_posterior_f_mean, precomputed_posterior_f_var = posterior.predict_f(
    Xnew, full_cov=True
)

np.testing.assert_allclose(precomputed_posterior_f_mean, posterior_f_mean)
np.testing.assert_allclose(precomputed_posterior_f_var, posterior_f_var)

# %% [markdown]
# We now demonstrate how to use these new types of inducing variables with the `SVGP` model class. First, let's create some toy data:

# %%
X = np.linspace(-2, 2, 510)
Xnew = np.linspace(-4, 4, 501)


def f(x):
    return np.cos(2 * np.pi * x / 4 * 2)


F = f(X)
Fnew = f(Xnew)
noise_scale = 0.1
np.random.seed(1)
Y = F + np.random.randn(*F.shape) * noise_scale

data = (X.reshape(-1, 1), Y.reshape(-1, 1))

# %%
plt.figure()
plt.plot(X, F, label="f(x)")
plt.plot(X, Y, ".", label="observations")
plt.legend()
plt.show()

# %% [markdown]
# Setting up an SVGP model with variational Fourier feature inducing variables is as simple as replacing the `inducing_variable` argument:

# %%
Mfreq = 9
m = gpflow.models.SVGP(
    kernel=gpflow.kernels.Matern32(),
    likelihood=gpflow.likelihoods.Gaussian(variance=noise_scale ** 2),
    inducing_variable=FourierFeatures1D(-4.5, 4.5, Mfreq),
    num_data=len(X),
    whiten=False,
)
gpflow.set_trainable(m.kernel, False)
gpflow.set_trainable(m.likelihood, False)
gpflow.set_trainable(
    m.inducing_variable, True
)  # whether to optimize bounds [a, b]

# %%
opt = gpflow.optimizers.Scipy()
opt.minimize(
    m.training_loss_closure(data),
    m.trainable_variables,
    options=dict(maxiter=reduce_in_tests(5000)),
)

gpflow.utilities.print_summary(m, fmt="notebook")

# %% [markdown]
# For comparison we also construct an SVGP model using inducing points and an exact GPR model:

# %%
m_ip = gpflow.models.SVGP(
    kernel=gpflow.kernels.Matern32(),
    likelihood=gpflow.likelihoods.Gaussian(variance=noise_scale ** 2),
    inducing_variable=np.linspace(-2, 2, Mfreq * 2 - 1)[:, None],
    num_data=len(X),
    whiten=False,
)
gpflow.set_trainable(m_ip.kernel, False)
gpflow.set_trainable(m_ip.likelihood, False)
gpflow.set_trainable(
    m_ip.inducing_variable, True
)  # whether to optimize inducing point locations

# %%
opt = gpflow.optimizers.Scipy()
opt.minimize(
    m_ip.training_loss_closure(data),
    m_ip.trainable_variables,
    options=dict(maxiter=reduce_in_tests(5000)),
)

gpflow.utilities.print_summary(m_ip, fmt="notebook")

# %%
m_ref = gpflow.models.GPR(
    (X.reshape(-1, 1), Y.reshape(-1, 1)), kernel=gpflow.kernels.Matern32()
)
m_ref.likelihood.variance = np.array(noise_scale ** 2).astype(np.float64)
gpflow.set_trainable(m_ref.kernel, False)
gpflow.set_trainable(m_ref.likelihood, False)

# Because we fixed the kernel and likelihood hyperparameters, we don't need to optimize anything.

gpflow.utilities.print_summary(m_ref, fmt="notebook")

# %%
exact_gpr_lml = m_ref.log_marginal_likelihood().numpy().item()
print("LML (exact GPR) =", exact_gpr_lml)
ip_svgp_elbo = m_ip.elbo(data).numpy().item()
print("ELBO (SVGP, inducing points) =", ip_svgp_elbo)
vff_svgp_elbo = m.elbo(data).numpy().item()
print("ELBO (SVGP, Fourier features) =", vff_svgp_elbo)

# %%
def plot_gp(m, Xnew, name=""):
    Fmean, Fvar = m.predict_f(Xnew[:, None])
    Fmean = Fmean.numpy().squeeze()
    Fvar = Fvar.numpy().squeeze()
    (p,) = plt.plot(Xnew, Fmean, label=name)
    plt.fill_between(
        Xnew,
        Fmean - 2 * np.sqrt(Fvar),
        Fmean + 2 * np.sqrt(Fvar),
        alpha=0.3,
        color=p.get_color(),
    )


def plot_data():
    plt.plot(Xnew, Fnew, label="f(x)")
    plt.plot(X, Y, ".", label="observations")


plt.figure(figsize=(15, 10))
plot_data()
plot_gp(m, Xnew, "VFF [ELBO={:.3}]".format(vff_svgp_elbo))
plot_gp(m_ip, Xnew, "inducing points [ELBO={:.3}]".format(ip_svgp_elbo))
plot_gp(m_ref, Xnew, "exact [LML={:.3}]".format(exact_gpr_lml))
plt.legend(loc="best")
plt.show()


# doc/sphinx/notebooks/advanced/varying_noise.pct.py

# ---
# jupyter:
#   jupytext:
#     formats: ipynb,.pct.py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.13.8
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Gaussian process regression with varying output noise

# %% [markdown]
# This notebook shows how to construct a Gaussian process model where different noise is assumed for different data points. The model is:
#
# \begin{align}
# f(\cdot) &\sim \mathcal{GP}\big(0, k(\cdot, \cdot)\big) \\
# y_i | f, x_i &\sim \mathcal N\big(y_i; f(x_i), \sigma^2_i\big)
# \end{align}
#
# We'll demonstrate three methods for specifying the data-point specific noise:
# * First we'll show how to fit the noise variance to a simple function.
# * In the second demonstration we'll assume that the data comes from two different groups and show how to learn separate noise variance for the groups.
# * Third we'll assume you have multiple samples at each location $x$ and show how to use the empirical variance.

# %%
import os
import warnings

warnings.simplefilter("ignore")
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"

# %%
from typing import Optional, Tuple

import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf

import gpflow as gf
from gpflow.ci_utils import reduce_in_tests
from gpflow.experimental.check_shapes import inherit_check_shapes
from gpflow.utilities import print_summary

# %matplotlib inline

gf.config.set_default_summary_fmt("notebook")

optimizer_config = dict(maxiter=reduce_in_tests(1000))
n_data = reduce_in_tests(300)
X_plot = np.linspace(0.0, 1.0, reduce_in_tests(101))[:, None]


# %% [markdown]
# To help us later we'll first define a small function for plotting data with predictions:

# %%
def plot_distribution(
    X: gf.base.AnyNDArray,
    Y: gf.base.AnyNDArray,
    X_plot: Optional[gf.base.AnyNDArray] = None,
    mean_plot: Optional[gf.base.AnyNDArray] = None,
    var_plot: Optional[gf.base.AnyNDArray] = None,
    X_err: Optional[gf.base.AnyNDArray] = None,
    mean_err: Optional[gf.base.AnyNDArray] = None,
    var_err: Optional[gf.base.AnyNDArray] = None,
) -> None:
    plt.figure(figsize=(15, 5))

    X = X.squeeze(axis=-1)
    Y = Y.squeeze(axis=-1)
    plt.scatter(X, Y, color="gray", alpha=0.8)

    def get_confidence_bounds(
        mean: gf.base.AnyNDArray, var: gf.base.AnyNDArray
    ) -> Tuple[gf.base.AnyNDArray, gf.base.AnyNDArray]:
        std = np.sqrt(var)
        return mean - 1.96 * std, mean + 1.96 * std

    if X_plot is not None:
        assert mean_plot is not None
        assert var_plot is not None
        X_plot = X_plot.squeeze(axis=-1)
        mean_plot = mean_plot.squeeze(axis=-1)
        var_plot = var_plot.squeeze(axis=-1)
        lower_plot, upper_plot = get_confidence_bounds(mean_plot, var_plot)
        plt.fill_between(
            X_plot, lower_plot, upper_plot, color="silver", alpha=0.25
        )
        plt.plot(X_plot, lower_plot, color="silver")
        plt.plot(X_plot, upper_plot, color="silver")
        plt.plot(X_plot, mean_plot, color="black")

    if X_err is not None:
        assert mean_err is not None
        assert var_err is not None
        lower_err, upper_err = get_confidence_bounds(mean_err, var_err)
        plt.vlines(X_err, lower_err, upper_err, color="black")

    plt.xlim([0, 1])
    plt.show()


# %% [markdown]
# ## Demo 1: known noise variances
# ### Generate data
# We create a utility function to generate synthetic data, including noise that varies amongst the data:

# %%
def generate_data() -> Tuple[gf.base.AnyNDArray, gf.base.AnyNDArray]:
    rng = np.random.default_rng(42)  # for reproducibility
    X = rng.uniform(0.0, 1.0, (n_data, 1))
    signal = (X - 0.5) ** 2 + 0.05
    noise_scale = 0.5 * signal
    noise = noise_scale * rng.standard_normal((n_data, 1))
    Y = signal + noise
    return X, Y


X, Y = generate_data()

# %% [markdown]
# The data alone looks like:

# %%
plot_distribution(X, Y)

# %% [markdown]
# ### Try a naive fit
# If we try to fit a naive GPR model to this data, which assumes a single shared noise variance value for all data points - as specified by the likelihood variance parameter - we get:

# %%
model = gf.models.GPR(
    data=(X, Y),
    kernel=gf.kernels.SquaredExponential(),
)
gf.optimizers.Scipy().minimize(
    model.training_loss, model.trainable_variables, options=optimizer_config
)
print_summary(model)

# %%
mean, var = model.predict_y(X_plot)
plot_distribution(X, Y, X_plot, mean.numpy(), var.numpy())

# %% [markdown]
# Notice how this naive model underestimates the noise near the center of the figure, but overestimates at for small and large values of x?

# %% [markdown]
# ### Fit a polynomial to the noise scale
#
# To fit a model with varying noise, instead of using the default single shared noise variance, we can create a Gaussian `Likelihood` with an input dependent (polynomial) `Function` for the scale of the noise,
# then pass that likelihood to our model.

# %%
model = gf.models.GPR(
    data=(X, Y),
    kernel=gf.kernels.SquaredExponential(),
    likelihood=gf.likelihoods.Gaussian(scale=gf.functions.Polynomial(degree=2)),
)
gf.optimizers.Scipy().minimize(
    model.training_loss, model.trainable_variables, options=optimizer_config
)
print_summary(model)

# %%
mean, var = model.predict_y(X_plot)
plot_distribution(X, Y, X_plot, mean.numpy(), var.numpy())


# %% [markdown]
# ## Demo 2: grouped noise variances
#
# In this demo, we won't assume that the noise variances is a function of $x$, but we will assume that they're known to be in two groups. This example represents a case where we might know that data has been collected by two instruments with different fidelity, but we do not know what those fidelities are.
#
# Of course it would be straightforward to add more groups. We'll stick with two for simplicity.

# %% [markdown]
# ### Generate data

# %%
def get_group(X: gf.base.AnyNDArray) -> gf.base.AnyNDArray:
    return np.where((((0.2 < X) & (X < 0.4)) | ((0.7 < X) & (X < 0.8))), 0, 1)


def generate_grouped_data() -> Tuple[gf.base.AnyNDArray, gf.base.AnyNDArray]:
    rng = np.random.default_rng(42)  # for reproducibility
    X = rng.uniform(0.0, 1.0, (n_data, 1))
    signal = np.sin(6 * X)
    noise_scale = 0.1 + 0.4 * get_group(X)
    noise = noise_scale * rng.standard_normal((n_data, 1))
    Y = signal + noise
    return X, Y


X, Y = generate_grouped_data()

# %% [markdown]
# And here's a plot of the raw data:

# %%
plot_distribution(X, Y)

# %% [markdown]
# ### Fit a naive model
#
# Again we'll start by fitting a naive GPR, to demonstrate the problem.

# %%
model = gf.models.GPR(
    data=(X, Y),
    kernel=gf.kernels.SquaredExponential(),
)
gf.optimizers.Scipy().minimize(
    model.training_loss, model.trainable_variables, options=optimizer_config
)
print_summary(model)

# %%
mean, var = model.predict_y(X_plot)
plot_distribution(X, Y, X_plot, mean.numpy(), var.numpy())

# %% [markdown]
# ### Data structure
#
# To model the different noise groups, we need to let the model know which group each data point belongs to. We'll do that by appending a column with the group to the $x$ data:

# %%
X_and_group = np.hstack([X, get_group(X)])
X_plot_and_group = np.hstack([X_plot, get_group(X_plot)])

# %% [markdown]
# ### Use multiple functions for the noise variance
#
# To model this we will create two noise variance functions, each of which are just a constant, and then switch between them depending on the group labels.
#
# Notice that we initialize the constant functions to a positive value. They would have defaulted to `0.0`, but that wouldn't work for a variance, which must be strictly positive.

# %%
model = gf.models.GPR(
    data=(X_and_group, Y),
    kernel=gf.kernels.SquaredExponential(active_dims=[0]),
    likelihood=gf.likelihoods.Gaussian(
        variance=gf.functions.SwitchedFunction(
            [
                gf.functions.Constant(1.0),
                gf.functions.Constant(1.0),
            ]
        )
    ),
)
gf.optimizers.Scipy().minimize(
    model.training_loss, model.trainable_variables, options=optimizer_config
)
print_summary(model)

# %%
mean, var = model.predict_y(X_plot_and_group)
plot_distribution(X, Y, X_plot, mean.numpy(), var.numpy())

# %% [markdown]
# ## Demo 3: Empirical noise variance
#
# In this demo we will assume that you have multiple measurements at each $x$ location, and want to use the empirical variance at each location.

# %% [markdown]
# ### Generate data

# %%
n_data = reduce_in_tests(20)
n_repeats = reduce_in_tests(10)


def generate_empiricial_noise_data() -> Tuple[
    gf.base.AnyNDArray, gf.base.AnyNDArray
]:
    rng = np.random.default_rng(42)  # for reproducibility
    X = rng.uniform(0.0, 1.0, (n_data, 1))
    signal = np.sin(6 * X)
    noise_scale = rng.uniform(0.1, 1.0, (n_data, 1))
    noise = noise_scale * rng.standard_normal((n_data, n_repeats))
    Y = signal + noise
    return X, Y


X, Y = generate_empiricial_noise_data()
Y_mean = np.mean(Y, axis=-1, keepdims=True)
Y_var = np.var(Y, axis=-1, keepdims=True)

# %% [markdown]
# For the sake of plotting we'll create flat lists of (x, y) pairs:

# %%
X_flat = np.broadcast_to(X, Y.shape).reshape((-1, 1))
Y_flat = Y.reshape((-1, 1))
plot_distribution(X_flat, Y_flat)

# %% [markdown]
# ### Data structure
#
# Our GPs don't like it when multiple data points occupy the same $x$ location. So we'll reduce this dataset to the means and the variances at each $x$, and use a custom function to model inject the empirical variances into the model.

# %%
Y_mean = np.mean(Y, axis=-1, keepdims=True)

# %% [markdown]
# ### Fit a naive model
#
# Again we'll start by fitting a naive GPR, to demonstrate the problem.
#
# Notice that this case is somewhat different from what we did above. Here we are not modelling $Y$, but the mean of $Y$, so the noise variance is not our uncertainty of $Y$, but our uncertainty of the mean of $Y$. Furthermore, while the naive model assumes a constant noise variance, the custom noise variance function we will develop below will only model the noise variance at the points where we have data, thus we cannot plot a continuous uncertainty. Instead, the shaded gray area is the confidence interval of $f$, and we use vertical black lines to indicate the confidence interval of the $Y$ mean.

# %%
model = gf.models.GPR(
    data=(X, Y_mean),
    kernel=gf.kernels.SquaredExponential(),
)
gf.optimizers.Scipy().minimize(
    model.training_loss, model.trainable_variables, options=optimizer_config
)
print_summary(model)

# %%
f_mean, f_var = model.predict_f(X_plot)
y_mean_mean, y_mean_var = model.predict_y(X)

plot_distribution(
    X_flat,
    Y_flat,
    X_plot,
    f_mean.numpy(),
    f_var.numpy(),
    X,
    y_mean_mean.numpy(),
    y_mean_var.numpy(),
)


# %% [markdown]
# ### Create custom function for the noise variance
#
# We're modelling the `Y_mean` and its standard error is `Y_var / n_repeats`. We will create a simple function that ignores its input and returns these values as a constant. This will obviously only work for the `X` that corresponds to the `Y` we computed the variance of - that is good enough for model training, but notice that it will not allow us to use `predict_y`.

# %%
class FixedVarianceOfMean(gf.functions.Function):
    def __init__(self, Y: gf.base.AnyNDArray):
        n_repeats = Y.shape[-1]
        self.var_mean = np.var(Y, axis=-1, keepdims=True) / n_repeats

    @inherit_check_shapes
    def __call__(self, X: gf.base.TensorType) -> tf.Tensor:
        return self.var_mean


# %% [markdown]
# Now we can plug that into our model:

# %%
model = gf.models.GPR(
    data=(X, Y_mean),
    kernel=gf.kernels.SquaredExponential(active_dims=[0]),
    likelihood=gf.likelihoods.Gaussian(variance=FixedVarianceOfMean(Y)),
)
gf.optimizers.Scipy().minimize(
    model.training_loss, model.trainable_variables, options=optimizer_config
)
print_summary(model)

# %%
f_mean, f_var = model.predict_f(X_plot)
y_mean_mean, y_mean_var = model.predict_y(X)

plot_distribution(
    X_flat,
    Y_flat,
    X_plot,
    f_mean.numpy(),
    f_var.numpy(),
    X,
    y_mean_mean.numpy(),
    y_mean_var.numpy(),
)

# %% [markdown]
# You may notice that the data points fall outside our confidence interval. Again, this is because the confidence interval is for the mean of $Y$, and not the $Y$ points themselves.

# %% [markdown]
# ## Further reading
# To model the variance using a second GP, see the [Heteroskedastic regression notebook](heteroskedastic.ipynb).


# doc/sphinx/notebooks/tailor/external-mean-function.pct.py

# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,.pct.py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.3.3
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Custom mean functions: metalearning with GPs
# One of the advantages of Gaussian process is their flexibility as a modeling tool. For instance, if the modeler knows that there is an underlying trend in the data, they can specify a mean function that captures this trend.
#
# In this notebook, we illustrate how to use GPflow to construct a custom neural network mean function for GPs that can capture complex trends. We look at this functionality in the context of metalearning, where a number of metatasks are available at train time and the user wants to adapt a flexible model to new tasks at test time.
#
# For an in-depth discussion on this topic, see *(Fortuin and Rätsch, 2019)*. This notebook reproduces section 4.2 of this paper.

# %%
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf

import gpflow
from gpflow.ci_utils import reduce_in_tests
from gpflow.kernels import RBF
from gpflow.models import GPR

# for reproducibility of this notebook:
np.random.seed(1)
tf.random.set_seed(2)

# %matplotlib inline

# %% [markdown]
# ## Generate the tasks
# To generate the meta and test tasks, we sample from a Gaussian process with an Squared Exponential covariance function and a sinusoidal mean function. Each task is a realization of this process.


# %%
def generate_data(num_functions=10, N=500):
    """
    For each function, sample the value at `N` equally spaced
    points in the [−5, 5] interval (Fortuin and Rätsch, 2019).

    Returns:
        Tuple of np.arrays of size (N, 1) and (N, num_functions).
    """
    jitter = 1e-6
    Xs = np.linspace(-5.0, 5.0, N)[:, None]
    kernel = RBF(lengthscales=1.0)
    cov = kernel(Xs)
    L = np.linalg.cholesky(cov + np.eye(N) * jitter)
    epsilon = np.random.randn(N, num_functions)
    F = np.sin(Xs) + np.matmul(L, epsilon)
    return Xs, F


# %% [markdown]
# We generate 10 tasks for illustration.

# %%
Xs, F = generate_data(10)

# %%
_ = plt.plot(Xs, F)

# %% [markdown]
# We generate the meta and test tasks.


# %%
def generate_meta_and_test_tasks(num_datapoints, num_meta, num_test, N):
    """Generates meta-task datasets {D_i} = {x_i, y_i} and target task training
    and test data {\tilde{x}, \tilde{y}} (Fortuin and Rätsch, 2019).

    Args:
        num_datapoints: The number of training points, \tilde{n} in Table S1.
        num_meta: The number of meta-tasks ("sampled functions").
        num_test: The number of test tasks ("unseen functions").
        N: Number of sampled data points per function.

    Returns:
        A tuple (meta, test) where
        - meta: List of num_meta pairs of arrays (X, Y) of size (n, 1) each.
        - test: List of num_test pairs of pairs of arrays of sizes
                (((num_datapoints, 1), (num_datapoints, 1)),
                 ((N - num_datapoints, 1), (N - num_datapoints, 1))).
    """
    Xs, F = generate_data(num_functions=num_meta + num_test, N=N)
    meta = []
    sd = 1e-1  # Standard deviation for normally-distributed observation noise.
    for i in range(num_meta):
        # We always use all data points of the curve to train the mean function,
        # i.e. n_i = N.
        noise = sd * np.random.randn(N, 1)
        Y = F[:, i][:, None] + noise
        meta.append((Xs, Y))
    test = []
    rand_indices = np.random.permutation(N)
    train_i = sorted(rand_indices[:num_datapoints])
    test_i = sorted(rand_indices[num_datapoints:])
    for i in range(num_test):
        # Form target training set, \tilde{D}_{train} (see Figure 1).
        noise = sd * np.random.randn(num_datapoints, 1)
        train_X = Xs[train_i]
        train_Y = F[train_i, num_meta + i][:, None] + noise
        # Form target test set, \tilde{D}_{test}, which is disjoint from the
        # target traininig set as it is in the original implementation.
        test_X = Xs[test_i]
        test_Y = F[test_i, num_meta + i][:, None]
        # Form target tasks' datasets \tilde{D} as a pair of pairs,
        # (\tilde{X}, \tilde{Y}) and (\tilde{X}*, \tilde{Y}*).
        test.append(((train_X, train_Y), (test_X, test_Y)))
    return meta, test


# %%
num_datapoints = 5
# Although the original experiment uses (1000, 200, 50) for the following
# parameters, we will use (50, 10, 500) for computational reasons.
num_meta_tasks = reduce_in_tests(50)
num_test_tasks = reduce_in_tests(10)
num_data_per_task = 500
num_iter = reduce_in_tests(5)
meta, test = generate_meta_and_test_tasks(
    num_datapoints, num_meta_tasks, num_test_tasks, num_data_per_task
)

# %% [markdown]
# ## Create the mean function
# We will use a Keras model Deep Neural Network as mean function.

# %%
from gpflow.config import default_float


def build_mean_function():
    tf.keras.backend.set_floatx("float64")
    assert default_float() == np.float64

    inputs = tf.keras.layers.Input(shape=(1,))
    x = tf.keras.layers.Dense(64, activation="relu")(inputs)
    x = tf.keras.layers.Dense(64, activation="relu")(x)
    outputs = tf.keras.layers.Dense(1)(x)
    return tf.keras.Model(inputs=inputs, outputs=outputs)


# %% [markdown]
# ## Build the GP metamodel
# Metalearning boils down to learning a good prior that can generalize to new tasks with a small number of data points. This framework is prevalent in GP modeling, where we usually maximize the marginal likelihood to learn a good set of hyperparameters that specify the GP prior.
#
# We perform the same optimization here, while sharing the hyperparameters across all the metatasks. For simplicity, we fix the kernel and likelihood parameters and learn those only for the mean function. Hence, our metalearning procedure is to cycle through the metatasks continuously, optimizing their marginal likelihood until a convergence criteria is reached (here, we just implement a fixed number of iterations over the tasks).
#
# To begin this process, first we create a utility function that takes in a task (X, Y) and a mean function and outputs a GP model.

# %%
from gpflow import set_trainable


def build_model(data, mean_function):
    model = GPR(data, kernel=RBF(), mean_function=mean_function)
    set_trainable(model.kernel, False)
    model.likelihood.variance.assign(1e-2)
    set_trainable(model.likelihood, False)
    return model


# %%
def create_optimization_step(optimizer, model: gpflow.models.GPR):
    @tf.function
    def optimization_step():
        with tf.GradientTape(watch_accessed_variables=False) as tape:
            tape.watch(model.trainable_variables)
            objective = model.training_loss()
        grads = tape.gradient(objective, model.trainable_variables)
        optimizer.apply_gradients(zip(grads, model.trainable_variables))
        return objective

    return optimization_step


def run_adam(model, iterations):
    """
    Utility function running the Adam optimizer

    :param model: GPflow model
    :param interations: number of iterations
    """
    # Create an Adam Optimizer action
    logf = []
    adam = tf.optimizers.Adam()
    optimization_step = create_optimization_step(adam, model)
    for step in range(iterations):
        loss = optimization_step()
        if step % 10 == 0:
            # print("num variables", len(model.trainable_variables))
            logf.append(loss.numpy())
    return logf


# %% [markdown]
# Next, we define the training loop for metalearning.

# %%
import time


def train_loop(meta_tasks):
    """
    Metalearning training loop. Trained for 100 epochs in original experiment.

    :param meta_tasks: list of metatasks.
    :param num_iter: number of iterations of tasks set
    :returns: a mean function object
    """
    # Initialize mean function
    mean_function = build_mean_function()
    # Iterate for several passes over the tasks set
    for iteration in range(num_iter):
        ts = time.time()
        print("Currently in meta-iteration/epoch {}".format(iteration))
        # Iterate over tasks
        for i, task in enumerate(meta_tasks):
            data = task  # (X, Y)
            model = build_model(data, mean_function=mean_function)
            run_adam(model, reduce_in_tests(100))

        print(">>>> iteration took {:.2f} s".format(time.time() - ts))
    return mean_function


mean_function_optimal = train_loop(meta)

# %% [markdown]
# Finally, we use the optimized mean function for all of the test tasks. **NOTE:** We do not do any further optimization for the hyperparameters in this step.

# %%
test_models = [build_model(data, mean_function_optimal) for (data, _) in test]

# %% [markdown]
# ## Assess the model
# We assess the performance of this procedure on the test tasks. For this, we use the mean squared error as a performance metric.


# %%
def mean_squared_error(y, y_pred):
    return np.mean((y - y_pred) ** 2)


# %%
mean_squared_errors = []
for i, test_task in enumerate(test):
    plt.figure(figsize=(8, 4))
    (train_X, train_Y), (Xs, F) = test_task
    pred_mean, pred_var = test_models[i].predict_f(Xs)
    plt.plot(Xs, pred_mean, label="Prediction mean", color="blue", linewidth=2)
    plt.fill_between(
        Xs.squeeze(1),
        tf.squeeze(pred_mean - pred_var),
        tf.squeeze(pred_mean + pred_var),
        alpha=0.25,
        facecolor="blue",
        label="Prediction variance",
    )
    plt.plot(train_X, train_Y, "ko", label="Training points")
    plt.plot(Xs, F, "ko", label="Ground truth", linewidth=2, markersize=1)
    mse = mean_squared_error(F, pred_mean)
    mean_squared_errors.append(mse)
    plt.title(f"Test Task {i + 1} | MSE = {mse:.2f}")
    plt.legend()
    plt.show()

# %%
mean_mse = np.mean(mean_squared_errors)
std_mse = np.std(mean_squared_errors) / np.sqrt(num_test_tasks)
print(
    f"The mean MSE over all {num_test_tasks} test tasks is {mean_mse:.2f} +/- {std_mse:.2f}"
)

# %% [markdown]
# We achieve comparable results to those reported in the paper.
#
# **NOTE:** We use only 50 metatasks and 10 test tasks over 5 epochs for scalability, whereas the paper uses 1,000 and 200 respectively over 100 epochs. To compensate, we sample 500 points per curve, whereas the paper samples only 50 points. Hence, there might be some discrepancies in the results.

# %% [markdown]
# ## References
# Fortuin, Vincent, and Gunnar Rätsch. "Deep Mean Functions for Meta-Learning in Gaussian Processes." arXiv preprint arXiv:1901.08098 (2019).


# doc/sphinx/notebooks/tailor/gp_nn.pct.py

# ---
# jupyter:
#   jupytext:
#     formats: ipynb,.pct.py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.3.3
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Mixing TensorFlow models with GPflow
#
# This notebook explores the combination of Keras TensorFlow neural networks with GPflow models.

# %%
from typing import Dict, Optional, Tuple

import numpy as np
import tensorflow as tf
import tensorflow_datasets as tfds
from scipy.cluster.vq import kmeans2

import gpflow
from gpflow.ci_utils import reduce_in_tests
from gpflow.utilities import to_default_float

iterations = reduce_in_tests(100)

# %% [markdown]
# ## Convolutional network inside a GPflow model

# %%
original_dataset, info = tfds.load(
    name="mnist", split=tfds.Split.TRAIN, with_info=True
)
total_num_data = info.splits["train"].num_examples
image_shape = info.features["image"].shape
image_size = tf.reduce_prod(image_shape)
batch_size = 32


def map_fn(input_slice: Dict[str, tf.Tensor]):
    updated = input_slice
    image = to_default_float(updated["image"]) / 255.0
    label = to_default_float(updated["label"])
    return tf.reshape(image, [-1, image_size]), label


autotune = tf.data.experimental.AUTOTUNE
dataset = (
    original_dataset.shuffle(1024)
    .batch(batch_size, drop_remainder=True)
    .map(map_fn, num_parallel_calls=autotune)
    .prefetch(autotune)
    .repeat()
)


# %% [markdown]
# Here we'll use the GPflow functionality, but put a non-GPflow model inside the kernel.\
# Vanilla ConvNet. This gets 97.3% accuracy on MNIST when used on its own (+ final linear layer) after 20K iterations

# %%
class KernelWithConvNN(gpflow.kernels.Kernel):
    def __init__(
        self,
        image_shape: Tuple,
        output_dim: int,
        base_kernel: gpflow.kernels.Kernel,
        batch_size: Optional[int] = None,
    ):
        super().__init__()
        with self.name_scope:
            self.base_kernel = base_kernel
            input_size = int(tf.reduce_prod(image_shape))
            input_shape = (input_size,)

            self.cnn = tf.keras.Sequential(
                [
                    tf.keras.layers.InputLayer(
                        input_shape=input_shape, batch_size=batch_size
                    ),
                    tf.keras.layers.Reshape(image_shape),
                    tf.keras.layers.Conv2D(
                        filters=32,
                        kernel_size=image_shape[:-1],
                        padding="same",
                        activation="relu",
                    ),
                    tf.keras.layers.MaxPool2D(pool_size=(2, 2), strides=2),
                    tf.keras.layers.Conv2D(
                        filters=64,
                        kernel_size=(5, 5),
                        padding="same",
                        activation="relu",
                    ),
                    tf.keras.layers.MaxPool2D(pool_size=(2, 2), strides=2),
                    tf.keras.layers.Flatten(),
                    tf.keras.layers.Dense(output_dim, activation="relu"),
                    tf.keras.layers.Lambda(to_default_float),
                ]
            )

            self.cnn.build()

    def K(
        self, a_input: tf.Tensor, b_input: Optional[tf.Tensor] = None
    ) -> tf.Tensor:
        transformed_a = self.cnn(a_input)
        transformed_b = self.cnn(b_input) if b_input is not None else b_input
        return self.base_kernel.K(transformed_a, transformed_b)

    def K_diag(self, a_input: tf.Tensor) -> tf.Tensor:
        transformed_a = self.cnn(a_input)
        return self.base_kernel.K_diag(transformed_a)


# %% [markdown]
# $K_{uf}$ is in ConvNN output space, therefore we need to update `Kuf` multidispatch.

# %%
class KernelSpaceInducingPoints(gpflow.inducing_variables.InducingPoints):
    pass


@gpflow.covariances.Kuu.register(KernelSpaceInducingPoints, KernelWithConvNN)
def Kuu(inducing_variable, kernel, jitter=None):
    func = gpflow.covariances.Kuu.dispatch(
        gpflow.inducing_variables.InducingPoints, gpflow.kernels.Kernel
    )
    return func(inducing_variable, kernel.base_kernel, jitter=jitter)


@gpflow.covariances.Kuf.register(
    KernelSpaceInducingPoints, KernelWithConvNN, object
)
def Kuf(inducing_variable, kernel, a_input):
    return kernel.base_kernel(inducing_variable.Z, kernel.cnn(a_input))


# %% [markdown]
# Now we are ready to create and initialize the model:

# %%
num_mnist_classes = 10
output_dim = 5
num_inducing_points = 100
images_subset, labels_subset = next(iter(dataset.batch(32)))
images_subset = tf.reshape(images_subset, [-1, image_size])
labels_subset = tf.reshape(labels_subset, [-1, 1])

kernel = KernelWithConvNN(
    image_shape,
    output_dim,
    gpflow.kernels.SquaredExponential(),
    batch_size=batch_size,
)

likelihood = gpflow.likelihoods.MultiClass(num_mnist_classes)

inducing_variable_kmeans = kmeans2(
    images_subset.numpy(), num_inducing_points, minit="points"
)[0]
inducing_variable_cnn = kernel.cnn(inducing_variable_kmeans)
inducing_variable = KernelSpaceInducingPoints(inducing_variable_cnn)

model = gpflow.models.SVGP(
    kernel,
    likelihood,
    inducing_variable=inducing_variable,
    num_data=total_num_data,
    num_latent_gps=num_mnist_classes,
)

# %% [markdown]
# And start optimization:

# %%
data_iterator = iter(dataset)
adam_opt = tf.optimizers.Adam(0.001)

training_loss = model.training_loss_closure(data_iterator)


@tf.function
def optimization_step():
    adam_opt.minimize(training_loss, var_list=model.trainable_variables)


for _ in range(iterations):
    optimization_step()

# %% [markdown]
# Let's do predictions after training. Don't expect that we will get a good accuracy, because we haven't run training for long enough.

# %%
m, v = model.predict_y(images_subset)
preds = np.argmax(m, 1).reshape(labels_subset.numpy().shape)
correct = preds == labels_subset.numpy().astype(int)
acc = np.average(correct.astype(float)) * 100.0

print("Accuracy is {:.4f}%".format(acc))


# doc/sphinx/notebooks/tailor/kernel_design.pct.py

# ---
# jupyter:
#   jupytext:
#     formats: ipynb,.pct.py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.3.3
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Kernel Design
#
# It's easy to make new kernels in GPflow. To demonstrate, we'll have a look at the Brownian motion kernel, whose function is
# \begin{equation}
# k(x, x') = \sigma^2 \text{min}(x, x')
# \end{equation}
# where $\sigma^2$ is a variance parameter.

# %%
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf

import gpflow
from gpflow.utilities import positive, print_summary

plt.style.use("ggplot")
# %matplotlib inline

# %% [markdown]
# To make this new kernel class, we inherit from the base class `gpflow.kernels.Kernel` and implement the three functions below. **NOTE:** Depending on the kernel to be implemented, other classes can be more adequate. For example, if the kernel to be implemented is isotropic stationary, you can immediately subclass `gpflow.kernels.IsotropicStationary` (at which point you
# only have to override `K_r` or `K_r2`; see the `IsotropicStationary` class docstring). Stationary but anisotropic kernels should subclass `gpflow.kernels.AnisotropicStationary` and override `K_d`.
#
# #### `__init__`
# In this simple example, the constructor takes no argument (though it could, if that was convenient, for example to pass in an initial value for `variance`). It *must* call the constructor of the superclass with appropriate arguments. Brownian motion is only defined in one dimension, and we'll assume that the `active_dims` are `[0]`, for simplicity.
#
# We've added a parameter to the kernel using the `Parameter` class. Using this class lets the parameter be used in computing the kernel function, and it will automatically be recognised for optimization (or MCMC). Here, the variance parameter is initialized at 1, and constrained to be positive.
#
# #### `K`
# This is where you implement the kernel function itself. This takes two arguments, `X` and `X2`. By convention, we make the second argument optional (it defaults to `None`).
#
# Inside `K`, all the computation must be done with TensorFlow - here we've used `tf.minimum`. When GPflow executes the `K` function, `X` and `X2` will be TensorFlow tensors, and parameters such as `self.variance` behave like TensorFlow tensors as well.
#
# #### `K_diag`
# This convenience function allows GPflow to save memory at predict time. It's simply the diagonal of the `K` function, in the case where `X2` is `None`. It must return a one-dimensional vector, so we use TensorFlow's reshape command.

# %%
class Brownian(gpflow.kernels.Kernel):
    def __init__(self):
        super().__init__(active_dims=[0])
        self.variance = gpflow.Parameter(1.0, transform=positive())

    def K(self, X, X2=None):
        if X2 is None:
            X2 = X
        return self.variance * tf.minimum(
            X, tf.transpose(X2)
        )  # this returns a 2D tensor

    def K_diag(self, X):
        return self.variance * tf.reshape(X, (-1,))  # this returns a 1D tensor


k_brownian = Brownian()
print_summary(k_brownian, fmt="notebook")

# %% [markdown]
# We can now evaluate our new kernel function and draw samples from a Gaussian process with this covariance:

# %%
np.random.seed(23)  # for reproducibility


def plotkernelsample(k, ax, xmin=0, xmax=3):
    xx = np.linspace(xmin, xmax, 300)[:, None]
    K = k(xx)
    ax.plot(xx, np.random.multivariate_normal(np.zeros(300), K, 5).T)
    ax.set_title("Samples " + k.__class__.__name__)


def plotkernelfunction(k, ax, xmin=0, xmax=3, other=0):
    xx = np.linspace(xmin, xmax, 100)[:, None]
    ax.plot(xx, k(xx, np.zeros((1, 1)) + other))
    ax.set_title(k.__class__.__name__ + " k(x, %.1f)" % other)


f, axes = plt.subplots(1, 2, figsize=(12, 4), sharex=True)
plotkernelfunction(k_brownian, axes[0], other=2.0)
plotkernelsample(k_brownian, axes[1])

# %% [markdown]
# ## Using the kernel in a model
#
# Because we've inherited from the `Kernel` base class, this new kernel has all the properties needed to be used in GPflow. It also has some convenience features such as allowing the user to call
#
# `k(X, X2)`
#
# which computes the kernel matrix.
#
# To show that this kernel works, let's use it inside GP regression. We'll see that Brownian motion has quite interesting properties. To add a little flexibility, we'll add a `Constant` kernel to our `Brownian` kernel, and the `GPR` class will handle the noise.

# %%
np.random.seed(42)
X = np.random.rand(5, 1)
Y = np.sin(X * 6) + np.random.randn(*X.shape) * 0.001

k1 = Brownian()
k2 = gpflow.kernels.Constant()
k = k1 + k2

m = gpflow.models.GPR((X, Y), kernel=k)
# m.likelihood.variance.assign(1e-6)

opt = gpflow.optimizers.Scipy()
opt.minimize(m.training_loss, variables=m.trainable_variables)
print_summary(m, fmt="notebook")

xx = np.linspace(0, 1.1, 100).reshape(100, 1)
mean, var = m.predict_y(xx)
plt.plot(X, Y, "kx", mew=2)
(line,) = plt.plot(xx, mean, lw=2)
_ = plt.fill_between(
    xx[:, 0],
    mean[:, 0] - 2 * np.sqrt(var[:, 0]),
    mean[:, 0] + 2 * np.sqrt(var[:, 0]),
    color=line.get_color(),
    alpha=0.2,
)

# %% [markdown]
# ## See also
#
# For more details on how to manipulate existing kernels (or the one you just created!), we refer to the [Manipulating kernels](../advanced/kernels.ipynb) notebook.


# doc/sphinx/notebooks/tailor/mdn_plotting.py

import numpy as np
from scipy.stats import norm


def make_grid(xx, yy):
    """
    Returns two n-by-n matrices. The first one contains all the x values
    and the second all the y values of a cartesian product between `xx` and `yy`.
    """
    n = len(xx)
    xx, yy = np.meshgrid(xx, yy)
    grid = np.array([xx.ravel(), yy.ravel()]).T
    x = grid[:, 0].reshape(n, n)
    y = grid[:, 1].reshape(n, n)
    return x, y


def plot(model, X, Y, axes, cmap, N_plot=100):
    xx = np.linspace(X.min() - 1, X.max() + 1, N_plot)[:, None]
    yy = np.linspace(Y.min() - 1, Y.max() + 1, N_plot)
    pis, mus, sigmas = [v.numpy() for v in model.eval_network(xx)]

    probs = norm.pdf(
        yy[:, None, None], loc=mus[None, :, :], scale=sigmas[None, :, :]
    )
    probs = np.sum(probs * pis[None, :, :], axis=-1)
    plot_x, plot_y = make_grid(xx, yy)
    axes[0].set_title("Posterior density and trainings data")
    axes[0].contourf(
        plot_x, plot_y, np.log(probs), 500, cmap=cmap, vmin=-5, vmax=5
    )
    axes[0].plot(X, Y, "ro", alpha=0.2, ms=3, label="data")
    axes[0].legend(loc=4)
    axes[1].set_title(
        r"$\mu_m(x)$ and their relative contribution shown by size"
    )
    axes[1].scatter(
        np.repeat(xx.flatten(), repeats=mus.shape[1]),
        mus.flatten(),
        s=pis.flatten() * 20,
    )


# doc/sphinx/notebooks/tailor/mixture_density_network.pct.py

# ---
# jupyter:
#   jupytext:
#     formats: ipynb,.pct.py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.14.1
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Mixture Density Networks in GPflow
#
# In this notebook we explain how to implement a Mixture Density Network (MDN) [1] using GPflow. In theory, this is similar to [this blog post](http://blog.otoro.net/2015/11/24/mixture-density-networks-with-tensorflow/) from 2015, but instead of using TensorFlow directly we'll use GPflow. GPflow is typically used for building Gaussian Process-based models, but the framework contains many useful methods and classes that can be used to quickly prototype a wide variety of ML algorithms. Excellent for doing research!
#
# We start by explaining why MDNs can be useful. We then examine a GPflow implementation of the model and use it for a couple of toy experiments.
#

# %% [markdown]
# ## Conditional Density Estimation models
# Imagine we are interested in performing regression on the following dataset.

# %%
import matplotlib.pyplot as plt
import numpy as np

# %matplotlib inline

np.random.seed(1)  # for reproducibility of this notebook

CMAP = plt.get_cmap("Blues")

# %%
N = 200
NOISE_STD = 5.0e-2


def sinusoidal_data(N, noise):
    Y = np.linspace(-2, 2, N)[:, None]
    X = np.sin(4 * Y) * 2.0 + Y * 0.5
    X += np.random.randn(N, 1) * noise
    Y += np.random.randn(N, 1) * noise
    return X, Y


X, Y = data = sinusoidal_data(N, NOISE_STD)

plt.plot(X, Y, "ro", alpha=0.3)
plt.xlabel("$x$")
_ = plt.ylabel("$y$")

# %% [markdown]
# At first sight, this dataset doesn't seem overly complex. Both input and output have a single dimension, and the data has a clear sinusoidal pattern. However, notice that a single input $x$ can correspond to multiple output values $y$, so for example $x=0$ can yield any of the values $\{-1.5, -3/4, 0, 0.8, 1.5\}$. Typical regression algorithms such as Linear Regression, Gaussian Process regression and Multilayer Perceptrons (MLPs) struggle as they can only predict one output value for every input.
#
# To model this dataset we can use a Conditional Density Estimation (CDE) model. CDE models infer $p(f(x)|x)$ instead of just calculating the expectation $E[f(x) | x]$. Modeling the complete distribution $p(f(x)|x)$ is typically harder but it reveals more interesting properties, such as the modes, outlier boundaries and samples. A real-world example might be modeling taxi drop-offs, conditioned on the pick-up location. We would expect a taxi drop-off location to be multi-modal as passengers need to go to different destinations (airport/city center/suburbs and so on) and the density depends on the starting location [2].

# %% [markdown]
# ## Mixture Density Network models
#
# Mixture Density Networks (MDNs) are a parametric class of models that allow for conditional density estimation. They consist of two parts: a neural net and a Mixture of Gaussians (MoG). The neural net is responsible for producing the characteristics of the MoG. In practice, given that the MoG consists of $M$ Gaussians, the neural net will output a collection of $M$ means, variances and weights $\{\mu_m, \sigma_m^2, \pi_m\}_{m=1}^M$. These means, variances and weights are used to define the conditional probability distribution function:
# \begin{equation}
# p(Y = y\,|\,X = x) = \sum_{m=1}^{M} \pi_{m}(x)\,\mathcal{N}\big(y\, \left|\,\mu_{m}(x), \sigma_{m}^2(x)\big)\right.
# \end{equation}
#
# Each of the parameters  $\pi_{m}(x), \mu_{m}(x), \sigma_{m}(x)$ of the distribution are determined by the neural net, as a function of the input $x$.
#
# We train the MDN's neural net by optimizing the model's likelihood:
# \begin{equation}
# \mathcal{L} \triangleq \text{argmax}_{\Theta} \prod_{n=1}^N p(Y = y_n | X = x_n)
# \end{equation}
#
# where $\Theta$ collects the neural net's weights and biases and $\{x_n, y_n\}_{n=1}^N$ represents our training dataset.

# %% [markdown]
# ## A GPflow MDN implementation
#
# GPflow doesn't reinvent the wheel; most of what follows is just plain Python/TensorFlow code. We choose to use GPflow, however, because it provides us with functionality to easily define a model. Once we have a GPflow model, we can specify its objective function, parameters and dataset. This extra layer of abstraction makes interacting with the model much easier, for example optimizing or performing inference.
#
# We begin by importing the required packages from TensorFlow and GPflow.

# %%
from typing import Callable, Optional, Tuple

import tensorflow as tf

from gpflow.base import Parameter
from gpflow.models import BayesianModel, ExternalDataTrainingLossMixin

# %% [markdown]
# Next, we create a `MDN` class that inherits from GPflow's `Model` class. We need to do the following:
# 1. Store each of the feature and target matrices (X, Y) as a `DataHolder` object.
# 2. Define our model's parameters using GPflow's `Parameter` and `ParamList` objects.
# 3. Define the objective function using the `_build_likelihood` method. When we optimize the model the negative of this function will be minimized.

# %%
class MDN(BayesianModel, ExternalDataTrainingLossMixin):
    def __init__(
        self,
        num_mixtures: Optional[int] = 5,
        inner_dims: Optional[list] = [10, 10],
        activation: Optional[
            Callable[[tf.Tensor], tf.Tensor]
        ] = tf.keras.activations.relu,
    ):
        super().__init__()

        # `self.dims` collects the neural net's input, hidden and output dimensions.
        # The number of output dims `self.dims[-1]` equals `num_mixtures` means +
        # `num _mixtures` variances + `num_mixtures` weights, a total of
        # 3 times `num_mixtures` variables.
        self.dims = [1] + list(inner_dims) + [3 * num_mixtures]
        self.activation = activation
        self._create_network()

    def _create_network(self):
        self.Ws, self.bs = [], []

        for dim_in, dim_out in zip(self.dims[:-1], self.dims[1:]):
            init_xavier_std = (2.0 / (dim_in + dim_out)) ** 0.5
            self.Ws.append(
                Parameter(np.random.randn(dim_in, dim_out) * init_xavier_std)
            )
            self.bs.append(Parameter(np.zeros(dim_out)))

    def eval_network(self, X):
        for i, (W, b) in enumerate(zip(self.Ws, self.bs)):
            X = tf.matmul(X, W) + b
            if i < len(self.bs) - 1:
                X = self.activation(X)

        pis, mus, sigmas = tf.split(X, 3, axis=1)
        pis = tf.nn.softmax(pis)  # make sure they normalize to 1
        sigmas = tf.exp(sigmas)  # make sure std. dev. are positive

        return pis, mus, sigmas

    def maximum_log_likelihood_objective(
        self, data: Tuple[tf.Tensor, tf.Tensor]
    ):
        x, y = data
        pis, mus, sigmas = self.eval_network(x)
        Z = (2 * np.pi) ** 0.5 * sigmas
        log_probs_mog = (
            (-0.5 * (mus - y) ** 2 / sigmas ** 2)
            - tf.math.log(Z)
            + tf.math.log(pis)
        )
        log_probs = tf.reduce_logsumexp(log_probs_mog, axis=1)
        return tf.reduce_sum(log_probs)


# %% [markdown]
# ### Notes
# - Given we are dealing with a MoG, the neural net output must comply with the following restrictions:
#   \begin{equation}
#   \sum_{m=1}^{M} \pi_{m}(x) = 1, \pi_m \ge 0\ \text{and}\ \sigma_m\ \forall\ m
#   \end{equation}
#   We achieve this by applying the `softmax` operator to the $\pi$'s and by taking the `exp` to the $\sigma$'s.
#
# - We use the "Xavier" initialization for the neural net's weights. (Glorot and Bengio, 2010).
#
# - Instead of calculating the pdf of the Gaussians, we work with the pdf `log` and use `tf.reduce_logsumexp`. This is mainly for numerical stability.

# %% [markdown]
# ## Experiment 1: The sinusoidal dataset
#
# Let's see how our model works in practice with the sinusoidal dataset presented earlier. We do this by initializing a new instance of our MDN model, and then specifying the dataset $(X, Y)$, the number of hidden units of the MDN's neural net, and the number of mixture components $M$.

# %%
model = MDN(inner_dims=[100, 100], num_mixtures=5)

from gpflow.utilities import print_summary

print_summary(model)

# %% [markdown]
# The objective function for MDN instances is the `maximum_log_likelihood_objective`, which we use for optimization of the parameters. GPflow ensures that only the variables stored in `Parameter` objects are optimized. For the MDN, the only parameters are the weights and the biases of the neural net.
#
# We use the `Scipy` optimizer, which is a wrapper around SciPy's L-BFGS optimization algorithm. Note that GPflow supports other TensorFlow optimizers such as `Adam`, `Adagrad`, and `Adadelta` as well.

# %%
from gpflow.ci_utils import reduce_in_tests
from gpflow.optimizers import Scipy

Scipy().minimize(
    model.training_loss_closure(data, compile=True),
    model.trainable_variables,
    options=dict(maxiter=reduce_in_tests(1500)),
)

print("Final Likelihood", model.maximum_log_likelihood_objective(data).numpy())

# %% [markdown]
# To evaluate the validity of our model, we draw the posterior density. We also plot $\mu(x)$ of the optimized neural net. Remember that for every $x$ the neural net outputs $M$ means $\mu_m(x)$. These determine the location of the Gaussians. We plot all $M$ means and use their corresponding mixture weight $\pi_m(X)$ to determine their size. Larger dots will have more impact in the Gaussian ensemble.

# %%
try:
    from mdn_plotting import plot
except:
    # VS CODE's root directory is GPflow's top-level directory
    from doc.sphinx.notebooks.tailor.mdn_plotting import plot

fig, axes = plt.subplots(1, 2, figsize=(12, 6))
for a in axes:
    a.set_xlim(-4, 4)
    a.set_ylim(-3, 3)
plot(model, X, Y, axes, cmap=CMAP)

# %% [markdown]
# ## Experiment 2: The half moon dataset
#
# The half moon dataset is available in the `scikit-learn` package.

# %%
from sklearn.datasets import make_moons


def moon_data(N, noise):
    data, _ = make_moons(n_samples=N, shuffle=True, noise=noise)
    X, Y = data[:, 0].reshape(-1, 1), data[:, 1].reshape(-1, 1)
    return X, Y


# %%
X, Y = data = moon_data(N, NOISE_STD)
plt.plot(X, Y, "ro", alpha=0.3)
plt.xlabel("$x$")
_ = plt.ylabel("$y$")

# %% [markdown]
# The only difference in the MDN's setup is that we lower the number of mixture components.

# %%
model = MDN(inner_dims=[100, 100], num_mixtures=5)

# %%
Scipy().minimize(
    model.training_loss_closure(data, compile=True),
    model.trainable_variables,
    options=dict(maxiter=reduce_in_tests(int(10e3))),
)

print("Final Likelihood", model.maximum_log_likelihood_objective(data).numpy())

# %%
fig, axes = plt.subplots(1, 2, figsize=(12, 6))
for a in axes:
    a.set_xlim(-2, 3)
    a.set_ylim(-1.5, 2)
plot(model, X, Y, axes, cmap=CMAP)

# %% [markdown]
# ## References
#
# [1] Bishop, Christopher M. Mixture density networks. Technical Report NCRG/4288, Aston University, Birmingham, UK, 1994.
#
# [2] Dutordoir, Vincent, et al. "Gaussian Process Conditional Density Estimation." Advances in Neural Information Processing Systems. 2018.


# doc/sphinx/notebooks/theory/FITCvsVFE.pct.py

# ---
# jupyter:
#   jupytext:
#     formats: ipynb,.pct.py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.3.3
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Comparing FITC approximation to VFE approximation
# This notebook examines why we prefer the Variational Free Energy (VFE) objective to the Fully Independent Training Conditional (FITC) approximation for our sparse approximations.
#
# %%
import matplotlib.pyplot as plt
from FITCvsVFE import (
    getTrainingTestData,
    plotComparisonFigure,
    plotPredictions,
    printModelParameters,
    repeatMinimization,
    stretch,
)

import gpflow
from gpflow.ci_utils import reduce_in_tests

# %matplotlib inline


# logging.disable(logging.WARN)  # do not clutter up the notebook with optimization warnings

# %% [markdown]
# First, we load the training data and plot it together with the exact GP solution (using the `GPR` model):

# %%
# Load the training data:
Xtrain, Ytrain, Xtest, Ytest = getTrainingTestData()


def getKernel():
    return gpflow.kernels.SquaredExponential()


# Run exact inference on training data:
exact_model = gpflow.models.GPR((Xtrain, Ytrain), kernel=getKernel())

opt = gpflow.optimizers.Scipy()
opt.minimize(
    exact_model.training_loss,
    exact_model.trainable_variables,
    method="L-BFGS-B",
    options=dict(maxiter=reduce_in_tests(20000)),
    tol=1e-11,
)

print("Exact model parameters:")
printModelParameters(exact_model)

figA, ax = plt.subplots(1, 1)
ax.plot(Xtrain, Ytrain, "ro")
plotPredictions(ax, exact_model, color="g")


# %%
def initializeHyperparametersFromExactSolution(sparse_model):
    sparse_model.likelihood.variance.assign(exact_model.likelihood.variance)
    sparse_model.kernel.variance.assign(exact_model.kernel.variance)
    sparse_model.kernel.lengthscales.assign(exact_model.kernel.lengthscales)


# %% [markdown]
# We now construct two sparse model using the VFE (`SGPR` model) and FITC (`GPRFITC` model) optimization objectives, with the inducing points being initialized on top of the training inputs, and the model hyperparameters (kernel variance and lengthscales, and likelihood variance) being initialized to the values obtained in the optimization of the exact `GPR` model:

# %%
# Train VFE model initialized from the perfect solution.
VFEmodel = gpflow.models.SGPR(
    (Xtrain, Ytrain), kernel=getKernel(), inducing_variable=Xtrain.copy()
)

initializeHyperparametersFromExactSolution(VFEmodel)

VFEcb = repeatMinimization(
    VFEmodel, Xtest, Ytest
)  # optimize with several restarts
print("Sparse model parameters after VFE optimization:")
printModelParameters(VFEmodel)

# %%
# Train FITC model initialized from the perfect solution.
FITCmodel = gpflow.models.GPRFITC(
    (Xtrain, Ytrain), kernel=getKernel(), inducing_variable=Xtrain.copy()
)

initializeHyperparametersFromExactSolution(FITCmodel)

FITCcb = repeatMinimization(
    FITCmodel, Xtest, Ytest
)  # optimize with several restarts
print("Sparse model parameters after FITC optimization:")
printModelParameters(FITCmodel)

# %% [markdown]
# Plotting a comparison of the two algorithms, we see that VFE stays at the optimum of exact GPR, whereas the FITC approximation eventually ends up with several inducing points on top of each other, and a worse fit:

# %%
figB, axes = plt.subplots(3, 2, figsize=(20, 16))

# VFE optimization finishes after 10 iterations, so we stretch out the training and test
# log-likelihood traces to make them comparable against FITC:
VFEiters = FITCcb.n_iters
VFElog_likelihoods = stretch(len(VFEiters), VFEcb.log_likelihoods)
VFEhold_out_likelihood = stretch(len(VFEiters), VFEcb.hold_out_likelihood)

axes[0, 0].set_title("VFE", loc="center", fontdict={"fontsize": 22})
plotComparisonFigure(
    Xtrain,
    VFEmodel,
    exact_model,
    axes[0, 0],
    axes[1, 0],
    axes[2, 0],
    VFEiters,
    VFElog_likelihoods,
    VFEhold_out_likelihood,
)

axes[0, 1].set_title("FITC", loc="center", fontdict={"fontsize": 22})
plotComparisonFigure(
    Xtrain,
    FITCmodel,
    exact_model,
    axes[0, 1],
    axes[1, 1],
    axes[2, 1],
    FITCcb.n_iters,
    FITCcb.log_likelihoods,
    FITCcb.hold_out_likelihood,
)

# %% [markdown]
# A more detailed discussion of the comparison between these sparse approximations can be found in [Understanding Probabilistic Sparse Gaussian Process Approximations](http://papers.nips.cc/paper/6477-understanding-probabilistic-sparse-gaussian-process-approximations) by Bauer, van der Wilk, and Rasmussen (2017).


# doc/sphinx/notebooks/theory/FITCvsVFE.py

import numpy as np
import tensorflow as tf

import gpflow
from gpflow.ci_utils import reduce_in_tests

nRepeats = reduce_in_tests(50)

predict_limits = [-4.0, 4.0]
inducing_points_limits = [-1.0, 9]
hold_out_limits = [0.20, 0.60]
optimization_limits = [18.0, 25.0]


def readCsvFile(fileName):
    return np.loadtxt(fileName).reshape(-1, 1)


def getTrainingTestData():
    overallX = readCsvFile("data/snelson_train_inputs.dat")
    overallY = readCsvFile("data/snelson_train_outputs.dat")

    trainIndices = []
    testIndices = []

    nPoints = overallX.shape[0]

    for index in range(nPoints):
        if index % 4 == 0:
            trainIndices.append(index)
        else:
            testIndices.append(index)

    Xtrain = overallX[trainIndices, :]
    Xtest = overallX[testIndices, :]
    Ytrain = overallY[trainIndices, :]
    Ytest = overallY[testIndices, :]

    return Xtrain, Ytrain, Xtest, Ytest


def getLogPredictiveDensities(targetValues, means, variances):
    assert targetValues.shape == means.shape
    assert variances.shape == means.shape

    deltas = targetValues - means
    mahalanobisTerms = -0.5 * deltas ** 2 / variances
    normalizationTerms = -0.5 * np.log(variances) - 0.5 * np.log(2.0 * np.pi)
    return mahalanobisTerms + normalizationTerms


def getKernel():
    return gpflow.kernels.SquaredExponential()


def getRegressionModel(X, Y):
    m = gpflow.models.GPR(X, Y, kern=getKernel())
    m.likelihood.variance = 1.0
    m.kern.lengthscales = 1.0
    m.kern.variance = 1.0
    return m


def getSparseModel(X, Y, isFITC=False):
    if isFITC:
        m = gpflow.models.GPRFITC(
            (X, Y), kernel=getKernel(), inducing_variable=X.copy()
        )
    else:
        m = gpflow.models.SGPR(
            (X, Y), kernel=getKernel(), inducing_variable=X.copy()
        )
    return m


def printModelParameters(model):
    print(
        "  Likelihood variance = {:.5g}".format(
            model.likelihood.variance.numpy()
        )
    )
    print(
        "  Kernel variance     = {:.5g}".format(model.kernel.variance.numpy())
    )
    print(
        "  Kernel lengthscale  = {:.5g}".format(
            model.kernel.lengthscales.numpy()
        )
    )


def plotPredictions(ax, model, color, label=None):
    Xtest = np.sort(readCsvFile("data/snelson_test_inputs.dat"))
    predMean, predVar = model.predict_y(Xtest)
    ax.plot(Xtest, predMean, color, label=label)
    ax.plot(Xtest, predMean + 2.0 * np.sqrt(predVar), color)
    ax.plot(Xtest, predMean - 2.0 * np.sqrt(predVar), color)


def repeatMinimization(model, Xtest, Ytest):
    callback = Callback(model, Xtest, Ytest)

    opt = gpflow.optimizers.Scipy()
    # print("Optimising for {} repetitions".format(nRepeats))
    for repeatIndex in range(nRepeats):
        # print(repeatIndex)
        opt.minimize(
            model.training_loss,
            model.trainable_variables,
            method="L-BFGS-B",
            tol=1e-11,
            options=dict(disp=False, maxiter=reduce_in_tests(2000)),
            step_callback=callback,
            compile=True,
        )
    return callback


def trainSparseModel(Xtrain, Ytrain, exact_model, isFITC, Xtest, Ytest):
    sparse_model = getSparseModel(Xtrain, Ytrain, isFITC)
    sparse_model.likelihood.variance = exact_model.likelihood.variance.numpy()
    sparse_model.kern.lengthscales = exact_model.kern.lengthscales.numpy()
    sparse_model.kern.variance = exact_model.kern.variance.numpy()
    return sparse_model, repeatMinimization(sparse_model, Xtest, Ytest)


def plotComparisonFigure(
    Xtrain,
    sparse_model,
    exact_model,
    ax_predictions,
    ax_inducing_points,
    ax_optimization,
    iterations,
    log_likelihoods,
    hold_out_likelihood,
    title=None,
):
    plotPredictions(ax_predictions, exact_model, "g", label="Exact")
    plotPredictions(ax_predictions, sparse_model, "b", label="Approximate")
    ax_predictions.legend(loc=9)
    ax_predictions.plot(
        sparse_model.inducing_variable.Z.numpy(),
        -1.0 * np.ones(Xtrain.shape),
        "ko",
    )
    ax_predictions.set_ylim(predict_limits)
    ax_inducing_points.plot(
        Xtrain, sparse_model.inducing_variable.Z.numpy(), "bo"
    )
    xs = np.linspace(
        ax_inducing_points.get_xlim()[0], ax_inducing_points.get_xlim()[1], 200
    )
    ax_inducing_points.plot(xs, xs, "g")
    ax_inducing_points.set_xlabel("Optimal inducing point position")
    ax_inducing_points.set_ylabel("Learnt inducing point position")
    ax_inducing_points.set_ylim(inducing_points_limits)
    ax_optimization.plot(iterations, -1.0 * np.array(log_likelihoods), "g-")
    ax_optimization.set_ylim(optimization_limits)
    ax2 = ax_optimization.twinx()
    ax2.plot(iterations, -1.0 * np.array(hold_out_likelihood), "b-")
    ax_optimization.set_xlabel("Minimization iterations")
    ax_optimization.set_ylabel("Minimization objective", color="g")
    ax2.set_ylim(hold_out_limits)
    ax2.set_ylabel("Hold out negative log likelihood", color="b")


class Callback:
    def __init__(self, model, Xtest, Ytest, holdout_interval=10):
        self.model = model
        self.holdout_interval = holdout_interval
        self.Xtest = Xtest
        self.Ytest = Ytest
        self.log_likelihoods = []
        self.hold_out_likelihood = []
        self.n_iters = []
        self.counter = 0

    def __call__(self, step, variables, values):
        # step will reset to zero between calls to minimize(), whereas counter will keep increasing

        if (self.counter <= 10) or (self.counter % self.holdout_interval) == 0:
            for var, val in zip(variables, values):
                var.assign(val)

            self.n_iters.append(self.counter)
            self.log_likelihoods.append(
                self.model.maximum_log_likelihood_objective().numpy()
            )

            predictive_mean, predictive_variance = self.model.predict_y(
                self.Xtest
            )
            test_log_likelihood = tf.reduce_mean(
                getLogPredictiveDensities(
                    self.Ytest, predictive_mean, predictive_variance
                )
            )
            self.hold_out_likelihood.append(test_log_likelihood.numpy())

        self.counter += 1


def stretch(lenNIters, initialValues):
    stretched = np.ones(lenNIters) * initialValues[-1]
    stretched[0 : len(initialValues)] = initialValues
    return stretched


def snelsonDemo():
    from IPython import embed
    from matplotlib import pyplot as plt

    Xtrain, Ytrain, Xtest, Ytest = getTrainingTestData()

    # run exact inference on training data.
    exact_model = getRegressionModel(Xtrain, Ytrain)
    opt = gpflow.train.ScipyOptimizer()
    opt.minimize(exact_model, maxiter=reduce_in_tests(2000000))

    figA, axes = plt.subplots(1, 1)
    inds = np.argsort(Xtrain.flatten())
    axes.plot(Xtrain[inds, :], Ytrain[inds, :], "ro")
    plotPredictions(axes, exact_model, "g", None)

    figB, axes = plt.subplots(3, 2)

    # run sparse model on training data initialized from exact optimal solution.
    VFEmodel, VFEcb = trainSparseModel(
        Xtrain, Ytrain, exact_model, False, Xtest, Ytest
    )
    FITCmodel, FITCcb = trainSparseModel(
        Xtrain, Ytrain, exact_model, True, Xtest, Ytest
    )

    print("Exact model parameters \n")
    printModelParameters(exact_model)
    print("Sparse model parameters for VFE optimization \n")
    printModelParameters(VFEmodel)
    print("Sparse model parameters for FITC optimization \n")
    printModelParameters(FITCmodel)

    VFEiters = FITCcb.n_iters
    VFElog_likelihoods = stretch(len(VFEiters), VFEcb.log_likelihoods)
    VFEhold_out_likelihood = stretch(len(VFEiters), VFEcb.hold_out_likelihood)

    plotComparisonFigure(
        Xtrain,
        VFEmodel,
        exact_model,
        axes[0, 0],
        axes[1, 0],
        axes[2, 0],
        VFEiters,
        VFElog_likelihoods,
        VFEhold_out_likelihood,
        "VFE",
    )
    plotComparisonFigure(
        Xtrain,
        FITCmodel,
        exact_model,
        axes[0, 1],
        axes[1, 1],
        axes[2, 1],
        FITCcb.n_iters,
        FITCcb.log_likelihoods,
        FITCcb.hold_out_likelihood,
        "FITC",
    )

    axes[0, 0].set_title("VFE", loc="center", fontdict={"fontsize": 22})
    axes[0, 1].set_title("FITC", loc="center", fontdict={"fontsize": 22})

    embed()


if __name__ == "__main__":
    snelsonDemo()


# doc/sphinx/notebooks/theory/Sanity_check.pct.py

# ---
# jupyter:
#   jupytext:
#     formats: ipynb,.pct.py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.3.3
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %% [markdown]
# ## Sanity checking when model behaviours should overlap
#
# Many of the model classes in GPflow have overlapping behaviour in special cases. In this notebook, we fit some approximations to a model with a Gaussian likelihood, and make sure they're all the same.
#
# The models are:
#  - `GPR`: Full Gaussian process regression.
#
#  - `VGP`: A Gaussian approximation with Variational Bayes.
#    Approximating a Gaussian posterior with a Gaussian should be exact.
#
#  - `SVGP`: a sparse GP, with a Gaussian approximation. The inducing points are set to be at the data points, so again, should be exact.
#
#  - `SVGP` (with whitened representation): As above, but with a rotation applied to whiten the representation of the process.
#
#  - `SGPR`: A sparse GP with a *collapsed* posterior (Titsias 2009). Again, the inducing points are fixed to the data points.
#
#  - `GPRFITC`: The FITC approximation. Again, the inducing points are fixed to the data points.
#
# In all cases the parameters are estimated by the method of maximum likelihood (or approximate maximum likelihood, as appropriate). The parameter estimates should all be the same.

# %%
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf

import gpflow
from gpflow import set_trainable
from gpflow.ci_utils import reduce_in_tests
from gpflow.config import default_float
from gpflow.models import (
    maximum_log_likelihood_objective,
    training_loss_closure,
)

# %matplotlib inline
plt.rcParams["figure.figsize"] = (12, 6)

# %%
np.random.seed(0)
X = np.random.rand(20, 1) * 10
Y = np.sin(X) + 0.9 * np.cos(X * 1.6) + np.random.randn(*X.shape) * 0.4
Xtest = np.random.rand(10, 1) * 10
_ = plt.plot(X, Y, "kx", mew=2)

# %%
data = (
    tf.convert_to_tensor(X, dtype=default_float()),
    tf.convert_to_tensor(Y, dtype=default_float()),
)
inducing_variable = tf.convert_to_tensor(X, dtype=default_float())

m1 = gpflow.models.GPR(data, kernel=gpflow.kernels.SquaredExponential())
m2 = gpflow.models.VGP(
    data,
    kernel=gpflow.kernels.SquaredExponential(),
    likelihood=gpflow.likelihoods.Gaussian(),
)
m3 = gpflow.models.SVGP(
    gpflow.kernels.SquaredExponential(),
    gpflow.likelihoods.Gaussian(),
    inducing_variable,
    q_diag=False,
)
set_trainable(m3.inducing_variable, False)

m4 = gpflow.models.SVGP(
    gpflow.kernels.SquaredExponential(),
    gpflow.likelihoods.Gaussian(),
    inducing_variable,
    q_diag=False,
    whiten=True,
)
set_trainable(m4.inducing_variable, False)

m5 = gpflow.models.SGPR(
    data,
    kernel=gpflow.kernels.SquaredExponential(),
    inducing_variable=inducing_variable,
)
set_trainable(m5.inducing_variable, False)

m6 = gpflow.models.GPRFITC(
    data,
    kernel=gpflow.kernels.SquaredExponential(),
    inducing_variable=inducing_variable,
)
set_trainable(m6.inducing_variable, False)

models = [m1, m2, m3, m4, m5, m6]

# %% [markdown]
# Now, we optimize the models. For `GPR`, `SVGP`, and `GPRFITC`, this simply optimizes the hyperparameters (since the inducing points are fixed). For the variational models, this jointly maximises the lower bound to the marginal likelihood (Evidence Lower Bound, ELBO) with respect to the variational parameters and the kernel and likelihood hyperparameters.

# %%
for m in models:
    opt = gpflow.optimizers.Scipy()
    loss_closure = training_loss_closure(m, data)
    opt.minimize(
        loss_closure,
        variables=m.trainable_variables,
        options=dict(maxiter=reduce_in_tests(1000)),
        compile=True,
    )


# %% [markdown]
# If everything worked as planned, the models should have the same:
#
#  - prediction functions
#  - log (marginal) likelihood
#  - kernel parameters
#
# For the variational models, where we use a ELBO in place of the likelihood, the ELBO should be tight to the likelihood in the cases studied here.

# %%
def plot(m, color, ax):
    xx = np.linspace(-1, 11, 100)[:, None]
    mu, var = m.predict_y(xx)
    ax.plot(xx, mu, color, lw=2)
    ax.fill_between(
        xx[:, 0],
        mu[:, 0] - 2 * np.sqrt(var[:, 0]),
        mu[:, 0] + 2 * np.sqrt(var[:, 0]),
        color=color,
        alpha=0.2,
    )
    ax.plot(X, Y, "kx", mew=2)
    ax.set_xlim(-1, 11)


f, ax = plt.subplots(3, 2, sharex=True, sharey=True, figsize=(12, 9))
plot(m1, "C0", ax[0, 0])
plot(m2, "C1", ax[1, 0])
plot(m3, "C2", ax[0, 1])
plot(m4, "C3", ax[1, 1])
plot(m5, "C4", ax[2, 0])
plot(m6, "C5", ax[2, 1])

# %% [markdown]
# Here are the kernels and likelihoods, which show the fitted kernel parameters and noise variance:

# %%
for m in models:
    print(m.__class__.__name__)
    print(f"  kernel lengthscale  = {m.kernel.lengthscales.numpy():.5g}")
    print(f"  kernel variance     = {m.kernel.variance.numpy():.5}")
    print(f"  likelihood variance = {m.likelihood.variance.numpy():.5}")

# %% [markdown]
# Here are the likelihoods (or ELBOs):

# %%
for m in models:
    print(
        f"{m.__class__.__name__:30}  {maximum_log_likelihood_objective(m, data)}"
    )


# doc/sphinx/notebooks/theory/cglb.pct.py

# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,.pct.py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.3.3
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %%
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
from cglb import load_snelson_data, plot_prediction

from gpflow.kernels import SquaredExponential
from gpflow.models import CGLB, GPR, SGPR
from gpflow.optimizers import Scipy

# %%
# %matplotlib inline

# %% [markdown]
# # Conjugate Gradient Lower Bound
#
# First, we load the `snelson1d` training dataset.

# %%
data = load_snelson_data()

# %% [markdown]
# The CGLB model introduces less bias in comparison to SGPR model.
# We can show empirically that CGLB has a lower bias by plotting the objective landscape with respect to different values of the lengthscale hyperparameters.

# %%

x, y = data
n = x.shape[0]
m = 10

iv_indices = np.random.choice(range(n), size=m, replace=False)
iv = x[iv_indices, :]

noise = 0.1
gpr = GPR(data, kernel=SquaredExponential(), noise_variance=noise)
cglb = CGLB(
    data,
    kernel=SquaredExponential(),
    noise_variance=noise,
    inducing_variable=iv,
)
sgpr = SGPR(
    data,
    kernel=SquaredExponential(),
    noise_variance=noise,
    inducing_variable=iv,
)


def loss_with_changed_parameter(model, parameter, value: float):
    original = parameter.numpy()
    parameter.assign(value)
    loss = model.training_loss()
    parameter.assign(original)
    return loss.numpy()


ls = np.linspace(0.01, 3, 100)
losses_fn = lambda m: [
    loss_with_changed_parameter(m, m.kernel.lengthscales, l) for l in ls
]
gpr_obj = losses_fn(gpr)
sgpr_obj = losses_fn(sgpr)
cglb_obj = losses_fn(cglb)

plt.plot(ls, gpr_obj, label="GPR")
plt.plot(ls, sgpr_obj, label="SGPR")
plt.plot(ls, cglb_obj, label="CGLB")
plt.xlabel("lengthscale")
plt.ylabel("Objective")
plt.legend()
plt.show()


# %% [markdown]
# The interface of CGLB model is similar to SGPR, however, it has extra options related to the conjugate gradient which is involved in the computation of $ yᵀK⁻¹y $.
# * `cg_tolerance`. This is the expected maximum error in the solution produced by the conjugate gradient,
# i.e. the ϵ threshold of the stopping criteria in the CG loop: $ rᵀQ⁻¹r < 0.5ϵ $.
# The default value for `cg_tolerance` is $ 1.0 $. It is a good working value for model training, which has been confirmed in practice. The `predict_f`, `predict_log_density`, and `predict_y` methods have their own `cg_tolerance` options as for predictions you might want to tighten up the CG solution closer to exact.
# * `max_cg_iters`. The maximum number of CG iterations.
# * `restart_cg_step`. The frequency with wich the CG resets the internal state to the initial position using current solution vector `v`.
# * `v_grad_optimization`. CGLB introduces auxiliary parameter `v`, and by default optimal `v` is found with the CG. However you can include `v` into the list of trainable model parameters.

# %%

cglb = CGLB(
    data,
    kernel=SquaredExponential(),
    noise_variance=noise,
    inducing_variable=iv,
    cg_tolerance=1.0,
    max_cg_iters=50,
    restart_cg_iters=50,
)
opt = Scipy()

# %% [markdown]
# We train the model as usual. Variables do not include the $ v $ auxiliary vector.

# %%
variables = cglb.trainable_variables
_ = opt.minimize(
    cglb.training_loss_closure(compile=False),
    variables,
    compile=False,
    options=dict(maxiter=100),
)

# %% [markdown]
# Below we compare prediction results for different CG tolerances. The `cg_tolerance=None` means that no CG is run to tune the $ v $ vector, and `cg_tolerance=0.01` is much lower value than the one used at the model optimization.

# %% [markdown]
# The 1D dataset which we use is too simple and it would be hard to see the difference in mean prediction of the model. As a demonstration we reset the $ v $ auxiliary vector to zero, so that predictions are identical to SGPR.

# %%
v = tf.zeros_like(cglb.aux_vec)
cglb.aux_vec.assign(v)

# %%
k = 100
xnew = np.linspace(x.min() - 1.0, x.max() + 1.0, k).reshape(-1, 1)

pred_no_tol = cglb.predict_y(xnew, cg_tolerance=None)
pred_tol = cglb.predict_y(xnew, cg_tolerance=0.01)

fig, axes = plt.subplots(2, 1, figsize=(7, 5), sharex=True)
fig.subplots_adjust(hspace=0.0)

x, y = x.squeeze(), y.squeeze()
xnew = xnew.squeeze()

mu_no_tol = pred_no_tol[0].numpy().squeeze()
std_no_tol = np.sqrt(pred_no_tol[1].numpy().squeeze())
plot_prediction(
    axes[0],
    x,
    y,
    xnew,
    mu_no_tol,
    std_no_tol,
    "tab:blue",
    "CGLB, use cached $v=0$",
)


mu_tol = pred_tol[0].numpy().squeeze()
std_tol = np.sqrt(pred_tol[1].numpy().squeeze())
plot_prediction(
    axes[1], x, y, xnew, mu_tol, std_tol, "tab:green", "CGLB, tol=0.01"
)

axes[0].legend()
axes[1].legend()
axes[0].set_xlim([-1, 7])
axes[1].set_xlim([-1, 7])
axes[0].set_ylim([-2.5, 2.5])
axes[1].set_ylim([-2.5, 2.5])
axes[0].set_ylabel("y")
axes[1].set_ylabel("y")
axes[1].set_xlabel("x")
plt.show()


# doc/sphinx/notebooks/theory/cglb.py

from pathlib import Path

import numpy as np


def load_snelson_data():
    curdir = Path(__file__).parent.resolve()
    train_inputs = "snelson_train_inputs.dat"
    train_outputs = "snelson_train_outputs.dat"
    datapath = lambda name: str(Path(curdir, "data", name).resolve())
    xfile = datapath(train_inputs)
    yfile = datapath(train_outputs)
    x = np.loadtxt(xfile).reshape(-1, 1)
    y = np.loadtxt(yfile).reshape(-1, 1)
    return (x, y)


def plot_prediction(ax, x, y, xnew, loc, scale, color, label):
    for k in (1, 2):
        lb = (loc - k * scale).squeeze()
        ub = (loc + k * scale).squeeze()
        alpha = 1 - 0.4 * k
        ax.fill_between(xnew, lb, ub, color=color, alpha=alpha)

    ax.plot(xnew, loc, color=color, label=label)
    ax.scatter(x, y, color="gray", s=8, alpha=0.7)


# doc/sphinx/notebooks/theory/upper_bound.pct.py

# ---
# jupyter:
#   jupytext:
#     formats: ipynb,.pct.py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.3.3
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Discussion of the GP marginal likelihood upper bound
#
# See the [`gp_upper` repository](https://github.com/markvdw/gp_upper) by Mark van der Wilk for code to tighten the upper bound through optimization, and a more comprehensive discussion.

# %%
import matplotlib.pyplot as plt

# %matplotlib inline
plt.rcParams["figure.figsize"] = (12, 6)

import logging

import numpy as np
import tensorflow as tf

import gpflow
from gpflow import set_trainable
from gpflow.ci_utils import reduce_in_tests
from gpflow.utilities import print_summary

logging.disable(logging.WARNING)

np.random.seed(1)

from FITCvsVFE import getTrainingTestData

# %%
X, Y, Xt, Yt = getTrainingTestData()


# %%
def plot_model(m, name=""):
    pX = np.linspace(-3, 9, 100)[:, None]
    pY, pYv = m.predict_y(pX)
    plt.plot(X, Y, "x")
    plt.plot(pX, pY)
    if not isinstance(m, gpflow.models.GPR):
        Z = m.inducing_variable.Z.numpy()
        plt.plot(Z, np.zeros_like(Z), "o")
    two_sigma = (2.0 * pYv ** 0.5)[:, 0]
    plt.fill_between(
        pX[:, 0], pY[:, 0] - two_sigma, pY[:, 0] + two_sigma, alpha=0.15
    )
    lml = m.maximum_log_likelihood_objective().numpy()
    plt.title("%s (lml = %f)" % (name, lml))
    return lml


# %% [markdown]
# ## Full model

# %%
gpr = gpflow.models.GPR((X, Y), gpflow.kernels.SquaredExponential())
gpflow.optimizers.Scipy().minimize(
    gpr.training_loss,
    gpr.trainable_variables,
    options=dict(maxiter=reduce_in_tests(1000)),
)
full_lml = plot_model(gpr)

# %% [markdown]
# ## Upper bounds for sparse variational models
# As a first investigation, we compute the upper bound for models trained using the sparse variational GP approximation.

# %%
Ms = np.arange(4, reduce_in_tests(20, test_n=6), 1)
vfe_lml = []
vupper_lml = []
vfe_hyps = []
for M in Ms:
    Zinit = X[:M, :].copy()
    vfe = gpflow.models.SGPR(
        (X, Y), gpflow.kernels.SquaredExponential(), inducing_variable=Zinit
    )
    gpflow.optimizers.Scipy().minimize(
        vfe.training_loss,
        vfe.trainable_variables,
        options=dict(disp=False, maxiter=reduce_in_tests(1000), compile=True),
    )

    vfe_lml.append(vfe.elbo().numpy())
    vupper_lml.append(vfe.upper_bound().numpy())
    vfe_hyps.append([(p.name, p.numpy()) for p in vfe.trainable_parameters])
    print("%i" % M, end=" ")

# %%
plt.title("LML bounds for models trained with SGPR")
plt.plot(Ms, vfe_lml, label="lower")
plt.plot(Ms, vupper_lml, label="upper")
plt.axhline(full_lml, label="full", alpha=0.3)
plt.xlabel("Number of inducing points")
plt.ylabel("LML estimate")
_ = plt.legend()

# %% [markdown]
# We see that the lower bound increases as more inducing points are added. Note that the upper bound does _not_ monotonically decrease! This is because as we train the sparse model, we also get better estimates of the hyperparameters. The upper bound will be different for this different setting of the hyperparameters, and is sometimes looser. The upper bound also converges to the true lml slower than the lower bound.

# %% [markdown]
# ### Upper bounds for fixed hyperparameters
# Here, we train sparse models with the hyperparameters fixed to the optimal value found previously.

# %%
fMs = np.arange(3, reduce_in_tests(20, test_n=5), 1)
fvfe_lml = []  # Fixed vfe lml
fvupper_lml = []  # Fixed upper lml

init_params = gpflow.utilities.parameter_dict(vfe)

# cannot copy this due to shape mismatch with different numbers of inducing points between models:
del init_params[".inducing_variable.Z"]

for M in fMs:
    Zinit = vfe.inducing_variable.Z.numpy()[:M, :]
    Zinit = np.vstack(
        (Zinit, X[np.random.permutation(len(X))[: (M - len(Zinit))], :].copy())
    )

    vfe = gpflow.models.SGPR(
        (X, Y), gpflow.kernels.SquaredExponential(), inducing_variable=Zinit
    )

    # copy hyperparameters (omitting inducing_variable.Z) from optimized model:
    gpflow.utilities.multiple_assign(vfe, init_params)

    set_trainable(vfe.kernel, False)
    set_trainable(vfe.likelihood, False)

    gpflow.optimizers.Scipy().minimize(
        vfe.training_loss,
        vfe.trainable_variables,
        options=dict(disp=False, maxiter=reduce_in_tests(1000)),
        compile=True,
    )

    fvfe_lml.append(vfe.elbo().numpy())
    fvupper_lml.append(vfe.upper_bound().numpy())
    print("%i" % M, end=" ")

# %%
plt.plot(fMs, fvfe_lml, label="lower")
plt.plot(fMs, fvupper_lml, label="upper")
plt.axhline(full_lml, label="full", alpha=0.3)
plt.xlabel("Number of inducing points")
plt.ylabel("LML estimate")
_ = plt.legend()

# %%
assert np.all(np.array(fvupper_lml) - np.array(fvfe_lml) > 0.0)

# %% [markdown]
# Now, as the hyperparameters are fixed, the bound _does_ monotonically decrease. We chose the optimal hyperparameters here, but the picture should be the same for any hyperparameter setting. This shows that we increasingly get a better estimate of the marginal likelihood as we add more inducing points.

# %% [markdown]
# ### A tight estimate bound does not imply a converged model

# %%
single_inducing_point = X[:1, :].copy()
vfe = gpflow.models.SGPR(
    (X, Y),
    gpflow.kernels.SquaredExponential(),
    inducing_variable=single_inducing_point,
)
objective = tf.function(vfe.training_loss)
gpflow.optimizers.Scipy().minimize(
    objective,
    vfe.trainable_variables,
    options=dict(maxiter=reduce_in_tests(1000)),
    compile=False,
)
# Note that we need to set compile=False here due to a discrepancy in compiling with tf.function
# see https://github.com/GPflow/GPflow/issues/1260

print("Lower bound: %f" % vfe.elbo().numpy())
print("Upper bound: %f" % vfe.upper_bound().numpy())

# %% [markdown]
# In this case we show that for the hyperparameter setting, the bound is very tight. However, this does _not_ imply that we have enough inducing points, but simply that we have correctly identified the marginal likelihood for this particular hyperparameter setting. In this specific case, where we used a single inducing point, the model collapses to not using the GP at all (lengthscale is really long to model only the mean). The rest of the variance is explained by noise. This GP can be perfectly approximated with a single inducing point.

# %%
plot_model(vfe)

# %%
print_summary(vfe, fmt="notebook")

# %% [markdown]
# This can be diagnosed by showing that there are other hyperparameter settings with higher upper bounds. This indicates that there might be better hyperparameter settings, but we cannot identify them due to the lack of inducing points. An example of this can be seen in the previous section.


# doc/sphinx/conf.py

# Configuration file for the Sphinx documentation builder.
#
# This file only contains a selection of the most common options. For a full
# list see the documentation:
# https://www.sphinx-doc.org/en/master/usage/configuration.html
from pathlib import Path
from typing import Sequence

# -- Path setup --------------------------------------------------------------

# If extensions (or modules to document with autodoc) are in another directory,
# add these directories to sys.path here. If the directory is relative to the
# documentation root, use os.path.abspath to make it absolute, like shown here.
#
# import os
# import sys
# sys.path.insert(0, os.path.abspath('.'))


# -- Project information -----------------------------------------------------

project = "GPflow"
copyright = "2022, The GPflow Contributors"
author = "The GPflow Contributors"

# The full version, including alpha/beta/rc tags
release = "2.6.3"


# -- General configuration ---------------------------------------------------

# Add any Sphinx extension module names here, as strings. They can be
# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
# ones.
extensions = [
    "sphinx.ext.autodoc",
    "sphinx.ext.autosummary",
    "sphinx.ext.todo",
    "sphinx.ext.mathjax",
    "sphinx.ext.viewcode",
    "sphinxcontrib.bibtex",
    "nbsphinx",
    "sphinx_autodoc_typehints",
    "IPython.sphinxext.ipython_console_highlighting",
]


# Configuration for highlighting in code examples:
highlight_options = {"stripall": True}

set_type_checking_flag = True
typehints_fully_qualified = False
always_document_param_types = True
# autoclass_content = 'both'

# Add any paths that contain templates here, relative to this directory.
templates_path = ["_templates"]

# List of patterns, relative to source directory, that match files and
# directories to ignore when looking for source files.
# This pattern also affects html_static_path and html_extra_path.
exclude_patterns: Sequence[str] = []


bibtex_bibfiles = ["refs.bib"]


# -- Options for HTML output -------------------------------------------------

# The theme to use for HTML and HTML Help pages.  See the documentation for
# a list of builtin themes.
#
html_theme = "pydata_sphinx_theme"
html_logo = "_static/gpflow_logo.svg"
html_css_files = ["pydata-custom.css"]

# theme-specific options. see theme docs for more info
html_theme_options = {
    "show_prev_next": False,
    "github_url": "https://github.com/GPflow/GPflow",
    "switcher": {
        "json_url": "https://gpflow.github.io/GPflow/versions.json",
        "version_match": Path("build_version.txt").read_text().strip(),
    },
    "navbar_end": ["version-switcher"],
}

# If True, show link to rst source on rendered HTML pages
html_show_sourcelink = False  # Remove 'view source code' from top of page (for html, not python)

# Add any paths that contain custom static files (such as style sheets) here,
# relative to this directory. They are copied after the builtin static files,
# so a file named "default.css" will overwrite the builtin "default.css".
html_static_path = ["_static"]


# doc/sphinx/notebooks/getting_started/large_data.pct.py
# ---
# jupyter:
#   jupytext:
#     formats: py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.14.1
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %%
# remove-cell
# pylint: disable=line-too-long,redefined-outer-name

# %% [markdown]
# # Large Data with SGPR
#
# Making predictions with Gaussian Processes is $O(N^3)$, which can be prohibitive for large $N$. In this chapter we introduce the Sparse Gaussian Process Regression (SGPR), which tries to solve this problem. The SGPR tries to approximate the $N$ actual data points with $M$ "inducing variables". With inducing variables the SGPR can make predicitions in $O(NM^2)$, which can make a big difference if $M << N$.
#
# Note: For large datasets the "SVGP" model may also be relevant. See our section on [Classification and other data distributions](classification_and_other_data_distributions.ipynb).

# %% [markdown]
# As usual we will start with our imports:

# %%
# hide: begin
import os
import warnings

warnings.simplefilter("ignore")
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
# hide: end

import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp

import gpflow

# hide: begin
# %matplotlib inline
plt.rcParams["figure.figsize"] = (12, 5)
# hide: end

# %% [markdown]
# Although the sparse models are meant for large datasets, smaller datasets can be easier to understand, so we will reuse the same data we used in our [Basic Usage](basic_usage.ipynb) chapter:

# %%
# hide: begin
# fmt: off
# hide: end
X = np.array(
    [
        [0.865], [0.666], [0.804], [0.771], [0.147], [0.866], [0.007], [0.026],
        [0.171], [0.889], [0.243], [0.028],
    ]
)
Y = np.array(
    [
        [1.57], [3.48], [3.12], [3.91], [3.07], [1.35], [3.80], [3.82], [3.49],
        [1.30], [4.00], [3.82],
    ]
)
# hide: begin
# fmt: on
# hide: end

_ = plt.plot(X, Y, "kx")


# %% [markdown]
# And a helper function for training and plotting our models:

# %%
def plot_model(model: gpflow.models.GPModel) -> None:
    X, Y = model.data
    opt = gpflow.optimizers.Scipy()
    opt.minimize(model.training_loss, model.trainable_variables)
    gpflow.utilities.print_summary(model, "notebook")

    Xplot = np.linspace(0.0, 1.0, 200)[:, None]

    y_mean, y_var = model.predict_y(Xplot, full_cov=False)
    y_lower = y_mean - 1.96 * np.sqrt(y_var)
    y_upper = y_mean + 1.96 * np.sqrt(y_var)

    _, ax = plt.subplots(nrows=1, ncols=1)
    ax.plot(X, Y, "kx", mew=2)
    (mean_line,) = ax.plot(Xplot, y_mean, "-")
    color = mean_line.get_color()
    ax.plot(Xplot, y_lower, lw=0.1, color=color)
    ax.plot(Xplot, y_upper, lw=0.1, color=color)
    ax.fill_between(
        Xplot[:, 0], y_lower[:, 0], y_upper[:, 0], color=color, alpha=0.1
    )

    # Also plot the inducing variables if possible:
    iv = getattr(model, "inducing_variable", None)
    if iv is not None:
        ax.scatter(iv.Z, np.zeros_like(iv.Z), marker="^")


# %% [markdown]
# As a reference, let us try fitting our usual GPR. Notice that this works because our dataset here still is of a quite managable size. In practice you should only use an SGPR when your dataset is too large for a GPR:

# %%
model = gpflow.models.GPR(
    (X, Y),
    kernel=gpflow.kernels.SquaredExponential(),
)
plot_model(model)

# %% [markdown]
# Let us try fitting an SGPR. As mentioned above the SGPR is based on "inducing points" / "inducing variables". The SGPR gains a speed-up by only capturing the shape of $f$ at a small(ish) number of inducing points. We need tell the model where to put these. In this example we will use four inducing points, initially spaced evenly:

# %%
inducing_points = np.array([[0.125], [0.375], [0.625], [0.875]])
model = gpflow.models.SGPR(
    (X, Y),
    kernel=gpflow.kernels.SquaredExponential(),
    inducing_variable=inducing_points,
)
plot_model(model)

# %% [markdown]
# Here we get an excellent fit, even though the original data has 12 points, and we are approximating it with only four inducing points.
#
# The blue triangles represent the positions of the inducing points.
#
# You may notice that even though we initialised our inducing variables to be evenly spaced, they are spaced irregularly when plotting the model. The inducing variables are represent by a GPflow Parameter, and is optimised during model training, along with the other parameters.

# %% [markdown]
# For illustrative purposes, let us see what happens if we don't use enough inducing points:

# %%
inducing_points = np.array([[0.25], [0.75]])
model = gpflow.models.SGPR(
    (X, Y),
    kernel=gpflow.kernels.SquaredExponential(),
    inducing_variable=inducing_points,
)
plot_model(model)

# %% [markdown]
# The two inducing points above are not enough to capture the data well, so we get an overly simplistic $f$, and the model compensates with an overly high variance.

# %% [markdown]
# ## Picking initial inducing points
#
# When you use a sparse model you will need to pick initial inducing points, and choosing these correctly can have a large impact on performance. Above we used a simple grid. This often works well when your data is 1D, but we do not recommend using a grid for higher-dimensional data. In this section we discuss some alternatives.
#
# First, let us declare some 2D data:

# %%
# hide: begin
# fmt: off
# hide: end
X = np.array(
    [
        [0.70, 0.70], [0.53, 0.81], [0.78, 0.36], [0.83, 0.09], [0.71, 0.55],
        [0.66, 0.75], [0.87, 0.50], [0.63, 0.65], [0.37, 0.90], [0.82, 0.11],
        [0.58, 0.61], [0.93, 0.21], [0.98, 0.18], [0.85, 0.27], [0.64, 0.77],
        [0.49, 0.73], [0.13, 0.82], [0.93, 0.08], [0.65, 0.71], [0.54, 0.83],
        [0.85, 0.20], [0.90, 0.07], [0.00, 0.84], [0.64, 0.81], [0.62, 0.70],
    ]
)
Y = np.array(
    [
        [0.83], [0.82], [0.60], [0.31], [0.73], [0.85], [0.70], [0.77], [0.77],
        [0.34], [0.73], [0.45], [0.39], [0.51], [0.85], [0.76], [0.46], [0.27],
        [0.82], [0.85], [0.44], [0.26], [0.28], [0.87], [0.81],
    ]
)
# hide: begin
# fmt: on
# hide: end

# %% [markdown]
# And write a helper-function to plot a model:

# %%
def plot_2d_model(model: gpflow.models.GPModel) -> None:
    n_grid = 50
    _, (ax_mean, ax_std) = plt.subplots(nrows=1, ncols=2, figsize=(11, 5))
    Xplots = np.linspace(0.0, 1.0, n_grid)
    Xplot1, Xplot2 = np.meshgrid(Xplots, Xplots)
    Xplot = np.stack([Xplot1, Xplot2], axis=-1)
    Xplot = Xplot.reshape([n_grid ** 2, 2])

    iv = getattr(model, "inducing_variable", None)
    # Do not optimize inducing variables, so that we can better see the impact their choice has. When solving
    # a real problem you should generally optimise your inducing points.
    if iv is not None:
        gpflow.set_trainable(iv, False)
    opt = gpflow.optimizers.Scipy()
    opt.minimize(model.training_loss, model.trainable_variables)

    y_mean, y_var = model.predict_y(Xplot)
    y_mean = y_mean.numpy()
    y_std = tf.sqrt(y_var).numpy()

    ax_mean.pcolor(Xplot1, Xplot2, y_mean.reshape(Xplot1.shape))
    ax_std.pcolor(Xplot1, Xplot2, y_std.reshape(Xplot1.shape))
    ax_mean.scatter(X[:, 0], X[:, 1], s=50, c="black")
    ax_std.scatter(X[:, 0], X[:, 1], s=50, c="black")

    # Also plot the inducing variables if possible:
    if iv is not None:
        ax_mean.scatter(iv.Z[:, 0], iv.Z[:, 1], marker="x", color="red")


# %% [markdown]
# Let us first make a plot of a dense model, for comparison:

# %%
model = gpflow.models.GPR(
    (X, Y),
    kernel=gpflow.kernels.RBF(),
)
plot_2d_model(model)

# %% [markdown]
# ### Random data samples
#
# A simple approach to selecting inducing points, is to randomly select a subset of your data:

# %%
rng = np.random.default_rng(1234)
n_inducing = 4
inducing_variable = rng.choice(X, size=n_inducing, replace=False)

model = gpflow.models.SGPR(
    (X, Y),
    kernel=gpflow.kernels.RBF(),
    inducing_variable=inducing_variable,
)
plot_2d_model(model)

# %% [markdown]
# ### k-means
#
# Another good approach is to use the [k-means algoritm](https://en.wikipedia.org/wiki/K-means_clustering), which tries to find $k$ centers of your data:

# %%
from scipy.cluster.vq import kmeans

n_inducing = 4
inducing_variable, _ = kmeans(X, n_inducing)

model = gpflow.models.SGPR(
    (X, Y),
    kernel=gpflow.kernels.RBF(),
    inducing_variable=inducing_variable,
)
plot_2d_model(model)

# %% [markdown]
# ### Uncorrelated samples
#
# It is a good idea to pick inducing points that are close to your data, which the above two algoritms does well. However they can struggle with some corner cases, such as if you do not (yet) have access to your data, or if you need to pick more inducing points than you have data. An approach that does not select points close to your data, but that can be more robust in some circumstances, is to use a [low-discrepancy sequence](https://en.wikipedia.org/wiki/Low-discrepancy_sequence):

# %%
n_dim = X.shape[-1]
n_inducing = 4
inducing_variable = tfp.mcmc.sample_halton_sequence(
    n_dim, n_inducing, seed=1234
)

model = gpflow.models.SGPR(
    (X, Y),
    kernel=gpflow.kernels.RBF(),
    inducing_variable=inducing_variable,
)
plot_2d_model(model)

# %% [markdown]
# ### Advanced initialisation
#
# If none of the above approaches work well for you, we recommend that advanced users to look at the [RobustGP repository](https://github.com/markvdw/RobustGP).
