import streamlit as st
from PIL import Image
import numpy as np
import main
import pandas as pd
import cv2

def get_image_np(uploaded_file):
    image = Image.open(uploaded_file)
    return np.array(image)

def show_predicted_caption(image_np, top_k=1):
    image = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
    captions = predict_caption(image, model, text_embeddings, valid_df['caption'].values, n=top_k)
    return captions

# Load the model and text embeddings
valid_df = pd.read_csv('testing_df.csv')
model, text_embeddings = main.get_text_embeddings(valid_df)

# App code
st.title("Medical Radiology Report Generator")

st.header("Personal Information")
first_name = st.text_input("First Name", "John")
last_name = st.text_input("Last Name", "Doe")
age = st.number_input("Age", min_value=0, max_value=120, value=25, step=1)
gender = st.selectbox("Gender", ["Male", "Female", "Other"])

st.write("Upload Scan to get Radiological Report:")
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])
if uploaded_file is not None:
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Image", use_column_width=True)
    st.write("")

    if st.button("Generate Caption"):
        with st.spinner("Generating caption..."):
            image_np = np.array(image)
            caption = show_predicted_caption(image_np)[0]

            st.success(f"Caption: {caption}")

            # Generate the radiology report
            radiology_report = generate_radiology_report(f"Write Complete Radiology Report for this: {caption}")

            # Add personal information to the radiology report
            radiology_report_with_personal_info = f"Patient Name: {first_name} {last_name}\nAge: {age}\nGender: {gender}\n\n{radiology_report}"

            container = st.container()
            with container:
                st.header("Radiology Report")
                st.write(radiology_report_with_personal_info)
                st.markdown(download_link(save_as_docx(radiology_report_with_personal_info, "radiology_report.docx"), "radiology_report.docx", "Download Report as DOCX"), unsafe_allow_html=True)

            # Advanced Feedback System
            st.header("Advanced Feedback System")
            feedback_options = ["Better", "Satisfied", "Worse"]
            feedback = st.radio("Rate the generated report:", feedback_options)

            top_k = 1
            while feedback == "Worse":
                top_k += 1
                with st.spinner("Regenerating report..."):
                    new_caption = show_predicted_caption(image_np, top_k=top_k)[-1]
                    radiology_report = generate_radiology_report(f"Write Complete Radiology Report for this: {new_caption}")

                    radiology_report_with_personal_info = f"Patient Name: {first_name} {last_name}\nAge: {age}\nGender: {gender}\n\n{radiology_report}"

                with container:
                    container.empty()
                    st.header("Radiology Report")
                    st.write(radiology_report_with_personal_info)
                    st.markdown(download_link(save_as_docx(radiology_report_with_personal_info, "radiology_report.docx"), "radiology_report.docx", "Download Report as DOCX"), unsafe_allow_html=True)

                feedback = st.radio("Rate the generated report:", feedback_options)