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.set_option('server.maxUploadSize', 300)

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

# Function to install dependencies
@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)

# Function to run inference
@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)

# Function to zip and download files
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")

# Main Streamlit code
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()

                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")

                        video_path = os.path.join(temp_dicom_dir_path, "images/summary/aaa.mp4")
                        image_path = os.path.join(temp_dicom_dir_path, "images/summary/diameter_graph.png")
                        largest_slice = os.path.join(temp_dicom_dir_path, "images/summary/out.png")

                        if os.path.exists(video_path) and os.path.exists(image_path) and os.path.exists(largest_slice):
                            zip_filename = os.path.join(temp_dir, "outputs.zip")
                            zip_and_download(temp_dicom_dir_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)
                        else:
                            st.error("Output files not found.")
                    else:
                        st.warning("No subdirectories found inside 'outputs'.")
                else:
                    st.error("The 'outputs' directory does not exist.")
        except Exception as e:
            st.error(f"Error: {str(e)}")

if __name__ == "__main__":
    main()