# -*- coding: utf-8 -*-
"""
Created on Wed Mar 01 17:08:10 2017

@author: Shreyas_V
"""
import Brain_pipeline
import Metrics
import numpy as np

#For testing a model

def predict_labels(test_images, model, mu, sigma):
    '''
    INPUT: a numpy array of 4x240x240 elements and a keras model already trained over various images
    OUTPUT: a numpy array of 240x240 label elements
    '''
    predicted_images = []   # list to maintain predicted labels
    for i in test_images:
        msk = ((i[0]+i[1]+i[2]+i[3])!=0.)
        patches = Brain_pipeline.test_patches(i, mu, sigma)
        if patches.shape[0] == 0:
            predicted_images.append(np.full((240, 240), 0))
            continue
        print "running..."
        predicted_slice = model.predict_classes(patches)
        predicted_slice = Brain_pipeline.reconstruct_labels(msk, predicted_slice)
        predicted_images.append(predicted_slice)
    return np.array(predicted_images)
    
def get_metrics(test_images, gt, msk):
    DSC = []
    acc = []
    DSC_core = []
    PPV = []
    for i, j in zip(test_images, gt):
        DSC.append(Metrics.DSC(i, j))
        acc.append(Metrics.accuracy(i, j, msk))
        DSC_core.append(Metrics.DSC_core_tumor(i, j))
        PPV.append(Metrics.PPV(i, j))
    return DSC, acc, DSC_core, PPV

def test_slices(test_images, gt, model, mu=0, sigma=1):
    pred = predict_labels(test_images, model, mu, sigma)
    msk = []
    for i in test_images:
        msk.append((i[0]+i[1]+i[2]+i[3])!=0.)
    return pred, get_metrics(pred, gt, msk)