#   Copyright 2024 - present The PyMC Developers
#
#   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 __future__ import annotations

import warnings

from dataclasses import field
from typing import Any, overload

import numpy as np
import pytensor
import scipy.linalg

from scipy.sparse import issparse

from pymc.pytensorf import floatX
from pymc.step_methods.state import (
    DataClassState,
    RandomGeneratorState,
    WithSamplingState,
    dataclass_state,
)
from pymc.util import RandomGenerator, get_random_generator

__all__ = [
    "QuadPotentialDiag",
    "QuadPotentialDiagAdapt",
    "QuadPotentialFull",
    "QuadPotentialFullAdapt",
    "QuadPotentialFullInv",
    "isquadpotential",
    "quad_potential",
]


def quad_potential(C, is_cov, rng=None):
    """
    Compute a QuadPotential object from a scaling matrix.

    Parameters
    ----------
    C: arraylike, 0 <= ndim <= 2
        scaling matrix for the potential
        vector treated as diagonal matrix.
    is_cov: Boolean
        whether C is provided as a covariance matrix or hessian
    rng: RandomGenerator
        An object that can produce be used to produce the step method's
        :py:class:`~numpy.random.Generator` object. Refer to
        :py:func:`pymc.util.get_random_generator` for more information.

    Returns
    -------
    q: Quadpotential
    """
    if issparse(C):
        if not chol_available:
            raise ImportError("Sparse mass matrices require scikits.sparse")
        elif is_cov:
            return QuadPotentialSparse(C, rng=rng)
        else:
            raise ValueError("Sparse precision matrices are not supported")

    partial_check_positive_definite(C)
    if C.ndim == 1:
        if is_cov:
            return QuadPotentialDiag(C, rng=rng)
        else:
            return QuadPotentialDiag(1.0 / C, rng=rng)
    else:
        if is_cov:
            return QuadPotentialFull(C, rng=rng)
        else:
            return QuadPotentialFullInv(C, rng=rng)


def partial_check_positive_definite(C):
    """Make a simple but partial check for Positive Definiteness."""
    if C.ndim == 1:
        d = C
    else:
        d = np.diag(C)
    (i,) = np.nonzero(np.logical_or(np.isnan(d), d <= 0))

    if len(i):
        raise PositiveDefiniteError("Simple check failed. Diagonal contains negatives", i)


class PositiveDefiniteError(ValueError):
    def __init__(self, msg, idx):
        super().__init__(msg)
        self.idx = idx
        self.msg = msg

    def __str__(self):
        return f"Scaling is not positive definite: {self.msg}. Check indexes {self.idx}."


@dataclass_state
class PotentialState(DataClassState):
    rng: RandomGeneratorState


class QuadPotential(WithSamplingState):
    dtype: np.dtype

    _state_class = PotentialState

    def __init__(self, rng=None):
        self.rng = get_random_generator(rng)

    @overload
    def velocity(self, x: np.ndarray, out: None) -> np.ndarray: ...

    @overload
    def velocity(self, x: np.ndarray, out: np.ndarray) -> None: ...

    def velocity(self, x: np.ndarray, out: np.ndarray | None = None) -> np.ndarray | None:
        """Compute the current velocity at a position in parameter space."""
        raise NotImplementedError("Abstract method")

    def energy(self, x, velocity=None):
        raise NotImplementedError("Abstract method")

    def random(self, x):
        raise NotImplementedError("Abstract method")

    def velocity_energy(self, x, v_out):
        raise NotImplementedError("Abstract method")

    def update(self, sample, grad, tune):
        """Inform the potential about a new sample during tuning.

        This can be used by adaptive potentials to change the
        mass matrix.
        """
        pass

    def raise_ok(self, map_info=None):
        """Check if the mass matrix is ok, and raise ValueError if not.

        Parameters
        ----------
        map_info: List of (name, shape, dtype)
            List tuples with variable name, shape, and dtype.

        Raises
        ------
        ValueError if any standard deviations are 0 or infinite

        Returns
        -------
        None
        """
        return None

    def reset(self):
        pass

    def stats(self):
        return {"largest_eigval": np.nan, "smallest_eigval": np.nan}

    def set_rng(self, rng: RandomGenerator):
        self.rng = get_random_generator(rng, copy=False)


def isquadpotential(value):
    """Check whether an object might be a QuadPotential object."""
    return isinstance(value, QuadPotential)


@dataclass_state
class QuadPotentialDiagAdaptState(PotentialState):
    _var: np.ndarray
    _stds: np.ndarray
    _inv_stds: np.ndarray
    _foreground_var: WeightedVarianceState
    _background_var: WeightedVarianceState
    _n_samples: int
    adaptation_window: int
    _mass_trace: list[np.ndarray] | None

    dtype: Any = field(metadata={"frozen": True})
    _n: int = field(metadata={"frozen": True})
    _discard_window: int = field(metadata={"frozen": True})
    _early_update: int = field(metadata={"frozen": True})
    _initial_mean: np.ndarray = field(metadata={"frozen": True})
    _initial_diag: np.ndarray = field(metadata={"frozen": True})
    _initial_weight: np.ndarray = field(metadata={"frozen": True})
    adaptation_window_multiplier: float = field(metadata={"frozen": True})
    _store_mass_matrix_trace: bool = field(metadata={"frozen": True})


class QuadPotentialDiagAdapt(QuadPotential):
    """Adapt a diagonal mass matrix from the sample variances."""

    _state_class = QuadPotentialDiagAdaptState

    def __init__(
        self,
        n,
        initial_mean,
        initial_diag=None,
        initial_weight=0,
        adaptation_window=101,
        adaptation_window_multiplier=1,
        dtype=None,
        discard_window=50,
        early_update=False,
        store_mass_matrix_trace=False,
        rng=None,
    ):
        """Set up a diagonal mass matrix.

        Parameters
        ----------
        n : int
            The number of parameters.
        initial_mean : np.ndarray
            An initial guess for the posterior mean of each parameter.
        initial_diag : np.ndarray
            An estimate of the posterior variance of each parameter.
        initial_weight : int
            How much weight the initial guess has compared to new samples during tuning.
            Measured in equivalent number of samples.
        adaptation_window : int
            The size of the adaptation window during tuning. It specifies how many samples
            are used to estimate the mass matrix in each section of the adaptation.
        adaptation_window_multiplier : float
            The factor with which we increase the adaptation window after each adaptation
            window.
        dtype : np.dtype
            The dtype used to store the mass matrix
        discard_window : int
            The number of initial samples that are just discarded and not used to estimate
            the mass matrix.
        early_update : bool
            Whether to update the mass matrix live during the first adaptation window.
        store_mass_matrix_trace : bool
            If true, store the mass matrix at each step of the adaptation. Only for debugging
            purposes.
        rng : Generator | int | None
            Numpy random number generator
        """
        if initial_diag is not None and initial_diag.ndim != 1:
            raise ValueError("Initial diagonal must be one-dimensional.")
        if initial_mean.ndim != 1:
            raise ValueError("Initial mean must be one-dimensional.")
        if initial_diag is not None and len(initial_diag) != n:
            raise ValueError(f"Wrong shape for initial_diag: expected {n} got {len(initial_diag)}")
        if len(initial_mean) != n:
            raise ValueError(f"Wrong shape for initial_mean: expected {n} got {len(initial_mean)}")

        if dtype is None:
            dtype = pytensor.config.floatX

        if initial_diag is None:
            initial_diag = np.ones(n, dtype=dtype)
            initial_weight = 1

        self.dtype = dtype
        self._n = n

        self._discard_window = discard_window
        self._early_update = early_update

        self._initial_mean = initial_mean
        self._initial_diag = initial_diag
        self._initial_weight = initial_weight
        self.adaptation_window = adaptation_window
        self.adaptation_window_multiplier = float(adaptation_window_multiplier)

        self._store_mass_matrix_trace = store_mass_matrix_trace
        self._mass_trace = []

        super().__init__(rng=rng)

        self.reset()

    def reset(self):
        self._var = np.array(self._initial_diag, dtype=self.dtype, copy=True)
        self._var_pytensor = pytensor.shared(self._var)
        self._stds = np.sqrt(self._initial_diag)
        self._inv_stds = floatX(1.0) / self._stds
        self._foreground_var = _WeightedVariance(
            self._n, self._initial_mean, self._initial_diag, self._initial_weight, self.dtype
        )
        self._background_var = _WeightedVariance(self._n, dtype=self.dtype)
        self._n_samples = 0

    def velocity(self, x, out=None):
        """Compute the current velocity at a position in parameter space."""
        return np.multiply(self._var, x, out=out)

    def energy(self, x, velocity=None):
        """Compute kinetic energy at a position in parameter space."""
        if velocity is not None:
            return 0.5 * x.dot(velocity)
        return 0.5 * x.dot(self._var * x)

    def velocity_energy(self, x, v_out):
        """Compute velocity and return kinetic energy at a position in parameter space."""
        self.velocity(x, out=v_out)
        return 0.5 * np.dot(x, v_out)

    def random(self):
        """Draw random value from QuadPotential."""
        vals = self.rng.normal(size=self._n).astype(self.dtype)
        return self._inv_stds * vals

    def _update_from_weightvar(self, weightvar):
        weightvar.current_variance(out=self._var)
        self._var = np.clip(self._var, 1e-12, 1e12)
        np.sqrt(self._var, out=self._stds)
        np.divide(1, self._stds, out=self._inv_stds)
        self._var_pytensor.set_value(self._var)

    def update(self, sample, grad, tune):
        """Inform the potential about a new sample during tuning."""
        if self._store_mass_matrix_trace:
            self._mass_trace.append(self._stds.copy())

        if not tune:
            return

        if self._n_samples > self._discard_window:
            self._foreground_var.add_sample(sample)
            self._background_var.add_sample(sample)

        if self._early_update or self._n_samples > self.adaptation_window:
            self._update_from_weightvar(self._foreground_var)

        if self._n_samples > 0 and self._n_samples % self.adaptation_window == 0:
            self._foreground_var = self._background_var
            self._background_var = _WeightedVariance(self._n, dtype=self.dtype)
            self.adaptation_window = int(self.adaptation_window * self.adaptation_window_multiplier)

        self._n_samples += 1

    def raise_ok(self, map_info):
        """Check if the mass matrix is ok, and raise ValueError if not.

        Parameters
        ----------
        map_info: List of (name, shape, dtype)
            List tuples with variable name, shape, and dtype.

        Raises
        ------
        ValueError if any standard deviations are 0 or infinite

        Returns
        -------
        None
        """
        if np.any(self._stds == 0):
            errmsg = ["Mass matrix contains zeros on the diagonal. "]
            last_idx = 0
            for name, shape, size, dtype in map_info:
                end = last_idx + size
                index = np.where(self._stds[last_idx:end] == 0)[0]
                errmsg.append(f"The derivative of RV `{name}`.ravel()[{index}] is zero.")
                last_idx += end

            raise ValueError("\n".join(errmsg))

        if np.any(~np.isfinite(self._stds)):
            errmsg = ["Mass matrix contains non-finite values on the diagonal. "]

            last_idx = 0
            for name, shape, size, dtype in map_info:
                end = last_idx + size
                index = np.where(~np.isfinite(self._stds[last_idx:end]))[0]
                errmsg.append(f"The derivative of RV `{name}`.ravel()[{index}] is non-finite.")
                last_idx = end
            raise ValueError("\n".join(errmsg))


@dataclass_state
class WeightedVarianceState(DataClassState):
    n_samples: int
    mean: np.ndarray
    raw_var: np.ndarray

    _dtype: Any = field(metadata={"frozen": True})


class _WeightedVariance(WithSamplingState):
    """Online algorithm for computing mean of variance."""

    _state_class = WeightedVarianceState

    def __init__(
        self, nelem, initial_mean=None, initial_variance=None, initial_weight=0, dtype="d"
    ):
        self._dtype = dtype
        self.n_samples = float(initial_weight)
        if initial_mean is None:
            self.mean = np.zeros(nelem, dtype="d")
        else:
            self.mean = np.array(initial_mean, dtype="d", copy=True)
        if initial_variance is None:
            self.raw_var = np.zeros(nelem, dtype="d")
        else:
            self.raw_var = np.array(initial_variance, dtype="d", copy=True)

        self.raw_var[:] *= self.n_samples

        if self.raw_var.shape != (nelem,):
            raise ValueError("Invalid shape for initial variance.")
        if self.mean.shape != (nelem,):
            raise ValueError("Invalid shape for initial mean.")

    def add_sample(self, x):
        x = np.asarray(x)
        self.n_samples += 1
        old_diff = x - self.mean
        self.mean[:] += old_diff / self.n_samples
        new_diff = x - self.mean
        self.raw_var[:] += old_diff * new_diff

    def current_variance(self, out=None):
        if self.n_samples == 0:
            raise ValueError("Can not compute variance without samples.")
        if out is not None:
            return np.divide(self.raw_var, self.n_samples, out=out)
        else:
            return (self.raw_var / self.n_samples).astype(self._dtype)

    def current_mean(self):
        return self.mean.copy(dtype=self._dtype)


@dataclass_state
class ExpWeightedVarianceState(DataClassState):
    _alpha: float
    _mean: np.ndarray
    _var: np.ndarray


class _ExpWeightedVariance(WithSamplingState):
    _state_class = ExpWeightedVarianceState

    def __init__(self, n_vars, *, init_mean, init_var, alpha):
        self._variance = init_var
        self._mean = init_mean
        self._alpha = alpha

    def add_sample(self, value):
        alpha = self._alpha
        delta = value - self._mean
        self._mean[...] += alpha * delta
        self._variance[...] = (1 - alpha) * (self._variance + alpha * delta**2)

    def current_variance(self, out=None):
        if out is None:
            out = np.empty_like(self._variance)
        np.copyto(out, self._variance)
        return out

    def current_mean(self, out=None):
        if out is None:
            out = np.empty_like(self._mean)
        np.copyto(out, self._mean)
        return out


@dataclass_state
class QuadPotentialDiagAdaptExpState(QuadPotentialDiagAdaptState):
    _alpha: float
    _stop_adaptation: float
    _variance_estimator: ExpWeightedVarianceState | None
    _variance_estimator_grad: ExpWeightedVarianceState | None


class QuadPotentialDiagAdaptExp(QuadPotentialDiagAdapt):
    _state_class = QuadPotentialDiagAdaptExpState

    def __init__(self, *args, alpha, use_grads=False, stop_adaptation=None, rng=None, **kwargs):
        """Set up a diagonal mass matrix.

        Parameters
        ----------
        n : int
            The number of parameters.
        initial_mean : np.ndarray
            An initial guess for the posterior mean of each parameter.
        initial_diag : np.ndarray
            An estimate of the posterior variance of each parameter.
        alpha : float
            Decay rate of the exponential weighted variance.
        use_grads : bool
            Use gradients, not only samples to estimate the mass matrix.
        stop_adaptation : int
            Stop the mass matrix adaptation after this many samples.
        dtype : np.dtype
            The dtype used to store the mass matrix
        discard_window : int
            The number of initial samples that are just discarded and not used to estimate
            the mass matrix.
        store_mass_matrix_trace : bool
            If true, store the mass matrix at each step of the adaptation. Only for debugging
            purposes.
        rng: RandomGenerator
            An object that can produce be used to produce the step method's
            :py:class:`~numpy.random.Generator` object. Refer to
            :py:func:`pymc.util.get_random_generator` for more information.
        """
        if len(args) > 3:
            raise ValueError("Unsupported arguments to QuadPotentialDiagAdaptExp")

        super().__init__(*args, rng=rng, **kwargs)
        self._alpha = alpha
        self._use_grads = use_grads

        if stop_adaptation is None:
            stop_adaptation = np.inf
        self._stop_adaptation = stop_adaptation
        self._variance_estimator = None
        self._variance_estimator_grad = None

    def update(self, sample, grad, tune):
        if tune and self._n_samples < self._stop_adaptation:
            if self._n_samples > self._discard_window:
                self._variance_estimator.add_sample(sample)
                if self._use_grads:
                    self._variance_estimator_grad.add_sample(grad)
            elif self._n_samples == self._discard_window:
                self._variance_estimator = _ExpWeightedVariance(
                    self._n,
                    init_mean=sample.copy(),
                    init_var=np.zeros_like(sample),
                    alpha=self._alpha,
                )
                if self._use_grads:
                    self._variance_estimator_grad = _ExpWeightedVariance(
                        self._n,
                        init_mean=grad.copy(),
                        init_var=np.zeros_like(grad),
                        alpha=self._alpha,
                    )

            if self._n_samples > 2 * self._discard_window:
                if self._use_grads:
                    self._update_from_variances(
                        self._variance_estimator, self._variance_estimator_grad
                    )
                else:
                    self._update_from_weightvar(self._variance_estimator)

            self._n_samples += 1

        if self._store_mass_matrix_trace:
            self._mass_trace.append(self._stds.copy())

    def _update_from_variances(self, var_estimator, inv_var_estimator):
        var = var_estimator.current_variance()
        inv_var = inv_var_estimator.current_variance()
        updated = np.sqrt(var / inv_var)
        self._var[:] = updated
        np.sqrt(updated, out=self._stds)
        np.divide(1, self._stds, out=self._inv_stds)


class QuadPotentialDiag(QuadPotential):
    """Quad potential using a diagonal covariance matrix."""

    def __init__(self, v, dtype=None, rng=None):
        """Use a vector to represent a diagonal matrix for a covariance matrix.

        Parameters
        ----------
        v: vector, 0 <= ndim <= 1
           Diagonal of covariance matrix for the potential vector
        dtype :
            The dtype to assign to the resulting momentum
        rng : RandomGenerator
            An object that can produce be used to produce the step method's
            :py:class:`~numpy.random.Generator` object. Refer to
            :py:func:`pymc.util.get_random_generator` for more information.
        """
        if dtype is None:
            dtype = pytensor.config.floatX
        self.dtype = dtype
        v = v.astype(self.dtype)
        s = v**0.5

        self.s = s
        self.inv_s = 1.0 / s
        self.v = v
        self.rng = get_random_generator(rng)

    def velocity(self, x, out=None):
        """Compute the current velocity at a position in parameter space."""
        if out is not None:
            np.multiply(x, self.v, out=out)
            return
        return self.v * x

    def random(self):
        """Draw random value from QuadPotential."""
        return floatX(self.rng.normal(size=self.s.shape)) * self.inv_s

    def energy(self, x, velocity=None):
        """Compute kinetic energy at a position in parameter space."""
        if velocity is not None:
            return 0.5 * np.dot(x, velocity)
        return 0.5 * x.dot(self.v * x)

    def velocity_energy(self, x, v_out):
        """Compute velocity and return kinetic energy at a position in parameter space."""
        np.multiply(x, self.v, out=v_out)
        return 0.5 * np.dot(x, v_out)


class QuadPotentialFullInv(QuadPotential):
    """QuadPotential object for Hamiltonian calculations using inverse of covariance matrix."""

    def __init__(self, A, dtype=None, rng=None):
        """Compute the lower cholesky decomposition of the potential.

        Parameters
        ----------
        A: matrix, ndim = 2
           Inverse of covariance matrix for the potential vector
        dtype :
            The dtype to assign to the resulting momentum
        rng : RandomGenerator
            An object that can produce be used to produce the step method's
            :py:class:`~numpy.random.Generator` object. Refer to
            :py:func:`pymc.util.get_random_generator` for more information.
        """
        if dtype is None:
            dtype = pytensor.config.floatX
        self.dtype = dtype
        self.L = floatX(scipy.linalg.cholesky(A, lower=True))
        self.rng = get_random_generator(rng)

    def velocity(self, x, out=None):
        """Compute the current velocity at a position in parameter space."""
        vel = scipy.linalg.cho_solve((self.L, True), x)
        if out is None:
            return vel
        out[:] = vel

    def random(self):
        """Draw random value from QuadPotential."""
        n = floatX(self.rng.normal(size=self.L.shape[0]))
        return np.dot(self.L, n)

    def energy(self, x, velocity=None):
        """Compute kinetic energy at a position in parameter space."""
        if velocity is None:
            velocity = self.velocity(x)
        return 0.5 * x.dot(velocity)

    def velocity_energy(self, x, v_out):
        """Compute velocity and return kinetic energy at a position in parameter space."""
        self.velocity(x, out=v_out)
        return 0.5 * np.dot(x, v_out)


class QuadPotentialFull(QuadPotential):
    """Basic QuadPotential object for Hamiltonian calculations."""

    def __init__(self, cov, dtype=None, rng=None):
        """Compute the lower cholesky decomposition of the potential.

        Parameters
        ----------
        A: matrix, ndim = 2
            scaling matrix for the potential vector
        dtype :
            The dtype to assign to the resulting momentum
        rng : RandomGenerator
            An object that can produce be used to produce the step method's
            :py:class:`~numpy.random.Generator` object. Refer to
            :py:func:`pymc.util.get_random_generator` for more information.
        """
        if dtype is None:
            dtype = pytensor.config.floatX
        self.dtype = dtype
        self._cov = np.array(cov, dtype=self.dtype, copy=True)
        self._chol = scipy.linalg.cholesky(self._cov, lower=True)
        self._n = len(self._cov)
        self.rng = get_random_generator(rng)

    def velocity(self, x, out=None):
        """Compute the current velocity at a position in parameter space."""
        return np.dot(self._cov, x, out=out)

    def random(self):
        """Draw random value from QuadPotential."""
        vals = self.rng.normal(size=self._n).astype(self.dtype)
        return scipy.linalg.solve_triangular(self._chol.T, vals, overwrite_b=True)

    def energy(self, x, velocity=None):
        """Compute kinetic energy at a position in parameter space."""
        if velocity is None:
            velocity = self.velocity(x)
        return 0.5 * np.dot(x, velocity)

    def velocity_energy(self, x, v_out):
        """Compute velocity and return kinetic energy at a position in parameter space."""
        self.velocity(x, out=v_out)
        return self.energy(x, v_out)

    __call__ = random


@dataclass_state
class QuadPotentialFullAdaptState(PotentialState):
    _previous_update: int
    _cov: np.ndarray
    _chol: np.ndarray
    _chol_error: scipy.linalg.LinAlgError | ValueError | None = None
    _foreground_cov: WeightedCovarianceState
    _background_cov: WeightedCovarianceState
    _n_samples: int
    adaptation_window: int

    dtype: Any = field(metadata={"frozen": True})
    _n: int = field(metadata={"frozen": True})
    _update_window: int = field(metadata={"frozen": True})
    _initial_mean: np.ndarray = field(metadata={"frozen": True})
    _initial_cov: np.ndarray = field(metadata={"frozen": True})
    _initial_weight: np.ndarray = field(metadata={"frozen": True})
    adaptation_window_multiplier: float = field(metadata={"frozen": True})


class QuadPotentialFullAdapt(QuadPotentialFull):
    """Adapt a dense mass matrix using the sample covariances."""

    _state_class = QuadPotentialFullAdaptState

    def __init__(
        self,
        n,
        initial_mean,
        initial_cov=None,
        initial_weight=0,
        adaptation_window=101,
        adaptation_window_multiplier=2,
        update_window=1,
        dtype=None,
        rng=None,
    ):
        warnings.warn("QuadPotentialFullAdapt is an experimental feature")

        if initial_cov is not None and initial_cov.ndim != 2:
            raise ValueError("Initial covariance must be two-dimensional.")
        if initial_mean.ndim != 1:
            raise ValueError("Initial mean must be one-dimensional.")
        if initial_cov is not None and initial_cov.shape != (n, n):
            raise ValueError(f"Wrong shape for initial_cov: expected {n} got {initial_cov.shape}")
        if len(initial_mean) != n:
            raise ValueError(f"Wrong shape for initial_mean: expected {n} got {len(initial_mean)}")

        if dtype is None:
            dtype = pytensor.config.floatX

        if initial_cov is None:
            initial_cov = np.eye(n, dtype=dtype)
            initial_weight = 1

        self.dtype = dtype
        self._n = n
        self._initial_mean = initial_mean
        self._initial_cov = initial_cov
        self._initial_weight = initial_weight

        self.adaptation_window = int(adaptation_window)
        self.adaptation_window_multiplier = float(adaptation_window_multiplier)
        self._update_window = int(update_window)

        self.rng = get_random_generator(rng)

        self.reset()

    def reset(self):
        self._previous_update = 0
        self._cov = np.array(self._initial_cov, dtype=self.dtype, copy=True)
        self._chol = scipy.linalg.cholesky(self._cov, lower=True)
        self._chol_error = None
        self._foreground_cov = _WeightedCovariance(
            self._n, self._initial_mean, self._initial_cov, self._initial_weight, self.dtype
        )
        self._background_cov = _WeightedCovariance(self._n, dtype=self.dtype)
        self._n_samples = 0

    def _update_from_weightvar(self, weightvar):
        weightvar.current_covariance(out=self._cov)
        try:
            self._chol = scipy.linalg.cholesky(self._cov, lower=True)
        except (scipy.linalg.LinAlgError, ValueError) as error:
            self._chol_error = error

    def update(self, sample, grad, tune):
        if not tune:
            return

        # Steps since previous update
        delta = self._n_samples - self._previous_update

        self._foreground_cov.add_sample(sample)
        self._background_cov.add_sample(sample)

        # Update the covariance matrix and recompute the Cholesky factorization
        # every "update_window" steps
        if (delta + 1) % self._update_window == 0:
            self._update_from_weightvar(self._foreground_cov)

        # Reset the background covariance if we are at the end of the adaptation
        # window.
        if delta >= self.adaptation_window:
            self._foreground_cov = self._background_cov
            self._background_cov = _WeightedCovariance(self._n, dtype=self.dtype)

            self._previous_update = self._n_samples
            self.adaptation_window = int(self.adaptation_window * self.adaptation_window_multiplier)

        self._n_samples += 1

    def raise_ok(self, vmap):
        if self._chol_error is not None:
            raise ValueError(str(self._chol_error))


@dataclass_state
class WeightedCovarianceState(DataClassState):
    n_samples: float
    mean: np.ndarray
    raw_cov: np.ndarray

    _dtype: Any = field(metadata={"frozen": True})


class _WeightedCovariance(WithSamplingState):
    """Online algorithm for computing mean and covariance.

    This implements the `Welford's algorithm
    <https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance>`_ based
    on the implementation in `the Stan math library
    <https://github.com/stan-dev/math>`_.

    """

    _state_class = WeightedCovarianceState

    def __init__(
        self,
        nelem,
        initial_mean=None,
        initial_covariance=None,
        initial_weight=0,
        dtype="d",
    ):
        self._dtype = dtype
        self.n_samples = float(initial_weight)
        if initial_mean is None:
            self.mean = np.zeros(nelem, dtype="d")
        else:
            self.mean = np.array(initial_mean, dtype="d", copy=True)
        if initial_covariance is None:
            self.raw_cov = np.eye(nelem, dtype="d")
        else:
            self.raw_cov = np.array(initial_covariance, dtype="d", copy=True)

        self.raw_cov[:] *= self.n_samples

        if self.raw_cov.shape != (nelem, nelem):
            raise ValueError("Invalid shape for initial covariance.")
        if self.mean.shape != (nelem,):
            raise ValueError("Invalid shape for initial mean.")

    def add_sample(self, x):
        x = np.asarray(x)
        self.n_samples += 1
        old_diff = x - self.mean
        self.mean[:] += old_diff / self.n_samples
        new_diff = x - self.mean
        self.raw_cov[:] += new_diff[:, None] * old_diff[None, :]

    def current_covariance(self, out=None):
        if self.n_samples == 0:
            raise ValueError("Can not compute covariance without samples.")
        if out is not None:
            return np.divide(self.raw_cov, self.n_samples - 1, out=out)
        else:
            return (self.raw_cov / (self.n_samples - 1)).astype(self._dtype)

    def current_mean(self):
        return np.array(self.mean, dtype=self._dtype)


try:
    import sksparse.cholmod as cholmod

    chol_available = True
except ImportError:
    chol_available = False

if chol_available:
    __all__ += ["QuadPotentialSparse"]

    import pytensor.sparse

    class QuadPotentialSparse(QuadPotential):
        def __init__(self, A, rng=None):
            """Compute a sparse cholesky decomposition of the potential.

            Parameters
            ----------
            A: matrix, ndim = 2
                scaling matrix for the potential vector
            rng : RandomGenerator
                An object that can produce be used to produce the step method's
                :py:class:`~numpy.random.Generator` object. Refer to
                :py:func:`pymc.util.get_random_generator` for more information.
            """
            self.A = A
            self.size = A.shape[0]
            self.factor = factor = cholmod.cholesky(A)
            self.d_sqrt = np.sqrt(factor.D())
            self.rng = get_random_generator(rng)

        def velocity(self, x):
            """Compute the current velocity at a position in parameter space."""
            A = pytensor.sparse.as_sparse(self.A)
            return pytensor.sparse.dot(A, x)

        def random(self):
            """Draw random value from QuadPotential."""
            n = floatX(self.rng.normal(size=self.size))
            n /= self.d_sqrt
            n = self.factor.solve_Lt(n)
            n = self.factor.apply_Pt(n)
            return n

        def energy(self, x):
            """Compute kinetic energy at a position in parameter space."""
            return 0.5 * x.T.dot(self.velocity(x))
