# %% Importing packages
import numpy as np
import cv2 as cv
import matplotlib.pyplot as plt
from matplotlib import patches
plt.rcParams['figure.figsize'] = [5, 10]
import os
import tensorflow as tf
from scipy.spatial import distance
import time
from glob import glob
from joblib import Parallel, delayed
from natsort import natsorted
plt.rcParams['figure.figsize'] = [50, 150]

# %% Citations
#############################################################
#############################################################



# %% Defining Functions
#############################################################
#############################################################

def process_file(file,
                 original_seg_directory,
                 high_res_seg_directory,
                 save_directory):

    try:
        file_id = file.split('-')[0] + '_'
        high_res_file = glob(high_res_seg_directory + file_id + '*.png')
        assert len(high_res_file) > 0, \
            f'could not find the pair to {file}, skipping file!'

        original_image = cv.imread(file,cv.IMREAD_UNCHANGED)
        original_seg = original_image[:,:,3]
        high_res_seg = cv.imread(high_res_file[0],cv.IMREAD_UNCHANGED)
        # 1: background, 2: muscle nuclei, 3: connective nuclei
        # 4: muscle, 5: connective, 6: blood

        original_seg[high_res_seg==1] = 0
        original_seg[high_res_seg==6] = 7

        original_image[:,:,3] = original_seg

        new_name = file_id + 'ML5.png'

        os.chdir(save_directory)
        cv.imwrite(new_name,original_image)
        os.chdir(original_seg_directory)


    except AssertionError as e:
        print(e)


#############################################################
#############################################################

original_seg_directory = '/media/briancottle/Samsung_T5/ML_Dataset_3'
high_res_seg_directory = '/home/briancottle/Research/Semantic_Segmentation/'\
                        'High_Res_Ilastik_Segmentation_V10/'\
                        'Ilastik_Segmentation_V10/'

save_directory = '/media/briancottle/Samsung_T5/ML_Dataset_5'

if not os.path.isdir(save_directory):
    os.mkdir(save_directory)

os.chdir(original_seg_directory)
original_files = glob('*.png')

contains_names_vascular = Parallel(
    n_jobs=22, verbose=5)(delayed(process_file)
    (file = name,
    original_seg_directory = original_seg_directory,
    high_res_seg_directory = high_res_seg_directory,
    save_directory = save_directory) for name in original_files
    )