# SCRIPT TO OBTAIN THE AVERAGE, NORMALIZED DLL4 INTENSITY IN IMAGES OF THE RETINA AS A FUNCTION OF THE DISTANCE FROM THE SPROUTING FRONT

from __future__ import division
from skimage import io
from pylab import np, plt
from skimage.io import imread, imshow, imsave
from skimage.filters import rank, gaussian, threshold_otsu
from skimage.morphology import disk, remove_small_objects, skeletonize, label
import scipy.misc
import seaborn as sns
import pandas as pd
import os
from fnmatch import fnmatch
import csv
import time
import scipy.misc
from shutil import copyfile
from scipy.stats import binned_statistic
from joblib import Parallel, delayed
import multiprocessing


def get_genotype(fn):
    if fn.endswith('- WT'):
        return 'WT'
    elif fn.endswith('- KO'):
        return 'KO'
    else:
        return 'None'


inputDir = '/Users/silvanus/Dropbox/angiogenesis/Yap_Taz/Dll4_analysis/data'
folder_list = []
for path, subdirs, files in os.walk(inputDir):
    folder_list.append(path)
folder_list = folder_list[1:]

folder_name = folder_list[0]

all_results = []

bin_edges = range(0, 500, 30)
bin_mps = (np.array(bin_edges[1:]) + np.array(bin_edges[:-1])) / 2

for folder_name in folder_list:
    # load the images (some of the masks are inverted, and some of the Dll4 images are 16 bit instead of 8bit - need to correct for that)
    im_vasc_mask = imread(folder_name + '/vasc_mask.tif') > 0
    if np.mean(im_vasc_mask) > 0.5:
        im_vasc_mask = ~im_vasc_mask
    im_artery_mask = imread(folder_name + '/artery_mask.tif') > 0
    if np.mean(im_artery_mask) > 0.5:
        im_artery_mask = ~im_artery_mask
    im_vasc_mask_wo_artery = im_vasc_mask * (~im_artery_mask)
    im_Dll4 = imread(folder_name + '/Dll4.tif')
    if len(np.shape(im_Dll4)) == 3:
        im_Dll4 = im_Dll4[:, :, 0]
    if np.max(im_Dll4) > 255:
        im_Dll4 = (im_Dll4 / 256.).astype('uint8')
    im_Dll4 = 255 - im_Dll4

    im_front = imread(folder_name + '/front.tif')
    Dll4_artery_mean = np.mean(im_artery_mask.astype('uint8') * im_Dll4)

    # calculate the distance from the front
    front_distance = scipy.ndimage.morphology.distance_transform_edt(im_front[:, :, 0] > 0) * 0.3125849

    # make masked vectors
    Dll4_outside_arteries = im_Dll4.flatten()[im_vasc_mask_wo_artery.flatten()]
    Dll4_in_arteries = im_Dll4.flatten()[im_artery_mask.flatten()]
    Dll4_artery_mean = np.mean(Dll4_in_arteries)
    Dll4_outside_vasculature = im_Dll4.flatten()[~im_vasc_mask.flatten()]
    dist_outside_arteries = front_distance.flatten()[im_vasc_mask_wo_artery.flatten()]
    # Dll4_artery_mean = np.mean(Dll4_outside_vasculature)

    mean_intensity, _, _ = binned_statistic(x=dist_outside_arteries, values=Dll4_outside_arteries / Dll4_artery_mean,
                                            bins=bin_edges)
    std_intensity, _, _ = binned_statistic(x=dist_outside_arteries, values=Dll4_outside_arteries / Dll4_artery_mean,
                                           bins=bin_edges, statistic=np.std)
    count_intensity, _, _ = binned_statistic(x=dist_outside_arteries, values=Dll4_outside_arteries / Dll4_artery_mean,
                                             bins=bin_edges, statistic=len)
    sem_intensity = std_intensity / np.sqrt(count_intensity)

    all_results.append({'folder_name': folder_name, 'mean_intensity': mean_intensity, 'std_intensity': std_intensity,
                        'count_intensity': count_intensity, 'sem_intensity': sem_intensity,
                        'genotype': get_genotype(folder_name)})

    print 'done: ' + folder_name

# plot all together
colorCodes = {'WT': 'red', 'KO': 'green'}
plt.figure()
for retina in all_results:
    plt.plot(bin_mps, retina['mean_intensity'], label=retina['folder_name'][retina['folder_name'].rfind('/') + 1:],
             color=colorCodes[retina['genotype']])
plt.legend()

# calculate and plot means
intensities_WT = []
intensities_KO = []

for retina in all_results:
    if retina['genotype'] == 'KO':
        intensities_KO.append(retina['mean_intensity'])
    elif retina['genotype'] == 'WT':
        intensities_WT.append(retina['mean_intensity'])

mean_intensities_KO = np.mean(intensities_KO, axis=0)
mean_intensities_WT = np.mean(intensities_WT, axis=0)
std_intensities_KO = np.std(intensities_KO, axis=0)
std_intensities_WT = np.std(intensities_WT, axis=0)

sns.set_style('whitegrid')
plt.figure()
sns.tsplot(intensities_WT, time=bin_mps, c='green')
sns.tsplot(intensities_KO, time=bin_mps, c='red')
plt.xlabel('distance to front [um]')
plt.ylabel('mean Dll4 / mean Dll4 in arteries')
plt.ylim([0, 4.0])
plt.xlim([0, 400])
plt.legend()
plt.savefig('outside_normalisation_SEM.pdf')

plt.figure()
plt.errorbar(x=bin_mps, y=mean_intensities_KO, yerr=std_intensities_KO, label='KO', color='red')
plt.errorbar(bin_mps, mean_intensities_WT, yerr=std_intensities_KO, label='WT', color='green')
plt.xlabel('distance to front [um]')
plt.ylabel('mean Dll4 / mean Dll4 in arteries')
plt.ylim([0, 4.0])
plt.xlim([0, 400])
plt.legend()
plt.savefig('outside_normalisation_STD.pdf')

sns.set_style('white')
plt.xlabel('distance to front [um]')
plt.ylabel('mean Dll4 / mean Dll4 in arteries')
plt.ylim([0, 4.0])
plt.xlim([0, 400])
plt.savefig('results/Dll4_KO_vs_WT_outside_normalization_all.pdf')
plt.savefig('results/KO_vs_WT_all.pdf')

im_vasc_mask_pup2 = imread(
    '/Users/silvanus/Dropbox/angiogenesis/YT mutant mice - Dll4 staining analysis/data/pup2B_Ib4_mask.tif') == 0
im_artery_mask_pup2 = imread(
    '/Users/silvanus/Dropbox/angiogenesis/YT mutant mice - Dll4 staining analysis/data/pup2B_artery_mask.tif') == 0
im_vasc_mask_wo_artery_pup2 = im_vasc_mask_pup2 * (~im_artery_mask_pup2)

im_Dll4_pup2 = 255 - (
imread('/Users/silvanus/Dropbox/angiogenesis/YT mutant mice - Dll4 staining analysis/data/pup2B_Dll4.tif')[:, :,
0] / 256).astype('uint8')
im_front_pup2 = imread(
    '/Users/silvanus/Dropbox/angiogenesis/YT mutant mice - Dll4 staining analysis/data/pup2B_front.tif')
Dll4_artery_mean_pup2 = np.mean(im_artery_mask_pup2.astype('uint8') * im_Dll4_pup2)

# calculate the distance from the front
front_distance_pup1 = scipy.ndimage.morphology.distance_transform_edt(im_front_pup1[:, :, 0] > 0) * 0.3125849
front_distance_pup2 = scipy.ndimage.morphology.distance_transform_edt(im_front_pup2[:, :, 0] > 0) * 0.3125849

# create vectors
Dll4_outside_arteries_pup1 = im_Dll4_pup1.flatten()[im_vasc_mask_wo_artery_pup1.flatten()]
Dll4_in_arteries_pup1 = im_Dll4_pup1.flatten()[im_artery_mask_pup1.flatten()]
Dll4_artery_mean_pup1 = np.mean(Dll4_in_arteries_pup1)
Dll4_outside_vasculature_pup1 = im_Dll4_pup1.flatten()[~im_vasc_mask_pup1.flatten()]
dist_outside_arteries_pup1 = front_distance_pup1.flatten()[im_vasc_mask_wo_artery_pup1.flatten()]

Dll4_outside_arteries_pup2 = im_Dll4_pup2.flatten()[im_vasc_mask_wo_artery_pup2.flatten()]
Dll4_in_arteries_pup2 = im_Dll4_pup2.flatten()[im_artery_mask_pup2.flatten()]
Dll4_artery_mean_pup2 = np.mean(Dll4_in_arteries_pup2)
Dll4_outside_vasculature_pup2 = im_Dll4_pup2.flatten()[~im_vasc_mask_pup2.flatten()]
dist_outside_arteries_pup2 = front_distance_pup2.flatten()[im_vasc_mask_wo_artery_pup2.flatten()]

bin_edges = range(0, 500, 15)
bin_mps = (np.array(bin_edges[1:]) + np.array(bin_edges[:-1])) / 2

mean_orientation_pup1, _, _ = binned_statistic(x=dist_outside_arteries_pup1,
                                               values=Dll4_outside_arteries_pup1 / Dll4_artery_mean_pup1,
                                               bins=bin_edges)
std_orientation_pup1, _, _ = binned_statistic(x=dist_outside_arteries_pup1,
                                              values=Dll4_outside_arteries_pup1 / Dll4_artery_mean_pup1,
                                              bins=bin_edges, statistic=np.std)
count_orientation_pup1, _, _ = binned_statistic(x=dist_outside_arteries_pup1,
                                                values=Dll4_outside_arteries_pup1 / Dll4_artery_mean_pup1,
                                                bins=bin_edges, statistic=len)
sem_orientation_pup1 = std_orientation_pup1 / np.sqrt(count_orientation_pup1)

mean_orientation_pup2, _, _ = binned_statistic(x=dist_outside_arteries_pup2,
                                               values=Dll4_outside_arteries_pup2 / Dll4_artery_mean_pup2,
                                               bins=bin_edges)
std_orientation_pup2, _, _ = binned_statistic(x=dist_outside_arteries_pup2,
                                              values=Dll4_outside_arteries_pup2 / Dll4_artery_mean_pup2,
                                              bins=bin_edges, statistic=np.std)
count_orientation_pup2, _, _ = binned_statistic(x=dist_outside_arteries_pup2,
                                                values=Dll4_outside_arteries_pup2 / Dll4_artery_mean_pup2,
                                                bins=bin_edges, statistic=len)
sem_orientation_pup2 = std_orientation_pup2 / np.sqrt(count_orientation_pup2)

sns.set_context('poster')
plt.figure()
plt.plot(bin_mps, mean_orientation_pup1, c='g', label='pup1B (KO)')
plt.plot(bin_mps, mean_orientation_pup2, c='r', label='pup2B (WT)')
plt.legend()
# plt.errorbar(x=bin_mps, y=mean_orientation_wt, yerr=std_orientation_wt, c='g', label='wt')
plt.xlabel('distance to front [um]')
plt.ylabel('mean Dll4 / mean Dll4 in arteries')
plt.ylim([0, 1.0])
plt.xlim([0, 400])
plt.savefig('results/Dll4_KO_vs_WT.pdf')
