import tensorflow as tf

import config

from model_utils import calculate_conv_output_size

n_x = config.IMAGE_PXL_SIZE_X

n_y = config.IMAGE_PXL_SIZE_Y

n_z = config.SLICES

# This handles padding in both convolution and pooling layers

strides = [[2, 2, 2],

           [2, 3, 3],

           [1, 1, 1],

           [2, 2, 2],

           [1, 1, 1],

           [1, 1, 1],

           [2, 2, 2]]

filters = [[4, 5, 5],

            [3, 4, 4],

            [3, 3, 3],

            [3, 3, 3],

            [3, 3, 3],

            [3, 3, 3],

            [3, 3, 3]]

padding_types = ['VALID'] * 7

# Default network config used with more slices

# and larger convololution stride on first layer

default_config = {

    'weights': [

        # Convolution layers

        ('wc1', tf.truncated_normal([4, 5, 5, config.NUM_CHANNELS, 16], stddev=0.01)),

        ('wc2', tf.truncated_normal([3, 3, 3, 16, 64], stddev=0.01)),

        ('wc3', tf.truncated_normal([3, 3, 3, 64, 64], stddev=0.01)),

        ('wc4', tf.truncated_normal([3, 3, 3, 64, 32], stddev=0.01)),

        # Fully connected layers

        ('wd1', tf.truncated_normal([calculate_conv_output_size(n_x, n_y, n_z, 

                                                                strides, 

                                                                filters,

                                                                padding_types, 

                                                                32), 

                                    100], stddev=0.01)),

        ('wd2', tf.truncated_normal([100, 50], stddev=0.01)),

        ('wout', tf.truncated_normal([50, config.N_CLASSES], stddev=0.01))

    ],

    'biases': (

        # Convolution layers

        ('bc1', tf.zeros([16])),

        ('bc2', tf.constant(1.0, shape=[64])),

        ('bc3', tf.zeros([64])),

        ('bc4', tf.constant(1.0, shape=[32])),

        # Fully connected layers

        ('bd1', tf.constant(1.0, shape=[100])),

        ('bd2', tf.constant(1.0, shape=[50])),

        ('bout', tf.constant(1.0, shape=[config.N_CLASSES]))

    ),

    'pool_strides': [

        [1, 2, 3, 3, 1],

        [1, 2, 2, 2, 1],

        [],

        [1, 2, 2, 2, 1],

    ],

    'pool_windows': [

        [1, 3, 4, 4, 1],

        [1, 3, 3, 3, 1],

        [],

        [1, 3, 3, 3, 1],

    ],

    'strides': [

        [1, 2, 2, 2, 1],

        [1, 1, 1, 1, 1],

        [1, 1, 1, 1, 1],

        [1, 1, 1, 1, 1],

    ]

}