import numpy as np
from scipy import signal
import math
# import huggingface_hub  # for loading model
import streamlit as st
# from transformers import AutoModel
# from transformers import TFAutoModel
# Needed for importing torch to use in the transformers model
# import torch
# import tensorflow
import matplotlib.pyplot as plt
# HELLO HUGGING FACE


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


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


########################################################################################################################
# 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 [3, 5, 10, 20]]

# Provide User Options
# 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)
density_2 = st.selectbox("Density 2:", density_options)

# 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)


def generate_unit_cell(shape, density, thickness):
    return globals()[shape](int(thickness), float(density), 28)

if st.button("Generate Endpoint Images"):
    plt.figure(1)
    st.header("Endpoints to be generated:")
    plt.subplot(1, 2, 1), plt.imshow(generate_unit_cell(shape_1, density_1, thickness_1), cmap='gray', vmin=0, vmax=1)
    plt.subplot(1, 2, 2), plt.imshow(generate_unit_cell(shape_2, density_2, thickness_2), cmap='gray', vmin=0, vmax=1)
    plt.figure(1)
    st.pyplot(plt.figure(1))

# Load the models from existing huggingface model
# Load the encoder model
encoder_model_boxes = huggingface_hub.from_pretrained_keras("cmudrc/2d-lattice-encoder")
# encoder_model = TFAutoModel.from_pretrained("cmudrc/2d-lattice-encoder")
# Load the decoder model
decoder_model_boxes = huggingface_hub.from_pretrained_keras("cmudrc/2d-lattice-decoder")
# decoder_model = TFAutoModel.from_pretrained("cmudrc/2d-lattice-decoder")