import numpy as np
import pandas as pd
import json
# 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
from Data_Generation.Piecewise_Box_Functions import basic_box_array, back_slash_array, forward_slash_array, hamburger_array, hot_dog_array
# from Piecewise_Box_Functions import basic_box_array, back_slash_array, forward_slash_array, hamburger_array, hot_dog_array
import matplotlib.pyplot as plt
from json import JSONEncoder



########################################################################################################################
# Make the data using all the code in Shape_Generation_Functions.py
def make_boxes(image_size, densities):
    """
    :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
    :return: [list[tuple]] - [Array, Density, Thickness of each strut type]
    """

    matrix = []

    # Establish the maximum thickness for each type of strut
    max_vert = int(np.ceil(1 / 2 * image_size) - 2)
    max_diag = int(image_size - 3)
    max_basic = int(np.ceil(1 / 2 * image_size) - 1)

    # Adds different density values
    for i in range(len(densities)):
        for j in range(1, max_basic):  # basic box loop, always want a border
            basic_box_thickness = j
            array_1 = basic_box_array(image_size, basic_box_thickness)
            if np.unique([array_1]).all() > 0:  # Checks if there is a solid figure
                break

            for k in range(0, max_vert):
                hamburger_box_thickness = k
                array_2 = hamburger_array(image_size, hamburger_box_thickness) + array_1
                array_2 = np.array(array_2 > 0, dtype=int)  # Keep all values 0/1
                if np.unique([array_2]).all() > 0:
                    break

                for l in range(0, max_vert):
                    hot_dog_box_thickness = l
                    array_3 = hot_dog_array(image_size, hot_dog_box_thickness) + array_2
                    array_3 = np.array(array_3 > 0, dtype=int)
                    if np.unique([array_3]).all() > 0:
                        break

                    for m in range(0, max_diag):
                        forward_slash_box_thickness = m
                        array_4 = forward_slash_array(image_size, forward_slash_box_thickness) + array_3
                        array_4 = np.array(array_4 > 0, dtype=int)
                        if np.unique([array_4]).all() > 0:
                            break

                        for n in range(0, max_diag):
                            back_slash_box_thickness = n
                            array_5 = back_slash_array(image_size, back_slash_box_thickness) + array_4
                            array_5 = np.array(array_5 > 0, dtype=int)
                            if np.unique([array_5]).all() > 0:
                                break
                            the_tuple = (array_5*densities[i], densities[i], basic_box_thickness,
                                         forward_slash_box_thickness, back_slash_box_thickness,
                                         hot_dog_box_thickness, hamburger_box_thickness)
                            matrix.append(the_tuple)

    return matrix


########################################################################################################################
image_size = 9
densities = [0.2, 0.4, 0.6, 0.8, 1]
boxes = make_boxes(image_size, densities)

box_arrays, box_density, basic_box_thickness, forward_slash_box_thickness, back_slash_box_thickness,hot_dog_box_thickness, hamburger_box_thickness\
    = list(zip(*boxes))[0], list(zip(*boxes))[1], list(zip(*boxes))[2], list(zip(*boxes))[3], list(zip(*boxes))[4], list(zip(*boxes))[5], list(zip(*boxes))[6]
print(type(box_arrays[0]))


class NumpyArrayEncoder(JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        return JSONEncoder.default(self, obj)


box_arrays = [json.dumps(x, cls=NumpyArrayEncoder) for x in box_arrays] # find argument for json dumps for numpy

# Create a dataframe to convert the data to a csv file
dataframe = (pd.DataFrame((box_arrays, box_density, basic_box_thickness, forward_slash_box_thickness, back_slash_box_thickness,hot_dog_box_thickness, hamburger_box_thickness)).T).astype(str)

# Rename the columns to the desired outputs
dataframe = dataframe.rename(columns={0: "Array", 1: "Density", 2:"Basic Box Thickness", 3:"Forward Slash Strut Thickness", 4:"Back Slash Strut Thickness", 5:"Vertical Strut Thickness", 6:"Horizontal Strut Thickness"})

csv = dataframe.to_csv('2D_Lattice.csv')


df = pd.read_csv('2D_Lattice.csv')
# print(df)
row = 0
box = df.iloc[row, 1]
print(box)
array = np.array(json.loads(box))
plt.imshow(array)
plt.show()

# import matplotlib.pyplot as plt
# plt.imshow(np.array(box))
# plt.show()
# print(df)