"""
Simulates the initial state discrimination experiment using different
methods, to compare the resulting error rates.
"""

import torch

from perm_hmm.util import num_to_data
from perm_hmm.postprocessing import ExactPostprocessor, EmpiricalPostprocessor
from perm_hmm.classifiers.perm_classifier import PermClassifier


class HMMSimulator(object):
    """
    Runs an experiment where data is generated by an HMM, then classified by
    a classifier.

    Instances of this class have the following attributes:

    ``phmm``:
        The :py:class:`~perm_hmm.models.hmms.PermutedDiscreteHMM` whose
        misclassification rates will be computed.
    """

    def __init__(self, phmm):
        """
        Initializes the experiment.

        :param perm_hmm.models.hmms.PermutedDiscreteHMM phmm:
            the model whose
            misclassification rate will be computed.
        """
        self.phmm = phmm
        """:py:class:`PermutedDiscreteHMM`
        The model whose misclassification rates we wish to analyze.
        """

    def all_classifications(self, num_steps, classifier=None, perm_policy=None, verbosity=0):
        """
        Computes the data required to compute the exact misclassification rate for the given classifier.

        This method always calls ``perm_policy.reset()`` if ``perm_policy`` is
        not ``None``.

        :param num_steps: Number of steps, int.
        :param classifier: Defaults to
            :py:class:`~perm_hmm.classifiers.perm_classifier.PermClassifier`,
            initialized with the hmm ``self.phmm``.
        :param perm_policy: Defaults to None. If specified, will call
            ``perm_policy.get_perms`` to compute the permutations.
        :param verbosity: If ``verbosity == 0``, only the
            :py:class:`~perm_hmm.postprocessing.ExactPostprocessor` needed to
            compute the misclassification rates is returned.

            If ``verbosity == 1``, this method returns a tuple, with the
            postprocessor as the first element, and a dictionary with keys:

                b"posterior_log_initial_state_dist":
                    The posterior log initial state distribution used to compute
                    the classifications.

                b"perms":
                    Only present if ``perm_policy`` is not ``None``. The
                    permutations computed from ``perm_policy.get_perms()``.

            If ``verbosity > 1``, this method returns a tuple, with the
            postprocessor as the first element, and a dictionary with keys as in
            the case with ``verbosity == 1`` and in addition,

                b"history":
                    Whatever is stored in ``perm_policy.calc_history`` after
                    calling ``perm_policy.get_perms``.

            Note that if ``verbosity > 1``, the simulator calls
            ``perm_policy.reset(save_history=True)`` before calling
            ``perm_policy.get_perms()``.
        """
        base = len(self.phmm.observation_dist.enumerate_support())
        data = torch.stack(
            [num_to_data(num, num_steps, base) for num in range(base**num_steps)]
        ).float()
        if verbosity > 1:
            save_history = True
        else:
            save_history = False
        if classifier is None:
            classifier = PermClassifier(self.phmm)
        if perm_policy is not None:
            perm_policy.reset(save_history=save_history)
            perms = perm_policy.get_perms(data)
            if save_history:
                history = perm_policy.calc_history
            classi_result = classifier.classify(data, perms=perms, verbosity=verbosity)
        else:
            perms = None
            classi_result = classifier.classify(data, verbosity=verbosity)
        if verbosity:
            classifications, classi_dict = classi_result
            if perm_policy is not None:
                classi_dict[b"perms"] = perms
                if save_history:
                    classi_dict[b"history"] = history
        else:
            classifications = classi_result
        lp = self.phmm.log_prob(data, perms)
        dist = self.phmm.posterior_log_initial_state_dist(data, perms)
        log_joint = dist.T + lp
        ep = ExactPostprocessor(
            log_joint,
            classifications,
        )
        if verbosity:
            return ep, classi_dict
        return ep

    def simulate(self, num_steps, num_samples, classifier=None, perm_policy=None, verbosity=0):
        """
        Computes the data required to compute the misclassification rates
        of the given classifier.

        This method always calls ``perm_policy.reset()`` if ``perm_policy`` is
        not ``None``.

        :param num_steps: Number of steps, int.
        :param num_samples: number of samples to draw from the hmm, int
        :param classifier: Defaults to
            :py:class:`~perm_hmm.classifiers.perm_classifier.PermClassifier`,
            initialized with the hmm ``self.phmm``.
        :param perm_policy: Defaults to None. If specified, will call
            ``self.hmm.sample(perm_policy=perm_policy)``.
        :param verbosity: If ``verbosity == 0``, only the
            :py:class:`~perm_hmm.postprocessing.EmpiricalPostprocessor` needed
            to compute the misclassification rates is returned.

            If ``verbosity == 1``, this method returns a tuple, with the
            postprocessor as the first element, and a dictionary with keys:

                b"posterior_log_initial_state_dist":
                    The posterior log initial state distribution used to compute
                    the classifications.

                b"perms":
                    Only present if ``perm_policy`` is not ``None``. The
                    permutations computed from ``perm_policy.get_perms()``.

            If ``verbosity > 1``, this method returns a tuple, with the
            postprocessor as the first element, and a dictionary with keys as in
            the case with ``verbosity == 1`` and in addition,

                b"history":
                    Whatever is stored in ``perm_policy.calc_history`` after
                    calling ``perm_policy.get_perms``.

            Note that if ``verbosity > 1``, the simulator calls
            ``perm_policy.reset(save_history=True)`` before calling
            ``perm_policy.get_perms()``.
        """
        if verbosity > 1:
            save_history = True
        else:
            save_history = False
        if perm_policy is not None:
            perm_policy.reset(save_history=save_history)
        output = self.phmm.sample((num_samples, num_steps), perm_policy=perm_policy)
        if perm_policy is not None:
            perms = perm_policy.perm_history
        else:
            perms = None
        history = None
        if save_history:
            if perm_policy is not None:
                history = perm_policy.calc_history
        data = output.observations
        if classifier is None:
            classifier = PermClassifier(self.phmm)
        if perms is not None:
            classi_result = classifier.classify(data, perms=perms, verbosity=verbosity)
        else:
            classi_result = classifier.classify(data, verbosity=verbosity)
        if verbosity:
            classifications, classi_dict = classi_result
            classi_dict[b"data"] = data
            if perm_policy is not None:
                classi_dict[b"perms"] = perms
                if history is not None:
                    classi_dict[b"history"] = history
        else:
            classifications = classi_result
        ep = EmpiricalPostprocessor(
            output.states[..., 0],
            classifications,
        )
        if verbosity:
            return ep, classi_dict
        return ep
