import numpy as np
from scipy import signal
import math
import matplotlib.pyplot as plt
from huggingface_hub import from_pretrained_keras
import streamlit as st
from elasticity import elasticity
import io
from voxel_to_SDF_to_STL import voxel_to_sdf, sdf_to_stl, single_body_check
from PIL import Image
import time

# Needed in requirements.txt for importing to use in the transformers model
import tensorflow


# HELLO HUGGING FACE

########################################################################################################################
# Define the piecewise functions to create each of the possible shapes
def basic_box_array(image_size):
    A = np.zeros((int(image_size), int(image_size)))  # Initializes A matrix with 0 values
    # Creates the outside edges of the box
    for i in range(image_size):
        for j in range(image_size):
            if i == 0 or j == 0 or i == image_size - 1 or j == image_size - 1:
                A[i][j] = 1
    return A


def back_slash_array(image_size):
    A = np.zeros((int(image_size), int(image_size)))  # Initializes A matrix with 0 values
    for i in range(image_size):
        for j in range(image_size):
            if i == j:
                A[i][j] = 1
    return A


def forward_slash_array(image_size):
    A = np.zeros((int(image_size), int(image_size)))  # Initializes A matrix with 0 values
    for i in range(image_size):
        for j in range(image_size):
            if i == (image_size - 1) - j:
                A[i][j] = 1
    return A


def hot_dog_array(image_size):
    # Places pixels down the vertical axis to split the box
    A = np.zeros((int(image_size), int(image_size)))  # Initializes A matrix with 0 values
    for i in range(image_size):
        for j in range(image_size):
            if j == math.floor((image_size - 1) / 2) or j == math.ceil((image_size - 1) / 2):
                A[i][j] = 1
    return A


def hamburger_array(image_size):
    # Places pixels across the horizontal axis to split the box
    A = np.zeros((int(image_size), int(image_size)))  # Initializes A matrix with 0 values
    for i in range(image_size):
        for j in range(image_size):
            if i == math.floor((image_size - 1) / 2) or i == math.ceil((image_size - 1) / 2):
                A[i][j] = 1
    return A


def center_array(image_size):
    A = np.zeros((int(image_size), int(image_size)))  # Initializes A matrix with 0 values
    for i in range(image_size):
        for j in range(image_size):
            if i == math.floor((image_size - 1) / 2) and j == math.ceil((image_size - 1) / 2):
                A[i][j] = 1
            if i == math.floor((image_size - 1) / 2) and j == math.floor((image_size - 1) / 2):
                A[i][j] = 1
            if j == math.ceil((image_size - 1) / 2) and i == math.ceil((image_size - 1) / 2):
                A[i][j] = 1
            if j == math.floor((image_size - 1) / 2) and i == math.ceil((image_size - 1) / 2):
                A[i][j] = 1
    return A


def update_array(array_original, array_new, image_size):
    A = array_original
    for i in range(image_size):
        for j in range(image_size):
            if array_new[i][j] == 1:
                A[i][j] = 1
    return A


def add_pixels(array_original, additional_pixels, image_size):
    # Adds pixels to the thickness of each component of the box
    A = array_original
    A_updated = np.zeros((int(image_size), int(image_size)))  # Initializes A matrix with 0 values
    for dens in range(additional_pixels):
        for i in range(1, image_size - 1):
            for j in range(1, image_size - 1):
                if A[i - 1][j] + A[i + 1][j] + A[i][j - 1] + A[i][j + 1] > 0:
                    A_updated[i][j] = 1
        A = update_array(A, A_updated, image_size)
    return A


########################################################################################################################
# Create the desired shape using the density and thickness
def basic_box(additional_pixels, density, image_size):
    A = basic_box_array(image_size)  # Creates the outside edges of the box
    # Increase the thickness of each part of the box
    A = add_pixels(A, additional_pixels, image_size)
    return A * density


def horizontal_vertical_box_split(additional_pixels, density, image_size):
    A = basic_box_array(image_size)  # Creates the outside edges of the box
    # Place pixels across the horizontal and vertical axes to split the box
    A = update_array(A, hot_dog_array(image_size), image_size)
    A = update_array(A, hamburger_array(image_size), image_size)
    # Increase the thickness of each part of the box
    A = add_pixels(A, additional_pixels, image_size)
    return A * density


def diagonal_box_split(additional_pixels, density, image_size):
    A = basic_box_array(image_size)  # Creates the outside edges of the box

    # Add pixels along the diagonals of the box
    A = update_array(A, back_slash_array(image_size), image_size)
    A = update_array(A, forward_slash_array(image_size), image_size)

    # Adds pixels to the thickness of each component of the box
    # Increase the thickness of each part of the box
    A = add_pixels(A, additional_pixels, image_size)
    return A * density


def back_slash_box(additional_pixels, density, image_size):
    A = basic_box_array(image_size)  # Initializes A matrix with 0 values
    A = update_array(A, back_slash_array(image_size), image_size)
    A = add_pixels(A, additional_pixels, image_size)
    return A * density


def forward_slash_box(additional_pixels, density, image_size):
    A = basic_box_array(image_size)  # Initializes A matrix with 0 values
    A = update_array(A, forward_slash_array(image_size), image_size)
    A = add_pixels(A, additional_pixels, image_size)
    return A * density


def hot_dog_box(additional_pixels, density, image_size):
    A = basic_box_array(image_size)  # Initializes A matrix with 0 values
    A = update_array(A, hot_dog_array(image_size), image_size)
    A = add_pixels(A, additional_pixels, image_size)
    return A * density


def hamburger_box(additional_pixels, density, image_size):
    A = basic_box_array(image_size)  # Initializes A matrix with 0 values
    A = update_array(A, hamburger_array(image_size), image_size)
    A = add_pixels(A, additional_pixels, image_size)
    return A * density


def x_plus_box(additional_pixels, density, image_size):
    A = basic_box_array(image_size)  # Initializes A matrix with 0 values
    A = update_array(A, hot_dog_array(image_size), image_size)
    A = update_array(A, hamburger_array(image_size), image_size)
    A = update_array(A, forward_slash_array(image_size), image_size)
    A = update_array(A, back_slash_array(image_size), image_size)
    A = add_pixels(A, additional_pixels, image_size)
    return A * density


def forward_slash_plus_box(additional_pixels, density, image_size):
    A = basic_box_array(image_size)  # Initializes A matrix with 0 values
    A = update_array(A, hot_dog_array(image_size), image_size)
    A = update_array(A, hamburger_array(image_size), image_size)
    A = update_array(A, forward_slash_array(image_size), image_size)
    # A = update_array(A, back_slash_array(image_size), image_size)
    A = add_pixels(A, additional_pixels, image_size)
    return A * density


def back_slash_plus_box(additional_pixels, density, image_size):
    A = basic_box_array(image_size)  # Initializes A matrix with 0 values
    A = update_array(A, hot_dog_array(image_size), image_size)
    A = update_array(A, hamburger_array(image_size), image_size)
    # A = update_array(A, forward_slash_array(image_size), image_size)
    A = update_array(A, back_slash_array(image_size), image_size)
    A = add_pixels(A, additional_pixels, image_size)
    return A * density


def x_hot_dog_box(additional_pixels, density, image_size):
    A = basic_box_array(image_size)  # Initializes A matrix with 0 values
    A = update_array(A, hot_dog_array(image_size), image_size)
    # A = update_array(A, hamburger_array(image_size), image_size)
    A = update_array(A, forward_slash_array(image_size), image_size)
    A = update_array(A, back_slash_array(image_size), image_size)
    A = add_pixels(A, additional_pixels, image_size)
    return A * density


def x_hamburger_box(additional_pixels, density, image_size):
    A = basic_box_array(image_size)  # Initializes A matrix with 0 values
    # A = update_array(A, hot_dog_array(image_size), image_size)
    A = update_array(A, hamburger_array(image_size), image_size)
    A = update_array(A, forward_slash_array(image_size), image_size)
    A = update_array(A, back_slash_array(image_size), image_size)
    A = add_pixels(A, additional_pixels, image_size)
    return A * density


def center_box(additional_pixels, density, image_size):
    A = basic_box_array(image_size)  # Initializes A matrix with 0 values
    A = update_array(A, center_array(image_size), image_size)
    A = add_pixels(A, additional_pixels, image_size)
    return A * density


########################################################################################################################
# The function to add thickness to struts in an array
def add_thickness(array_original, thickness: int) -> np.ndarray:
    """
    :param array_original: [ndarray] - an array with thickness 1 of any shape type
    :param thickness: [int] - the number of pixels to be activated surrounding the base shape
    :return: [ndarray] - the output is a unit cell that has been convolved to expand the number of pixels activated
    based on the desired thickness. The activated pixels are 1 (white) and the deactivated pixels are 0 (black)
    """
    A = array_original
    if thickness == 0:  # want an array of all 0's for thickness = 0
        A[A > 0] = 0
    else:
        filter_size = 2*thickness - 1 # the size of the filter needs to extend far enough to reach the base shape
        filter = np.zeros((filter_size, filter_size))
        filter[np.floor((filter_size - 1) / 2).astype(int), :] = filter[:, np.floor((filter_size - 1) / 2).astype(int)] =1
        filter[np.ceil((filter_size - 1) / 2).astype(int), :] = filter[:, np.ceil((filter_size - 1) / 2).astype(int)] = 1
        # The filter is made into a '+' shape using these functions
        convolution = signal.convolve2d(A, filter, mode='same')
        A = np.where(convolution <= 1, convolution, 1)
    return A


# The function to efficiently combine arrays in a list
def combine_arrays(arrays):
    output_array = np.sum(arrays, axis=0)  # Add the list of arrays
    output_array = np.array(output_array > 0, dtype=int)  # Convert all values in array to 1
    return output_array


########################################################################################################################
# Explain the App
st.header("Multi-Lattice Generator Through a VAE Model")
st.write("Shape: the type of shape the lattice will have")
st.write("Density: the pixel intensity of each activated pixel")
st.write("Thickness: the additional pixels added to the base shape")
st.write("Interpolation Length: the number of internal interpolation points that will exist in the interpolation")
########################################################################################################################
# Provide the Options for users to select from
shape_options = ("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")
density_options = ["{:.2f}".format(x) for x in np.linspace(0.1, 1, 10)]
thickness_options = [str(int(x)) for x in np.linspace(0, 10, 11)]
interpolation_options = [str(int(x)) for x in np.linspace(2, 20, 19)]

# Provide User Options
st.header("Option 1: Perform a Linear Interpolation")
# Select Shapes
shape_1 = st.selectbox("Shape 1", shape_options)
shape_2 = st.selectbox("Shape 2", shape_options)

# Select Density
density_1 = st.selectbox("Density 1:", density_options, index=len(density_options)-1)
density_2 = st.selectbox("Density 2:", density_options, index=len(density_options)-1)

# Select Thickness
thickness_1 = st.selectbox("Thickness 1", thickness_options)
thickness_2 = st.selectbox("Thickness 2", thickness_options)

# Select Interpolation Length
interp_length = st.selectbox("Interpolation Length", interpolation_options, index=2)


# Define the function to generate unit cells based on user inputs
def generate_unit_cell(shape, density, thickness):
    return globals()[shape](int(thickness), float(density), 28)


def display_arrays(array1, array2, label_1, label_2):
    # A Function to plot two arrays side by side in streamlit
    # Create two columns
    col1, col2 = st.columns(2)

    # Populate the first column with array1
    col1.header(label_1)
    col1.write(array1)

    # Populate the second column with array2
    col2.header(label_2)
    col2.write(array2)


# Generate the endpoints
number_1 = generate_unit_cell(shape_1, density_1, thickness_1)
number_2 = generate_unit_cell(shape_2, density_2, thickness_2)

# Calculate the elasticity for the shapes:
elasticity_1 = elasticity(number_1)
elasticity_2 = elasticity(number_2)

# Display the endpoints to the user
if st.button("Generate Endpoint Images and Elasticity Tensors"):
    plt.figure(1)
    st.header("Endpoints to be generated:")
    plt.subplot(1, 2, 1), plt.imshow(number_1, cmap='gray', vmin=0, vmax=1), plt.title("Shape 1:")
    display_arrays(elasticity_1, elasticity_2, "Elasticity Tensor of Shape 1", "Elasticity Tensor of Shape 2")
    plt.subplot(1, 2, 2), plt.imshow(number_2, cmap='gray', vmin=0, vmax=1), plt.title("Shape 2:")
    plt.figure(1)
    st.pyplot(plt.figure(1))
########################################################################################################################
# Load the models from existing huggingface model
# Load the encoder model
encoder_model_boxes = from_pretrained_keras("cmudrc/2d-lattice-encoder")

# Load the decoder model
decoder_model_boxes = from_pretrained_keras("cmudrc/2d-lattice-decoder")


########################################################################################################################
# Encode the Desired Endpoints
# resize the array to match the prediction size requirement
number_1_expand = np.expand_dims(np.expand_dims(number_1, axis=2), axis=0)
number_2_expand = np.expand_dims(np.expand_dims(number_2, axis=2), axis=0)

# Determine the latent point that will represent our desired number
latent_point_1 = encoder_model_boxes.predict(number_1_expand)[0]
latent_point_2 = encoder_model_boxes.predict(number_2_expand)[0]

latent_dimensionality = len(latent_point_1)  # define the dimensionality of the latent space
########################################################################################################################
# Establish the Framework for a LINEAR Interpolation
num_interp = int(interp_length)   # the number of images to be pictured
latent_matrix = []  # This will contain the latent points of the interpolation
for column in range(latent_dimensionality):
    new_column = np.linspace(latent_point_1[column], latent_point_2[column], num_interp)
    latent_matrix.append(new_column)
latent_matrix = np.array(latent_matrix).T  # Transposes the matrix so that each row can be easily indexed


########################################################################################################################
# Create a gif from an interpolation
def interpolate_gif(decoder, latent_endpoint_1, latent_endpoint_2, n=100):
    # z = np.stack([latent_endpoint_1 + (latent_endpoint_2 - latent_endpoint_1) * t for t in np.linspace(0, 1, n)])

    # interpolate_list = decoder.predict(z)
    # interpolate_list = (interpolate_list * 255).astype(np.uint8)

    # images_list = [Image.fromarray(img.reshape(28, 28)).resize((256, 256)) for img in interpolate_list]
    # images_list = images_list + images_list[::-1]  # loop back beginning
    predicted_interps = []
    interp_latent = np.linspace(latent_endpoint_1, latent_endpoint_2, n)

    figure = np.zeros((28, 28 * n))
    for i in range(n):
        generated_image = decoder.predict(np.array([interp_latent[i]]))[0]
        figure[0:28, i * 28:(i + 1) * 28, ] = generated_image[:, :, -1]
        predicted_interps.append(generated_image[:, :, -1])


    # Regular Save for GIF
    # images_list[0].save(f'{filename}.gif',save_all=True,append_images=images_list[1:],loop=1)

    images_list = [Image.fromarray(img.reshape(28, 28)).resize((256, 256)) for img in predicted_interps]
    images_list = images_list + images_list[::-1]  # loop back beginning

    # Create a BytesIO object to hold the GIF data
    gif_bytes = io.BytesIO()
    images_list[0].save(
        gif_bytes,
        format='GIF',
        save_all=True,
        append_images=images_list[1:],
        loop=0, duration=100)  # Set loop to 0 for infinite looping

    # Reset the BytesIO object to the beginning
    gif_bytes.seek(0)

    st.video(gif_bytes)
    # return gif_bytes


########################################################################################################################
# Create an STL file from an interpolation
def convert_to_2_5d_sdf(interpolation, voxel_threshold, pixel_thickness):
    # Thresholding determines the distance from the SDF that is used, the threshold provided is a divisor

    # 1. Convert the interpolation into a 3D structure
    interpolation_3d = [interpolation] * pixel_thickness

    # 2. Convert the voxels into an SDF
    sdf = voxel_to_sdf(interpolation_3d, voxel_threshold)
    return sdf

def convert_sdf_to_stl(sdf, threshold_divisor):
    # 3. Check if the SDF is a single body, then convert into an STL
    if single_body_check(sdf, threshold_divisor):
        # Thresholding determines the distance from the SDF that is used, the thresdhold provided is a divisor
        stl = sdf_to_stl(sdf, threshold_divisor)
        return stl


########################################################################################################################
# Plotting the Interpolation in 2D Using Chosen Points
if st.checkbox("Generate Linear Interpolation"):
    # Generate the set of latent points in the interpolation
    linear_interp_latent = np.linspace(latent_point_1, latent_point_2, num_interp)
    linear_predicted_interps = []
    figure_2 = np.zeros((28, 28 * num_interp))

    # Predict the image for each latent point
    for i in range(num_interp):
        generated_image = decoder_model_boxes.predict(np.array([linear_interp_latent[i]]))[0]
        figure_2[0:28, i * 28:(i + 1) * 28, ] = generated_image[:, :, -1]
        linear_predicted_interps.append(generated_image[:, :, -1])

    st.image(figure_2, width=600)

    # Code to display a gif
    # interpolate_gif(decoder_model_boxes, latent_point_1, latent_point_2)

    # Code for generating the STL file
    st.subheader("Create an STL file from the extruded image!")

    if st.checkbox("Select to begin model generation"):
        # Creating an STL file of the linear interpolation

        pixel_thickness_input = st.number_input("(1) Select a pixel thickness for the 3D model: ", min_value=1, value=28)
        # Set the image threshold for binarization
        voxel_threshold_input = st.slider("(2) Select a value to threshold the image (Recommend <= 0.1) "
                                          "Higher values will result in less defined shapes: ",
                                          min_value=0.0001, max_value=0.999, value=0.1, key='voxel_threshold')

        # Create the SDF File
        linear_sdf = convert_to_2_5d_sdf(figure_2, voxel_threshold_input, pixel_thickness_input)

        # Set the threshold for the Mesh
        threshold_divisor_input = st.slider("(3) Choose a threshold divisor for the SDF: ", min_value=0.0,
                                            max_value=5.0,
                                            value=3.0, key="divisor_threshold")

        linear_sdf_min = np.min(linear_sdf)
        linear_sdf_max = np.max(linear_sdf)
        st.info("Lower SDF Thresholds will result in smoother, but less accurate shapes. Higher thresholds will result in more "
                "rugged shapes, but they are more accurate. Suggested value for threshold is less than: " +
                str((linear_sdf_max - abs(linear_sdf_min)) / 2))

        if st.button("Generate STL Model"):
            # Generate the STL File
            linear_stl = convert_sdf_to_stl(linear_sdf, threshold_divisor=threshold_divisor_input)

            # Download the STL File
            with open(linear_stl, 'rb') as file:
                st.download_button(
                    label='Download STL',
                    data=file,
                    file_name='linear_interpolation.stl',
                    key='stl-download'
                    )

########################################################################################################################
# Provide User Options
st.header("Option 2: Perform a Mesh Interpolation")
st.write("The four corners of this mesh are defined using the shapes in both Option 1 and Option 2")
# Select Shapes
shape_3 = st.selectbox("Shape 3", shape_options)
shape_4 = st.selectbox("Shape 4", shape_options)

# Select Density
density_3 = st.selectbox("Density 3:", density_options, index=len(density_options)-1)
density_4 = st.selectbox("Density 4:", density_options, index=len(density_options)-1)

# Select Thickness
thickness_3 = st.selectbox("Thickness 3", thickness_options)
thickness_4 = st.selectbox("Thickness 4", thickness_options)

# Generate the endpoints
number_3 = generate_unit_cell(shape_3, density_3, thickness_3)
number_4 = generate_unit_cell(shape_4, density_4, thickness_4)

# Display the endpoints to the user
if st.button("Generate Endpoint Images for Mesh and Elasticity Tensors"):
    plt.figure(1)
    st.header("Endpoints to be generated:")
    elasticity_3 = elasticity(number_3)
    elasticity_4 = elasticity(number_4)
    display_arrays(elasticity_1, elasticity_2, "Elasticity Tensor of Shape 1", "Elasticity Tensor of Shape 2")
    display_arrays(elasticity_3, elasticity_4, "Elasticity Tensor of Shape 3", "Elasticity Tensor of Shape 4")
    plt.subplot(2, 2, 1), plt.imshow(number_1, cmap='gray', vmin=0, vmax=1)
    plt.subplot(2, 2, 2), plt.imshow(number_2, cmap='gray', vmin=0, vmax=1)
    plt.subplot(2, 2, 3), plt.imshow(number_3, cmap='gray', vmin=0, vmax=1)
    plt.subplot(2, 2, 4), plt.imshow(number_4, cmap='gray', vmin=0, vmax=1)
    plt.figure(1)
    st.pyplot(plt.figure(1))
########################################################################################################################
# Encode the Desired Endpoints
# resize the array to match the prediction size requirement
number_3_expand = np.expand_dims(np.expand_dims(number_3, axis=2), axis=0)
number_4_expand = np.expand_dims(np.expand_dims(number_4, axis=2), axis=0)

# Determine the latent point that will represent our desired number
latent_point_3 = encoder_model_boxes.predict(number_3_expand)[0]
latent_point_4 = encoder_model_boxes.predict(number_4_expand)[0]

latent_dimensionality = len(latent_point_1)  # define the dimensionality of the latent space
########################################################################################################################
# Plot a Mesh Gridded Interpolation
if st.checkbox("Generate Mesh Interpolation"):
    latent_matrix_2 = []  # This will contain the latent points of the interpolation
    for column in range(latent_dimensionality):
        new_column = np.linspace(latent_point_3[column], latent_point_4[column], num_interp)
        latent_matrix_2.append(new_column)
    latent_matrix_2 = np.array(latent_matrix_2).T  # Transposes the matrix so that each row can be easily indexed

    mesh = []  # This will create a mesh by interpolating between the two interpolations
    for column in range(num_interp):
        row = np.linspace(latent_matrix[column], latent_matrix_2[column], num_interp)
        mesh.append(row)

    mesh = np.transpose(mesh, axes=(1, 0, 2))  # Transpose the array so it matches the original interpolation
    generator_model = decoder_model_boxes

    figure_3 = np.zeros((28 * num_interp, 28 * num_interp))

    mesh_predicted_interps = []
    for i in range(num_interp):
        for j in range(num_interp):
            generated_image = generator_model.predict(np.array([mesh[i][j]]))[0]
            figure_3[i * 28:(i + 1) * 28, j * 28:(j + 1) * 28, ] = generated_image[:, :, -1]
            mesh_predicted_interps.append(generated_image[:, :, -1])

    st.image(figure_3, width=600)

    # Code for generating the STL file
    st.subheader("Create an STL file from the extruded image!")

    if st.checkbox("Select to begin model generation"):
        # Creating an STL file of the linear interpolation

        mesh_pixel_thickness_input = st.number_input("(1) Select a pixel thickness for the 3D model: ", min_value=1,
                                                value=28)
        # Set the image threshold for binarization
        mesh_voxel_threshold_input = st.slider("(2) Select a value to threshold the image (Recommend <= 0.1) "
                                          "Higher values will result in less defined shapes: ",
                                          min_value=0.0001, max_value=0.999, value=0.1, key='voxel_threshold')

        # Create the SDF File
        mesh_sdf = convert_to_2_5d_sdf(figure_3, mesh_voxel_threshold_input, mesh_pixel_thickness_input)

        # Set the threshold for the Mesh
        mesh_threshold_divisor_input = st.slider("(3) Choose a threshold divisor for the SDF: ", min_value=0.0,
                                            max_value=5.0,
                                            value=3.0, key="divisor_threshold")

        mesh_sdf_min = np.min(mesh_sdf)
        mesh_sdf_max = np.max(mesh_sdf)
        st.info(
            "Lower SDF Thresholds will result in smoother, but less accurate shapes. Higher thresholds will result in more "
            "rugged shapes, but they are more accurate. Suggested value for threshold is less than: " +
            str((mesh_sdf_max - abs(mesh_sdf_min)) / 2))

        if st.button("Generate STL Model"):
            # Generate the STL File
            linear_stl = convert_sdf_to_stl(mesh_sdf, threshold_divisor=mesh_threshold_divisor_input)

            # Download the STL File
            with open(linear_stl, 'rb') as file:
                st.download_button(
                    label='Download STL',
                    data=file,
                    file_name='interpolation.stl',
                    key='stl-download'
                )