import streamlit as st
import pandas as pd
import numpy as np
import os
from PIL import Image
import torch.nn as nn
import pandas as pd
import torch
from torch.optim import Adam
from torch.utils.data import Dataset, DataLoader


# Load the dataset
st.session_state.train_df = pd.read_csv("data/asl_data/sign_mnist_train.csv")
st.session_state.valid_df = pd.read_csv("data/asl_data/sign_mnist_valid.csv")

# Streamlit app
st.title("ASL Recognition App")

# Tabs: Dataset, Model, Prediction
tab1, tab2, tab3 = st.tabs(["Dataset", "Model", "Prediction"])

# Dataset Tab
with tab1:
    st.header("Dataset Overview")
    st.write("Displaying sample images from the training dataset.")
    
    # Convert CSV image data to images (Placeholder code for now)
    num_samples = 10
    train_data = st.session_state.train_df.values[:, 1:]
    labels = st.session_state.train_df.values[:, 0]
    
    # Display sample images
    col1, col2, col3, col4, col5 = st.columns(5)
    for i in range(num_samples):
        img = train_data[i].reshape(28, 28)  # Assuming images are 28x28 grayscale
        img = Image.fromarray(img.astype('uint8'))
        with [col1, col2, col3, col4, col5][i % 5]:
            st.image(img, caption=f"Label: {labels[i]}", use_column_width=True)

# Model Tab
with tab2:
    st.header("Model Training")
    st.write("Click the button to train the model.")
    if st.button("Train Model"):
        st.write("Training started... (Placeholder for model training)")

# Prediction Tab
with tab3:
    st.header("Make a Prediction")
    
    # File uploader
    uploaded_file = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
    
    # Camera input
    camera_file = st.camera_input("Take a picture")
    
    if uploaded_file is not None:
        image = Image.open(uploaded_file)
        st.image(image, caption="Uploaded Image", use_column_width=True)
        st.write("(Placeholder for image preprocessing and prediction)")
    
    elif camera_file is not None:
        image = Image.open(camera_file)
        st.image(image, caption="Captured Image", use_column_width=True)
        st.write("(Placeholder for image preprocessing and prediction)")