import numpy as np


class QuadraticConvexTransport:
    """T(x)=grad(.5*x^T A*x+b^T x), with A symmetric positive definite."""
    def __init__(self, A, b=None):
        self.A = np.asarray(A, dtype=float)
        self.b = np.zeros(self.A.shape[0]) if b is None else np.asarray(b, dtype=float)
        if not np.allclose(self.A, self.A.T):
            raise ValueError("A must be symmetric")
        if np.min(np.linalg.eigvalsh(self.A)) <= 0:
            raise ValueError("potential is not strongly convex")

    def map(self, x):
        return np.asarray(x) @ self.A.T + self.b

    def potential(self, x):
        x = np.asarray(x)
        return 0.5 * np.einsum("...i,ij,...j->...", x, self.A, x) + x @ self.b

    def logdet_jacobian(self, x):
        return np.full(len(np.atleast_2d(x)), np.linalg.slogdet(self.A)[1])

    def gaussian_kl_to_standard(self, mean, covariance):
        """Exact KL(T#N(mean,cov)||N(0,I))."""
        mean, covariance = np.asarray(mean), np.asarray(covariance)
        m = self.A @ mean + self.b
        c = self.A @ covariance @ self.A.T
        d = len(mean)
        return 0.5 * (np.trace(c) + m @ m - d - np.linalg.slogdet(c)[1])


def change_of_variables_kl(samples, transport, log_source_density):
    x = np.asarray(samples)
    tx = transport.map(x)
    return float(np.mean(log_source_density(x) - transport.logdet_jacobian(x)
                         - log_source_density(tx)))


def standard_normal_logpdf(x):
    x = np.asarray(x)
    return -0.5 * (np.sum(x * x, axis=-1) + x.shape[-1] * np.log(2 * np.pi))