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

"""

Created on Wed Sep  2 15:54:42 2020

@author: Billy

"""

import CellDetector

import DataPreparation

import unet

import os

#Example Script using all of the created classes:

#assemble all svs into a folder and save the name of the folder as svs_loc

svs_loc = "C:\\Users\\Billy\\Downloads\\Data"

#save the model checkpoint folder, so that it can be loaded when the unet is initialised

model_loc = 'N:\\8_(384, 384, 3)\\cp.ckpt'

#initialise your data prep object with the svs_loc

cd = DataPreparation.DataPreparation(svs_loc)

#initialise your unet model with the location of your .cpkt file, 

#which will be on the same level as all of the wieghts and model parameters

unet = unet.uNet_segmentor(model_loc,window = 384, first_conv=8)

#this method automatically crops and covnerts all svs into png images

cd.AutocropAll()

#this will create subfolders, with converted, cropped version of the svs files.

#These subfolders will contain png files corresponding to images from each svs file.

#loop through these subfolders and analyse them, like so:

#firstly the pngs were saved in a sibling folder to svs_loc.

#this process saves the locations of the pngs as png_loc:

dir_ = os.path.split(svs_loc)[0]

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

#using this location, create a list of subfolders:

subfolders = [x[0] for x in os.walk(png_loc)][1:]

# The next part makes predictions for every image that was generated by cd.AutocropAll(),

#and generates a dataset for each subfolder.

#

#It is useful at this stage to go into each subfolder and check the cropping is correct on all photos. Ideally,

#the png images would be cropped such that only the epithelium is contained. Using your operating system's

#photo manipulation app, rotate/straighten the photo and crop, to cut as much stroma and background out as possible.

for subfolder in subfolders:

    subfolder_loc = os.path.join(dir_,subfolder)

    for file in os.listdir(subfolder):

        file_loc = os.path.join(subfolder_loc,file)

        unet.predict(file_loc)

    cd = CellDetector.CellDetector(subfolder_loc, cell_only = True, isOpt = True)

    cd.predictAll(subfolder_loc, isOpt=True, cell_only=False)

    data = cd.dataCleaner(subfolder_loc, cell_only=True)

# if you wish to train the model, you need to provide the location of the training images 

# and their corresponding ground truths, like so:

image_loc = "N:\\Bill_Mcgough\\Correct Labelling\\Labels\\ImageData"

label_loc = "N:\\Bill_Mcgough\\Correct Labelling\\Labels\\MatlabLabelData"

unet.learn_from_matlab(image_loc, label_loc, photo_per_mask=65, batch = 10, epochs = 5000, colour = True)