from model_definition.baseline import baseline_config

from model_definition.additional_layers import additional_layers_config

from model_definition.default import default_config

class ModelConfig(object):

    def __init__(self, conv_layers_num, fc_layers_num, config_dict):

        self._weights = config_dict['weights']

        self._biases = config_dict['biases']

        self._fc_layers_with_dropout = []

        self._conv_layers_num = conv_layers_num

        self._fc_layers_num = fc_layers_num

        self._config_dict = config_dict

    def get_fc_weights(self):

        return self._weights[self._conv_layers_num:]

    def get_fc_biases(self):

        return self._biases[self._conv_layers_num:]

    def get_conv_weights(self):

        return self._weights[0:self._conv_layers_num]

    def get_conv_biases(self):

        return self._biases[0:self._conv_layers_num]

    def get_strides(self):

        return self._config_dict['strides']

    def get_pool_strides(self):

        return self._config_dict['pool_strides']

    def get_pool_windows(self):

        return self._config_dict['pool_windows']

    def has_fc_dropout(self, index):

        return index in self._fc_layers_with_dropout

    def has_dropout_after_convolutions(self):

        return False

    def with_l2_norm(self):

        return False

class BaselineConfig(ModelConfig):

    def __init__(self, conv_layers_num=3, fc_layers_num=2, 

                 config_dict=baseline_config):

        super(BaselineConfig, self).__init__(conv_layers_num,

                                             fc_layers_num,

                                             config_dict)

class NoRegularizationConfig(ModelConfig):

    def __init__(self, conv_layers_num=4, fc_layers_num=3, 

                 config_dict=additional_layers_config):

        super(NoRegularizationConfig, self).__init__(

            conv_layers_num, fc_layers_num, config_dict)

# Three different regularization options used for

# building the network configuration

class OneDropoutRegularizationConfig(NoRegularizationConfig):

    def __init__(self, conv_layers_num=4, fc_layers_num=3, 

                 config_dict=additional_layers_config):

        super(OneDropoutRegularizationConfig, self).__init__(

            conv_layers_num, fc_layers_num, config_dict)

        # Dropout on second fully connected layers, zero based indexing

        self._fc_layers_with_dropout = [1]

class DropoutAfterConvolutionsConfig(OneDropoutRegularizationConfig):

    def __init__(self, conv_layers_num=4, fc_layers_num=3, 

                 config_dict=additional_layers_config):

        super(DropoutAfterConvolutionsConfig, self).__init__(

            conv_layers_num, fc_layers_num, config_dict)

    def has_dropout_after_convolutions(self):

        return True

class DropoutsWithL2RegularizationConfig(DropoutAfterConvolutionsConfig):

    def __init__(self, conv_layers_num=4, fc_layers_num=3, 

                 config_dict=additional_layers_config):

        super(DropoutsWithL2RegularizationConfig, self).__init__(

            conv_layers_num, fc_layers_num, config_dict)

    def with_l2_norm(self):

        return True

# With regularization - 2 dropouts and L2 norm

# filters and strides adopted to handle more slices with the same

# network depth

class DefaultConfig(ModelConfig):

    def __init__(self, conv_layers_num=4, fc_layers_num=3, 

                 config_dict=default_config):

        super(DefaultConfig, self).__init__(conv_layers_num,

                                            fc_layers_num,

                                            config_dict)

        self._fc_layers_with_dropout = [1]

    def has_dropout_after_convolutions(self):

        return True

    def with_l2_norm(self):

        return True