#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Created on Thu Jul  9 08:30:11 2020

@author: Billy
"""

import time
import os
import glob
import math
import skimage.measure as measure
from skimage.segmentation import watershed
import pandas as pd
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
import cv2
import random



class CellDetector:
    
    #initialise the Cell Detector
    #This class converts AI outputs into Datasets
    #one must initialise this class with the folder of the AI outputs
    def __init__(self, SaveLoc,outputs = 11, save_prefix = 'predicted'):
        assert type(SaveLoc) == type('')
        
        if not os.path.exists(SaveLoc):
            raise ValueError("Save directory is illegitimate:", SaveLoc,". Please enter the full filepath of an existing directory containing CellSegmentor's outputs.\n")
        else:
            os.chdir(SaveLoc)
            self.net_images = glob.glob("*.png") + glob.glob("*.tif") 
            if len(self.net_images) == 0:
                raise ValueError("Directory", SaveLoc," exists, but there are no Neural Network outputs in this location.")
            self.net_images = [k for k in self.net_images if save_prefix in k]
            if len(self.net_images) == 0:
                raise ValueError("Directory", SaveLoc," exists, but there are no valid images in this directory. Ensure \'", save_prefix,"\' appears in the save name for all Neural Network outputs. This should occur by default; please don't change file names.\n")
        
        print("Found the following outputs:", self.net_images,"\n")
        self.outputs = outputs
        self.save_fold = SaveLoc
        
        
        
    #this function detects all of the cells within a layer of an AI output image and return a dataframe of those cells' properties
    #this function recievees the file location of an image as input.
    #
    #layers: 0 = basal, 1= prickle, 2 = superficial
    def layerDetection(self, imdir, layer = 1, cell_only=True, by_nucleus=False,
                       plot =[], subplot_shape=(0,0), param_open =1, 
                       param_close = 0, param_erosion=3):
        
        assert type(imdir) == type('')
        
        if subplot_shape == (0,0):
            subplot_shape = (int(len(plot)/2)+1, 2)
        elif not(subplot_shape[0]*subplot_shape[1] == len(plot) or subplot_shape[0]*subplot_shape[1] == len(plot)-1):
            print(subplot_shape)
            print(subplot_shape[0]*subplot_shape[1])
            print(len(plot))
            print(subplot_shape[0]*subplot_shape[1] != len(plot) or subplot_shape[0]*subplot_shape[1] != len(plot)-1)
            print("Input Subplot shape is invalid. Automatically being changed to:",(int(len(plot)/2), 2))
            subplot_shape = (int(len(plot)/2)+1, 2)
        
        if not os.path.exists(imdir):
            raise ValueError("Image directory is illegitimate:", imdir,"\n")
            
        dicter = {}
        dicter[0] = 'Basal'
        dicter[1] = 'Prickle'
        dicter[2] = 'Superficial'
        
        if not (layer in dicter):
            raise ValueError("'Layer' arg must be an int between 0-2, corresponding to the layers like such:",dicter)
            
        parent_dir, file = os.path.split(imdir)
        imdata = np.array(Image.open(imdir).convert('L')).astype('uint8')
        
        if layer == 0:
            image = np.where(np.logical_or(imdata==255,imdata==230),255,0).astype('uint8')
            nuclei = np.where(imdata==255,255,0).astype('uint8')
        elif layer ==1:
            image = np.where(np.logical_or(imdata==204,imdata==179),255,0).astype('uint8')
            nuclei = np.where(imdata==204,255,0).astype('uint8')
        else:
            image = np.where(np.logical_or(imdata==153,imdata==128),255,0).astype('uint8')
            nuclei = np.where(imdata==128,255,0).astype('uint8')
    
        kernel = np.ones((3,3),np.uint8)
        opening = cv2.morphologyEx(image,cv2.MORPH_CLOSE,kernel, iterations = param_open)
        erosion = cv2.erode(opening, kernel,iterations = param_erosion)
        closing = cv2.morphologyEx(erosion,cv2.MORPH_OPEN,kernel, iterations = param_close)
        
        if not by_nucleus:
            ret, markers = cv2.connectedComponents(closing)
        else:
            nuclei = cv2.morphologyEx(nuclei,cv2.MORPH_OPEN,kernel, iterations = 1)
            ret, markers = cv2.connectedComponents(nuclei)

        # Marker labelling
        markers = markers.astype(np.int32)
        
        labels = watershed(image, markers, mask=cv2.morphologyEx(image,cv2.MORPH_OPEN,kernel, iterations = 1))
        
        plot_dict = {'AI':imdata,'Watershed':labels,'Original':np.asarray(Image.open(os.path.join(parent_dir,file[10:]))), 'Markers':markers, 'Nuclei':nuclei, 'Binary':image}
        
        if len(plot)==0:
            pass
        elif len(plot) == 1:
            if plot[0] in plot_dict:
                plt.plot(plot_dict[plot[0]])
                plt.title(plot[0])
            else:
                print("Your desired plot",plot[0],"is not available. Chose from this list:",plot_dict.keys())
        else:
            
            index = 0
            while not (plot[index] in plot_dict.keys()):
                index+=1
                
            ax1 = plt.subplot(subplot_shape[0],subplot_shape[1], 1)
            ax1.title.set_text(plot[index])
            ax1.imshow(plot_dict[plot[index]])
                
            for i in range(len(plot)):
                i = i+index
                try:
                    plt.subplot(subplot_shape[0],subplot_shape[1],i+1, sharex=ax1, sharey=ax1)
                    plt.imshow(plot_dict[plot[i]])
                    plt.title(plot[i])
                except:
                    print(plot[i],"is not available to plot, so has been skipped. Only choose values from:",plot_dict.keys())
        
        table = measure.regionprops(labels)
        num_cells = len(table)

        if not cell_only:
            table_nuclei = measure.regionprops(watershed(markers, markers, mask=markers))
            table.extend(table_nuclei)

        lister =[]
        if layer == 0:
            layer='Basal'
        elif layer == 1:
            layer = 'Prickle'
        else:
            layer = 'Superficial'
        
        for i,cell in enumerate(table):
            entry ={}
            if i+1>num_cells:
                entry['Type'] = 'nucleus'
            else:
                entry['Type'] = 'cell'
            circ = (4*np.pi*cell['Area'])/(cell['Perimeter']**2)
            if circ >100:
                circ = np.nan
            area = (cell['convex_area']/2)+cell['perimeter']
            entry['Origin_Image'] = os.path.splitext(file)[0]
            entry['Layer'] = layer
            entry['Identifier'] = cell['label']
            entry['Area'] = area
            if circ == np.nan:
                entry['Perimeter'] = cell['perimeter']*1.1
            else:
                entry['Perimeter'] = math.sqrt((2*area)/circ)
            entry['Centroid_y'] = cell['centroid'][0]
            entry['Centroid_x'] = cell['centroid'][1]
            entry['Solidity'] = cell['solidity']
            entry['Major axis diameter'] = cell['major_axis_length']/2
            entry['Minor axis diameter'] = cell['minor_axis_length']/2
            entry['Circularity'] = circ
            lister.append(entry)
            
        return lister
    
    
    #extracts all of the cell data from each layer in an image and saves a dataframe of that data
    def predictOne(self, imdir = None,isRandom = False,cell_only=True, isOpt=True):
    
        if isOpt: optimise = [(0,3,0),(1,0,1),(1,2,1),(2,0,2),(1,3,2),(4,1,3),(1,3,0)]
        if imdir == None:
            if isRandom:
                img_choice = random.randint(0,len(self.net_images)-1)
            else:
                choice_dict = {}
                for i,file in enumerate(self.net_images):
                    choice_dict[i] = file
                ask_str = "Choose one of the following images:\n"+ str(choice_dict)+ " \n"
                img_choice = int(input(ask_str))
                
            imdir = os.path.join(self.save_fold, self.net_images[img_choice])
        
        data = []
        log = []
        if isOpt:
            for j in range(3):
                data_store = {}
                arr_store = {}
                if j ==0:
                    area_boundary = 100
                else:
                    area_boundary = 250
                for arangement in optimise:
                    opening,erosion,closing = arangement
                    datum = cD.layerDetection(imdir, layer=j, cell_only=cell_only, by_nucleus=False,
                                              param_open=opening, param_close = closing, param_erosion = erosion)
                    key = len(list(filter(lambda d: d['Area'] > area_boundary, datum)))
                    
                    while key in arr_store.keys():
                        key = key+0.01
                    data_store[key] =datum
                    arr_store[key] = str(arangement)+" False"
                    
                    if j ==0:
                        datum = cD.layerDetection(imdir, layer=j, cell_only=cell_only, by_nucleus=True,
                                                  param_open=opening, param_close = closing, param_erosion = erosion)
                        key = len(list(filter(lambda d: d['Area'] > area_boundary, datum)))
                        while key in arr_store.keys():
                            key = key+0.01
                        data_store[key] = datum
                        arr_store[key] = str(arangement)+" True"
                key = np.max(list(arr_store.keys()))
                log.append("Layer "+str(j)+", arrangment: "+arr_store[key])
                data.extend(data_store[key])
        else:
            for i in range(len(self.net_images)): 
                imdir = os.path.join(self.save_fold, self.net_images[i])
                for j in range(3):
                    data.extend(cD.layerDetection(imdir, layer=j, cell_only=cell_only, by_nucleus=False))
            
        df = pd.DataFrame(data)
        save = os.path.join(self.save_fold, "nc_measurements_"+os.path.splitext(os.path.split(imdir)[1])[0]+'.xlsx')
        df.to_excel(save)
        return df
    
    
    #extracts all of the cell data from each layer in each image in a folder and saves a dataframe of that data
    #this function recieves the file location of a folder of images as input
    def predictAll(self, save_loc=None,cell_only=True, isOpt=False):
        #The typical most-succesful configurations of image-tuning parameters
        if isOpt: optimise = [(0,3,0),(1,0,1),(1,2,1),(2,0,2),(1,3,2),(4,1,3),(1,3,0)]
            
        if save_loc == None:
            save_loc = self.save_fold
            
        if not os.path.exists(save_loc):
            raise ValueError("Save directory is illegitimate:", save_loc,". Please enter the full filepath of an existing directory containing CellSegmentor's outputs.\n")
        log = []
        if isOpt:
            for i in range(len(self.net_images)): 
                data = []
                imdir = os.path.join(self.save_fold, self.net_images[i])
                print("Scraping cell data from:",imdir)
                for j in range(3):
                    data_store = {}
                    arr_store = {}
                    if j ==0:
                        area_boundary = 70
                    else:
                        area_boundary = 250
                    for arangement in optimise:
                        opening,erosion,closing = arangement
                        datum = cD.layerDetection(imdir, layer=j, cell_only=cell_only, by_nucleus=False,
                                                  param_open=opening, param_close = closing, param_erosion = erosion)
                        key = len(list(filter(lambda d: d['Area'] > area_boundary, datum)))
                        
                        while key in arr_store.keys():
                            key = key+0.01
                        data_store[key] =datum
                        arr_store[key] = str(arangement)+" False"
                        
                        if j ==0:
                            datum = cD.layerDetection(imdir, layer=j, cell_only=cell_only, by_nucleus=True,
                                                      param_open=opening, param_close = closing, param_erosion = erosion)
                            key = len(list(filter(lambda d: d['Area'] > area_boundary, datum)))
                            while key in arr_store.keys():
                                key = key+0.01
                            data_store[key] = datum
                            arr_store[key] = str(arangement)+" True"
                    key = np.max(list(arr_store.keys()))
                    log.append("Layer "+str(j)+", arrangment: "+arr_store[key])
                    data.extend(data_store[key])
                    
                df = pd.DataFrame(data)
                save = os.path.join(save_loc, "nc_measurements_"+os.path.splitext(self.net_images[i])[0]+'.xlsx')
                df.to_excel(save)
                print("Successfully scraped data. Saving to... ",save)
        else:
            for i in range(len(self.net_images)): 
                data = []
                imdir = os.path.join(self.save_fold, self.net_images[i])
                for j in range(3):
                    data.extend(cD.layerDetection(imdir, layer=j, cell_only=cell_only, by_nucleus=False))
                df = pd.DataFrame(data)
                save = os.path.join(save_loc, "nc_measurements_"+os.path.splitext(self.net_images[i])[0]+'.xlsx')
                df.to_excel(save)
                print("Successfully scraped data. Saving to... ",save)
        print(log)
        return df
    
    
    #aims to clean data, by destroying anaomalous results in a given dataset. This function recieves the file location of the dataset as input.
    # anomalous results are defined by having highly atypical cell area, perimeter and circularity characterisitics.
    def dataCleaner(self,save_loc = None,filename=None,cell_only=True, basal_circularity_bounds=(0.5,1.05), basal_area_bounds = (40,350),
                    prickle_circularity_bounds = (0.5,1.05),prickle_area_bounds=(60,1400), superficial_circularity_bounds=(0.5,1.05),
                    superficial_area_bounds=(60,1500), solidity_bound = 0.7):
        
            
        if save_loc == None:
            save_loc = self.save_fold
            
        if not os.path.exists(save_loc):
            raise ValueError("Save directory is illegitimate:", save_loc,". Please enter the full filepath of an existing directory containing CellSegmentor's outputs.\n")
        
        if filename != None:
            open_ = os.path.join(save_loc, filename)
            non_cleaned = [open_]
        else:
            os.chdir(save_loc)
            non_cleaned = glob.glob("*xlsx*")
            non_cleaned = [k for k in non_cleaned if 'nc_measurement' in k]
            
        
        print(non_cleaned)
        layers = ['Basal','Prickle','Superficial']
        cyto_bounds =[[basal_circularity_bounds,basal_area_bounds],
                       [prickle_circularity_bounds, prickle_area_bounds],
                       [superficial_circularity_bounds, superficial_area_bounds]]
        
        nuc_bounds = [[basal_circularity_bounds,(basal_area_bounds[0]/1.5,basal_area_bounds[1])],
                       [prickle_circularity_bounds, (prickle_area_bounds[0]/1.5, prickle_area_bounds[1]/1.5)],
                       [superficial_circularity_bounds, (superficial_area_bounds[0]*1.5,superficial_area_bounds[1]/1.5) ]]
        
        for image in non_cleaned:
            non_clean =  os.path.join(save_loc,image)
            df = pd.read_excel( non_clean,index=False)
            save_name = os.path.join(save_loc,"clean_data_"+os.path.splitext(image)[0].replace("nc_measurements_","")+'.xlsx')
            cell_store = {'Basal':0, 'Prickle':0, 'Superficial':0}
            for i,layer in enumerate(layers):
                cells = df[(df['Layer']==layer) & (df['Type'] == 'cell')]
                cells.reset_index(drop=True, inplace=True)
                cells = cells[(cells['Circularity']>cyto_bounds[i][0][0]) & (cells['Circularity']<cyto_bounds[i][0][1])]
                cells = cells[(cells['Area']>cyto_bounds[i][1][0]) & (cells['Area']<cyto_bounds[i][1][1])]
                cells = cells[(cells['Circularity']>cyto_bounds[i][0][0]*1.1) & (cells['Area']<cyto_bounds[i][1][1]*0.9)]
                cells = cells[(cells['Circularity']<cyto_bounds[i][0][1]*0.9) & (cells['Area']>cyto_bounds[i][1][0]*1.1)]
                cells = cells[((cells['Solidity']>solidity_bound) | (cells['Area']> 2*cyto_bounds[i][1][0]))]
                cell_store[layer] = cells
                
            if not cell_only:
                nuclei_store = {'Basal':0, 'Prickle':0, 'Superficial':0}
                for i,layer in enumerate(layers):
                    nuclei = df[(df['Layer']==layer) & (df['Type'] == 'nucleus')]
                    nuclei.reset_index(drop=True, inplace=True) 
                    nuclei = nuclei[(nuclei['Circularity']>nuc_bounds[i][0][0]) & (nuclei['Circularity']<nuc_bounds[i][0][1])]
                    nuclei = nuclei[(nuclei['Area']>nuc_bounds[i][1][0]) & (nuclei['Area']<nuc_bounds[i][1][1])]
                    nuclei = nuclei[(nuclei['Circularity']>nuc_bounds[i][0][0]*1.1) & (nuclei['Area']<nuc_bounds[i][1][1]*0.9)]
                    nuclei = nuclei[(nuclei['Circularity']<nuc_bounds[i][0][1]*0.9) & (nuclei['Area']>nuc_bounds[i][1][0]*1.1)]
                    nuclei = nuclei[((nuclei['Solidity']>solidity_bound) | (nuclei['Area']> 2*cyto_bounds[i][1][0]))]
                    nuclei_store[layer] = nuclei
                frames = list(cell_store.values())+list(nuclei_store.values())
            else:
                frames = list(cell_store.values())
            output_df = pd.concat(frames)
            output_df.to_excel(save_name)
            os.remove(non_clean)
    
        
if __name__ == '__main__':
    output_loc = "C:\\Users\\Billy\\Downloads\\Prepared_SVS"
    picLoc = os.path.join(output_loc, 'predicted_cropped_Normal buccal mucosa_0_mag20.png')
    
    cD = CellDetector(output_loc)
    time.sleep(2)
    
    cD.predictOne(picLoc)

    cD.predictAll(output_loc, isOpt=True, cell_only=False)
    
    cD.layerDetection(picLoc,layer=0, plot=['AI','Watershed','Original','Binary'],
                      by_nucleus=False, cell_only=True,subplot_shape=(2,2),
                      param_open=0,  param_erosion = 3, param_close = 0)
    
    data = cD.dataCleaner(output_loc, cell_only=False)