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, 'numpy.swapaxes', 'np.swapaxes', True, 'import numpy as np\n'), (62, 'hmc.linalg.solve_psd', 'solve_psd', False, 'from hmc.linalg import solve_psd\n'), (63, 'hmc.integrators.fields.riemannian.grad_logdet', 'riemannian.grad_logdet', False, 'from hmc.integrators.fields import riemannian\n'), (71, 'hmc.integrators.fields.riemannian.velocity', 'riemannian.velocity', False, 'from hmc.integrators.fields import riemannian\n'), (72, 'hmc.integrators.fields.riemannian.force', 'riemannian.force', False, 'from hmc.integrators.fields import riemannian\n')]