import os

import xml.etree.ElementTree as ET

import cv2

# import numpy as np

data_dir = "Sag_1-491_5380x2jpgXmls_ResnetRecrop_Jun223020/"

sag_dir = data_dir[:-1] + "-scaledby150-cropped-version3/"

scale_factor = 0.5

grading = ['1', '2', '3', '4']

grading_LUT = {

            'normal': '1',

            'mild': '2',

            'moderate': '3',

            'severe': '4',

        }

sag_label_file = data_dir[:-1] + "-Jun242020-sagittal-label.txt"

sag_label = {}

with open(sag_label_file, 'r') as f:

    lines = f.readlines()

    for e in lines:

        tmp = e.strip().split(' ')

        tmp2 = []

        for i in range(1, len(tmp)):

            tmp2.append(tmp[i])

        sag_label[tmp[0]] = tmp2

print("sag_label: ", sag_label)

if not os.path.exists(sag_dir):

    os.makedirs(sag_dir)

for e in grading:

    if not os.path.exists(os.path.join(sag_dir, e)):

        os.makedirs(os.path.join(sag_dir, e))

def scale_crop(xmin, xmax, ymin, ymax, factor, img_shape):

    """

    the imgs top left corner is (0,0),

    x min-max is from point A,

    y min-max is from point B

    """

    cropped_w = xmax - xmin

    cropped_h = ymax - ymin

    xmin -= ((cropped_w * factor) // 2)

    ymin -= ((cropped_h * factor) // 2)

    xmax += ((cropped_w * factor) // 2)

    ymax += ((cropped_h * factor) // 2)

    # cv2 img shape information

    height = img_shape[0]

    width = img_shape[1]

    return (

            int(max(xmin, 0)),

            int(min(xmax, width)),

            int(max(ymin, 0)),

            int(min(ymax, height))

            )

for file in os.listdir(data_dir):

    if '.xml' in file:

        print("\nXML file: ", file)

        # Load an color image in grayscale, flag=0

        img = cv2.imread(

            os.path.join(data_dir, file.replace('xml', 'jpg')),

            0)

        # NOTE: some imgs are with ext of JPG uppercase

        # JPG == jpg == jpeg == JPEG in storage

        if img is None:

            img = cv2.imread(

                os.path.join(

                    data_dir,

                    file.replace('xml', 'JPG')

                ), 0)

        print("img.shape: ", img.shape)

        data = ET.parse(os.path.join(data_dir, file))

        root = data.getroot()

        crops = []

        for o in root.findall('object'):

            xmin = int(o.find('bndbox').find('xmin').text)

            xmax = int(o.find('bndbox').find('xmax').text)

            ymin = int(o.find('bndbox').find('ymin').text)

            ymax = int(o.find('bndbox').find('ymax').text)

            print(">>>>before scaling, xmin, xmax, ymin, ymax: ", xmin, xmax, ymin, ymax)

            (xmin, xmax, ymin, ymax) = scale_crop(xmin, xmax, ymin, ymax, scale_factor, img.shape)

            print("<<<<after scaling, xmin, xmax, ymin, ymax: ", xmin, xmax, ymin, ymax)

            cropped_img = img[ymin:ymax, xmin:xmax]

            crops.append((ymin, cropped_img))

        crops = sorted(crops, key=lambda x: x[0])

        for e in crops:

            print("=======crops item:=======")

            print(e[0])

            #  print("cropped img: ", e[1])

            print("img.shape: ", e[1].shape)

        original_labels = sag_label.get(file, '')

        print("original_labels: ", original_labels)

        labels = [grading_LUT.get(i, i) for i in original_labels]

        print("labels: ", labels)

        file_index = 0

        for i in range(0, len(labels)):

            print('labels[i]: ', labels[i])

            #  print('crops[i][1]: ', crops[i][1])

            file_name = os.path.join(

                    sag_dir,

                    labels[i],

                    file[:-3])

            unique_file_name = file_name + \

                    'file-' + str(file_index) + \

                    '-label-' + str(labels[i]) + \

                    ".png"

            print("file_name: ", file_name)

            print("unique_file_name: ", unique_file_name)

            # write to destination folder with the filename:

            # <original_filename>.file-<index>-label-<label>.png

            if not cv2.imwrite(unique_file_name, crops[i][1]):

                raise Exception("Could not write image")

            file_index += 1