import numpy as np

from Data_Generation.Shape_Generation_Functions import basic_box, diagonal_box_split, horizontal_vertical_box_split, \
    back_slash_box, forward_slash_box, back_slash_plus_box, forward_slash_plus_box, hot_dog_box, hamburger_box, \
    x_hamburger_box, x_hot_dog_box, x_plus_box

import matplotlib.pyplot as plt

from Data_Generation.Piecewise_Box_Functions import add_pixels


########################################################################################################################
# Make the data using all the code in Shape_Generation_Functions.py
def make_boxes(image_size, densities, shapes):
    """
    :param image_size: [int] - the pixel height and width of the generated arrays
    :param densities: [list] - of the values of each of the active pixels in each shape
    :param shapes: [list] - of the various shapes desired for the dataset
    :return: [list[tuple]] - [Array, Density, Thickness, Shape]
    """

    # matrix = []
    #
    # for function in shapes:  # Adds different types of shapes
    #
    #     # Adds different density values
    #     for i in range(len(densities)):
    #         # Loops through the possible thickness values
    #         for j in range(image_size):  # Adds additional Pixels
    #             thickness = j
    #             Array = (function(thickness, densities[i], image_size))
    #
    #             # Checks if there are any 0's left in the array to append
    #             if (np.where((Array == float(0)))[0] > 0).any():
    #                 the_tuple = (Array, str(function.__name__), densities[i], thickness)
    #                 matrix.append(the_tuple)
    #
    #             # Prevents solids shapes from being appended to the array
    #             else:
    #                 break
    matrix = []
    base_shapes = []
    density_1 = []
    for function in shapes:  # Create an array of the base shapes
        thickness = 0
        Array = function(thickness, 1, image_size)
        # density_1_tuple = np.array([Array, str(function.__name__), 1, thickness])  # Array, Shape, Density, Thickness
        # base_shapes.append(density_1_tuple)

        density_1 = np.append(density_1,(np.array([Array, str(function.__name__), 1, thickness])), axis=1)  # Array, Shape, Density, Thickness
        # Add one to the thickness of the previous array
        # for j in range(image_size):
        while (np.where((Array == float(0)))[0] > 0).any():
            # Checks if there are any 0's left in the array to append
            # if (np.where((Array == float(0)))[0] > 0).any():
                # density_1.append(density_1_tuple, axis=0)
            thickness += 1
            if np.shape(density_1) == (4,):
                Array = add_pixels(density_1[0], 1) # will add 1 pixel to each previous array, rather than adding multiple and having to loop

            else:
                print(np.shape(density_1))
                print(density_1[-1][0])
                Array = add_pixels(density_1[-1][0], 1)
            # print(np.shape(Array))
            density_1_tuple = np.array([Array, str(function.__name__), 1, thickness])
            # else:  # Prevents solids shapes from being appended to the array
            #     break
            density_1 = np.vstack((density_1, density_1_tuple))

    matrix = []
    # print(np.shape(density_1[0]))
    # print(density_1[:][0])
    for i in range(len(densities)):
        some = np.multiply(density_1[:][0],densities[i]) #,density_1[:1])
        # print(np.shape(some))
        matrix.append(tuple(some))


        # # Adds different density values
        # for i in range(len(densities)):
        #     # Loops through the possible thickness values
        #     for j in range(image_size):  # Adds additional Pixels
        #         thickness = j
        #         Array = (function(thickness, densities[i], image_size))
        #
        #         # Checks if there are any 0's left in the array to append
        #         if (np.where((Array == float(0)))[0] > 0).any():
        #             the_tuple = (Array, str(function.__name__), densities[i], thickness)
        #             matrix.append(the_tuple)
        #
        #         # Prevents solids shapes from being appended to the array
        #         else:
        #             break
    return matrix


########################################################################################################################
# Testing
image_size = 9
densities = [1]
shapes = [basic_box, diagonal_box_split, horizontal_vertical_box_split, back_slash_box, forward_slash_box,
                 back_slash_plus_box, forward_slash_plus_box, hot_dog_box, hamburger_box, x_hamburger_box,
                 x_hot_dog_box, x_plus_box]

boxes = make_boxes(image_size, densities, shapes)

# print(np.shape(boxes))
desired_label = 'basic_box'
desired_density = 1
desired_thickness = 0

box_arrays, box_shape, box_density, box_thickness,  = list(zip(*boxes))[0], list(zip(*boxes))[1], list(zip(*boxes))[2], list(zip(*boxes))[3]
# print(np.shape(box_arrays))
# print(np.shape(box_shape))
# print(np.shape(box_density))

indices = [i for i in range(len(box_arrays)) if box_shape[i] == desired_label and box_density[i] == desired_density and box_thickness[i] == desired_thickness]
plt.imshow(box_arrays[indices[0]])
plt.show()