# Author: @ar3792
"""
SMU-Net. Ensembled version of U-Net and Small-U-Net
"""
import os
import sys
import random
import warnings
import tensorflow as tf
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from tqdm import tqdm
from itertools import chain
from skimage.io import imread, imshow, imread_collection, concatenate_images
from skimage.transform import resize
from skimage.morphology import label
from keras.models import Model, load_model
from keras.layers import Input,add, Merge, concatenate,average
from keras.layers.core import Lambda
from keras.layers.convolutional import Conv2D, Conv2DTranspose
from keras.layers.pooling import MaxPooling2D
from keras.layers.merge import concatenate
from keras.callbacks import EarlyStopping, ModelCheckpoint
from keras import backend as K
from sklearn.model_selection import train_test_split

#model = get_basic_unet_model(model_input)
#Defining Image size and channels
IMG_WIDTH = 512
IMG_HEIGHT = 512
IMG_CHANNELS = 3
TRAIN_PATH = '../input/stage1_train/'
TEST_PATH = '../input/stage1_test/'

warnings.filterwarnings('ignore', category=UserWarning, module='skimage')
seed = 42
random.seed = seed
np.random.seed = seed


def load_train_data():
    X_train = np.zeros((len(train_ids), IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS), dtype=np.uint8)
    Y_train = np.zeros((len(train_ids), IMG_HEIGHT, IMG_WIDTH, 1), dtype=np.bool)
    print('Getting and resizing train images and masks ... ')
    sys.stdout.flush()
    for n, id_ in tqdm(enumerate(train_ids), total=len(train_ids)):
        path = TRAIN_PATH + id_
        img = imread(path + '/images/' + id_ + '.png')[:,:,:IMG_CHANNELS]
        img = resize(img, (IMG_HEIGHT, IMG_WIDTH), mode='constant', preserve_range=True)
        X_train[n] = img
        mask = np.zeros((IMG_HEIGHT, IMG_WIDTH, 1), dtype=np.bool)
        for mask_file in next(os.walk(path + '/masks/'))[2]:
            mask_ = imread(path + '/masks/' + mask_file)
            mask_ = np.expand_dims(resize(mask_, (IMG_HEIGHT, IMG_WIDTH), mode='constant', 
                                          preserve_range=True), axis=-1)
            mask = np.maximum(mask, mask_)
        Y_train[n] = mask
    return X_train,Y_train



# Iou metric Adapted and modified from https://www.kaggle.com/aglotero/another-iou-metric
def iou_metric(y_true_in, y_pred_in, print_table=False):
    labels = label(y_true_in > 0.0)
    y_pred = label(y_pred_in > 0.5)
    
    true_obj = len(np.unique(labels))
    pred_obj = len(np.unique(y_pred))

    intersection = np.histogram2d(labels.flatten(), y_pred.flatten(), bins=(true_obj, pred_obj))[0]

    # Compute areas (needed for finding the union between all objects)
    area_true = np.histogram(labels, bins = true_obj)[0]
    area_pred = np.histogram(y_pred, bins = pred_obj)[0]
    area_true = np.expand_dims(area_true, -1)
    area_pred = np.expand_dims(area_pred, 0)

    # Compute union
    union = area_true + area_pred - intersection

    # Exclude background from the analysis
    intersection = intersection[1:,1:]
    union = union[1:,1:]
    union[union == 0] = 1e-9

    # Compute the intersection over union
    iou = intersection / union

    # Precision helper function
    def precision_at(threshold, iou):
        matches = iou > threshold
        true_positives = np.sum(matches, axis=1) == 1   # Correct objects
        false_positives = np.sum(matches, axis=0) == 0  # Missed objects
        false_negatives = np.sum(matches, axis=1) == 0  # Extra objects
        tp, fp, fn = np.sum(true_positives), np.sum(false_positives), np.sum(false_negatives)
        return tp, fp, fn

    # Loop over IoU thresholds
    precision = []
    if print_table:
        print("Thresh\tTP\tFP\tFN\tPrec.")
    for t in np.arange(0.5, 1.0, 0.05):
        tp, fp, fn = precision_at(t, iou)
        if (tp + fp + fn) > 0:
            p = tp / (tp + fp + fn)
        else:
            p = 0
        if print_table:
            print("{:1.3f}\t{}\t{}\t{}\t{:1.3f}".format(t, tp, fp, fn, p))
        precison.append(p)
    
    if print_table:
        print("AP\t-\t-\t-\t{:1.3f}".format(np.mean(precision)))
    return np.mean(precision)

def iou_metric_batch(y_true_in, y_pred_in):
    batch_size = y_true_in.shape[0]
    metric = []
    for batch in range(batch_size):
        value = iou_metric(y_true_in[batch], y_pred_in[batch])
        metric.append(value)
    return np.array(np.mean(metric), dtype=np.float32)

def my_iou_metric(label, pred):
    metric_value = tf.py_func(iou_metric_batch, [label, pred], tf.float32)
    return metric_value

def get_basic_unet_model(model_input):
    #inputs = Input((IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS))
    s = Lambda(lambda x: x / 255) (model_input)

    c1 = Conv2D(8, (3, 3), activation='relu', padding='same') (s)
    c1 = Conv2D(8, (3, 3), activation='relu', padding='same') (c1)
    p1 = MaxPooling2D((2, 2)) (c1)

    c2 = Conv2D(16, (3, 3), activation='relu', padding='same') (p1)
    c2 = Conv2D(16, (3, 3), activation='relu', padding='same') (c2)
    p2 = MaxPooling2D((2, 2)) (c2)

    c3 = Conv2D(32, (3, 3), activation='relu', padding='same') (p2)
    c3 = Conv2D(32, (3, 3), activation='relu', padding='same') (c3)
    p3 = MaxPooling2D((2, 2)) (c3)

    c4 = Conv2D(64, (3, 3), activation='relu', padding='same') (p3)
    c4 = Conv2D(64, (3, 3), activation='relu', padding='same') (c4)
    p4 = MaxPooling2D(pool_size=(2, 2)) (c4)

    c5 = Conv2D(128, (3, 3), activation='relu', padding='same') (p4)
    c5 = Conv2D(128, (3, 3), activation='relu', padding='same') (c5)

    u6 = Conv2DTranspose(64, (2, 2), strides=(2, 2), padding='same') (c5)
    u6 = concatenate([u6, c4])
    c6 = Conv2D(64, (3, 3), activation='relu', padding='same') (u6)
    c6 = Conv2D(64, (3, 3), activation='relu', padding='same') (c6)

    u7 = Conv2DTranspose(32, (2, 2), strides=(2, 2), padding='same') (c6)
    u7 = concatenate([u7, c3])
    c7 = Conv2D(32, (3, 3), activation='relu', padding='same') (u7)
    c7 = Conv2D(32, (3, 3), activation='relu', padding='same') (c7)

    u8 = Conv2DTranspose(16, (2, 2), strides=(2, 2), padding='same') (c7)
    u8 = concatenate([u8, c2])
    c8 = Conv2D(16, (3, 3), activation='relu', padding='same') (u8)
    c8 = Conv2D(16, (3, 3), activation='relu', padding='same') (c8)

    u9 = Conv2DTranspose(8, (2, 2), strides=(2, 2), padding='same') (c8)
    u9 = concatenate([u9, c1], axis=3)
    c9 = Conv2D(8, (3, 3), activation='relu', padding='same') (u9)
    c9 = Conv2D(8, (3, 3), activation='relu', padding='same') (c9)

    outputs = Conv2D(1, (1, 1), activation='sigmoid') (c9)

    model = Model(inputs=[model_input], outputs=[outputs])
    model.compile(optimizer='adam', loss='binary_crossentropy', metrics=[my_iou_metric,'accuracy'])
    model.summary()
    return model

#Adapted from basic U-Net architecture. // Decreases number of hidden neurons to prevent overfitting and reduces number of conv layers
def get_small_unet_model(model_input):
    #inputs = Input((IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS))
    s = Lambda(lambda x: x / 255) (model_input)

    c1 = Conv2D(4, (3, 3), activation='relu', padding='same') (s)
    c1 = Conv2D(4, (3, 3), activation='relu', padding='same') (c1)
    p1 = MaxPooling2D((2, 2)) (c1)

    c2 = Conv2D(8, (3, 3), activation='relu', padding='same') (p1)
    c2 = Conv2D(8, (3, 3), activation='relu', padding='same') (c2)
    p2 = MaxPooling2D((2, 2)) (c2)

    c3 = Conv2D(16, (3, 3), activation='relu', padding='same') (p2)
    c3 = Conv2D(16, (3, 3), activation='relu', padding='same') (c3)
    p3 = MaxPooling2D((2, 2)) (c3)

    c4 = Conv2D(32, (3, 3), activation='relu', padding='same') (p3)
    c4 = Conv2D(32, (3, 3), activation='relu', padding='same') (c4)
    p4 = MaxPooling2D(pool_size=(2, 2)) (c4)

    c5 = Conv2D(64, (3, 3), activation='relu', padding='same') (p4)
    c5 = Conv2D(64, (3, 3), activation='relu', padding='same') (c5)

    u6 = Conv2DTranspose(32, (2, 2), strides=(2, 2), padding='same') (c5)
    u6 = concatenate([u6, c4])
    c6 = Conv2D(32, (3, 3), activation='relu', padding='same') (u6)
    c6 = Conv2D(32, (3, 3), activation='relu', padding='same') (c6)

    u7 = Conv2DTranspose(16, (2, 2), strides=(2, 2), padding='same') (c6)
    u7 = concatenate([u7, c3])
    c7 = Conv2D(16, (3, 3), activation='relu', padding='same') (u7)
    c7 = Conv2D(16, (3, 3), activation='relu', padding='same') (c7)

    u8 = Conv2DTranspose(8, (2, 2), strides=(2, 2), padding='same') (c7)
    u8 = concatenate([u8, c2])
    c8 = Conv2D(8, (3, 3), activation='relu', padding='same') (u8)
    c8 = Conv2D(8, (3, 3), activation='relu', padding='same') (c8)

    u9 = Conv2DTranspose(4, (2, 2), strides=(2, 2), padding='same') (c8)
    u9 = concatenate([u9, c1], axis=3)
    c9 = Conv2D(4, (3, 3), activation='relu', padding='same') (u9)
    c9 = Conv2D(4, (3, 3), activation='relu', padding='same') (c9)

    outputs = Conv2D(1, (1, 1), activation='sigmoid') (c9)

    model = Model(inputs=[model_input], outputs=[outputs])
    model.compile(optimizer='adam', loss='binary_crossentropy', metrics=[my_iou_metric,'accuracy',dice_coef])
    model.summary()
    return model

#Ensemble U-Net and Small-U-Net as SMU-Net
def get_smu_net_model(models,mode_input):
    outputs = concatenate([model.outputs[0] for model in models])
    y = Conv2D(1,(1,1),padding = 'same',activation = 'sigmoid',use_bias=True)(outputs)
    model = Model(model_input, y, name='ensemble')
    
    return model

#Define model
model_input = Input((IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS))

model_1 = get_basic_unet_model(model_input)
model_2 = get_small_unet_model(model_input)

#Training basic model and saving weights with early stopping and check pointing
earlystopper = EarlyStopping(patience=5, verbose=1)
checkpointer = ModelCheckpoint('model_1-dsbowl2018-1_50_epochs.h5', verbose=1, save_best_only=True)
results = model_1.fit(X_train, Y_train, validation_split=0.1, batch_size=4, epochs=50, 
                    callbacks=[earlystopper, checkpointer])

earlystopper = EarlyStopping(patience=5, verbose=1)
checkpointer = ModelCheckpoint('model_2-dsbowl2018-2_50_epochs.h5', verbose=1, save_best_only=True)
results = model_2.fit(X_train, Y_train, validation_split=0.1, batch_size=8, epochs=20, 
                    callbacks=[earlystopper, checkpointer])

#Concatenate models and get ensemble

models = [model_1,model_2]

ensemble_model_conv = get_smu_net_model(models,model_input)

#Make only the last Convolutional layer trainable
for layer in ensemble_model_conv.layers[0:-2]:
    layer.trainable = False
ensemble_model_conv.compile(optimizer='adam', loss='binary_crossentropy', metrics=[my_iou_metric,'accuracy',dice_coef])


#Training weights
earlystopper = EarlyStopping(patience=5, verbose=1)
checkpointer = ModelCheckpoint('ensemble-model-dsbowl2018-2_50_epochs.h5', verbose=1, save_best_only=True)
results = ensemble_model_conv.fit(X_train, Y_train, validation_split=0.1, batch_size=8, epochs=50, 
                    callbacks=[earlystopper, checkpointer])


#Load weights
#model_1.load_weights('model_1-dsbowl2018-1.h5') #Weights available on request

#Testing model 
#Load test data ( Matrices avaliable on request)
X_test = np.load('../../test_data/test_matrices/X_test.npy')
Y_test = np.load('../../test_data/test_matrices/Y_test.npy')

#Evaluate model 

ensemble_model_conv.evaluate(X_test*255.0,Y_test) #Multiplied by 255 to account for scale





