"""
Module for Gaussian measures on Hilbert spaces. 
"""

import numpy as np
from scipy.linalg import eigh
from scipy.sparse import diags
from scipy.stats import multivariate_normal
from pygeoinf.hilbert_space import (
    LinearOperator,
    DiagonalLinearOperator,
    EuclideanSpace,
)
from pygeoinf.direct_sum import (
    BlockDiagonalLinearOperator,
)


class GaussianMeasure:
    """
    Class for Gaussian measures on a Hilbert space.
    """

    def __init__(
        self,
        /,
        *,
        covariance=None,
        covariance_factor=None,
        expectation=None,
        sample=None,
        inverse_covariance=None,
        inverse_covariance_factor=None,
    ):
        """
        Args:
            domain (HilbertSpace): The Hilbert space on which the measure is defined.
            covariance (LinearOperator): A self-adjoint and non-negative linear operator on the domain.
                Default is None.
            covariance_factor (LinearOperator): A linear operator from Euclidean space to the domain of
                the measure such that C = LL*. Default is None.
            expectation (vector): The expectation value of the measure. Default is zero.
            sample (callable): A functor that returns a random sample from the measure. Default is None.
            inverse_covariance (LinearOperator | None): The inverse of the covariance. Default is None.
            inverse_covariance_factor (LinearOperator | None): The inverse of the covariance. Default is None.

        Notes:
            Either the covariance or covariance factor must be specified, but not both. In the covariance
            factor is provided, then a sampling method is generated by default. Otherwise one can optionally
            be provided.
        """

        if covariance is None and covariance_factor is None:
            raise ValueError(
                "Neither covariance or covariance factor has been provided"
            )

        self._covariance_factor = covariance_factor
        self._covariance = (
            covariance_factor @ covariance_factor.adjoint
            if covariance is None
            else covariance
        )
        self._domain = self._covariance.domain
        self._sample = sample if covariance_factor is None else self._sample_from_factor
        self._inverse_covariance_factor = inverse_covariance_factor
        if inverse_covariance_factor is not None:
            self._inverse_covariance = (
                inverse_covariance_factor.adjoint @ inverse_covariance_factor
            )
        elif inverse_covariance is not None:
            self._inverse_covariance = inverse_covariance
        else:
            self._inverse_covariance = None

        if expectation is None:
            self._expectation = self.domain.zero
        else:
            self._expectation = expectation

    @staticmethod
    def from_standard_deviation(domain, standard_deviation, /, *, expectation=None):
        """
        Forms a Gaussian measure on a Hilbert space with zero
        expectation and whose covariance is proportional to the
        identity operator.

        Args:
            domain (HilbertSpace): The Hilbert space on which the measure
                is defined.
            standard_devitation (float): The standard deviation by which the
                identity is scaled to form the covariance.
            expectation (vector): The expectation of the measure, Default is
                None which corresponds to the zero-vector.

        Notes:
            This measure is only well-defined for finite-dimensional spaces and not
            for finite-dimensional approximations to infinite-dimensional spaces.
        """
        covariance_factor = standard_deviation * domain.identity_operator()
        inverse_covariance_factor = (
            1 / standard_deviation
        ) * domain.identity_operator()
        return GaussianMeasure(
            covariance_factor=covariance_factor,
            inverse_covariance_factor=inverse_covariance_factor,
            expectation=expectation,
        )

    @staticmethod
    def from_standard_deviations(domain, standard_deviations, /, *, expectation=None):
        """
        Forms a Gaussian measure on a Hilbert space with zero
        expectation and whose covariance is diagonal within the
        Galerkin representation.

        Args: domain (HilbertSpace): The Hilbert space on which the measure
                is defined.
        standard_devitations (numpy vector): The diagonal values of the covariance
            within its Galerkin representation.
        expectation (vector): The expectation of the measure, Default is
                None which corresponds to the zero-vector.

        Raises:
            ValueError: If the dimension of the standard deviation vector does not
                match that of the Hilbert space.

        """
        if standard_deviations.size != domain.dim:
            raise ValueError(
                "Standard deviation vector does not have the correct length"
            )
        euclidean = EuclideanSpace(domain.dim)
        covariance_factor = DiagonalLinearOperator(
            euclidean, domain, standard_deviations, galerkin=True
        )
        return GaussianMeasure(
            covariance_factor=covariance_factor,
            inverse_covariance_factor=covariance_factor.inverse,
            expectation=expectation,
        )

    @staticmethod
    def from_covariance_matrix(domain, covariance_matrix, /, *, expectation=None):
        """
        Forms a Gaussian measure from the Galerking representation
        of the covariance matrix.
        """

        eigenvalues, U = eigh(covariance_matrix)
        if any(val < 0 for val in eigenvalues):
            raise ValueError("Covariance matrix is not non-negative")

        values = np.fromiter([np.sqrt(val) for val in eigenvalues], dtype=float)
        D = diags([values], [0])
        Di = diags([np.reciprocal(values)], [0])
        L = U @ D
        Li = Di @ U.T

        covariance_factor = LinearOperator.from_matrix(
            EuclideanSpace(domain.dim), domain, L
        )
        inverse_covariance_factor = LinearOperator.from_matrix(
            domain, EuclideanSpace(domain.dim), Li
        )

        return GaussianMeasure(
            covariance_factor=covariance_factor,
            inverse_covariance_factor=inverse_covariance_factor,
            expectation=expectation,
        )

    @staticmethod
    def from_samples(domain, samples):
        """
        Forms a Gaussian measure from a set of samples based on
        the sample expectation and sample variance.

        Args:
            domain (HilbertSpace): The space the measure is defined on.
            samples ([vectors]): A list of samples.

        Notes:
            :
        """
        assert all([domain.is_element(x) for x in samples])

        n = len(samples)
        expectation = domain.zero
        for x in samples:
            expectation = domain.axpy(1 / n, x, expectation)

        if n == 1:
            covariance = LinearOperator.self_adjoint(domain, lambda x: domain.zero)

            def sample():
                return expectation

        else:
            offsets = [domain.subtract(x, expectation) for x in samples]
            covariance = LinearOperator.self_adjoint_from_tensor_product(
                domain, offsets
            ) / (n - 1)

            def sample():
                x = expectation
                randoms = np.random.randn(len(offsets))
                for y, r in zip(offsets, randoms):
                    x = domain.axpy(r / np.sqrt(n - 1), y, x)
                return x

        return GaussianMeasure(
            covariance=covariance, expectation=expectation, sample=sample
        )

    @staticmethod
    def from_direct_sum(measures):
        """
        Forms the direct sum of a list of Gaussian measures.
        """

        expectation = [measure.expectation for measure in measures]

        covariance = BlockDiagonalLinearOperator(
            [measure.covariance for measure in measures]
        )

        inverse_covariance = (
            BlockDiagonalLinearOperator(
                [measure.inverse_covariance for measure in measures]
            )
            if all(measure.inverse_covariance_set for measure in measures)
            else None
        )

        def sample_impl():
            return [measure.sample() for measure in measures]

        sample = (
            sample_impl if all(measure.sample_set for measure in measures) else None
        )

        return GaussianMeasure(
            covariance=covariance,
            expectation=expectation,
            sample=sample,
            inverse_covariance=inverse_covariance,
        )

    @property
    def domain(self):
        """The Hilbert space the measure is defined on."""
        return self._domain

    @property
    def covariance(self):
        """The covariance operator as an instance of LinearOperator."""
        return self._covariance

    @property
    def inverse_covariance_set(self):
        """Returns true if inverse covariance is available"""
        return self._inverse_covariance is not None

    @property
    def inverse_covariance(self):
        """The covariance operator as an instance of LinearOperator."""
        if self._inverse_covariance is None:
            raise NotImplementedError("Inverse covariance is not set")
        else:
            return self._inverse_covariance

    @property
    def covariance_factor_set(self):
        """
        True if covariance factor set.
        """
        return self._covariance_factor is not None

    @property
    def covariance_factor(self):
        """
        Return the covariance factor if set.
        """
        if not self.covariance_factor_set:
            raise NotImplementedError("Covariance factor has not been set")
        return self._covariance_factor

    @property
    def inverse_covariance_factor_set(self):
        """
        True if inverse covariance factor set.
        """
        return self._inverse_covariance_factor is not None

    @property
    def inverse_covariance_factor(self):
        """
        Return the inverse covariance factor if set.
        """
        if not self.covariance_factor_set:
            raise NotImplementedError("Covariance factor has not been set")
        return self._inverse_covariance_factor

    @property
    def expectation(self):
        """The expectation of the measure."""
        return self._expectation

    @property
    def sample_set(self):
        """
        Returns true if sample method is available.
        """
        return self._sample is not None

    def sample(self):
        """
        Returns a random sample drawn from the measure.
        """
        if self._sample is None:
            raise NotImplementedError("Sample method is not set.")
        else:
            return self._sample()

    def samples(self, n):
        """
        Returns a list of n samples form the measure.
        """
        assert n >= 1
        return [self.sample() for _ in range(n)]

    def sample_expectation(self, n):
        """
        Returns the sample expectation using n > 1 samples.
        """
        assert n >= 1
        expectation = self.domain.zero
        for _ in range(n):
            sample = self.sample()
            expectation = self.domain.axpy(1 / n, sample, expectation)
        return expectation

    def sample_pointwise_variance(self, n):
        """
        Returns the sample variance using n > 1 samples.
        Available only if multiplication of vectors is defined
        on the domain of the measure.
        """
        if not self.domain.has_vector_multiply:
            raise NotImplementedError("Method not defined for this measure")

        samples = self.samples(n)
        expectation = self.expectation
        variance = self.domain.zero

        for sample in samples:
            diff = self.domain.subtract(sample, expectation)
            prod = self.domain.vector_multiply(diff, diff)
            variance = self.domain.axpy(1 / n, prod, variance)

        return variance

    def affine_mapping(self, /, *, operator=None, translation=None):
        """
        Returns the push forward of the measure under an affine mapping.

        Args:
            operator (LinearOperator): The operator part of the mapping.
            translation (vector): The translational part of the mapping.

        Returns:
            GaussianMeasure: The transformed measure defined on the
                codomain of the operator.

        Raises:
            ValueError: If the domain of the operator domain is not
                the domain of the measure.

        Notes:
            If operator is not set, it defaults to the identity.
            It translation is not set, it defaults to zero.
        """

        if operator is None:
            _operator = self.domain.identity_operator()
        else:
            _operator = operator

        if translation is None:
            _translation = _operator.codomain.zero
        else:
            _translation = translation

        expectation = _operator.codomain.add(_operator(self.expectation), _translation)

        if self.covariance_factor_set:

            covariance_factor = _operator @ self.covariance_factor

            return GaussianMeasure(
                covariance_factor=covariance_factor, expectation=expectation
            )

        else:

            covariance = _operator @ self.covariance @ _operator.adjoint

            def sample():
                return _operator.codomain.add(_operator(self.sample()), _translation)

            return GaussianMeasure(
                covariance=covariance, expectation=expectation, sample=sample
            )

    def as_multivariate_normal(self):
        """
        Returns the measure in the form of a scipy.stats multivariate_normal
        distribution. This method is available only if the measure is
        defined on a EuclideanSpace.
        """

        if not isinstance(self.domain, EuclideanSpace):
            raise NotImplementedError(
                "Method only defined for measures on Euclidean space"
            )

        return multivariate_normal(
            mean=self.expectation, cov=self.covariance.matrix(dense=True)
        )

    def low_rank_approximation(
        self,
        rank,
        /,
        *,
        power=0,
        method="fixed",
        rtol=1e-2,
        scaling=None,
    ):
        """
        Returns an approximation to the measure with the covariance operator
        replaced by a low-rank Cholesky decomposition along with the associated
        sampling method.
        """

        covariance = self.covariance

        if scaling is None:
            covariance_factor = covariance.random_cholesky(
                rank, power=power, method=method, rtol=rtol
            )

        else:
            covariance = scaling @ covariance @ scaling
            covariance_factor = covariance.random_cholesky(
                rank, power=power, method=method, rtol=rtol
            )
            covariance_factor = scaling.inverse @ covariance_factor

        return GaussianMeasure(
            covariance_factor=covariance_factor,
            expectation=self.expectation,
        )

    def _sample_from_factor(self):
        covariance_factor = self.covariance_factor
        value = covariance_factor(np.random.randn(self.covariance_factor.domain.dim))
        return (
            self.domain.add(value, self.expectation)
            if self.expectation is not None
            else value
        )

    def __neg__(self):
        """Negative of the measure."""

        if self.covariance_factor_set:
            return GaussianMeasure(
                covariance_factor=self.covariance_factor,
                expectation=self.domain.negative(self.expectation),
            )
        else:
            return GaussianMeasure(
                covariance=self.covariance,
                expectation=self.domain.negative(self.expectation),
                sample=lambda: self.domain.negative(self.sample()),
            )

    def __mul__(self, alpha):
        """Multiply the measure by a scalar."""
        if self.covariance_factor_set:
            return GaussianMeasure(
                covariance_factor=alpha * self.covariance_factor,
                expectation=self.domain.multiply(alpha, self.expectation),
            )
        return GaussianMeasure(
            covariance=alpha * alpha * self.covariance,
            expectation=self.domain.multiply(alpha, self.expectation),
            sample=lambda: self.domain.multiply(alpha, self.sample()),
        )

    def __rmul__(self, alpha):
        """Multiply the measure by a scalar."""
        return self * alpha

    def __add__(self, other):
        """Add two measures on the same domain."""
        return GaussianMeasure(
            covariance=self.covariance + other.covariance,
            expectation=self.domain.add(self.expectation, other.expectation),
            sample=lambda: self.domain.add(self.sample(), other.sample()),
        )

    def __sub__(self, other):
        """Subtract two measures on the same domain."""
        return GaussianMeasure(
            covariance=self.covariance + other.covariance,
            expectation=self.domain.subtract(self.expectation, other.expectation),
            sample=lambda: self.domain.subtract(self.sample(), other.sample()),
        )
