from typing import Callable

import numpy as np

from hmc.integrators.states.leapfrog_state import LeapfrogState
from hmc.integrators.fields import riemannian
from hmc.linalg import solve_psd


class RiemannianLeapfrogState(LeapfrogState):
    """The Riemannian leapfrog state uses the Fisher information matrix to provide
    a position-dependent Riemannian metric. As such, computing the gradients of
    the Hamiltonian requires higher derivatives of the metric, which vanish in
    the Euclidean case.

    """
    def __init__(self,
                 position: np.ndarray,
                 momentum: np.ndarray):
        super().__init__(position, momentum)
        self._jac_metric: np.ndarray
        self._grad_logdet_metric: np.ndarray

    @property
    def requires_update(self) -> bool:
        o = self.log_posterior is None or \
            self.grad_log_posterior is None or \
            self.metric is None or \
            self.inv_metric is None or \
            self.jac_metric is None or \
            self.grad_logdet_metric is None
        return o

    @property
    def jac_metric(self):
        return self._jac_metric

    @jac_metric.setter
    def jac_metric(self, value):
        self._jac_metric = value

    @jac_metric.deleter
    def jac_metric(self):
        del self._jac_metric

    @property
    def grad_logdet_metric(self):
        return self._grad_logdet_metric

    @grad_logdet_metric.setter
    def grad_logdet_metric(self, value):
        self._grad_logdet_metric = value

    @grad_logdet_metric.deleter
    def grad_logdet_metric(self):
        del self._grad_logdet_metric

    def update(self, auxiliaries: Callable):
        num_dims = len(self.position)
        log_posterior, grad_log_posterior, metric, jac_metric = auxiliaries(self.position)
        jac_metric = np.swapaxes(jac_metric, 0, -1)
        inv_metric, sqrtm_metric = solve_psd(metric, return_chol=True)
        grad_logdet_metric = riemannian.grad_logdet(inv_metric, jac_metric, num_dims)
        self.log_posterior = log_posterior
        self.grad_log_posterior = grad_log_posterior
        self.metric = metric
        self.sqrtm_metric = sqrtm_metric
        self.inv_metric = inv_metric
        self.jac_metric = jac_metric
        self.grad_logdet_metric = grad_logdet_metric
        self.velocity = riemannian.velocity(inv_metric, self.momentum)
        self.force = riemannian.force(self.velocity, grad_log_posterior, jac_metric, grad_logdet_metric)

    def clear(self):
        super().clear()
        del self.jac_metric
        del self.grad_logdet_metric
        del self.metric
        del self.inv_metric
        del self.logdet_metric
        del self.sqrtm_metric

[((61, 21, 61, 51), 'numpy.swapaxes', 'np.swapaxes', ({(61, 33, 61, 43): 'jac_metric', (61, 45, 61, 46): '0', (61, 48, 61, 50): '-1'}, {}), '(jac_metric, 0, -1)', True, 'import numpy as np\n'), ((62, 35, 62, 70), 'hmc.linalg.solve_psd', 'solve_psd', (), '', False, 'from hmc.linalg import solve_psd\n'), ((63, 29, 63, 85), 'hmc.integrators.fields.riemannian.grad_logdet', 'riemannian.grad_logdet', ({(63, 52, 63, 62): 'inv_metric', (63, 64, 63, 74): 'jac_metric', (63, 76, 63, 84): 'num_dims'}, {}), '(inv_metric, jac_metric, num_dims)', False, 'from hmc.integrators.fields import riemannian\n'), ((71, 24, 71, 70), 'hmc.integrators.fields.riemannian.velocity', 'riemannian.velocity', ({(71, 44, 71, 54): 'inv_metric', (71, 56, 71, 69): 'self.momentum'}, {}), '(inv_metric, self.momentum)', False, 'from hmc.integrators.fields import riemannian\n'), ((72, 21, 72, 104), 'hmc.integrators.fields.riemannian.force', 'riemannian.force', ({(72, 38, 72, 51): 'self.velocity', (72, 53, 72, 71): 'grad_log_posterior', (72, 73, 72, 83): 'jac_metric', (72, 85, 72, 103): 'grad_logdet_metric'}, {}), '(self.velocity, grad_log_posterior, jac_metric,\n    grad_logdet_metric)', False, 'from hmc.integrators.fields import riemannian\n')]