import tensorflow as tf
from training_utils.architecture.simple_mlp import simple_mlp


class BaseModel(object):
    """
    A base class for an input-output model that can be trained.
    """
    
    def __init__(self, params):
        self.p = params
        self.make_architecture()
        self.make_processing_functions()
        
    def make_architecture(self):
        """
        Create the NN architecture for the model.
        """
        self.arch = simple_mlp(num_inputs=self.p.model.num_inputs,
                               num_outputs=self.p.model.num_outputs,
                               params=self.p.model.arch)
    
    def compute_loss_function(self, raw_data, is_training=None, return_loss_components=False,
                              return_loss_components_and_output=False):
        """
        Compute the loss function for a given dataset.
        """
        # Create the NN inputs and labels
        processed_data = self.create_nn_inputs_and_outputs(raw_data, is_training=is_training)

        # Predict the NN output
        nn_output = self.predict_nn_output(processed_data['inputs'], is_training=is_training)
        
        # Compute the regularization loss, prediction loss and the total loss
        regularization_loss = 0.
        model_variables = self.get_trainable_vars()
        for model_variable in model_variables:
            regularization_loss += tf.nn.l2_loss(model_variable)
        regularization_loss = self.p.loss.regn * regularization_loss
        
        if self.p.loss.loss_type == 'mse':
            prediction_loss = tf.losses.mean_squared_error(nn_output, processed_data['labels'])
        elif self.p.loss.loss_type == 'l2_loss':
            prediction_loss = tf.nn.l2_loss(nn_output - processed_data['labels'])
        else:
            raise NotImplementedError
        
        total_loss = prediction_loss + regularization_loss
       
        if return_loss_components_and_output:
            return regularization_loss, prediction_loss, total_loss, nn_output
        elif return_loss_components:
            return regularization_loss, prediction_loss, total_loss
        else:
            return total_loss
    
    def get_trainable_vars(self):
        """
        Get a list of the trainable variables of the model.
        """
        return self.arch.variables
    
    def create_nn_inputs_and_outputs(self, raw_data, is_training=None):
        """
        Create the NN inputs and outputs from the raw data batch. All pre-processing should go here.
        """
        raise NotImplementedError
    
    def predict_nn_output(self, data, is_training=None):
        """
        Predict the NN output to a given input.
        """
        assert is_training is not None
        
        if is_training:
            # Use dropouts
            tf.keras.backend.set_learning_phase(1)
        else:
            # Do not use dropouts
            tf.keras.backend.set_learning_phase(0)
        
        return self.arch.predict_on_batch(data)

    def predict_nn_output_with_postprocessing(self, data, is_training=None):
        """
        Predict the NN output to a given input with an optional post processing function
        applied. By default there is no post processing function applied.
        """
        return self.predict_nn_output(data, is_training=is_training)

    def make_processing_functions(self):
        """
        Initialize the processing functions if required.
        """
        return

cef765be086761d8af8cffb97a8cb81d49c278b4
<mask token>


class BaseModel(object):
    <mask token>

    def __init__(self, params):
        self.p = params
        self.make_architecture()
        self.make_processing_functions()
    <mask token>

    def compute_loss_function(self, raw_data, is_training=None,
        return_loss_components=False, return_loss_components_and_output=False):
        """
        Compute the loss function for a given dataset.
        """
        processed_data = self.create_nn_inputs_and_outputs(raw_data,
            is_training=is_training)
        nn_output = self.predict_nn_output(processed_data['inputs'],
            is_training=is_training)
        regularization_loss = 0.0
        model_variables = self.get_trainable_vars()
        for model_variable in model_variables:
            regularization_loss += tf.nn.l2_loss(model_variable)
        regularization_loss = self.p.loss.regn * regularization_loss
        if self.p.loss.loss_type == 'mse':
            prediction_loss = tf.losses.mean_squared_error(nn_output,
                processed_data['labels'])
        elif self.p.loss.loss_type == 'l2_loss':
            prediction_loss = tf.nn.l2_loss(nn_output - processed_data[
                'labels'])
        else:
            raise NotImplementedError
        total_loss = prediction_loss + regularization_loss
        if return_loss_components_and_output:
            return regularization_loss, prediction_loss, total_loss, nn_output
        elif return_loss_components:
            return regularization_loss, prediction_loss, total_loss
        else:
            return total_loss

    def get_trainable_vars(self):
        """
        Get a list of the trainable variables of the model.
        """
        return self.arch.variables

    def create_nn_inputs_and_outputs(self, raw_data, is_training=None):
        """
        Create the NN inputs and outputs from the raw data batch. All pre-processing should go here.
        """
        raise NotImplementedError

    def predict_nn_output(self, data, is_training=None):
        """
        Predict the NN output to a given input.
        """
        assert is_training is not None
        if is_training:
            tf.keras.backend.set_learning_phase(1)
        else:
            tf.keras.backend.set_learning_phase(0)
        return self.arch.predict_on_batch(data)

    def predict_nn_output_with_postprocessing(self, data, is_training=None):
        """
        Predict the NN output to a given input with an optional post processing function
        applied. By default there is no post processing function applied.
        """
        return self.predict_nn_output(data, is_training=is_training)

    def make_processing_functions(self):
        """
        Initialize the processing functions if required.
        """
        return

<mask token>


class BaseModel(object):
    <mask token>

    def __init__(self, params):
        self.p = params
        self.make_architecture()
        self.make_processing_functions()

    def make_architecture(self):
        """
        Create the NN architecture for the model.
        """
        self.arch = simple_mlp(num_inputs=self.p.model.num_inputs,
            num_outputs=self.p.model.num_outputs, params=self.p.model.arch)

    def compute_loss_function(self, raw_data, is_training=None,
        return_loss_components=False, return_loss_components_and_output=False):
        """
        Compute the loss function for a given dataset.
        """
        processed_data = self.create_nn_inputs_and_outputs(raw_data,
            is_training=is_training)
        nn_output = self.predict_nn_output(processed_data['inputs'],
            is_training=is_training)
        regularization_loss = 0.0
        model_variables = self.get_trainable_vars()
        for model_variable in model_variables:
            regularization_loss += tf.nn.l2_loss(model_variable)
        regularization_loss = self.p.loss.regn * regularization_loss
        if self.p.loss.loss_type == 'mse':
            prediction_loss = tf.losses.mean_squared_error(nn_output,
                processed_data['labels'])
        elif self.p.loss.loss_type == 'l2_loss':
            prediction_loss = tf.nn.l2_loss(nn_output - processed_data[
                'labels'])
        else:
            raise NotImplementedError
        total_loss = prediction_loss + regularization_loss
        if return_loss_components_and_output:
            return regularization_loss, prediction_loss, total_loss, nn_output
        elif return_loss_components:
            return regularization_loss, prediction_loss, total_loss
        else:
            return total_loss

    def get_trainable_vars(self):
        """
        Get a list of the trainable variables of the model.
        """
        return self.arch.variables

    def create_nn_inputs_and_outputs(self, raw_data, is_training=None):
        """
        Create the NN inputs and outputs from the raw data batch. All pre-processing should go here.
        """
        raise NotImplementedError

    def predict_nn_output(self, data, is_training=None):
        """
        Predict the NN output to a given input.
        """
        assert is_training is not None
        if is_training:
            tf.keras.backend.set_learning_phase(1)
        else:
            tf.keras.backend.set_learning_phase(0)
        return self.arch.predict_on_batch(data)

    def predict_nn_output_with_postprocessing(self, data, is_training=None):
        """
        Predict the NN output to a given input with an optional post processing function
        applied. By default there is no post processing function applied.
        """
        return self.predict_nn_output(data, is_training=is_training)

    def make_processing_functions(self):
        """
        Initialize the processing functions if required.
        """
        return

<mask token>


class BaseModel(object):
    """
    A base class for an input-output model that can be trained.
    """

    def __init__(self, params):
        self.p = params
        self.make_architecture()
        self.make_processing_functions()

    def make_architecture(self):
        """
        Create the NN architecture for the model.
        """
        self.arch = simple_mlp(num_inputs=self.p.model.num_inputs,
            num_outputs=self.p.model.num_outputs, params=self.p.model.arch)

    def compute_loss_function(self, raw_data, is_training=None,
        return_loss_components=False, return_loss_components_and_output=False):
        """
        Compute the loss function for a given dataset.
        """
        processed_data = self.create_nn_inputs_and_outputs(raw_data,
            is_training=is_training)
        nn_output = self.predict_nn_output(processed_data['inputs'],
            is_training=is_training)
        regularization_loss = 0.0
        model_variables = self.get_trainable_vars()
        for model_variable in model_variables:
            regularization_loss += tf.nn.l2_loss(model_variable)
        regularization_loss = self.p.loss.regn * regularization_loss
        if self.p.loss.loss_type == 'mse':
            prediction_loss = tf.losses.mean_squared_error(nn_output,
                processed_data['labels'])
        elif self.p.loss.loss_type == 'l2_loss':
            prediction_loss = tf.nn.l2_loss(nn_output - processed_data[
                'labels'])
        else:
            raise NotImplementedError
        total_loss = prediction_loss + regularization_loss
        if return_loss_components_and_output:
            return regularization_loss, prediction_loss, total_loss, nn_output
        elif return_loss_components:
            return regularization_loss, prediction_loss, total_loss
        else:
            return total_loss

    def get_trainable_vars(self):
        """
        Get a list of the trainable variables of the model.
        """
        return self.arch.variables

    def create_nn_inputs_and_outputs(self, raw_data, is_training=None):
        """
        Create the NN inputs and outputs from the raw data batch. All pre-processing should go here.
        """
        raise NotImplementedError

    def predict_nn_output(self, data, is_training=None):
        """
        Predict the NN output to a given input.
        """
        assert is_training is not None
        if is_training:
            tf.keras.backend.set_learning_phase(1)
        else:
            tf.keras.backend.set_learning_phase(0)
        return self.arch.predict_on_batch(data)

    def predict_nn_output_with_postprocessing(self, data, is_training=None):
        """
        Predict the NN output to a given input with an optional post processing function
        applied. By default there is no post processing function applied.
        """
        return self.predict_nn_output(data, is_training=is_training)

    def make_processing_functions(self):
        """
        Initialize the processing functions if required.
        """
        return

import tensorflow as tf
from training_utils.architecture.simple_mlp import simple_mlp


class BaseModel(object):
    """
    A base class for an input-output model that can be trained.
    """

    def __init__(self, params):
        self.p = params
        self.make_architecture()
        self.make_processing_functions()

    def make_architecture(self):
        """
        Create the NN architecture for the model.
        """
        self.arch = simple_mlp(num_inputs=self.p.model.num_inputs,
            num_outputs=self.p.model.num_outputs, params=self.p.model.arch)

    def compute_loss_function(self, raw_data, is_training=None,
        return_loss_components=False, return_loss_components_and_output=False):
        """
        Compute the loss function for a given dataset.
        """
        processed_data = self.create_nn_inputs_and_outputs(raw_data,
            is_training=is_training)
        nn_output = self.predict_nn_output(processed_data['inputs'],
            is_training=is_training)
        regularization_loss = 0.0
        model_variables = self.get_trainable_vars()
        for model_variable in model_variables:
            regularization_loss += tf.nn.l2_loss(model_variable)
        regularization_loss = self.p.loss.regn * regularization_loss
        if self.p.loss.loss_type == 'mse':
            prediction_loss = tf.losses.mean_squared_error(nn_output,
                processed_data['labels'])
        elif self.p.loss.loss_type == 'l2_loss':
            prediction_loss = tf.nn.l2_loss(nn_output - processed_data[
                'labels'])
        else:
            raise NotImplementedError
        total_loss = prediction_loss + regularization_loss
        if return_loss_components_and_output:
            return regularization_loss, prediction_loss, total_loss, nn_output
        elif return_loss_components:
            return regularization_loss, prediction_loss, total_loss
        else:
            return total_loss

    def get_trainable_vars(self):
        """
        Get a list of the trainable variables of the model.
        """
        return self.arch.variables

    def create_nn_inputs_and_outputs(self, raw_data, is_training=None):
        """
        Create the NN inputs and outputs from the raw data batch. All pre-processing should go here.
        """
        raise NotImplementedError

    def predict_nn_output(self, data, is_training=None):
        """
        Predict the NN output to a given input.
        """
        assert is_training is not None
        if is_training:
            tf.keras.backend.set_learning_phase(1)
        else:
            tf.keras.backend.set_learning_phase(0)
        return self.arch.predict_on_batch(data)

    def predict_nn_output_with_postprocessing(self, data, is_training=None):
        """
        Predict the NN output to a given input with an optional post processing function
        applied. By default there is no post processing function applied.
        """
        return self.predict_nn_output(data, is_training=is_training)

    def make_processing_functions(self):
        """
        Initialize the processing functions if required.
        """
        return

import tensorflow as tf
from training_utils.architecture.simple_mlp import simple_mlp


class BaseModel(object):
    """
    A base class for an input-output model that can be trained.
    """
    
    def __init__(self, params):
        self.p = params
        self.make_architecture()
        self.make_processing_functions()
        
    def make_architecture(self):
        """
        Create the NN architecture for the model.
        """
        self.arch = simple_mlp(num_inputs=self.p.model.num_inputs,
                               num_outputs=self.p.model.num_outputs,
                               params=self.p.model.arch)
    
    def compute_loss_function(self, raw_data, is_training=None, return_loss_components=False,
                              return_loss_components_and_output=False):
        """
        Compute the loss function for a given dataset.
        """
        # Create the NN inputs and labels
        processed_data = self.create_nn_inputs_and_outputs(raw_data, is_training=is_training)

        # Predict the NN output
        nn_output = self.predict_nn_output(processed_data['inputs'], is_training=is_training)
        
        # Compute the regularization loss, prediction loss and the total loss
        regularization_loss = 0.
        model_variables = self.get_trainable_vars()
        for model_variable in model_variables:
            regularization_loss += tf.nn.l2_loss(model_variable)
        regularization_loss = self.p.loss.regn * regularization_loss
        
        if self.p.loss.loss_type == 'mse':
            prediction_loss = tf.losses.mean_squared_error(nn_output, processed_data['labels'])
        elif self.p.loss.loss_type == 'l2_loss':
            prediction_loss = tf.nn.l2_loss(nn_output - processed_data['labels'])
        else:
            raise NotImplementedError
        
        total_loss = prediction_loss + regularization_loss
       
        if return_loss_components_and_output:
            return regularization_loss, prediction_loss, total_loss, nn_output
        elif return_loss_components:
            return regularization_loss, prediction_loss, total_loss
        else:
            return total_loss
    
    def get_trainable_vars(self):
        """
        Get a list of the trainable variables of the model.
        """
        return self.arch.variables
    
    def create_nn_inputs_and_outputs(self, raw_data, is_training=None):
        """
        Create the NN inputs and outputs from the raw data batch. All pre-processing should go here.
        """
        raise NotImplementedError
    
    def predict_nn_output(self, data, is_training=None):
        """
        Predict the NN output to a given input.
        """
        assert is_training is not None
        
        if is_training:
            # Use dropouts
            tf.keras.backend.set_learning_phase(1)
        else:
            # Do not use dropouts
            tf.keras.backend.set_learning_phase(0)
        
        return self.arch.predict_on_batch(data)

    def predict_nn_output_with_postprocessing(self, data, is_training=None):
        """
        Predict the NN output to a given input with an optional post processing function
        applied. By default there is no post processing function applied.
        """
        return self.predict_nn_output(data, is_training=is_training)

    def make_processing_functions(self):
        """
        Initialize the processing functions if required.
        """
        return
