import streamlit as st
import torch
import bitsandbytes
import accelerate
import scipy
from PIL import Image
import torch.nn as nn
from transformers import Blip2Processor, Blip2ForConditionalGeneration, InstructBlipProcessor, InstructBlipForConditionalGeneration
from my_model.object_detection import detect_and_draw_objects
from my_model.captioner.image_captioning import get_caption
from my_model.utilities import free_gpu_resources


    

def answer_question(image, question, model, processor):
    
    
    image = Image.open(image)

    inputs = processor(image, question, return_tensors="pt").to("cuda", torch.float16)

    if isinstance(model, torch.nn.DataParallel):
    # Use the 'module' attribute to access the original model
        out = model.module.generate(**inputs, max_length=100, min_length=20)
    else:
       
        out = model.generate(**inputs, max_length=100, min_length=20)

    answer = processor.decode(out[0], skip_special_tokens=True).strip()
    return answer






    


# Set up the sidebar navigation
st.sidebar.title("Navigation")
selection = st.sidebar.radio("Go to", ["Home", "View PDF", "Run Inference"])

# Set up the main page content based on navigation selection
if selection == "Home":
    st.title("Welcome to LLM Architecture Assessment")
    st.write("Home page content goes here...")
    # You can include more content for the home page here

elif selection == "View PDF":
    st.title("View PDF")
    st.write("Click the link below to view the PDF.")
    # Example to display a link to a PDF
    st.download_button(
        label="Download PDF",
        data=open("path/to/your/pdf.pdf", "rb"),
        file_name="example.pdf",
        mime="application/octet-stream"
    )
    # You can include more content for the PDF page here

elif selection == "Run Inference":
    st.title("Run Inference")
    st.write("This page allows you to run the space for inference.")
    # You would include your inference code here
    # For example, if you have a form to collect user input for the model:
    user_input = st.text_input("Enter your text here...")
    if st.button("Run"):
        # Call your model inference function
        # result = run_inference(user_input)
        # st.write(result)
        pass  # Replace pass with your inference code

# Other pages and functionality would be added in a similar manner.







        


st.title("Image Question Answering")

# File uploader for the image
image = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])

# Text input for the question
question = st.text_input("Enter your question about the image:")

if st.button('Generate Caption'):
    free_gpu_resources()
    if image is not None:
        # Display the image
        st.image(image, use_column_width=True)
        caption = get_caption(image)
        st.write(caption)
        free_gpu_resources()

    else:
        st.write("Please upload an image and enter a question.")
        
if st.button("Get Answer"):
    if image is not None and question:
        # Display the image
        st.image(image, use_column_width=True)
        # Get and display the answer
        model, processor = load_caption_model()
        answer = answer_question(image, question, model, processor)
        st.write(answer)
    else:
        st.write("Please upload an image and enter a question.")






# Object Detection

# Object Detection UI in the sidebar
st.sidebar.title("Object Detection")
# Dropdown to select the model
detect_model = st.sidebar.selectbox("Choose a model for object detection:", ["detic", "yolov5"])
# Slider for threshold with default values based on the model
threshold = st.sidebar.slider("Select Detection Threshold", 0.1, 0.9, 0.2 if detect_model == "yolov5" else 0.4)
# Button to trigger object detection
detect_button = st.sidebar.button("Detect Objects")


def perform_object_detection(image, model_name, threshold):
    """
    Perform object detection on the given image using the specified model and threshold.

    Args:
    image (PIL.Image): The image on which to perform object detection.
    model_name (str): The name of the object detection model to use.
    threshold (float): The threshold for object detection.

    Returns:
    PIL.Image, str: The image with drawn bounding boxes and a string of detected objects.
    """


    
    # Perform object detection and draw bounding boxes

    processed_image, detected_objects = detect_and_draw_objects(image, model_name, threshold)

    return processed_image, detected_objects



# Check if the 'Detect Objects' button was clicked
if detect_button:
    if image is not None:
        # Open the uploaded image
        try:
            image = Image.open(image)

            # Display the original image
            st.image(image, use_column_width=True, caption="Original Image")

            # Perform object detection
            processed_image, detected_objects = perform_object_detection(image, detect_model, threshold)


            # Display the image with detected objects
            st.image(processed_image, use_column_width=True, caption="Image with Detected Objects")

            # Display the detected objects as text
            st.write(detected_objects)


        except Exception as e:
            st.error(f"Error loading image: {e}")

    else:
        st.write("Please upload an image for object detection.")