# -*- coding: utf-8 -*-
"""
Created on Sun Jul 23 16:04:33 2017

@author: david
"""

import numpy
import colorsys
import sklearn
import sklearn.metrics

#generate a 2D color class map using a stack of binary class images
def classcolor2(C):
    
    #determine the number of classes
    nc = C.shape[-1]

    #generate an RGB image
    RGB = numpy.zeros((C.shape[0], C.shape[1], 3), dtype=numpy.ubyte)
    
    #for each class
    for c in range(0, nc):
        hsv = (c * 1.0 / nc, 1, 1)
        color = numpy.asarray(colorsys.hsv_to_rgb(hsv[0], hsv[1], hsv[2])) * 255
        RGB[C[:, :, c], :] = color
    
    return RGB

#create a function that loads a set of class images as a stack of binary masks
#input: list of class image names
#output: X x Y x C stack of binary mask images
#example: image2class(("class_coll.bmp", "class_epith.bmp"))

#create a set of feature/target pairs for classification
#input: envi file object, stack of class masks, list of class names
#output: feature matrix (features x pixels), target matrix (1 x pixels)
#example: generate_training(("class_coll.bmp", "class_epith.bmp"), (1, 2))

#create a class mask stack from an X x Y x C probability image
#input: X x Y x C image giving the probability P(c |x,y)
#output: X x Y x C binary class image

#calculate an ROC curve given a probability image and mask of "True" values
def image2roc(P, t_vals, mask=[]):
    
    if not P.shape == t_vals.shape:
        print("ERROR: the probability and mask images must be the same shape")
        return
    
    #if a mask image isn't provided, create one for the entire image
    if mask == []:
        mask = numpy.ones(t_vals.shape, dtype=numpy.bool)
    
    #create masks for the positive and negative probability scores
    mask_p = t_vals
    mask_n = mask - mask * t_vals
    
    #calculate the indices for the positive and negative scores
    idx_p = numpy.nonzero(mask_p)
    idx_n = numpy.nonzero(mask_n)
    
    Pp = P[idx_p]
    Pn = P[idx_n]

    Lp = numpy.ones((Pp.shape), dtype=numpy.bool)
    Ln = numpy.zeros((Pn.shape), dtype=numpy.bool)
    
    scores = numpy.concatenate((Pp, Pn))
    labels = numpy.concatenate((Lp, Ln))
    
    return sklearn.metrics.roc_curve(labels, scores)