import numpy as np

from deepneuro.outputs.inference import ModelInference
from deepneuro.utilities.util import add_parameter


class GANInference(ModelInference):

    def load(self, kwargs):

        """ Parameters
            ----------
            output_num : int, optional
                The amount of batches generated by the GAN.

        """

        super(GANInference, self).load(kwargs)

        # Patching Parameters
        add_parameter(self, kwargs, 'output_num', 10)
        add_parameter(self, kwargs, 'sampling_std', 1)
        add_parameter(self, kwargs, 'sampling_mean', 0)

    def generate(self, model=None):

        self.generate_output_directory()

        if model is not None:
            self.model = model

        self.latent_size = self.model.latent_size

        for batch_idx in range(0, self.output_num, self.batch_size):

            minibatch_length = min(self.batch_size, self.output_num - batch_idx)

            input_data = {'input_data': np.random.normal(self.sampling_mean, self.sampling_std, [minibatch_length, self.latent_size]), 
                    'casename': [''] * minibatch_length, 
                    'input_data_augmentation_string': [str(i) for i in range(batch_idx, batch_idx + minibatch_length)],
                    'input_data_affine': [None] * minibatch_length}

            self.process_case(input_data)
            self.postprocess(input_data)      

        self.close_output()

        return_dict = {'data': self.return_objects, 'filenames': self.return_filenames}
        return return_dict

    def process_case(self, input_data):

        input_data = input_data['input_data']
        output_data = self.model.predict(input_data)

        self.return_objects.append(output_data)

        return output_data