import random

from learning.learning_utils.encoder_sampler_all_random import EncoderSamplerAllRandom
from learning.learning_utils.encoder_sampler_bfs_reuse import EncoderSamplerBFSReUse
from learning.learning_utils.encoder_sampler_reuse import EncoderSamplerReUse
from learning.learning_utils.encoder_sampler_forward_reuse import EncoderSamplerForwardReUse
from learning.learning_utils.encoder_sampler_same_policy import EncoderSamplerSamePolicy


class EncoderSamplerWrapper:
    """ Class for generating samples. A single sample is of the form {(x, a, x', y, s, s')} where
    x and x' are observations, a is an action, s and s' are states and y is a {0, 1} indicator.
    When state information is computable, s and s' represent the state of x and x' respectively. However,


    The samples are generated according to one of the styles:
    ALL_RANDOM: The imposter is generated by following a homing policy randomly uniformly and taking
                an action randomly uniformly.

    SAME_POLICY: The imposter is generated by following the same homing policy as the one that generated x
                but then taking an action randomly uniformly.

    REUSE:  We generate observed data by selecting a homing policy uniformly and then roll-in with it and then
            taking an action uniformly. We then generate negative samples for each observed data (x_i, a_i, x'_i) by
            sampling x'_j ~ uniform({x'_k}_k) and generating fake transitions (x_i, a_i, x'_j).

    FORWARD_REUSE: We generate observed data by selecting a homing policy uniformly and then roll-in with it and then
            taking an action uniformly. We then generate negative samples for each observed data (x_i, a_i, x'_i) by
            sampling (x_j, a_j) ~ uniform({(x'_k, a_k)}_k) and generating fake transitions (x_j, a_j, x'_i).

    BFS_REUSE: We generate observed data by iterating over all possible previous timestep homing policies and actions
            and then roll-in with that policy and taking the action. We then generate negative samples for each
            observed data (x_i, a_i, x'_i) by sampling x'_j ~ uniform({x'_k}_k) and generating fake transitions
            (x_i, a_i, x'_j).
    """

    def __init__(self, constants):
        self.sampling_style = constants["encoder_sampling_style"]
        self.data_aggregation = constants["data_aggregation"]
        self.use_selection_weights = constants["bias_homing_policy"]

        if self.sampling_style == "all_random":
            self.sampler = EncoderSamplerAllRandom()
        elif self.sampling_style == "same_policy":
            self.sampler = EncoderSamplerSamePolicy()
        elif self.sampling_style == "reuse":
            self.sampler = EncoderSamplerReUse()
        elif self.sampling_style == "forwardreuse":
            self.sampler = EncoderSamplerForwardReUse()
        elif self.sampling_style == "bfs_reuse":
            self.sampler = EncoderSamplerBFSReUse()
        else:
            raise AssertionError("Unhandled sampling style %r" % self.sampling_style)

    def gather_samples(self, env, actions, step, homing_policies, num_samples, dataset, selection_weights=None):

        if self.data_aggregation:
            new_dataset = dataset
        else:
            new_dataset = []

        if not self.use_selection_weights:
            # When not biasing the set of homing policies. Algorithm would be pick them uniformly.
            selection_weights = None

        if selection_weights is not None and step > 1:
            assert len(selection_weights) == len(homing_policies[step - 1]), \
                "Weights supplied should match the number of learned policies for the previous step"

        collected_dataset = self.sampler.gather_samples(num_samples, env, actions, step,
                                                        homing_policies, selection_weights)

        new_dataset.extend(collected_dataset)

        if self.data_aggregation:
            # When aggregating the data. Shuffle it so the samples from different time steps are mixed.
            random.shuffle(new_dataset)

        return new_dataset
