from abc import abstractmethod, ABC
from typing import Tuple
from numbers import Number
from .._state import State, ParameterSpaceState


class Perturbation(ABC):
    """Base class for a perturbation callable

    Its subclasses represent different types of perturbations. The core functions of
    each subclass are :meth:`perturb` and :meth:`log_prior_ratio` and are to be called
    by :meth:`bayesbay.BaseMarkovChain.advance_chain` on each iteration.

    Users don't instantiate these perturbation classes directly; the instances are
    generated by :class:`bayesbay.parameterization.Parameterization` and
    :class:`bayesbay.LogLikelihood` instead.
    """

    @abstractmethod
    def perturb(self, state: State) -> Tuple[State, Number]:
        r"""proposes a new state from the given state and calculates the log of
        the corresponding partial acceptance probability

        .. math::
            \underbrace{\alpha_{p}}_{\begin{array}{c} \text{Partial} \\ \text{acceptance} \\ \text{probability} \end{array}} =
            \underbrace{\frac{p\left({\bf m'}\right)}{p\left({\bf m}\right)}}_{\text{Prior ratio}}
            \underbrace{\frac{q\left({\bf m} \mid {\bf m'}\right)}{q\left({\bf m'} \mid {\bf m}\right)}}_{\text{Proposal ratio}}
            \underbrace{\lvert \mathbf{J} \rvert}_{\begin{array}{c} \text{Jacobian} \\ \text{determinant} \end{array}}

        where :math:`{\bf d}_{obs}` denotes the observed data and
        :math:`\mathbf{J}` the Jacobian of the transformation.

        Parameters
        ----------
        state : State
            the given current state

        Returns
        -------
        Tuple[State, Number]
            the proposed new state and
            :math:`\alpha_{p} = \log(
            \frac{p({\bf m'})}{p({\bf m})}
            \frac{q\left({\bf m}
            \mid {\bf m'}\right)}{q\left({\bf m'} \mid {\bf m}\right)}
            \lvert \mathbf{J} \rvert)`
        """
        raise NotImplementedError

    def __call__(self, state: State) -> Tuple[State, Number]:
        """same as :meth:`perturb`

        Parameters
        ----------
        state : State
            the given current state

        Returns
        -------
        Tuple[State, Number]
            proposed new state and the partial acceptance probability excluding log
            likelihood ratio for this perturbation
        """
        return self.perturb(state)

    @property
    def type(self) -> str:
        """type of current perturbation (i.e. class name)

        Returns
        -------
        str
            the name of this current class
        """
        return self.__class__.__name__

    @property
    def __name__(self) -> str:
        """same as :attr:`type`

        Returns
        -------
        str
            the name of this current class
        """
        return self.type

    def __repr__(self) -> str:
        return self.__name__


class ParamSpaceMixin:
    @abstractmethod
    def perturb_param_space_state(
        self, ps_state: State
    ) -> Tuple[ParameterSpaceState, Number]:
        r"""proposes a new parameter space state from the given state and calculates
        the log of the corresponding partial acceptance probability

        Parameters
        ----------
        state : ParameterSpaceState
            the given current parameter space state

        Returns
        -------
        Tuple[State, Number]
            the proposed new parameter space state and the partial acceptance
            probability
        """
        raise NotImplementedError
