from __future__ import absolute_import
from __future__ import print_function
import os

os.environ['KERAS_BACKEND'] = 'theano'
#os.environ['THEANO_FLAGS']='mode=FAST_RUN,device=cuda,floatX=float32,optimizer=None'

import keras.models as models
from keras.layers.core import Layer, Dense, Dropout, Activation, Flatten, Reshape, Permute
from keras.layers.convolutional import Convolution2D, MaxPooling2D, UpSampling2D, ZeroPadding2D
from keras.layers.normalization import BatchNormalization

from keras import backend as K
K.set_image_data_format('channels_first')

import numpy as np
import json
np.random.seed(7) # 0bserver07 for reproducibility

data_shape = 360*480

def create_encoding_layers():
    '''
    ZeroPadding2D(padding=(pad,pad)),
    Convolution2D(filter_size, kernel, kernel, border_mode='valid'),
    BatchNormalization(),
    Activation('relu'),
    MaxPooling2D(pool_size=(pool_size, pool_size)),

    ZeroPadding2D(padding=(pad,pad)),
    Convolution2D(128, kernel, kernel, border_mode='valid'),
    BatchNormalization(),
    Activation('relu'),
    MaxPooling2D(pool_size=(pool_size, pool_size)),

    ZeroPadding2D(padding=(pad,pad)),
    Convolution2D(256, kernel, kernel, border_mode='valid'),
    BatchNormalization(),
    Activation('relu'),
    MaxPooling2D(pool_size=(pool_size, pool_size)),

    ZeroPadding2D(padding=(pad,pad)),
    Convolution2D(512, kernel, kernel, border_mode='valid'),
    BatchNormalization(),
    Activation('relu'),
    '''

    kernel = 3
    filter_size = 64
    pad = 1
    pool_size = 2
    return [
        ZeroPadding2D(padding=(pad,pad)),
        Convolution2D(filter_size, (kernel, kernel), padding='valid'),
        BatchNormalization(),
        Activation('relu'),
        MaxPooling2D(pool_size=(pool_size, pool_size), data_format="channels_first"),

        ZeroPadding2D(padding=(pad,pad)),
        Convolution2D(128, (kernel, kernel), padding='valid'),
        BatchNormalization(),
        Activation('relu'),
        MaxPooling2D(pool_size=(pool_size, pool_size),  data_format="channels_first"),

        ZeroPadding2D(padding=(pad,pad)),
        Convolution2D(256, (kernel, kernel), padding='valid'),
        BatchNormalization(),
        Activation('relu'),
        MaxPooling2D(pool_size=(pool_size, pool_size)),

        ZeroPadding2D(padding=(pad,pad)),
        Convolution2D(512, (kernel, kernel), padding='valid'),
        BatchNormalization(),
        Activation('relu'),
    ]

def create_decoding_layers():
    '''
    ZeroPadding2D(padding=(pad,pad)),
    Convolution2D(512, kernel, kernel, border_mode='valid'),
    BatchNormalization(),

    UpSampling2D(size=(pool_size,pool_size)),
    ZeroPadding2D(padding=(pad,pad)),
    Convolution2D(256, kernel, kernel, border_mode='valid'),
    BatchNormalization(),

    UpSampling2D(size=(pool_size,pool_size)),
    ZeroPadding2D(padding=(pad,pad)),
    Convolution2D(128, kernel, kernel, border_mode='valid'),
    BatchNormalization(),

    UpSampling2D(size=(pool_size,pool_size)),
    ZeroPadding2D(padding=(pad,pad)),
    Convolution2D(filter_size, kernel, kernel, border_mode='valid'),
    BatchNormalization(),
    '''

    kernel = 3
    filter_size = 64
    pad = 1
    pool_size = 2
    return[
        ZeroPadding2D(padding=(pad,pad)),
        Convolution2D(512, (kernel, kernel), padding='valid'),
        BatchNormalization(),

        UpSampling2D(size=(pool_size,pool_size)),
        ZeroPadding2D(padding=(pad,pad)),
        Convolution2D(256, (kernel, kernel), padding='valid'),
        BatchNormalization(),

        UpSampling2D(size=(pool_size,pool_size)),
        ZeroPadding2D(padding=(pad,pad)),
        Convolution2D(128, (kernel, kernel), padding='valid'),
        BatchNormalization(),

        UpSampling2D(size=(pool_size,pool_size)),
        ZeroPadding2D(padding=(pad,pad)),
        Convolution2D(filter_size, (kernel, kernel), padding='valid'),
        BatchNormalization(),
    ]

segnet_basic = models.Sequential()

segnet_basic.add(Layer(input_shape=(3, 360, 480)))

segnet_basic.encoding_layers = create_encoding_layers()
for l in segnet_basic.encoding_layers:
    segnet_basic.add(l)

# Note: it this looks weird, that is because of adding Each Layer using that for loop
# instead of re-writting mode.add(somelayer+params) everytime.

segnet_basic.decoding_layers = create_decoding_layers()
for l in segnet_basic.decoding_layers:
    segnet_basic.add(l)

#segnet_basic.add(Convolution2D(12, 1, 1, border_mode='valid',))
segnet_basic.add(Convolution2D( 12 , (1, 1), padding='valid',))

segnet_basic.add(Reshape((12,data_shape), input_shape=(12,360,480)))
segnet_basic.add(Permute((2, 1)))
segnet_basic.add(Activation('softmax'))

# Save model to JSON

with open('segNet_basic_model.json', 'w') as outfile:
    outfile.write(json.dumps(json.loads(segnet_basic.to_json()), indent=2))