from util.image import Image
from scipy.misc import imread
import numpy as np
import config as cfg
import time

"""
Implements methods to classify images
"""


def classify(melanoma, ground, feature, classifier, block=True):
    seg = []
    tim = []
    dim = []

    for (melanoma_item, ground_item) in zip(melanoma, ground):
        print('Segmentating...')
        print('\t'+melanoma_item)
        img = Image(cfg.melanoma_path + melanoma_item, cfg.ground_path + ground_item, cfg.block)
        size = img.get_shape()
        portion = img.get_portion()

        dim.append(size)

        img_seg = np.zeros((size[0], size[1]))

        row = [portion, size[0] - portion]
        col = [portion, size[1] - portion]

        if block:
            st = time.time()
            seg.append(per_block(img, img_seg, row, col, feature, classifier))
            tim.append(time.time() - st)
        else:
            st = time.time()
            seg.append(per_pixel(img, img_seg, row, col, feature, classifier))
            tim.append(time.time() - st)

    return seg, tim, dim


def per_block(img, img_seg, row, col, feature, classifier):
    ratio = 4
    adv = 2 * ratio + 1

    for r in range(row[0], row[1], adv):
        for c in range(col[0], col[1], adv):
            blk = img.get_block([r, c])
            val = feature.features(blk)
            pre = classifier.predict([val])
            img_seg[r: r + adv, c: c + adv] = pre
    return img_seg


def per_pixel(img, img_seg, row, col, feature, classifier):
    for r in range(row[0], row[1]):
        for c in range(col[0], col[1]):
            blk = img.get_block([r, c])
            val = feature.features(blk)
            pre = classifier.predict([val])
            img_seg[r, c] = pre
    return img_seg


def local_error(confmat):
    """
    Calculates the accuracy by each image

    Parameters
    ----------
    confmat: list of lists
        The confusion matrix to calculate the accuracy, TP, FP, FN, TN

    Returns
    -------
        A list of list with 3 values (Sensitivity, Specificity, Accuracy)
    """
    local_err = []
    for mat in confmat:
        TP = mat[0]
        FP = mat[1]
        FN = mat[2]
        TN = mat[3]
        sensitivity = TP / (TP + FN)
        specificity = TN / (TN + FP)
        accuracy = (TP + TN) / (TP + FP + TN + FN)
        local_err.append([sensitivity, specificity, accuracy])
    return local_err


def total_error(local_acc):
    """
    Calculates the mean accuracy of a list of local accuracys

    Parameters
    ----------
    local_acc: list of lists
        The local accuracy of each image

    Returns
    -------
        3 values, (Sensitivity, Specificity, Accuracy)
    """
    acc = np.array(local_acc)
    sensitivity_mean = acc[:, 0].mean()
    sensitivity_std = acc[:, 0].std()

    specificity_mean = acc[:, 1].mean()
    specificity_std = acc[:, 1].std()

    accuracy_mean = acc[:, 2].mean()
    accuracy_std = acc[:, 2].std()

    return [sensitivity_mean, sensitivity_std], [specificity_mean, specificity_std], [accuracy_mean, accuracy_std]


def estimate_error(confmat):
    """
    Calculates the accuracy from a confusion matrix.

    Parameters
    ----------
    confmat: list of lists
        The confusion matrix to calculate the accuracy

    Returns
    -------
        3 values, (Sensitivity, Specificity, Accuracy)
    """
    cm = np.array(confmat)
    TP = cm[:, 0].sum()
    FP = cm[:, 1].sum()
    FN = cm[:, 2].sum()
    TN = cm[:, 3].sum()
    sensitivity = TP / (TP + FN)
    specificity = TN / (TN + FP)
    accuracy = (TP + TN) / (TP + FP + TN + FN)

    return sensitivity, specificity, accuracy


def confusion_matrix(seg, ground_list):
    """
    Calculates the confusion matrix

    Parameters
    ----------
    seg: list
        A list with the segmented images. Each item of the list is a 2D-array
    ground_list: list
        A list with all the ground_truth path of the segmented images.

    Returns
    -------
    A list of 1D-array. Each sublist contents the values TP, FP, FN, TN of a confusion matrix.
    """
    from skimage import io

    acc = []

    for s, g in zip(seg, ground_list):
        ground = imread(cfg.ground_path + g)

        m1 = ground.astype(int)
        m2 = s.astype(int)

        conf = np.zeros((4,), dtype=int) # TP, FP, FN, TN
        for row in range(m1.shape[0]):
            for col in range(m1.shape[1]):
                val1 = m1[row][col]
                val2 = m2[row][col]

                if val1 == val2:
                    if val1 == 255:
                        conf[0] += 1
                    elif val1 == 0:
                        conf[3] += 1
                elif val1 > val2: