import numpy as np

class _Topology:
    """Represents a topology for a HMM, imposing restrictions on the transition matrix and initial state distribution.

    Parameters
    ----------
    n_states: int
        Number of states in the HMM.

    random_state: numpy.random.RandomState
        A random state object for reproducible randomness.
    """

    def __init__(self, n_states: int, random_state: np.random.RandomState):
        self._n_states = n_states
        self._random_state = random_state

    def uniform_initial(self) -> np.ndarray:
        """Sets the initial state distribution as a discrete uniform distribution.

        Returns
        -------
        initial: :class:`numpy:numpy.ndarray` (float)
            The initial state distribution of shape `(n_states,)`.
        """
        return np.ones(self._n_states) / self._n_states

    def random_initial(self) -> np.ndarray:
        """Sets the initial state distribution by randomly sampling probabilities generated by a Dirichlet distribution.

        Returns
        -------
        initial: :class:`numpy:numpy.ndarray` (float)
            The initial state distribution of shape `(n_states,)`.
        """
        return self._random_state.dirichlet(np.ones(self._n_states), size=1).flatten()

    def uniform_transitions(self) -> np.ndarray:
        """Sets the transition matrix as uniform (equal probability of transitioning
        to all other possible states from each state) corresponding to the topology.

        Returns
        -------
        transitions: :class:`numpy:numpy.ndarray` (float)
            The uniform transition matrix of shape `(n_states, n_states)`.
        """
        raise NotImplementedError

    def random_transitions(self) -> np.ndarray:
        """Sets the transition matrix as random (random probability of transitioning
        to all other possible states from each state) by sampling probabilities
        from a Dirichlet distribution - according to the topology.

        Returns
        -------
        transitions: :class:`numpy:numpy.ndarray` (float)
            The random transition matrix of shape `(n_states, n_states)`.
        """
        raise NotImplementedError

    def validate_initial(self, initial: np.ndarray) -> None:
        """Validates an initial state distribution according to the topology's restrictions.

        Parameters
        ----------
        initial: numpy.ndarray (float)
            The initial state distribution to validate.
        """
        if not isinstance(initial, np.ndarray):
            raise TypeError('Initial state distribution must be a numpy.ndarray')
        if not initial.shape == (self._n_states,):
            raise ValueError('Initial state distribution must be of shape (n_states,)')
        if not np.isclose(initial.sum(), 1):
            raise ValueError('Initial state distribution must sum to one')

    def validate_transitions(self, transitions: np.ndarray) -> None:
        """Validates a transition matrix according to the topology's restrictions.

        Parameters
        ----------
        transitions: numpy.ndarray (float)
            The transition matrix to validate.
        """
        if not isinstance(transitions, np.ndarray):
            raise TypeError('Transition matrix must be a numpy.ndarray')
        if not transitions.shape == (self._n_states, self._n_states):
            raise ValueError('Transition matrix must be of shape (n_states, n_states)')
        if not np.allclose(transitions.sum(axis=1), np.ones(self._n_states)):
            raise ValueError('Transition probabilities out of each state must sum to one')