import streamlit as st
import pydicom
import matplotlib.pyplot as plt
import zipfile
import os
import subprocess
from datetime import datetime
import shutil
import moviepy.video.io.ImageSequenceClip
from io import BytesIO
from tkinter import Tcl
from PIL import Image

subprocess_executed = False

logo_image_path = '1.png'

st.image(logo_image_path, width=150)
st.title("Automated Abdominal Aortic Aneurysm Detection")

@st.cache_resource
def install_dependencies():
    command = "chmod +x install.sh"
    subprocess.run(command, shell=True, check=True)

    command = "./install.sh"
    subprocess.run(command, shell=True, check=True)

@st.cache_resource
def run_inference():
    command = "chmod +x inference.sh"
    subprocess.run(command, shell=True, check=True)

    command = "./inference.sh"
    subprocess.run(command, shell=True, check=True)

def zip_and_download(directory, zip_filename):
    zipf = zipfile.ZipFile(zip_filename, 'w', zipfile.ZIP_DEFLATED)
    for root, _, files in os.walk(directory):
        for file in files:
            zipf.write(os.path.join(root, file), os.path.relpath(os.path.join(root, file), directory))
    zipf.close()

    with open(zip_filename, 'rb') as f:
        bytes_data = f.read()
    st.download_button(label="Download Outputs", data=bytes_data, file_name=zip_filename, mime="application/zip")

def search_file(start_path, target_file):
    for root, _, files in os.walk(start_path):
        if target_file in files:
            return os.path.join(root, target_file)
    return None

def fetch_outputs():
    outputs_directory = "/home/user/app/C2C/outputs"

    if os.path.exists(outputs_directory):
        subdirectories = [subdir for subdir in os.listdir(outputs_directory)
                          if os.path.isdir(os.path.join(outputs_directory, subdir))]

        first_subdirectory = subdirectories[0] if subdirectories else None

        if first_subdirectory:
            subdirectory_path = os.path.join(outputs_directory, first_subdirectory)
            temp_dicom_dir_path = os.path.join(subdirectory_path, "temp_dicom_dir")

            for root, dirs, _ in os.walk(temp_dicom_dir_path):
                for subdir in dirs:
                    dicom_folder_path = os.path.join(root, subdir)
                    video_path = search_file(dicom_folder_path, "aaa.mp4")
                    if video_path:  # Use this folder if it contains the mp4 file
                        image_path = search_file(dicom_folder_path, "diameter_graph.png")
                        largest_slice = search_file(dicom_folder_path, "out.png")

                        if video_path and image_path and largest_slice:
                            zip_filename = os.path.join("/home/user/app/C2C/temp_dicom_dir", "outputs.zip")
                            zip_and_download(dicom_folder_path, zip_filename)

                            st.title("Largest Slice")
                            st.image(largest_slice, use_column_width=True)

                            st.title("Video")
                            st.video(video_path, format="video/mp4")

                            st.title("Diameter Graph")
                            st.image(image_path, use_column_width=True)
                            return
            st.error("Output files not found in any subdirectory")
        else:
            st.warning("Output files not found")
    else:
        st.error("Output files not found")

def main():
    st.write("Upload a ZIP file containing DICOM slices")
    uploaded_zip_file = st.file_uploader("Upload a .zip file", type=["zip"])

    if uploaded_zip_file is not None:
        try:
            install_dependencies()

            temp_dir = "/home/user/app/C2C/temp_dicom_dir"
            os.makedirs(temp_dir, exist_ok=True)

            with zipfile.ZipFile(uploaded_zip_file, "r") as zip_ref:
                zip_ref.extractall(temp_dir)

            st.success("Zip file uploaded successfully")

            if st.button("Analyze"):
                st.success("Analysis started (expected time: 5 mins)")
                run_inference()

            if st.button("Fetch Outputs"):
                fetch_outputs()
        except Exception as e:
            st.error(f"Error: {str(e)}")

if __name__ == "__main__":
    main()