import numpy as np
import skfmm

"""
    This function to detect the soma position is not the key focus of our work,
    we directly refer the traditional way to detect the location of soma from
    Rivulet2 https://github.com/RivuletStudio/rivuletpy
"""

def loadtiff3d(filepath): # reverse the order of input
    """Load a tiff file into 3D numpy array"""

    import tifffile as tiff
    a = tiff.imread(filepath)

    stack = []
    for sample in a:
        stack.append(np.rot90(np.fliplr(np.flipud(sample)))) # flipud: flips elements in up/down direction
                                                             # fliplf: flips elements in left/right direction
                                                             # rot90: rotate an array 90 degree
    out = np.dstack(stack)

    return out

def det_soma_pos(filepath):
    print("Detecting the soma position...")

    img = loadtiff3d(filepath)
    img = (img-img.min())/(img.max()-img.min())
    bimg = (img > 0.9).astype('int')
    dt = skfmm.distance(bimg, dx=1.1)
    soma_pos = np.asarray(np.unravel_index(dt.argmax(), dt.shape))
    
    return soma_pos
    
