# %% importing packages

import argparse
from typing import NamedTuple
import numpy as np
import cv2 as cv
import os
import matplotlib.pyplot as plt
import tqdm
from natsort import natsorted
from glob import glob
import magic
import re
from PIL import Image

plt.rcParams['figure.figsize'] = [50, 150]



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

# You should try and make this a "select this directory and pad/process all the 
# files in it this way" script. You currently have WAY too much repetition
# %%
TissueChainID = '26/01/'
cropping = 0

old_files = 0

fiduciary_files = 1
nodal_files = 1
seg_files = 1
high_res_files = 1

nodal_white = 1
fiduciary_white = 1

reduction_size = 4

if reduction_size == 4:
    reduction_name = 'QuarterScale'

if reduction_size == 8:
    reduction_name = 'EighthScale'



base_directory = '/var/confocaldata/HumanNodal/HeartData/'+ TissueChainID
JPG_directory = '/var/confocaldata/HumanNodal/HeartData/'+ TissueChainID +'JPG/'
jpg_file_names = glob(JPG_directory + '*.jpg')


if old_files:
    ML_directory = '/var/confocaldata/HumanNodal/HeartData/'+ TissueChainID +'uNet_Segmentations/'
    Nodal_Seg_directory = '/var/confocaldata/HumanNodal/HeartData/'+ TissueChainID +'Segmentations/Nodal Segmentation FullScale_NoPad/'
    if fiduciary_files:
        Fiduciary_Seg_directory = '/var/confocaldata/HumanNodal/HeartData/'+ TissueChainID +'Segmentations/Fiduciary Segmentation FullScale_NoPad/'
else:
    ML_directory = '/var/confocaldata/HumanNodal/HeartData/'+ TissueChainID +'uNet_Segmentations/'
    Nodal_Seg_directory = '/var/confocaldata/HumanNodal/HeartData/'+ TissueChainID +'Segmentations/Nodal Segmentation/'
    if fiduciary_files:
        Fiduciary_Seg_directory = '/var/confocaldata/HumanNodal/HeartData/'+ TissueChainID +'Segmentations/Fiduciary Segmentation/'

high_res_directory = '/var/confocaldata/HumanNodal/HeartData/'+ TissueChainID +'HighResSeg/'

os.chdir(ML_directory)
jpg_file_names = natsorted(jpg_file_names)
padding_size = 4000 # + 1536

# %% USE THIS SECTION FOR CROPPING THE SEGMENTATIONS AFTER THE UNET HAS AT IT
if cropping:
    for idx in tqdm.tqdm(range(len(jpg_file_names))):

        out_directory = './../Cropped_uNet_Segmentations/'
        
        # create the directory for saving if it doesn't already exist
        if not os.path.isdir(out_directory):
            os.mkdir(out_directory)

        os.chdir(out_directory)

        jpg_file = jpg_file_names[idx]
        id = jpg_file.split('/')[-1].split('.')[0]
        id = id.split('_')[0] + '_' + id.split('_')[1] + '_' + id.split('_')[2]
        ml_file = glob(ML_directory + f'{id}_*.png')[0]

        jpg_image1 = cv.imread(jpg_file)
        ml_image1 = cv.imread(ml_file)[:,:,0]
        [x,y,z] = jpg_image1.shape

        cropped_ml1 = ml_image1[padding_size:padding_size+x,
                                padding_size:padding_size+y]

        cv.imwrite(
            id + 
            f'_CroppedSeg.png',
            cropped_ml1
            )

# %%
ML_directory = '/var/confocaldata/HumanNodal/HeartData/'+ TissueChainID +'Cropped_uNet_Segmentations/'
ml_file_names = glob(ML_directory + '*.png')
all_image_sizes = []
for file_name in ml_file_names:
    header = magic.from_file(file_name)
    size = re.search('(\d+) x (\d+)',header).groups()
    sizes = [int(a) for a in size]
    all_image_sizes.append(sizes)

max_width = np.max(np.asarray(all_image_sizes)[:,0])
max_height = np.max(np.asarray(all_image_sizes)[:,1])

idx = 200
additional_padding = 4000
os.chdir(JPG_directory)
out_big_directory = base_directory + 'Padded_Images/'
out_small_directory = base_directory + 'Padded_Images_' + reduction_name + '/'
out_parent_list = [out_big_directory,out_small_directory]
out_list = []
if not os.path.isdir(out_big_directory):
    os.mkdir(out_big_directory)

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

for idx, out_directory in enumerate(out_parent_list):
    jpg_out = out_directory + 'JPG'
    seg_out = out_directory + 'Seg'
    nodal_out = out_directory + 'Nodal'
    fiduciary_out = out_directory + 'Fiduciary'
    high_res_out = out_directory + 'HighRes'
    out_list.append([jpg_out,seg_out,nodal_out,high_res_out,fiduciary_out])

    for directory in out_list[idx]:
        if not os.path.isdir(directory):
            os.mkdir(directory)



for idx in tqdm.tqdm(range(len(jpg_file_names))):

    jpg_file = jpg_file_names[idx]
    id = jpg_file.split('/')[-1].split('.')[0]
    id = id.split('_')[0] + '_' + id.split('_')[1] + '_' + id.split('_')[2]
    
    
    
    # Create a separate section for the nodal tissue stuff, as it looks like 
    # nodal segmentation will actually happen after the registration using the
    # segmentations
    # change this to _*.png if you are not using the FullScale_NoPad segmentations
    # will need to scale the segmentations for newer segmentations that haven't been 
    # performed using previously padded images. This section is below, starting with


    jpg_image = cv.imread(jpg_file)
    [height,width,z] = jpg_image.shape

    height_diff = max_height - height
    width_diff = max_width - width

    if height_diff%2 == 1:
        pad_top = np.floor(height_diff/2) + additional_padding
        pad_bottom = np.floor(height_diff/2) + additional_padding
        pad_bottom += 1
    else:
        pad_top = np.floor(height_diff/2) + additional_padding
        pad_bottom = np.floor(height_diff/2) + additional_padding

    if width_diff%2 == 1:
        pad_left = np.floor(width_diff/2) + additional_padding
        pad_right = np.floor(width_diff/2) + additional_padding
        pad_right += 1
    else:
        pad_left = np.floor(width_diff/2) + additional_padding
        pad_right = np.floor(width_diff/2) + additional_padding

    padded_jpg = cv.copyMakeBorder(jpg_image,
                                int(pad_top),
                                int(pad_bottom),
                                int(pad_left),
                                int(pad_right),
                                borderType=cv.cv2.BORDER_CONSTANT,
                                value=[255,255,255])

    os.chdir(out_list[0][0])
    cv.imwrite(
        id + 
        f'_Padded.png',
        padded_jpg
        )

    [pad_height,pad_width,z] = padded_jpg.shape
    

    width_small = int(pad_width/reduction_size)
    height_small = int(pad_height/reduction_size)
    jpg_small = cv.resize(padded_jpg,[width_small,height_small],cv.INTER_AREA)
    
    os.chdir(out_list[1][0])
    cv.imwrite(
        id + 
        f'_Padded_' + reduction_name + '.png',
        jpg_small
        )


    if seg_files:
        ml_file = glob(ML_directory + f'{id}_*.png')[0]
        ml_image = cv.imread(ml_file)[:,:,0]
        padded_seg = cv.copyMakeBorder(ml_image,
                                    int(pad_top),
                                    int(pad_bottom),
                                    int(pad_left),
                                    int(pad_right),
                                    borderType=cv.cv2.BORDER_CONSTANT,
                                    value=[0,0,0])
        os.chdir(out_list[0][1])
        cv.imwrite(
            id + 
            f'_Padded_Seg.png',
            padded_seg
            )

        seg_small = np.array(Image.fromarray(padded_seg).resize((width_small,height_small), Image.NEAREST))
        os.chdir(out_list[1][1])
        cv.imwrite(
            id + 
            f'_Padded_Seg_' + reduction_name + '.png',
            seg_small
            )


    if nodal_files:

        if old_files:
            nodal_file = glob(Nodal_Seg_directory + f'{id}-*.png')[0]
        else:
            nodal_file = glob(Nodal_Seg_directory + f'{id}_*.png')[0]

        nodal_image = cv.imread(nodal_file)[:,:,0]
        # be warry of this, you may need to use this later, though I'm not sure what
        # it was originally used for.
        if nodal_white:
            if sum(sum(nodal_image)) > 0:
                nodal_image = ~nodal_image
        # This accounts for the nodal segmentation images being a quarter the
        # original size, but you should make sure that you haven't already done the 
        # fullscale noPad stuff yet
        if ~old_files:
            nodal_image = np.array(Image.fromarray(nodal_image).resize((width,height), Image.NEAREST))

        padded_nodal = cv.copyMakeBorder(nodal_image,
                                    int(pad_top),
                                    int(pad_bottom),
                                    int(pad_left),
                                    int(pad_right),
                                    borderType=cv.cv2.BORDER_CONSTANT,
                                    value=[0,0,0])

        os.chdir(out_list[0][2])
        cv.imwrite(
            id + 
            f'_Padded_Nodal.png',
            padded_nodal
            )

        nodal_small = np.array(Image.fromarray(padded_nodal).resize((width_small,height_small), Image.NEAREST))
        os.chdir(out_list[1][2])
        cv.imwrite(
            id + 
            f'_Padded_Nodal_' + reduction_name + '.png',
            nodal_small
            )





    if high_res_files:
        high_res_file = glob(high_res_directory + f'{id}_*.png')[0]
        high_res_image = cv.imread(high_res_file)[:,:,0]

        padded_high_res = cv.copyMakeBorder(high_res_image,
                                    int(pad_top),
                                    int(pad_bottom),
                                    int(pad_left),
                                    int(pad_right),
                                    borderType=cv.cv2.BORDER_CONSTANT,
                                    value=[0,0,0])


        os.chdir(out_list[0][3])
        cv.imwrite(
            id + 
            f'_Padded_HighRes.png',
            padded_high_res
            )

        high_res_small = np.array(Image.fromarray(padded_high_res).resize((width_small,height_small), Image.NEAREST))
        os.chdir(out_list[1][3])
        cv.imwrite(
            id + 
            f'_Padded_HighRes_' + reduction_name + '.png',
            high_res_small
            )





    if fiduciary_files:
        if old_files:
            fiduciary_file = glob(Fiduciary_Seg_directory + f'{id}-*.png')[0]
        else:
            fiduciary_file = glob(Fiduciary_Seg_directory + f'{id}_*.png')[0]

        fiduciary_image = cv.imread(fiduciary_file)[:,:,0]

        if fiduciary_white:
            if sum(sum(fiduciary_image)) > 0:
                fiduciary_image = ~fiduciary_image

        if ~old_files:
            fiduciary_image = np.array(Image.fromarray(fiduciary_image).resize((width,height), Image.NEAREST))

        padded_fiduciary = cv.copyMakeBorder(fiduciary_image,
                                    int(pad_top),
                                    int(pad_bottom),
                                    int(pad_left),
                                    int(pad_right),
                                    borderType=cv.cv2.BORDER_CONSTANT,
                                    value=[0,0,0])
        os.chdir(out_list[0][4])
        cv.imwrite(
            id + 
            f'_Padded_Fiduciary.png',
            padded_fiduciary
            )

        fiduciary_small = np.array(Image.fromarray(padded_fiduciary).resize((width_small,height_small), Image.NEAREST))
        os.chdir(out_list[1][4])
        cv.imwrite(
            id + 
            f'_Padded_Fiduciary_' + reduction_name + '.png',
            fiduciary_small
            )

 # %%