import os
import sys
import streamlit as st
from streamlit_image_zoom import image_zoom
from PIL import Image
import pydicom
import cv2
import subprocess
import numpy as np
import glob
import io
############### Import PATH

script_dir = os.path.dirname(os.path.abspath(__file__))
yolov9 = os.path.join(script_dir, '..', 'yolov9')
sys.path.append(yolov9)


try:
    import torchmcubes
    import torch
    import torchvision
    import fpdf
except ImportError:
    subprocess.check_call(['pip', 'install', 'git+https://github.com/tatsy/torchmcubes.git'])
    subprocess.check_call(['pip', 'install','fpdf'])


from yolov9.detect_dual import predict_image
# predict_image("004f33259ee4aef671c2b95d54e4be68.png")

st.markdown("<h1 style='text-align: center;'>Chào mừng tới Chuẩn đoán ung thư phổi 🎈</h1>", unsafe_allow_html=True)

################## CONVERT DICOM TO PNG ##################
def convert_dcm_to_png(input_image_path, output_image_path='a.png'):
    ds = pydicom.dcmread(input_image_path)
    img = ds.pixel_array
    img = cv2.normalize(img, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
    cv2.imwrite(output_image_path, img)



############### DELETED EXISTING FILES ##################
def delete_images_in_folder(folder_path):
    image_extensions = ['*.jpg', '*.jpeg', '*.png', '*.bmp', '*.gif']
    for ext in image_extensions:
        files = glob.glob(os.path.join(folder_path, ext))
        for file in files:
            try:
                os.remove(file)
                print(f"Deleted: {file}")
            except Exception as e:
                print(f"Error deleting {file}: {e}")

with st.sidebar:
    st.markdown("## Truyền vào bản scans của bạn")
    uploaded_files = st.file_uploader("Choose scans...", type=["jpg", "jpeg", "png", "dicom"], accept_multiple_files=True)

status_images=False
col_1,col_3 ,col_2 = st.columns([7,1 ,7.5])
with col_1:
    with st.expander("Intructions"):
        st.markdown("Truyền ảnh vào sau đó ấn vào **Detect Lung Cancer** để xem chẩn đoán")
    if uploaded_files:
        for uploaded_file in uploaded_files:
            file_type = uploaded_file.name.split('.')[-1].lower()
            if file_type in ["jpg", "jpeg", "png"]:
                img = Image.open(uploaded_file)
                img.save('temp_image.png')
                st.markdown("<div style='display: flex; justify-content: center;'>", unsafe_allow_html=True)
                width, height = img.size
                image_zoom(img, mode="both", keep_aspect_ratio=True, zoom_factor=4.0, increment=0.2)
                st.markdown("</div>", unsafe_allow_html=True)

                btn_convert=st.button("Detect Lung Cancer")
                if btn_convert:
                    delete_images_in_folder("pages/output_yolov9")
                    predict_image("temp_image.png")
                    status_images=True
            elif file_type in ["dicom", "dcm"]:
                convert_dcm_to_png(uploaded_file)
                img = Image.open('a.png').convert('RGB')
                img.save('temp_image.png')
                st.markdown("<div style='display: flex; justify-content: center;'>", unsafe_allow_html=True)
                width, height = img.size
                image_zoom(img, mode="both",size=(width//5, height//5), keep_aspect_ratio=True, zoom_factor=4.0, increment=0.2)
                st.markdown("</div>", unsafe_allow_html=True)
                btn_convert = st.button("Detect Lung Cancer")
                if btn_convert:
                    delete_images_in_folder("pages/output_yolov9")
                    predict_image("temp_image.png")
                    status_images=True
    else:
        st.info("Yêu cầu một bản scans của ảnh phổi để xem ảnh chẩn đoán")

if "image_saved" not in st.session_state:
    st.session_state.image_saved = False

if status_images:
    with col_2:
        with st.expander("Lung Cancer Detected"):
            st.markdown("Truyền ảnh vào sau đó chọn **Detect Lung Cancer** để quan sát ảnh 3D.")

        uploaded_files = "pages/output_yolov9/temp_image.png"
        img = Image.open(uploaded_files)
        img.save('temp_image.png')

        st.markdown("<div style='display: flex; justify-content: center;'>", unsafe_allow_html=True)
        width, height = img.size
        image_zoom(img, mode="both", size=(int(width / 4.3), int(height / 4.3)),
                   keep_aspect_ratio=True, zoom_factor=4.0, increment=0.2)

        # Chuyển đổi ảnh sang định dạng byte để tải xuống
        img_byte_arr = io.BytesIO()
        img.save(img_byte_arr, format='PNG')
        img_byte_arr = img_byte_arr.getvalue()

        # Nút lưu ảnh
        btn_saved = st.download_button(
            label="Save Image",
            data=img_byte_arr,
            file_name="lung_cancer_detected.png",
            mime="image/png",
            key="save_image"
        )

        status_images=True