# score.py
import os
import settings
from tqdm import tqdm
from skimage.io import imread
import numpy as np
import metrics
from keras.models import Model, load_model
from skimage.morphology import label
from skimage.transform import resize

# Settings
max_number_of_samples = 50 # this sets how many samples of the full training set to look at, this is not random, just the first N samples
model_file = "model-dsbowl2018-3.h5"
X_train_file = "X_train_256_256.npy"

# Define the object by object mean IOU calculation for a given image
def object_mean_iou(y_labeled_true, y_labeled_pred):
    num_y_labeled_true = y_labeled_true.max()
    num_y_labeled_pred = y_labeled_pred.max()
    threshold_ious = []
    for threshold in [0.5, 0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95]:
        true_positives = 0
        false_negative = 0
        pred_obj_preload = []
        pred_obj_size = []
        for predicted_object_id in range(1, num_y_labeled_pred + 1):
            pred_obj = y_labeled_pred == predicted_object_id
            pred_obj_preload.append(pred_obj)
            pred_obj_size.append(np.count_nonzero(pred_obj))
        # a true positive for given threshold is when a _single_ object in the prediction corresponds to a given true object
        for true_object_id in range(1, num_y_labeled_true + 1):
            true_obj = y_labeled_true == true_object_id
            true_obj_size = np.count_nonzero(true_obj)
            matches = 0
            for predicted_object_id in range(1, num_y_labeled_pred + 1):
                # calculate the iou for this object and the true object
                this_pred_obj = pred_obj_preload[predicted_object_id-1]
                this_pred_obj_size = pred_obj_size[predicted_object_id-1]
                intersection = np.count_nonzero(true_obj & this_pred_obj)
                union = true_obj_size + this_pred_obj_size - intersection
                iou = intersection / union
                if iou > threshold:
                    matches += 1
            if matches == 1:
                true_positives += 1
            if matches == 0:
                false_negative += 1
        false_positive = num_y_labeled_pred - true_positives
        threshold_ious.append(true_positives / (true_positives + false_positive + false_negative))
    return sum(threshold_ious) / len(threshold_ious)

train_ids = next(os.walk(settings.TRAIN_PATH))[1][:max_number_of_samples]
# this is the 256x256 training data already loaded into a numpy array and resized
X_train = np.load(X_train_file)[:max_number_of_samples]
# this is the true sizes of the training images for upsampling the masks after prediction
sizes_train = np.load("sizes_train.npy")

# Get and resize train images and masks
Y_labeled_true = []
IMG_CHANNELS = settings.IMG_CHANNELS
print('Loading true masks ...')
for n, id_ in tqdm(enumerate(train_ids), total=len(train_ids)):
    path = settings.TRAIN_PATH + id_
    img = imread(path + '/images/' + id_ + '.png')[:,:,:IMG_CHANNELS]
    mask = np.zeros((img.shape[0], img.shape[1], 1), dtype=np.bool)
    next_object_number = 1
    for mask_file in next(os.walk(path + '/masks/'))[2]:
        mask_ = imread(path + '/masks/' + mask_file)
        mask_ = next_object_number * (np.expand_dims(mask_, axis=-1) // 255)
        mask = np.maximum(mask, mask_)
        next_object_number += 1
    Y_labeled_true.append(np.squeeze(mask, -1))

# Use the same labeling methodology here that you use before submitting
model = load_model(model_file, custom_objects={"mean_iou": metrics.mean_iou, "mean_iou2": metrics.mean_iou2})
preds_train = model.predict(X_train, verbose=1)
preds_train_t = (preds_train > 0.5).astype(np.uint8)

# Create list of upsampled test masks
preds_test_upsampled = []
for i in range(len(preds_train_t)):
    preds_test_upsampled.append(resize(np.squeeze(preds_train_t[i]),
                                       (sizes_train[i][0], sizes_train[i][1]),
                                       mode='constant', preserve_range=True))
Y_labeled_pred = [label(pred) for pred in preds_test_upsampled]

images_iou = []
for n in range(len(Y_labeled_true)):
    image_iou = object_mean_iou(Y_labeled_true[n], Y_labeled_pred[n])
    print("Image " + str(n) + " IOU: " + str(image_iou))
    images_iou.append(image_iou)
print("mIOU (LB) Score: " + str(sum(images_iou) / len(images_iou)))