from skimage import io
from skimage.morphology import disk
from datetime import datetime
from skimage.filters.rank import median
from skimage.color import rgb2gray
from scipy import ndimage as ndi
from skimage import (
    color, measure, morphology, segmentation
)
from skimage.filters import rank

import matplotlib.pyplot as plt
import numpy as np
import cv2


now = datetime.now()

# CHANGE DIRECTION WHERE IMAGE IS
filename = 'D:\Vasily\DIPLOM_THESIS/analysis/x100_2/control_2.tif'

# READ IMAGE
orig = io.imread(filename)

# TO GRAYSCALE
if len(orig.shape) == 3:  # if RGB
    img_gs = rgb2gray(orig)
else:
    img_gs = orig  # if not RGB

# CONVERT TO UINT8
img_gs = img_gs / img_gs.max()  # normalizes data in range 0 - 1
img_gs = 255 * img_gs
img_8 = img_gs.astype(np.uint8)

# PREPROCESSING (MEDIAN FILTERING)
img_med = median(img_8, disk(3))

# LOCAL OTSU
ret, img_bin = cv2.threshold(img_med, 100, 255, cv2.THRESH_BINARY)
radius = 15
selem = disk(radius)
local = rank.otsu(img_med, selem)
loc_otsu = img_med >= local

new_otsu = np.zeros(img_bin.shape)
new_otsu[loc_otsu == True] = 1
new_otsu = 255 * new_otsu
new_otsu = new_otsu.astype(np.uint8)

local_otsu = cv2.bitwise_and(new_otsu, img_bin)
img_bin = median(local_otsu, disk(3))

# CLEARING UNNECESSARY OBJECTS
remove_holes = morphology.remove_small_holes(img_bin, 5 ** 2)
remove_objects = morphology.remove_small_objects(remove_holes, 12.5 ** 2)
clear_border = segmentation.clear_border(remove_objects, buffer_size=0, bgval=0, mask=None, out=None)

# WATERSHED SEGMENTATION
distance = ndi.distance_transform_edt(clear_border)
# H-EXTREMA
h = 5
h_maxima = morphology.extrema.h_maxima(distance, h)
label_h_maxima = measure.label(h_maxima)
labels_h = segmentation.watershed(-distance, label_h_maxima, mask=clear_border)
color_labels_h = color.label2rgb(labels_h, bg_label=0)

# MORPHOLOGICAL MEASUREMENT
props = measure.regionprops(labels_h)
area_arr = np.zeros((len(props)))
ecc_arr = np.zeros((len(props)))
perim_arr = np.zeros((len(props)))
circ_arr = np.zeros((len(props)))

for i in range(0, len(props)):
    area_arr[i] = props[i].area * (0.028**2)
    ecc_arr[i] = props[i].eccentricity
    perim_arr[i] = props[i].perimeter * 0.028
    act_circularity = (4 * 3.14 * area_arr[i]) / (perim_arr[i] ** 2)
    circ_arr[i] = act_circularity

properties = ['area', 'eccentricity', 'perimeter']

print('Image name: ', filename.split('/')[-1])
print('Act Number: ', str(len(props)))
print('Total Area:', str(sum(area_arr)))
print('Act Mean Area: ', str(np.mean(area_arr)))
print('Act Std Area: ', str(np.std(area_arr)))
print('Act Circularity: ', str(np.mean(circ_arr)))
print('Act Eccentricity: ', str(np.mean(ecc_arr)))
print('Act Perimeter: ', str(np.mean(perim_arr)))
print('..................')

then = datetime.now()

print('TOTAL TIME: ', then - now)


fsize_title = 16
fsize_lbls = 13
fsize_text = 18
fsize = (9,9)

# COLORMAP MITOCHONDRIA COLORED BY AREA
object_areas = np.array([objf["area"] for objf in props])
object_areas = np.insert(-5000,1,object_areas)
colored_area_image = object_areas[labels_h]

fig, ax = plt.subplots(figsize=fsize)
im = plt.imshow(colored_area_image, cmap='seismic')
plt.title('Mitochondria colored by area', fontsize=fsize_title)
cbar = fig.colorbar(im, ax=ax, shrink=0.85)
cbar.ax.set_title("Area")
plt.axis("off")
plt.show()

# HISTOGRAM OF AREA
fsize_title = 16
fsize_lbls = 13
fsize_text = 18
fig, ax = plt.subplots()
plt.hist(object_areas[1:], int(len(object_areas)/5), facecolor='g', alpha=0.9)
plt.xlabel("Area [px²]", fontsize=fsize_lbls)
plt.ylabel("Number of objects [-]", fontsize=fsize_lbls)
plt.title("Area distribution \n" + 'N=' + str(len(props)), fontsize=fsize_title)
plt.text(.11, .89, '\u03BC' + '=' + str(np.round(np.mean(area_arr) / (0.028**2),2)) + ', ' + '\u03C3' + '=' + str(np.round(np.std(area_arr)/(0.028**2),2)), ha='left', va='top',fontsize=fsize_text, transform=ax.transAxes)
plt.grid(True)
plt.show()