import streamlit as st
from PIL import Image
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from io import BytesIO
import uuid
import gc

import sys
import os

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from segmentation_model import load_model,transform_image, run_inference, save_input_image, save_objects_and_metadata, extract_object
# from models.identification_model import identify_and_extract_objects
# from models.text_extraction_model import extract_text
# from models.summarization_model import summarize_text
# from utils.data_mapping import create_summary_table



model = load_model()

def resize_image(image, size=(800, 800)):
    return image.resize(size, Image.ANTIALIAS)

def display_masks(outputs, image, threshold=0.5):
    masks = outputs[0]['masks']
    scores = outputs[0]['scores']
    
    fig, ax = plt.subplots()
    ax.imshow(np.array(image))

    extracted_objects = []
    
    for i in range(len(scores)):
        if scores[i] > threshold:
            mask = masks[i].squeeze().cpu().numpy()
            mask = np.where(mask > 0.5, 1, 0).astype(np.uint8)

            object_img = extract_object(image,mask)
            extracted_objects.append(object_img)
            #Display the mask
            ax.imshow(mask, cmap='jet', alpha=0.5)  # Overlay mask on image
    
    st.pyplot(fig)

    return extracted_objects

    

st.title("Image Segmentation with Mask R-CNN")

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

if uploaded_file is not None:
    # Convert uploaded file to PIL Image
    image = uploaded_file
    st.image(image, caption='Uploaded Image.', use_column_width=True)
    image = Image.open(uploaded_file).convert('RGB')
    # Generate a unique master ID for the image
    master_id = str(uuid.uuid4())
    
    # Save the input image
    save_input_image(image, master_id)
    # Transform image
    image_tensor = transform_image(image)
    outputs = run_inference(model, image_tensor)

    extracted_objects = display_masks(outputs, image)
    
    if extracted_objects:
    # Save the extracted objects and their metadata
        metadata = save_objects_and_metadata(extracted_objects, master_id)
    
    # Display metadata as a JSON output
        st.write("Metadata for extracted objects:")
        st.json(metadata)
    
    # Display each extracted object
        st.write("Extracted Objects:")
        for i, obj_img in enumerate(extracted_objects):
         st.image(obj_img, caption=f'Object {i+1}', use_column_width=True)
    else:
        st.write("No objects were detected")


    # del extracted_objects
    # gc.collect()
    
    # Display results
    #display_masks(outputs, image) 





if uploaded_file is not None:
    image = Image.open(uploaded_file).convert("RGB")
    st.image(image, caption='Uploaded Image.', use_column_width=True)
    
    image_tensor = transform_image(image)
    outputs = run_inference(model, image_tensor)
    
    display_masks(outputs, image)



# def upload_image():
#     uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
#     if uploaded_file is not None:
#         image = Image.open(uploaded_file)
#         return image
#     return None


# # def display_segmentation(image):
# #     st.image(image, caption="Original Image", use_column_width=True)
    
#     # Transform and run inference
#     # image_tensor = transform_image(image)
#     # outputs = run_inference(image_tensor)

#     # # Save segmented objects
#     # output_dir = 'segmented_objects/'
#     # save_segmented_objects(image, outputs, output_dir)
    
#     # segmented_images = [Image.open(f"{output_dir}object_{i+1}.png") for i in range(len(outputs[0]['scores']))]
#     # for img in segmented_images:
#     #     st.image(img, caption="Segmented Object", use_column_width=True)




# def main():
#     st.title("Image Processing Pipeline")

# #     uploaded_file = st.file_uploader("Upload an image", type=["jpg", "png"])
# #     if uploaded_file:
# #         image_path = f"data/input_images/{uploaded_file.name}"
# #         image = Image.open(uploaded_file)
# #         image.save(image_path)  # Save the uploaded image for further processing
# #         st.image(image, caption="Uploaded Image")

# #         if st.button("Segment Image"):
# #             segmented = segment_image(image_path)
# #             st.image(segmented, caption="Segmented Image", use_column_width=True)

# #         if st.button("Identify and Extract Objects"):
# #             objects_data = identify_and_extract_objects(image_path)
# #             extracted_objects = []

# #             for obj_data in objects_data:
# #                 object_image = Image.open(obj_data['Image Path'])
# #                 text = extract_text(object_image)
# #                 summary = summarize_text(text)
# #                 obj_data['Text'] = text
# #                 obj_data['Summary'] = summary
# #                 extracted_objects.append(obj_data)

# #                 st.image(object_image, caption=f"Object {obj_data['ID']} - Label {obj_data['Label']}")

# #             summary_file = create_summary_table(extracted_objects)
# #             st.write(pd.DataFrame(extracted_objects))
# #             st.download_button(label="Download Summary Table", data=open(summary_file).read(), file_name="summary.csv")

# if __name__ == "__main__":
#     main()