import streamlit as st
import numpy as np
import matplotlib.pyplot as plt

# Define a function to calculate assembly A
def calculate_assembly(ai_values, ni_values, N_T):
    return sum([np.exp(ai) * ((ni - 1) / N_T) for ai, ni in zip(ai_values, ni_values)])

# Sample data for top ten geometric assemblies
example_assemblies = {
    "Lighthouse": {"ai": 2.5, "ni": 1},
    "Eyeglasses": {"ai": 1.8, "ni": 2},
    "Stool": {"ai": 1.2, "ni": 3},
    "Window": {"ai": 1.5, "ni": 4},
    "Hand": {"ai": 2.0, "ni": 5},
    "Dice": {"ai": 0.8, "ni": 6},
    "Heaven": {"ai": 2.3, "ni": 7},
    "Gate": {"ai": 1.7, "ni": 8},
    "Cat": {"ai": 2.1, "ni": 9},
    "Toes": {"ai": 1.0, "ni": 10},
}

# Streamlit App
st.title("The Assembly Equation Interactive Simulator")
st.write("""
The Assembly Equation is defined as:
\( A = \sum_{i=1}^{N} e^{a_i} \left(\frac{n_i - 1}{N_T}\right) \)

Where:
- \( A \) is the assembly of the ensemble
- \( a_i \) is the assembly index of object \( i \)
- \( n_i \) is the copy number of object \( i \)
- \( N \) is the total number of unique objects
- \( N_T \) is the total number of objects in the ensemble
""")

# Sidebar for inputs
st.sidebar.header("Input Parameters")
N = st.sidebar.number_input("Enter the total number of unique objects (N):", min_value=1, value=len(example_assemblies))
N_T = st.sidebar.number_input("Enter the total number of objects in the ensemble (N_T):", min_value=1, value=50)

# Inputs for each object
ai_values = []
ni_values = []
objects = list(example_assemblies.keys())

for i in range(N):
    obj = objects[i]
    ai = st.sidebar.number_input(f"Enter the assembly index of {obj} (a{i+1}):", value=example_assemblies[obj]["ai"], key=f"a{i+1}")
    ni = st.sidebar.number_input(f"Enter the copy number of {obj} (n{i+1}):", min_value=0, value=example_assemblies[obj]["ni"], key=f"n{i+1}")
    ai_values.append(ai)
    ni_values.append(ni)

# Button to calculate assembly
if st.sidebar.button('Calculate Assembly'):
    A = calculate_assembly(ai_values, ni_values, N_T)
    st.write("The assembly of the ensemble (A) is:", A)
    
    # Plotting the assembly index
    fig, ax = plt.subplots()
    ax.bar(objects, ai_values, color='skyblue')
    ax.set_xlabel('Objects')
    ax.set_ylabel('Assembly Index')
    ax.set_title('Assembly Index of Each Object')
    st.pyplot(fig)
else:
    st.write("Enter values and press 'Calculate Assembly' to see the result.")

# Showing the dictionary of examples
st.write("Example inputs for simple assemblies:")
st.json(example_assemblies)

# Run the app with streamlit run <this_script_name>.py