#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

r"""Linear Elliptical Slice Sampler.

References

.. [Gessner2020]
    A. Gessner, O. Kanjilal, and P. Hennig. Integrals over gaussians under
    linear domain constraints. AISTATS 2020.


This implementation is based (with multiple changes / optimiations) on
the following implementations based on the algorithm in [Gessner2020]_:
- https://github.com/alpiges/LinConGauss
- https://github.com/wjmaddox/pytorch_ess

The implementation here differentiates itself from the original implementations with:
1) Support for fixed feature equality constraints.
2) Support for non-standard Normal distributions.
3) Numerical stability improvements, especially relevant for high-dimensional cases.

Notably, this implementation does not rely on an adaptive `delta_theta` parameter in
order to determine if two neighboring constraint intersection angles `theta` lead to a
change in the feasibility of the sample. This both simplifies the implementation and
makes it more robust to numerical imprecisions when two constraint intersection angles
are close to each other.
"""

from __future__ import annotations

import math
from typing import List, Optional, Tuple, Union

import torch
from botorch.utils.sampling import PolytopeSampler
from linear_operator.operators import DiagLinearOperator, LinearOperator
from torch import Tensor

_twopi = 2.0 * math.pi


class LinearEllipticalSliceSampler(PolytopeSampler):
    r"""Linear Elliptical Slice Sampler.

    Ideas:
    - Add batch support, broadcasting over parallel chains.
    - Optimize computations if possible, potentially with torch.compile.
    - Extend fixed features constraint to general linear equality constraints.
    """

    def __init__(
        self,
        inequality_constraints: Optional[Tuple[Tensor, Tensor]] = None,
        bounds: Optional[Tensor] = None,
        interior_point: Optional[Tensor] = None,
        fixed_indices: Optional[Union[List[int], Tensor]] = None,
        mean: Optional[Tensor] = None,
        covariance_matrix: Optional[Union[Tensor, LinearOperator]] = None,
        covariance_root: Optional[Union[Tensor, LinearOperator]] = None,
        check_feasibility: bool = False,
        burnin: int = 0,
        thinning: int = 0,
    ) -> None:
        r"""Initialize LinearEllipticalSliceSampler.

        Args:
            inequality_constraints: Tensors `(A, b)` describing inequality constraints
                 `A @ x <= b`, where `A` is an `n_ineq_con x d`-dim Tensor and `b` is
                 an `n_ineq_con x 1`-dim Tensor, with `n_ineq_con` the number of
                 inequalities and `d` the dimension of the sample space. If omitted,
                 must provide `bounds` instead.
            bounds: A `2 x d`-dim tensor of box bounds. If omitted, must provide
                `inequality_constraints` instead.
            interior_point: A `d x 1`-dim Tensor presenting a point in the (relative)
                interior of the polytope. If omitted, an interior point is determined
                automatically by solving a Linear Program. Note: It is crucial that
                the point lie in the interior of the feasible set (rather than on the
                boundary), otherwise the sampler will produce invalid samples.
            fixed_indices: Integer list or `d`-dim Tensor representing the indices of
                dimensions that are constrained to be fixed to the values specified in
                the `interior_point`, which is required to be passed in conjunction with
                `fixed_indices`.
            mean: The `d x 1`-dim mean of the MVN distribution (if omitted, use zero).
            covariance_matrix: The `d x d`-dim covariance matrix of the MVN
                distribution (if omitted, use the identity).
            covariance_root: A `d x d`-dim root of the covariance matrix such that
                covariance_root @ covariance_root.T = covariance_matrix. NOTE: This
                matrix is assumed to be lower triangular. covariance_root can only be
                passed in conjunction with fixed_indices if covariance_root is a
                DiagLinearOperator. Otherwise the factorization would need to be re-
                computed, as we need to solve in `standardize`.
            check_feasibility: If True, raise an error if the sampling results in an
                infeasible sample. This creates some overhead and so is switched off
                by default.
            burnin: Number of samples to generate upon initialization to warm up the
                sampler.
            thinning: Number of samples to skip before returning a sample in `draw`.

        This sampler samples from a multivariante Normal `N(mean, covariance_matrix)`
        subject to linear domain constraints `A x <= b` (intersected with box bounds,
        if provided).
        """
        if interior_point is not None and interior_point.ndim == 1:
            interior_point = interior_point.unsqueeze(-1)

        if mean is not None and mean.ndim == 1:
            mean = mean.unsqueeze(-1)

        super().__init__(
            inequality_constraints=inequality_constraints,
            # TODO: Support equality constraints?
            interior_point=interior_point,
            bounds=bounds,
        )
        tkwargs = {"device": self.x0.device, "dtype": self.x0.dtype}
        if covariance_matrix is not None and covariance_root is not None:
            raise ValueError(
                "Provide either covariance_matrix or covariance_root, not both."
            )

        # can't unpack inequality constraints directly if bounds are passed
        A, b = self.A, self.b
        self._Az, self._bz = A, b
        self._is_fixed, self._not_fixed = None, None
        if fixed_indices is not None:
            mean, covariance_matrix, covariance_root = (
                self._fixed_features_initialization(
                    A=A,
                    b=b,
                    interior_point=interior_point,
                    fixed_indices=fixed_indices,
                    mean=mean,
                    covariance_matrix=covariance_matrix,
                    covariance_root=covariance_root,
                )
            )

        self._mean = mean
        # Have to delay factorization until after fixed features initialization.
        if covariance_matrix is not None:  # implies root is None
            covariance_root, info = torch.linalg.cholesky_ex(covariance_matrix)
            not_psd = torch.any(info)
            if not_psd:
                raise ValueError(
                    "Covariance matrix is not positive definite. "
                    "Currently only non-degenerate distributions are supported."
                )
        self._covariance_root = covariance_root

        # Rewrite the constraints as a system that constrains a standard Normal.
        self._standardization_initialization()

        # state of the sampler ("current point")
        self._x = self.x0.clone()
        self._z = self._transform(self._x)

        # We will need the following repeatedly, let's allocate them once
        self._zero = torch.zeros(1, **tkwargs)
        self._nan = torch.tensor(float("nan"), **tkwargs)
        self._full_angular_range = torch.tensor([0.0, _twopi], **tkwargs)
        self.check_feasibility = check_feasibility
        self._lifetime_samples = 0
        if burnin > 0:
            self.thinning = 0
            self.draw(burnin)
        self.thinning = thinning

    def _fixed_features_initialization(
        self,
        A: Tensor,
        b: Tensor,
        interior_point: Optional[Tensor],
        fixed_indices: Union[List[int], Tensor],
        mean: Optional[Tensor],
        covariance_matrix: Optional[Tensor],
        covariance_root: Optional[Tensor],
    ) -> Tuple[Optional[Tensor], Optional[Tensor]]:
        """Modifies the constraint system (A, b) due to fixed indices and assigns
        the modified constraints system to `self._Az`, `self._bz`. NOTE: Needs to be
        called prior to `self._standardization_initialization` in the constructor.
        covariance_root and fixed_indices can both not be None only if covariance_root
        is a DiagLinearOperator. Otherwise, the covariance matrix would need to be
        refactorized.

        Returns:
            Tuple of `mean` and `covariance_matrix` tensors of the non-fixed dimensions.
        """
        if interior_point is None:
            raise ValueError(
                "If `fixed_indices` are provided, an interior point must also be "
                "provided in order to infer feasible values of the fixed features."
            )

        root_is_diag = isinstance(covariance_root, DiagLinearOperator)
        if covariance_root is not None and not root_is_diag:
            root_is_diag = (covariance_root.diag().diag() == covariance_root).all()
            if root_is_diag:  # convert the diagonal root to a DiagLinearOperator
                covariance_root = DiagLinearOperator(covariance_root.diagonal())
            else:  # otherwise, fail
                raise ValueError(
                    "Provide either covariance_root or fixed_indices, not both."
                )
        d = interior_point.shape[0]
        is_fixed, not_fixed = get_index_tensors(fixed_indices=fixed_indices, d=d)
        self._is_fixed = is_fixed
        self._not_fixed = not_fixed
        # Transforming constraint system to incorporate fixed features:
        # A @ x - b = (A[:, fixed] @ x[fixed] + A[:, not fixed] @ x[not fixed]) - b
        #           = A[:, not fixed] @ x[not fixed] - (b - A[:, fixed] @ x[fixed])
        #           = Az @ z - bz
        self._Az = A[:, not_fixed]
        self._bz = b - A[:, is_fixed] @ interior_point[is_fixed]
        if mean is not None:
            mean = mean[not_fixed]
        if covariance_matrix is not None:  # subselect active dimensions
            covariance_matrix = covariance_matrix[
                not_fixed.unsqueeze(-1), not_fixed.unsqueeze(0)
            ]
        if root_is_diag:  # in the special case of diagonal root, can subselect
            covariance_root = DiagLinearOperator(covariance_root.diagonal()[not_fixed])

        return mean, covariance_matrix, covariance_root

    def _standardization_initialization(self) -> None:
        """For non-standard mean and covariance, we're going to rewrite the problem as
        sampling from a standard normal distribution subject to modified constraints.
            A @ x - b = A @ (covar_root @ z + mean) - b
                      = (A @ covar_root) @ z - (b - A @ mean)
                      = _Az @ z - _bz
        NOTE: We need to standardize bz before Az in the following, because it relies
        on the untransformed Az. We can't simply use A instead because Az might have
        been subject to the fixed features transformation.
        """
        if self._mean is not None:
            self._bz = self._bz - self._Az @ self._mean
        if self._covariance_root is not None:
            self._Az = self._Az @ self._covariance_root

    @property
    def lifetime_samples(self) -> int:
        """The total number of samples generated by the sampler during its lifetime."""
        return self._lifetime_samples

    def draw(self, n: int = 1) -> Tuple[Tensor, Tensor]:
        r"""Draw samples.

        Args:
            n: The number of samples.

        Returns:
            A `n x d`-dim tensor of `n` samples.
        """
        samples = []
        for _ in range(n):
            for _ in range(self.thinning):
                self.step()
            samples.append(self.step())
        return torch.cat(samples, dim=-1).transpose(-1, -2)

    def step(self) -> Tensor:
        r"""Take a step, return the new sample, update the internal state.

        Returns:
            A `d x 1`-dim sample from the domain.
        """
        nu = torch.randn_like(self._z)
        theta = self._draw_angle(nu=nu)
        z = self._get_cart_coords(nu=nu, theta=theta)
        self._z[:] = z
        x = self._untransform(z)
        self._x[:] = x
        self._lifetime_samples += 1
        if self.check_feasibility and (not self._is_feasible(self._x)):
            Axmb = self.A @ self._x - self.b
            violated_indices = Axmb > 0
            raise RuntimeError(
                "Sampling resulted in infeasible point. \n\t- Number "
                f"of violated constraints: {violated_indices.sum()}."
                f"\n\t- Magnitude of violations: {Axmb[violated_indices]}"
                "\n\t- If the error persists, please report this bug on GitHub."
            )
        return x

    def _draw_angle(self, nu: Tensor) -> Tensor:
        r"""Draw the rotation angle.

        Args:
            nu: A `d x 1`-dim tensor (the "new" direction, drawn from N(0, I)).

        Returns:
            A `1`-dim Tensor containing the rotation angle (radians).
        """
        rot_angle, rot_slices = self._find_rotated_intersections(nu)
        rot_lengths = rot_slices[:, 1] - rot_slices[:, 0]
        cum_lengths = torch.cumsum(rot_lengths, dim=0)
        cum_lengths = torch.cat((self._zero, cum_lengths), dim=0)
        rnd_angle = cum_lengths[-1] * torch.rand(
            1, device=cum_lengths.device, dtype=cum_lengths.dtype
        )
        idx = torch.searchsorted(cum_lengths, rnd_angle) - 1
        return (rot_slices[idx, 0] + rnd_angle + rot_angle) - cum_lengths[idx]

    def _get_cart_coords(self, nu: Tensor, theta: Tensor) -> Tensor:
        r"""Determine location on ellipsoid in cartesian coordinates.

        Args:
            nu: A `d x 1`-dim tensor (the "new" direction, drawn from N(0, I)).
            theta: A `k`-dim tensor of angles.

        Returns:
            A `d x k`-dim tensor of samples from the domain in cartesian coordinates.
        """
        return self._z * torch.cos(theta) + nu * torch.sin(theta)

    def _find_rotated_intersections(self, nu: Tensor) -> Tuple[Tensor, Tensor]:
        r"""Finds rotated intersections.

        Rotates the intersections by the rotation angle and makes sure that all
        angles lie in [0, 2*pi].

        Args:
            nu: A `d x 1`-dim tensor (the "new" direction, drawn from N(0, I)).

        Returns:
            A two-tuple containing rotation angle (scalar) and a
            `num_active / 2 x 2`-dim tensor of shifted angles.
        """
        slices = self._find_active_intersections(nu)
        rot_angle = slices[0]
        slices = (slices - rot_angle).reshape(-1, 2)
        # Ensuring that we don't sample within numerical precision of the boundaries
        # due to resulting instabilities in the constraint satisfaction.
        eps = 1e-6 if slices.dtype == torch.float32 else 1e-12
        eps = torch.tensor(eps, dtype=slices.dtype, device=slices.device)
        eps = eps.minimum(slices.diff(dim=-1).abs() / 4)
        slices = slices + torch.cat((eps, -eps), dim=-1)
        # NOTE: The remainder call relies on the epsilon contraction, since the
        # remainder of_twopi divided by _twopi is zero, not _twopi.
        return rot_angle, slices.remainder(_twopi)

    def _find_active_intersections(self, nu: Tensor) -> Tensor:
        """
        Find angles of those intersections that are at the boundary of the integration
        domain by adding and subtracting a small angle and evaluating on the ellipse
        to see if we are on the boundary of the integration domain.

        Args:
            nu: A `d x 1`-dim tensor (the "new" direction, drawn from N(0, I)).

        Returns:
            A `num_active`-dim tensor containing the angles of active intersection in
            increasing order so that activation happens in positive direction. If a
            slice crosses `theta=0`, the first angle is appended at the end of the
            tensor. Every element of the returned tensor defines a slice for elliptical
            slice sampling.
        """
        theta = self._find_intersection_angles(nu)
        theta_active, delta_active = self._active_theta_and_delta(
            nu=nu,
            theta=theta,
        )
        if theta_active.numel() == 0:
            theta_active = self._full_angular_range
            # TODO: What about `self.ellipse_in_domain = False` in the original code?
        elif delta_active[0] == -1:  # ensuring that the first interval is feasible

            theta_active = torch.cat((theta_active[1:], theta_active[:1]))

        return theta_active.view(-1)

    def _find_intersection_angles(self, nu: Tensor) -> Tensor:
        """Compute all of the up to 2*n_ineq_con intersections of the ellipse
        and the linear constraints.

        For background, see equation (2) in
        http://proceedings.mlr.press/v108/gessner20a/gessner20a.pdf

        Args:
            nu: A `d x 1`-dim tensor (the "new" direction, drawn from N(0, I)).

        Returns:
            An `M`-dim tensor, where `M <= 2 * n_ineq_con` (with `M = n_ineq_con`
            if all intermediate computations yield finite numbers).
        """
        # Compared to the implementation in https://github.com/alpiges/LinConGauss
        # we need to flip the sign of A b/c the original algorithm considers
        # A @ x + b >= 0 feasible, whereas we consider A @ x - b <= 0 feasible.
        g1 = -self._Az @ self._z
        g2 = -self._Az @ nu
        r = torch.sqrt(g1**2 + g2**2)
        phi = 2 * torch.atan(g2 / (r + g1)).squeeze()

        arg = -(self._bz / r).squeeze()
        # Write NaNs if there is no intersection
        arg = torch.where(torch.absolute(arg) <= 1, arg, self._nan)
        # Two solutions per linear constraint, shape of theta: (n_ineq_con, 2)
        acos_arg = torch.arccos(arg)
        theta = torch.stack((phi + acos_arg, phi - acos_arg), dim=-1)
        theta = theta[torch.isfinite(theta)]  # shape: `n_ineq_con - num_not_finite`
        theta = torch.where(theta < 0, theta + _twopi, theta)  # in [0, 2*pi]
        return torch.sort(theta).values

    def _active_theta_and_delta(self, nu: Tensor, theta: Tensor) -> Tensor:
        r"""Determine active indices.

        Args:
            nu: A `d x 1`-dim tensor (the "new" direction, drawn from N(0, I)).
            theta: A sorted `M`-dim tensor of intersection angles in [0, 2pi].

        Returns:
            A tuple of Tensors of active constraint intersection angles `theta_active`,
            and the change in the feasibility of the points on the ellipse on the left
            and right of the active intersection angles `delta_active`. `delta_active`
            is is negative if decreasing the angle renders the sample feasible, and
            positive if increasing the angle renders the sample feasible.
        """
        # In order to determine if an angle that gives rise to an intersection with a
        # constraint boundary leads to a change in the feasibility of the solution,
        # we evaluate the constraints on the midpoint of the intersection angles.
        # This gets rid of the `delta_theta` parameter in the original implementation,
        # which cannot be set universally since it can be both 1) too large, when
        # the distance in adjacent intersection angles is small, and 2) too small,
        # when it approaches the numerical precision limit.
        # The implementation below solves both problems and gets rid of the parameter.
        if len(theta) < 2:  # if we have no or only a tangential intersection
            theta_active = torch.tensor([], dtype=theta.dtype, device=theta.device)
            delta_active = torch.tensor([], dtype=int, device=theta.device)
            return theta_active, delta_active
        theta_mid = (theta[:-1] + theta[1:]) / 2  # midpoints of intersection angles
        last_mid = (theta[:1] + theta[-1:] + _twopi) / 2
        last_mid = last_mid.where(last_mid < _twopi, last_mid - _twopi)
        theta_mid = torch.cat((last_mid, theta_mid, last_mid), dim=0)
        samples_mid = self._get_cart_coords(nu=nu, theta=theta_mid)
        delta_feasibility = (
            self._is_feasible(samples_mid, transformed=True).to(dtype=int).diff()
        )
        active_indices = delta_feasibility.nonzero()
        return theta[active_indices], delta_feasibility[active_indices]

    def _is_feasible(self, points: Tensor, transformed: bool = False) -> Tensor:
        r"""Returns a Boolean tensor indicating whether the `points` are feasible,
        i.e. they satisfy `A @ points <= b`, where `(A, b)` are the tensors passed
        as the `inequality_constraints` to the constructor of the sampler.

        Args:
            points: A `d x M`-dim tensor of points.
            transformed: Wether points are assumed to be transformed by a change of
                basis, which means feasibility should be computed based on the
                transformed constraint system (_Az, _bz), instead of (A, b).

        Returns:
            An `M`-dim binary tensor where `True` indicates that the associated
            point is feasible.
        """
        A, b = (self._Az, self._bz) if transformed else (self.A, self.b)
        return (A @ points <= b).all(dim=0)

    def _transform(self, x: Tensor) -> Tensor:
        """Transforms the input so that it is equivalent to a standard Normal variable
        constrained with the modified system constraints (self._Az, self._bz).

        Args:
            x: The input tensor to be transformed, usually `d x 1`-dimensional.

        Returns:
            A `d x 1`-dimensional tensor of transformed values subject to the modified
            system of constraints.
        """
        if self._not_fixed is not None:
            x = x[self._not_fixed]
        return self._standardize(x)

    def _untransform(self, z: Tensor) -> Tensor:
        """The inverse transform of the `_transform`, i.e. maps `z` back to the original
        space where it is subject to the original constraint system (self.A, self.b).

        Args:
            z: The transformed tensor to be un-transformed, usually `d x 1`-dimensional.

        Returns:
            A `d x 1`-dimensional tensor of un-transformed values subject to the
            original system of constraints.
        """
        if self._is_fixed is None:
            return self._unstandardize(z)
        else:
            x = self._x.clone()  # _x already contains the fixed values
            x[self._not_fixed] = self._unstandardize(z)
            return x

    def _standardize(self, x: Tensor) -> Tensor:
        """_transform helper standardizing the input `x`, which is assumed to be a
        `d x 1`-dim Tensor, or a `len(self._not_fixed) x 1`-dim if there are fixed
        indices.
        """
        z = x
        if self._mean is not None:
            z = z - self._mean
        root = self._covariance_root
        if root is not None:
            z = torch.linalg.solve_triangular(root, z, upper=False)

        return z

    def _unstandardize(self, z: Tensor) -> Tensor:
        """_untransform helper un-standardizing the input `z`, which is assumed to be a
        `d x 1`-dim Tensor, or a `len(self._not_fixed) x 1`-dim if there are fixed
        indices.
        """
        x = z
        if self._covariance_root is not None:
            x = self._covariance_root @ x
        if self._mean is not None:
            x = x + self._mean
        return x


def get_index_tensors(
    fixed_indices: Union[List[int], Tensor], d: int
) -> Tuple[Tensor, Tensor]:
    """Converts `fixed_indices` to a `d`-dim integral Tensor that is True at indices
    that are contained in `fixed_indices` and False otherwise.

    Args:
        fixed_indices: A list or Tensoro of integer indices to fix.
        d: The dimensionality of the Tensors to be indexed.

    Returns:
        A Tuple of integral Tensors partitioning [1, d] into indices that are fixed
        (first tensor) and non-fixed (second tensor).
    """
    is_fixed = torch.as_tensor(fixed_indices)
    dtype, device = is_fixed.dtype, is_fixed.device
    dims = torch.arange(d, dtype=dtype, device=device)
    not_fixed = torch.tensor([i for i in dims if i not in is_fixed], dtype=dtype)
    return is_fixed, not_fixed
