import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

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 Entire Dataset
st.write("Click 'Generate Dataset' to generate the dataset based on the conditions set previously:")
if st.button('Generate Dataset'):  # Generate the samples
    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]

    # 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()
    st.write("Here is what the generated data looks like (double click on the 'Array' cells to view the full array):")
    st.write(dataframe)  # Display the data generated
    st.write("Click 'Download' to download a CSV file of the dataset:")
    st.download_button("Download Dataset", csv, file_name='2D_Lattice.csv')