# -*- coding: utf-8 -*-

"""

Created on Fri Jul 10 13:14:31 2020

@author: Billy

"""

import slideio

import glob

import os

import time

import numpy as np

import gc

import progressbar

from PIL import Image

from skimage import filters

class DataPreparation:

    #this function initialises a data cleaner

    #the job of this class is to generate x20 magnification .png images from svs slides,

    #whilst cutting out as much background and non-epithelial area as possible.

    #

    #this function needs the location of the folder that contains svs images as an input

    def __init__(self, svs_loc, png_loc = None):

        assert type(svs_loc) == type('')

        if not os.path.exists(svs_loc):

            raise ValueError("SVS directory is illegitimate:", svs_loc,". Please enter the full filepath of an existing directory containing .svs files.\n")

        else:

            os.chdir(svs_loc)

            self.svs_images = glob.glob("*.svs")

            if len(self.svs_images) == 0:

                raise ValueError("Directory", svs_loc," exists, but there are no .svs images in this location.")

            self.svs_loc = svs_loc

        print("Found the following outputs:", self.svs_images,"\n")

        if png_loc == None:

            parent_dir, dir_ = os.path.split(svs_loc)

            png_loc = os.path.join(parent_dir,"Prepared_SVS")

        if not os.path.exists(png_loc):

            os.mkdir(png_loc)

        self.png_fold = png_loc

    #this function receives the file location of a folder that contains svs images,

    #and generates a subling subfolder populated with png images.

    #

    #this function also saves a log file, to inform the user about the new image's geometric properties.

    def AutocropAll(self, svs_loc = None, png_loc= None, max_mag= 20):

        if svs_loc == None:

            svs_loc = self.svs_loc

        if png_loc == None:

            png_loc = self.png_fold

        if not os.path.exists(svs_loc):

            raise ValueError("SVS directory is illegitimate:", svs_loc,". Please enter the full filepath of an existing directory containing .svs files.\n")

        if not os.path.exists(png_loc):

            print("File Location", png_loc,"does not exist. Making this directory...")

            os.mkdir(png_loc)

            print("Successfully create .png save directory.")

        def consecutive(data, stepsize=1):

            arr_consec= np.split(data, np.where(np.diff(data) != stepsize)[0]+1)

            return max(arr_consec, key = len) 

        widgets = [

                'Cropping: ', progressbar.Percentage(),

                ' ', progressbar.AnimatedMarker(),

                ' ', progressbar.ETA(),

            ]

        bar = progressbar.ProgressBar(

        widgets=widgets,

        maxval=len(self.svs_images)).start()

        log_loc = os.path.join(png_loc, 'log.txt')

        with open(log_loc, 'w') as filetowrite:

            for i in range(len(self.svs_images)): 

                bar.update(i)

                information ={}

                file,svs = os.path.splitext(self.svs_images[i])

                pic1_loc = os.path.join(svs_loc, self.svs_images[i])

                slide= slideio.open_slide(pic1_loc, 'SVS')

                scene = slide.get_scene(0)

                mag = scene.magnification

                pixel_size = scene.resolution[0]

                _,_,width,height = scene.rect

                img_fold = os.path.join(png_loc, file)

                if not os.path.exists(img_fold):

                    os.makedirs(img_fold)

                while mag>max_mag:

                    width = int(np.round(width/2))

                    height = int(np.round(height/2))

                    mag = mag/2

                    pixel_size = pixel_size*2

                image= scene.read_block(scene.rect,(width,height))

                image_data_bw = image.min(axis=2)

                information['ImageName'] = file

                information['Magnification'] = mag

                information['ImagePixelHeight'] = height

                information['ImagePixelWidth'] = width

                information['PixelSizeMeters'] = pixel_size

                filetowrite.write(str(information))

                filetowrite.write(' \n ')

                filetowrite.write('#####')         

                filetowrite.write(' \n ')

                object_h = self.ObjectSplitter(image_data_bw, axis = 0)

                width,height=np.shape(image_data_bw)

                for j,indices_v in enumerate(object_h):

                    real_objects = self.ObjectSplitter(image_data_bw[0:height,indices_v[0]:indices_v[1]], axis = 1)

                    for k,indices_h in enumerate(real_objects):

                        full_path = os.path.join(img_fold, file+"_"+str(j)+"_mag"+str(int(mag))+".png")

                        if not os.path.exists(full_path):

                            self.BackgroundReducer(image[indices_h[0]:indices_h[1],indices_v[0]:indices_v[1]],  full_path)  

        bar.finish()

    #return the indices splitting pairs of an image array depending on the percentage that of pixel 'completion' along a given axis.

    #This is to say, this function generates indices that an imaged should be cropped between,

    #either vertically or horziontally, based upon the percentage of white background in the image

    #the default percentage is 2%

    def ObjectSplitter(self, image_arr,percentage_threshold=2, axis =0):

        val = filters.threshold_otsu(image_arr)

        data = np.sum(image_arr < val,axis=axis)

        n = data.shape[0]

        data = np.where(data<(percentage_threshold/100)*n,0,1)

        loc_run_start = np.empty(n, dtype=bool)

        loc_run_start[0] = True

        np.not_equal(data[:-1], data[1:], out=loc_run_start[1:])

        run_starts = np.nonzero(loc_run_start)[0].tolist()

        # find run values

        run_values = data[loc_run_start].tolist()

        # find run lengths

        run_lengths = np.diff(np.append(run_starts, n)).tolist()

        counter = 0

        for i in range(len(run_starts)):

            idx = i-counter

            if run_lengths[idx]<0.02*n:

                if idx==0:continue

                if run_lengths[idx-1]>0.05*n and run_values[idx-1]==1:

                    run_lengths[idx-1] += run_lengths[idx]

                    run_lengths.pop(idx)

                    run_starts.pop(idx)

                    run_values.pop(idx)

                    counter+=1

                    continue

                if idx>=len(run_starts)-1:continue

                if run_lengths[idx+1]>0.05*n and run_values[idx+1]==1:

                    run_lengths[idx+1] += run_lengths[idx]

                    run_lengths.pop(idx)

                    run_starts.pop(idx)

                    run_values.pop(idx)

                    counter+=1

                    continue

                run_lengths[idx-1] += run_lengths[idx]

                run_lengths.pop(idx)

                run_starts.pop(idx)

                run_values.pop(idx)

                counter+=1

        object_pairs = []

        for i in range(len(run_values)):

            if not run_values[i]==1:continue

            object_pairs.append((run_starts[i], run_starts[i]+run_lengths[i]))

        return object_pairs

    #this function uses the indices splitting pairs to split input images.

    def BackgroundReducer(self, image, png_save, true_boundary = 0.01):

        def consecutive(data, stepsize=1):

            arr_consec= np.split(data, np.where(np.diff(data) != stepsize)[0]+1)

            return max(arr_consec, key = len) 

        image_data_bw = image.min(axis=2)

        gc.collect()

        non_empty = np.where(image_data_bw<220,True, False)

        non_empty_columns = np.where(np.sum(non_empty,axis=0)>true_boundary*np.shape(non_empty)[1])

        non_empty_rows = np.where(np.sum(non_empty,axis=1)>true_boundary*np.shape(non_empty)[0])

        non_empty_cols_consec = consecutive(non_empty_columns)

        non_empty_rows_consec = consecutive(non_empty_rows) 

        try:

            cropBox = (np.min(non_empty_rows_consec), np.max(non_empty_rows_consec), np.min(non_empty_cols_consec), np.max(non_empty_cols_consec))

        except:

            print("Improper Object found. Moving on...")

            gc.collect()

            return

        Image.fromarray(image[cropBox[0]:cropBox[1]+1, cropBox[2]:cropBox[3]+1 , :]).save(png_save)

        gc.collect()

if __name__ == '__main__':

    a= DataPreparation("C:\\Users\\Billy\\Downloads\\Data")

    time.sleep(2)

    a.AutocropAll()