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

@author: david
"""

import numpy
import colorsys
from scipy import misc
from envi import envi

#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"))
def image2class(*masks):
    #get num of mask file names
    num_masks = len(masks)

    if num_masks == 0:
        print("ERROR: mask filenames not provided")
        print("Usage example: image2class(('class_coll.bmp', 'class_epith.bmp'))")
        return

    #load first mask to get dimensions -- assuming all masks have same dimensions
    mask = misc.imread(masks[0], flatten=True).astype(numpy.bool)

    mask_stack = numpy.zeros((mask.shape[0], mask.shape[1], num_masks), dtype=numpy.bool)
    mask_stack[:,:,0] = mask

    #load the rest of masks
    for i in range(1, num_masks):
        mask_stack[:,:,i] = misc.imread(masks[i], flatten=True).astype(numpy.bool)

    return mask_stack


#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))
def generate_training(filename, mask_stack):
    #create envi file object
    E = envi(filename)

    # get number of classes
    C = mask_stack.shape[2]

    #get number of annotated pixels
    num_pixels = 0
    for i in range(0,C):
        num_pixels += numpy.count_nonzero(mask_stack[:,:,i])

    feature_matrix = numpy.zeros((E.header.bands, num_pixels), dtype=E.header.data_type)
    target_matrix = numpy.zeros((1, num_pixels))

    #load masks and append the corresponding pixels to feature matrix and labels to target_matrix
    idx = 0
    for i in range(0,C):
        mask = E.loadmask(mask_stack[:,:,i])
        feature_matrix[:, idx:idx + mask.shape[1]] = mask
        target_matrix[idx:idx+mask.shape[1]] = (i+1)
        idx += mask.shape[1]

    return feature_matrix, target_matrix

#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
def prob2class(prob_image):
    class_image = numpy.zeros_like(prob_image)
    #get nonzero indices
    nnz_idx = numpy.transpose(numpy.nonzero(numpy.sum(prob_image, axis=0)))

    #set pixel corresponding to max probability to 1
    for idx in nnz_idx:
        idx_max_prob = numpy.argmax(prob_image[:, idx[0], idx[1]])
        class_image[idx_max_prob, idx[0], idx[1]] = 1

    return class_image


#create an ROC curve calculator
#input: X x Y x C image giving the probability P(c | x,y)
#output: ROC curve