import matplotlib.pyplot as plt
import numpy as np


import streamlit as st
from Data_Generation.Dataset_Generation_Functions import make_boxes
from Data_Generation.Piecewise_Box_Functions import basic_box_array, forward_slash_array, combine_arrays, add_thickness
########################################################################################################################
# User Inputs
image_size = st.slider('Select a value for the image size', min_value=9, max_value=28)
# st.write(x, 'squared is', x * x)

density_selection = st.slider('Select a value for the number of equally spaced density values (0, 1]', min_value=1, max_value=10)
########################################################################################################################
# Compute Example Shapes

densities = np.linspace(0, 1, num=density_selection+1)[1:]

sample_basic_box = basic_box_array(image_size, 1)
sample_forward_slash_box = forward_slash_array(image_size, 1)
sample_combined = combine_arrays([sample_forward_slash_box, sample_basic_box])
print(sample_combined)

sample_density = np.array([sample_combined * density_value for density_value in densities])

# copy = sample_combined
sample_thickness = []
# test = add_thickness(copy, 2)
# sample_thickness.append(test)
# print(sample_thickness)
for i in [1, 2, 3, 4]:
    copy = sample_combined
    print(i)
    test = add_thickness(copy, i)
    print(test)
    sample_thickness.append(test)
    # print(sample_thickness)
########################################################################################################################
# Output Example Shapes
st.write("Click 'Generate Samples' to show some density values that would exist in your dataset:")


# Show samples of various density values

if st.button('Generate Samples'):  # Generate the samples
    plt.figure(1)
    st.header("Sample Density Figures:")
    max_figures = min(density_selection, 5)
    for i in range(max_figures):
        plt.subplot(1, max_figures+1, i+1), plt.imshow(sample_density[i], cmap='gray', vmin=0, vmax=1)
        if i != 0:  # Show y-label for only first figure
            plt.tick_params(left=False, labelleft=False)
        plt.title("Density: " + str(round(densities[i], 4)), fontsize=6)
    plt.figure(1)
    # cax = plt.axes([0.85, 0.1, 0.075, 0.8])
    # plt.colorbar(cax=cax, shrink=0.1)
    st.pyplot(plt.figure(1))

    # Show samples of various thickness values
    st.header("Sample Thickness Figures:")
    plt.figure(2)
    for i in range(len(sample_thickness)):
        plt.subplot(1, 5, i+1), plt.imshow(sample_thickness[i], cmap='gray', vmin=0, vmax=1)
        if i != 0:  # Show y-label for only first figure
            plt.tick_params(left=False, labelleft=False)
        plt.title("Thickness: " + str(i+1), fontsize=6)
    plt.figure(2)
    # cax = plt.axes([0.85, 0.1, 0.075, 0.8])
    # plt.colorbar(cax=cax, shrink=0.1)
    st.pyplot(plt.figure(2))



########################################################################################################################
# Output Generated Examples

    # plt.figure(2)
    # for j in range(5):  # shows  5 random images to the users to view samples of the dataset
    #     i = np.random.randint(0, len(result))
    #     plt.subplot(550 + 1 + j)
    #     plt.imshow(result[i], cmap='gray', vmin=0, vmax=1)
    # plt.figure(2)
    # st.pyplot(plt.figure(2))


'''
# Testing
image_size = 100
densities = [1]

boxes = make_boxes(image_size, densities)

desired_density = 1
# desired_thickness = 0

desired_basic_box_thickness = 1
desired_forward_slash_box_thickness = 2
desired_back_slash_box_thickness = 0
desired_hot_dog_box_thickness = 0
desired_hamburger_box_thickness = 0


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(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_density[i] == desired_density
           and basic_box_thickness[i] == desired_basic_box_thickness
           and forward_slash_box_thickness[i] == desired_forward_slash_box_thickness
           and back_slash_box_thickness[i] == desired_back_slash_box_thickness
           and hot_dog_box_thickness[i] == desired_hot_dog_box_thickness
           and hamburger_box_thickness[i] == desired_hamburger_box_thickness]
plt.imshow(box_arrays[indices[0]], cmap='gray', vmin=0, vmax=1)
plt.show()
'''



'''

# food = load_dataset("cmudrc/2d-lattices", split="train[:15]")  # Loads the training data samples
food = load_dataset("cmudrc/2d-lattices", split="train+test")  # Loads all of the data, for use after training

# checks to see if the dataset has been assigned a class label
# if type(food.features["label"]) != 'datasets.features.features.ClassLabel': # Cast to ClassLabel
#     food = food.class_encode_column('label')
print(food)
desired_label = 'x_plus_box'
desired_thickness = 3
desired_density = 1

data_frame = pd.DataFrame(food)
# print(data_frame)

shape_rows = data_frame['Shape'] == desired_label
# print(shape_rows)

thickness_rows = data_frame['Thickness'] == desired_thickness
# print(thickness_rows)

density_rows = data_frame['Density'] == desired_density
# print(density_rows)

desired_output = data_frame.loc[shape_rows & thickness_rows & density_rows].iloc[0]['Array']
print(desired_output)
print(type(desired_output))


example_point = numpy.array(json.loads(desired_output))

plt.imshow(example_point)
plt.show()


all_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]

base_shapes = [basic_box, back_slash_box, forward_slash_box, hot_dog_box, hamburger_box]
image_size = 256
density = [1]

boxes = make_boxes(image_size, density, all_shapes)


box_arrays, box_shape, box_density, box_thickness,  = list(zip(*boxes))[0], list(zip(*boxes))[1], list(zip(*boxes))[2], list(zip(*boxes))[3]

# indices_1 = [i for i in range(len(boxes)) if boxes[1][i] == str(base_shapes[0]) and boxes[2][i] == density[0] and boxes[3][i] == desired_thickness]
indices_1 = [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]
print(indices_1)
# indices_1 = random.randint(0, len(box_arrays))


# plt.imshow(box_arrays[indices_1])
plt.imshow(box_arrays[indices_1[0]])
plt.show()
'''

'''trainer.push_to_hub()''' # Need to figure out how to push the model to the hub