# -*- coding: utf-8 -*-
"""
Created on Sun Jul 23 13:52:22 2017

@author: david
"""
import numpy

#sift a 2D hyperspectral image into a PxB matrix where P is the number of pixels and B is the number of bands
def sift2(I, mask = []):
    
    #get the shape of the input array
    S = I.shape
    
    #convert that array into a 1D matrix
    M = numpy.reshape(I, (S[0], S[1] * S[2]))
    
    #gif no mask is provided, just return all pixels
    if mask == []:
        return M

    #if a mask is provided, only return pixels corresponding to that mask
    flatmask = numpy.reshape(mask, (S[1] * S[2]))
    i = numpy.flatnonzero(flatmask)                     #get the nonzero indices
    return M[:, i]                                      #return pixels corresponding to the masked values

def unsift2(M, mask):
    
    #get the size of the input matrix
    S = M.shape
    
    #count the number of nonzero values in the mask
    nnz = numpy.count_nonzero(mask)
    
    #the number of masked values should be the same as the number of pixels in the input matrix
    if len(S) == 1:
    	if not S[0] == nnz:
    		print("ERROR: expected " + str(nnz) + " pixels based on the mask but there are " + str(S[0]) + " in the matrix.")
    elif not S[1] == nnz:
        print("ERROR: expected " + str(nnz) + " pixels based on the mask but there are " + str(S[1]) + " in the matrix.")
    
    
    i = numpy.nonzero(mask)
    
    if len(S) == 1:
    	I = numpy.zeros((1, mask.shape[0], mask.shape[1]), dtype=M.dtype)
    else:
    	I = numpy.zeros((M.shape[0], mask.shape[0], mask.shape[1]), dtype=M.dtype)
    I[:, i[0], i[1]] = M
    return I

#create a function that sifts a color image
#input: image name, mask