from __future__ import print_function

import cv2
import numpy as np
import math
from keras.models import Model
from keras.layers import Input, merge, Convolution2D, MaxPooling2D, UpSampling2D
from keras.optimizers import Adam, SGD
from keras.callbacks import ModelCheckpoint, LearningRateScheduler
from keras import backend as K
from keras.preprocessing.image import ImageDataGenerator
from sklearn.cross_validation import KFold
from sklearn.utils import shuffle
from data import load_train_data, load_test_data
from keras.regularizers import l2, activity_l2
from keras.callbacks import LearningRateScheduler

img_rows = 64
img_cols = 80

smooth = 1.
epochs = 1
# init_learning_rate = 0.1

# Some Model Assumptions.
# Reduction down to 64 by 80 from 420 by 580
#OLD
def dice_coef(y_true, y_pred):
    y_true_f = K.flatten(y_true)
    y_pred_f = K.flatten(y_pred)
    intersection = K.sum(y_true_f * y_pred_f)
    return (2. * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)


def dice_coef_loss(y_true, y_pred):
    return -new_dice_coef(y_true, y_pred)

#NEW
def new_dice_coef(y_true, y_pred):
    y_true_f = K.batch_flatten(y_true)
    y_pred_f = K.batch_flatten(y_pred)
    intersection = 2. * K.sum(y_true_f * y_pred_f, axis=1, keepdims=True) + smooth
    union = K.sum(y_true_f, axis=1, keepdims=True) + K.sum(y_pred_f, axis=1, keepdims=True) + smooth
    return K.mean(intersection / union)

def step_decay(epoch):
    initial_lrate = init_learning_rate
    drop = 0.5
    epochs_drop = 4.0
    lrate = initial_lrate * math.pow(drop, math.floor((1+epoch)/epochs_drop))
    return lrate

def get_unet():
    inputs = Input((1, img_rows, img_cols))
    conv1 = Convolution2D(32, 3, 3, activation='relu', border_mode='same', W_regularizer=l2(0.01))(inputs)
    conv1 = Convolution2D(32, 3, 3, activation='relu', border_mode='same', W_regularizer=l2(0.01))(conv1)
    pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)

    conv2 = Convolution2D(64, 3, 3, activation='relu', border_mode='same')(pool1)
    conv2 = Convolution2D(64, 3, 3, activation='relu', border_mode='same')(conv2)
    pool2 = MaxPooling2D(pool_size=(2, 2))(conv2)

    conv3 = Convolution2D(128, 3, 3, activation='relu', border_mode='same')(pool2)
    conv3 = Convolution2D(128, 3, 3, activation='relu', border_mode='same')(conv3)
    pool3 = MaxPooling2D(pool_size=(2, 2))(conv3)

    conv4 = Convolution2D(256, 3, 3, activation='relu', border_mode='same')(pool3)
    conv4 = Convolution2D(256, 3, 3, activation='relu', border_mode='same')(conv4)
    pool4 = MaxPooling2D(pool_size=(2, 2))(conv4)

    conv5 = Convolution2D(512, 3, 3, activation='relu', border_mode='same', W_regularizer=l2(0.01))(pool4)
    conv5 = Convolution2D(512, 3, 3, activation='relu', border_mode='same', W_regularizer=l2(0.01))(conv5)

    up6 = merge([UpSampling2D(size=(2, 2))(conv5), conv4], mode='concat', concat_axis=1)
    conv6 = Convolution2D(256, 3, 3, activation='relu', border_mode='same')(up6)
    conv6 = Convolution2D(256, 3, 3, activation='relu', border_mode='same')(conv6)

    up7 = merge([UpSampling2D(size=(2, 2))(conv6), conv3], mode='concat', concat_axis=1)
    conv7 = Convolution2D(128, 3, 3, activation='relu', border_mode='same')(up7)
    conv7 = Convolution2D(128, 3, 3, activation='relu', border_mode='same')(conv7)

    up8 = merge([UpSampling2D(size=(2, 2))(conv7), conv2], mode='concat', concat_axis=1)
    conv8 = Convolution2D(64, 3, 3, activation='relu', border_mode='same')(up8)
    conv8 = Convolution2D(64, 3, 3, activation='relu', border_mode='same')(conv8)

    up9 = merge([UpSampling2D(size=(2, 2))(conv8), conv1], mode='concat', concat_axis=1)
    conv9 = Convolution2D(32, 3, 3, activation='relu', border_mode='same', W_regularizer=l2(0.01))(up9)
    conv9 = Convolution2D(32, 3, 3, activation='relu', border_mode='same', W_regularizer=l2(0.01))(conv9)

    conv10 = Convolution2D(1, 1, 1, activation='sigmoid')(conv9)

    model = Model(input=inputs, output=conv10)
    # sgd = SGD(lr=1e-3, decay=1e-6, momentum=0.9, nesterov=True)
    #With: model.compile(optimizer=sgd, loss='categorical_crossentropy', metrics=[dice_coef])
    # model.compile(optimizer=sgd, loss=dice_coef_loss, metrics=[dice_coef])
    model.compile(optimizer=Adam(lr=1e-5), loss=dice_coef_loss, metrics=[new_dice_coef])


    return model

# Preprocess Images, Make any sort of Poolling etc Changes here. 
def preprocess(imgs):
    imgs_p = np.ndarray((imgs.shape[0], imgs.shape[1], img_rows, img_cols), dtype=np.uint8)
    for i in range(imgs.shape[0]):
        imgs_p[i, 0] = cv2.resize(imgs[i, 0], (img_cols, img_rows), interpolation=cv2.INTER_CUBIC)
    return imgs_p


def train_and_predict():
    print('-'*30)
    print('Loading and Preprocessingng train data...')
    print('-'*30)
    imgs_train, imgs_mask_train = load_train_data()

    imgs_train = preprocess(imgs_train)
    imgs_mask_train = preprocess(imgs_mask_train)

    imgs_train = imgs_train.astype('float32')
    mean = np.mean(imgs_train)  # mean for data centering
    std = np.std(imgs_train)  # std for data normalization

    imgs_train -= mean
    imgs_train /= std

    imgs_mask_train = imgs_mask_train.astype('float32')
    imgs_mask_train /= 255.  # scale masks to [0, 1]

    print('-'*30)
    print('Shuffling Training Data...')
    print('-'*30)
    imgs_train, imgs_mask_train = shuffle(imgs_train, imgs_mask_train, random_state=0)

    print('-'*30)
    print('Creating and compiling model...')
    print('-'*30)
    model = get_unet()
    model_checkpoint = ModelCheckpoint('unet.hdf5', monitor='loss', save_best_only=True)

    print('-'*30)
    print('Image Preprocessing...')
    print('-'*30)
    #ImageDataGen
    datagen = ImageDataGenerator(shear_range=0.0,rotation_range=5,height_shift_range=0.01,width_shift_range=0.01)#,fill_mode = "constant", rotation_range=270,height_shift_range=0.1,width_shift_range=0.1,shear_range=0.0,zoom_range=0.2,horizontal_flip=True,vertical_flip=True)
    # datagen.fit(imgs_train, imgs_mask_train, augment=True)

    print('-'*30)
    print('Fitting model...')
    print('-'*30)
    #Step Decay
    # lrate = LearningRateScheduler(step_decay)

    #New Fit
    model.load_weights('unet.hdf5')

    model.fit_generator(datagen.flow(imgs_train, imgs_mask_train,shuffle=True,batch_size=32),samples_per_epoch=imgs_train.shape[0], nb_epoch=epochs,verbose=1, callbacks=[model_checkpoint]) #, save_to_dir='rawtest/'
    #Old Fit
    #model.fit(imgs_train, imgs_mask_train, batch_size=32, nb_epoch=epochs, verbose=1, shuffle=True,validation_split=0.1,callbacks=[model_checkpoint])
    print('-'*30)
    print('Loading and preprocessing test data...')
    print('-'*30)
    imgs_test, imgs_id_test = load_test_data()
    imgs_test = preprocess(imgs_test)

    imgs_test = imgs_test.astype('float32')
    imgs_test -= mean
    imgs_test /= std

    print('-'*30)
    print('Loading saved weights...')
    print('-'*30)
    model.load_weights('unet.hdf5')

    print('-'*30)
    print('Predicting masks on test data...')
    print('-'*30)

    imgs_mask_test = model.predict(imgs_test, verbose=1)
    np.save('imgs_mask_test.npy', imgs_mask_test)


if __name__ == '__main__':
    train_and_predict()
