import streamlit as st
import numpy as np
import cv2 as cv
import supervision as sv
from PIL import Image

class Upload_View:
    def __init__(self, app, model):
        self.app = app
        self.model = model
    def show(self):
        col1_back, col2_back = st.columns([0.2, 0.8])
        with col1_back:
            if st.button("Back", key='upload_back', icon = ':material/arrow_back:', type = 'primary'):
                self.app.change_page("Main")
        st.markdown("<h1 style='text-align: center;'>🧠 Upload MRI or CT scan image</h1>", unsafe_allow_html=True)
        st.divider()
        col1_img, col2_img = st.columns(2,gap='medium')
        with col1_img:
            st.header('Input')
            uploaded_file = st.file_uploader("Choose an image...", type=["png", "jpg", "jpeg"])
            if uploaded_file is not None:
                st.divider()
                st.image(uploaded_file, caption="Uploaded Image", use_container_width=True)
        with col2_img:
            st.header('Output')
            if uploaded_file is not None:
                # Convert uploaded image to OpenCV format
                image = Image.open(uploaded_file).convert("RGB")
                image_np = np.array(image)  # Convert to NumPy array (H, W, C)
                image_bgr = cv.cvtColor(image_np, cv.COLOR_RGB2BGR)  # Convert RGB to BGR for OpenCV

                # Run YOLO inference
                with st.spinner("Making the predictions... ⏳"):
                    results = self.model(image_bgr)[0]

                # Convert YOLO results into Supervision Detections format
                detections = sv.Detections.from_ultralytics(results)

                # Generate labels with confidence scores
                labels = [
                    f"{results.names[class_id]}: {confidence*100:.2f}%"
                    for class_id, confidence in zip(detections.class_id, detections.confidence)
                ]

                # Annotate image with bounding boxes & labels
                box_annotator = sv.BoxAnnotator()
                label_annotator = sv.LabelAnnotator()

                annotated_image = box_annotator.annotate(image_np, detections)
                annotated_image = label_annotator.annotate(annotated_image, detections, labels=labels)

                # Add tumor count and confidence description
                num_tumors = len(detections)  # Total number of detected tumors

                # Display the annotated image
                st.image(annotated_image, caption="Tumor Detection Result", use_container_width=True)

                # Display tumor count and confidence information using Streamlit text element
                st.subheader(f"Number of Tumors Detected: {num_tumors}")

                # Convert annotated image to bytes for download
                annotated_pil = Image.fromarray(annotated_image)
                img_bytes = annotated_pil.convert("RGB")
                from io import BytesIO
                buf = BytesIO()
                img_bytes.save(buf, format="PNG")
                byte_im = buf.getvalue()

                # Add download button
                st.download_button(
                    label="📥 Download Annotated Image",
                    data=byte_im,
                    file_name="tumor_prediction.png",
                    mime="image/png"
                )