# app.py
import streamlit as st
from generate import generate_image
import os

# Set up page config
st.set_page_config(page_title="Visual Reconstruction from Brain", layout="centered")

# Configure Streamlit to use the correct host
import streamlit.web.cli as stcli
import sys
sys.argv = ["streamlit", "run", "scripts/app.py", "--server.address", "10.192.12.247", "--server.port", "8501", "--browser.serverAddress", "10.192.12.247"]

st.title("🧠 Imagine an Image!")

# Subject selection
sub = st.selectbox("Select Subject", options=[1, 2, 5, 7], index=0)

# Image ID input
image_id = st.number_input("Enter Image ID", min_value=0, step=1)

original_path = f'data/nsddata_stimuli/test_images/{image_id}.png'
if os.path.exists(original_path):
    st.image(original_path, caption="Original Image", use_column_width=True)
else:
    st.warning("Original image not found.")
# Text prompt
annot = st.text_input("Describe what you imagined", placeholder="e.g., a dog under a tree")

# Parameters
strength = st.slider("Diffusion Strength", 0.0, 1.0, 0.75, 0.05)
mixing = st.slider("Mixing Strength", 0.0, 1.0, 0.4, 0.05)

# Submit button
if st.button("Reconstruct Image"):
    with st.spinner("Reconstructing... please wait"):
        try:
            original_path, imagined_path = generate_image(sub, image_id, annot, strength, mixing)
            # if os.path.exists(original_path):
            #     st.image(original_path, caption="Original Image", use_column_width=True)
            # else:
                # st.warning("Original image not found.")
            if os.path.exists(imagined_path):
                st.image(imagined_path, caption="Imagined Reconstruction", use_column_width=True)
            else:
                st.warning("Imagined image not found.")
        except Exception as e:
            st.error(f"⚠️ Error during generation: {e}")

# Optional: For cloud users
st.markdown("---")
# st.markdown("🔗 Access the app at: http://10.192.12.247:8501")