# -*- coding: utf-8 -*-
"""
Created on Sun Mar 12 21:54:40 2017

@author: david
"""

import numpy
import scipy.ndimage
import progressbar
import glob

def st2(I, s=1, dtype=numpy.float32):   
    
    #calculates the 2D structure tensor for an image using a gaussian window with standard deviation s
    
    #calculate the gradient
    dI = numpy.gradient(I)
    
    #calculate the dimensions of the tensor field
    field_dims = [dI[0].shape[0], dI[0].shape[1], 3]
    
    #allocate space for the tensor field
    Tg = numpy.zeros(field_dims, dtype=dtype)
    
    #calculate the gradient components of the tensor
    ti = 0
    for i in range(2):
        for j in range(i + 1):
            Tg[:, :, ti] = dI[j] * dI[i]
            ti = ti + 1
    
    #blur the tensor field
    T = numpy.zeros(field_dims, dtype=dtype)
    
    for i in range(3):
        T[:, :, i] = scipy.ndimage.filters.gaussian_filter(Tg[:, :, i], [s, s])

    
    return T

def st3(I, s=1):
    #calculate the structure tensor field for the 3D input image I given the window size s in 3D
    #check the format for the window size
    if type(s) is not list:
        s = [s] * 3
    elif len(s) == 1:
        s = s * 3
    elif len(s) == 2:
        s.insert(1, s[0])
        
    print("\nCalculating gradient...")
    dI = numpy.gradient(I)
    #calculate the dimensions of the tensor field
    field_dims = [dI[0].shape[0], dI[0].shape[1], dI[0].shape[2], 6]
    
    #allocate space for the tensor field
    Tg = numpy.zeros(field_dims, dtype=numpy.float32)
    
    #calculate the gradient components of the tensor
    ti = 0
    print("Calculating tensor components...")
    bar = progressbar.ProgressBar()
    bar.max_value = 6
    for i in range(3):
        for j in range(i + 1):
            Tg[:, :, :, ti] = dI[j] * dI[i]
            ti = ti + 1
            bar.update(ti)
    
    #blur the tensor field
    T = numpy.zeros(field_dims, dtype=numpy.float32)
    
    print("\nConvolving tensor field...")
    bar = progressbar.ProgressBar()
    bar.max_value = 6
    for i in range(6):
        T[:, :, :, i] = scipy.ndimage.filters.gaussian_filter(Tg[:, :, :, i], s)
        bar.update(i+1)
        
    return T

def st(I, s=1):
    if I.ndim == 3:
        return st3(I, s)
    elif I.ndim == 2:
        return st2(I, s)
    else:
        print("Image must be 2D or 3D")
    return
        
   
    
def sym2mat(T):
    #Calculate the full symmetric matrix from a list of lower triangular elements.
    #The lower triangular components in the input field T are assumed to be the
    #   final dimension of the input matrix.
    
    #       | 1  2  4  7  |
    #       | 0  3  5  8  |
    #       | 0  0  6  9  |
    #       | 0  0  0  10 |
   
    in_s = T.shape
    
    #get the number of tensor elements
    n = in_s[T.ndim - 1]
    
    #calculate the dimension of the symmetric matrix
    d = int(0.5 * (numpy.sqrt(8. * n + 1.) - 1.))
    
    #calculate the dimensions of the output field
    out_s = list(in_s)[:-1] + [d] + [d]

    #allocate space for the output field
    R = numpy.zeros(out_s)
    
    ni = 0
    for i in range(d):
        for j in range(i + 1):
            R[..., i, j] = T[..., ni]
            if i != j:
                R[..., j, i] = T[..., ni]
            ni = ni + 1
    
    return R   

def st2vec(S, vector='largest'):
    #Create a color image from a 2D or 3D structure tensor slice
    
    if(S.ndim != 3):
        print("ERROR: a 2D slice is expected")
        return
      
    #convert the field to a full rank-2 tensor
    T = sym2mat(S);
    del(S)
    
    #calculate the eigenvectors and eigenvalues
    l, v = numpy.linalg.eig(T)
    
    #get the dimension of the tensor field
    d = T.shape[2]
    
    #allocate space for the vector field
    V = numpy.zeros([T.shape[0], T.shape[1], 3])
    
    idx = l.argsort()
    
    for di in range(d):
        if vector == 'smallest':
            b = idx[:, :, 0] == di
        elif vector == 'largest':
            b = idx[:, :, d-1] == di
        else:
            b = idx[:, :, 1] == di
        V[b, 0:d] = v[b, :, di]
    
    return V

def loadstack(filemask):
    #Load an image stack as a 3D grayscale array
    
    #get a list of all files matching the given mask
    files = [file for file in glob.glob(filemask)]
    
    #calculate the size of the output stack
    I = scipy.misc.imread(files[0])
    X = I.shape[0]
    Y = I.shape[1]
    Z = len(files)
    
    #allocate space for the image stack
    M = numpy.zeros([X, Y, Z]).astype('float32')
    
    #create a progress bar
    bar = progressbar.ProgressBar()
    bar.max_value = Z
    
    #for each file
    for z in range(Z):
        #load the file and save it to the image stack
        M[:, :, z] = scipy.misc.imread(files[z], flatten="True").astype('float32')
        bar.update(z+1)
    return M

def anisotropy(S):

    Sf = sym2mat(S)
    
    #calculate the eigenvectors and eigenvalues
    l, v = numpy.linalg.eig(Sf)
    
    #store the sorted eigenvalues
    ls = numpy.sort(l)
    l0 = ls[:, :, 0]
    l1 = ls[:, :, 1]
    l2 = ls[:, :, 2]
    
    #calculate the linear anisotropy
    Cl = (l2 - l1)/(l2 + l1 + l0)
    
    #calculate the planar anisotropy
    Cp = 2 * (l1 - l0) / (l2 + l1 + l0)
    
    #calculate the spherical anisotropy
    Cs = 3 * l0 / (l2 + l1 + l0)
    
    #calculate the fractional anisotropy
    l_hat = (l0 + l1 + l2)/3
    fa_num = (l2 - l_hat) ** 2 + (l1 - l_hat) ** 2 + (l0 - l_hat) ** 2;
    fa_den = l0 ** 2 + l1 ** 2 + l2 ** 2
    FA = numpy.sqrt(3./2.) * numpy.sqrt(fa_num) / numpy.sqrt(fa_den)
    
    return FA, Cl, Cp, Cs

def st2amira(filename, T):
    #generates a tensor field that can be imported into Amira
    
    #   0    dx dx   ----> 0
    #   1    dx dy   ----> 1
    #   2    dy dy   ----> 3
    #   3    dx dz   ----> 2
    #   4    dy dz   ----> 4
    #   5    dz dz   ----> 5
    
    #swap the 2nd and 3rd tensor components
    A = numpy.copy(T)
    A[..., 3] = T[..., 2]
    A[..., 2] = T[..., 3]
    
    #roll the tensor axis so that it is the leading component
    #A = numpy.rollaxis(A, A.ndim - 1)
    A.tofile(filename)
    print("\n", A.shape)

def resample3(T, s=2):
    #resample a tensor field by an integer factor s
    #This function first convolves the field with a box filter and then
    #   re-samples to create a smaller field
    
    #check the format for the window size
    if type(s) is not list:
        s = [s] * 3
    elif len(s) == 1:
        s = s * 3
    elif len(s) == 2:
        s.insert(1, s[0])
    s = numpy.array(s)
    
    bar = progressbar.ProgressBar()
    bar.max_value = T.shape[3]
    
    #blur with a uniform box filter of size r
    for t in range(T.shape[3]):
        T[..., t] = scipy.ndimage.filters.uniform_filter(T[..., t], 2 * s)
        bar.update(t+1)
        
    #resample at a rate of r
    R = T[::s[0], ::s[1], ::s[2], :]
    return R

def color3(prefix, T, vector='largest', aniso=True):
    #Saves a stack of color images corresponding to the eigenvector and optionally scaled by anisotropy
    
    bar = progressbar.ProgressBar()
    bar.max_value = T.shape[2]
    
    #for each z-axis slice
    for z in range(T.shape[2]):
        S = T[:, :, z, :]                           #get the slice
        V = st2vec(S, vector='smallest')   #calculate the vector
        C = numpy.absolute(V)                       #calculate the absolute value
        
        if aniso == True:                              #if the image is scaled by anisotropy
            FA, Cl, Cp, Cs = anisotropy(S)          #calculate the anisotropy of the slice
            if vector == 'largest':
                A = Cl
            elif vector == 'smallest':
                A = Cp
        else:                                       #otherwise just scale by 1
            A = numpy.ones(T.shape[0], T.shape[1])
        image = C * numpy.expand_dims(A, 3)
        
        filename = prefix + str(z).zfill(4) + ".bmp"
        scipy.misc.imsave(filename, image)
        bar.update(z + 1)