import os
import streamlit as st
import tensorflow as tf
from PIL import Image
import numpy as np
from huggingface_hub import login, hf_hub_download

# Authenticate with Hugging Face token (if available)
hf_token = os.environ.get("HF_TOKEN") 
if hf_token:
    login(token=hf_token)

# Download and load the model from the Hugging Face Hub
repo_id = os.environ.get("MODEL_ID", "willco-afk/tree-test-x")  # Get repo ID from secret or default
filename = "your_trained_model.keras"  # Updated filename
cache_dir = "./models"  # Local directory to cache the model
os.makedirs(cache_dir, exist_ok=True)  
model_path = hf_hub_download(repo_id=repo_id, filename=filename, cache_dir=cache_dir)

# Load the model
model = tf.keras.models.load_model(model_path) 

# Streamlit UI
tab1, tab2 = st.tabs(["Christmas Tree Classifier", "Sample Image Links"])

# Tab 1: Christmas Tree Classifier
with tab1:
    st.title("Christmas Tree Classifier")
    st.write("Upload an image of a Christmas tree to classify it:")

    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

    if uploaded_file is not None:
        # Display the uploaded image
        image = Image.open(uploaded_file)
        st.image(image, caption="Uploaded Image.", use_container_width=True)  # Updated to use_container_width
        st.write("")
        st.write("Classifying...")

        # Preprocess the image
        image = image.resize((224, 224))  # Resize to match your model's input size
        image_array = np.array(image) / 255.0  # Normalize pixel values
        image_array = np.expand_dims(image_array, axis=0)  # Add batch dimension

        # Make prediction
        prediction = model.predict(image_array)
        
        # Get predicted class
        predicted_class = "Decorated" if prediction[0][0] >= 0.5 else "Undecorated"

        # Display the prediction
        st.write(f"Prediction: {predicted_class}")

# Tab 2: Sample Image Links
with tab2:
    st.title("Sample Image Links")
    
    # First header and paragraph with link
    st.header("View some of my decorated and undecorated tree samples for the Model here:")
    st.write("You can view the sample images for both decorated and undecorated trees in the following link:")
    st.write("[View sample images here](https://www.dropbox.com/scl/fo/cuzo12z39cxv6joz7gz2o/ACf5xSjT7nHqMRdgh21GYlc?rlkey=w10usqhkngf2uxwvllgnqb8tf&st=ld22fl4c&dl=0)")

    # Second header and paragraph with download link
    st.header("Download the tree sample pictures to test them on the model yourself here:")
    st.write("You can download the tree sample images by clicking on the link below to test them in the Christmas Tree Classifier:")
    st.write("[Download sample images here](https://www.dropbox.com/scl/fo/cuzo12z39cxv6joz7gz2o/ACf5xSjT7nHqMRdgh21GYlc?rlkey=w10usqhkngf2uxwvllgnqb8tf&st=ld22fl4c&dl=1)")