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


# Flag to track if subprocess commands have been executed
subprocess_executed = False

# Streamlit app title
st.title("DICOM Image Viewer")

# Upload a ZIP file containing DICOM slices
uploaded_zip_file = st.file_uploader("Upload a ZIP file containing DICOM slices", type=["zip"])
st.write("Upload a ZIP file containing DICOM slices to view the images.")

@st.cache_resource
def install_dependencies():
        command = "chmod +x install.sh"
        try:
            subprocess.run(command, shell=True, check=True)
            print("Script 'install.sh' has been made executable.")
        except subprocess.CalledProcessError as e:
            print(f"Error while making the script executable: {e}")

        command = "./install.sh"
        try:
            subprocess.run(command, shell=True, check=True)
            print("Script 'install.sh' has started running.")
            install_script_executed = True
        except subprocess.CalledProcessError as e:
            print(f"Error while running the script: {e}")

@st.cache_resource
def run_inference():
        command = "chmod +x inference.sh"
        try:
            subprocess.run(command, shell=True, check=True)
            print("Script 'inference.sh' has been made executable.")
        except subprocess.CalledProcessError as e:
            print(f"Error while making the script executable: {e}")

        command = "./inference.sh"
        try:
            subprocess.run(command, shell=True, check=True)
            print("Script 'inference.sh' has started running.")
        except subprocess.CalledProcessError as e:
            print(f"Error while running the script: {e}")

# Function to read and display the DICOM image
def convert_dicom_to_png(dicom_data, output_dir, index):
    # Convert DICOM to a PNG image
    plt.imshow(dicom_data.pixel_array, cmap=plt.cm.bone)
    plt.axis("off")
    plt.title("DICOM Image")
    plt.tight_layout()
    img_buffer = BytesIO()
    plt.savefig(img_buffer, format="png")
    img_buffer.seek(0)
    output_path = os.path.join(output_dir, f"image_{index:03d}.png")
    with open(output_path, "wb") as f:
        f.write(img_buffer.read())
    return output_path

if uploaded_zip_file is not None:
    try:
        install_dependencies()
            
        temp_dir = "/home/user/app/C2C/temp_dicom_dir"
        output_image_dir = "/home/user/app/C2C/images_for_video"
        
        os.makedirs(temp_dir, exist_ok=True)
        os.makedirs(output_image_dir, exist_ok=True)
        full_path = os.path.abspath(temp_dir)

        st.write(f"Temporary directory path: {full_path}")
        
        with zipfile.ZipFile(uploaded_zip_file, "r") as zip_ref:
            zip_ref.extractall(temp_dir)

        # Get a list of DICOM files in the directory
        dicom_files = [os.path.join(temp_dir, f) for f in os.listdir(temp_dir) if f.endswith(".dcm")]
        dicom_files.sort()  # Sort the files

        if len(dicom_files) == 0:
            st.error("No DICOM files found in the ZIP archive.")
        else:
            # Convert DICOM slices to PNG frames and save them to the output folder
            for i, dicom_file in enumerate(dicom_files):
                dicom_data = pydicom.dcmread(dicom_file)
                convert_dicom_to_png(dicom_data, output_image_dir, i)

        image_folder = output_image_dir
        fps = 20
        image_files = [
            os.path.join(image_folder, img)
            for img in Tcl().call("lsort", "-dict", os.listdir(image_folder))
            if img.endswith(".png")
        ]
        clip = moviepy.video.io.ImageSequenceClip.ImageSequenceClip(
            image_files, fps=fps
        )
        clip.write_videofile("aaa.mp4")

        st.video("aaa.mp4", format="video/mp4")
            
    except Exception as e:
        st.error(f"Error: {str(e)}")

if st.button("Analyze"):
    run_inference()

    # Assuming the 'outputs' directory is located at "/home/user/app/C2C/comp2comp/outputs"
    outputs_directory = "/home/user/app/C2C/outputs"

    # Get a list of subdirectories inside the 'outputs' directory
    subdirectories = [subdir for subdir in os.listdir(outputs_directory)
        if os.path.isdir(os.path.join(outputs_directory, subdir))]

    # Use the first (or only) subdirectory, or None if there are no subdirectories
    first_subdirectory = subdirectories[0] if subdirectories else None

    if first_subdirectory:
        # Construct the paths to the video and image files inside the first subdirectory
        subdirectory_path = os.path.join(outputs_directory, first_subdirectory)
        temp_dicom_dir_path = os.path.join(subdirectory_path, "temp_dicom_dir")

        # Get a list of subdirectories inside the 'temp_dicom_dir' directory
        dicom_subdirectories = [subdir for subdir in os.listdir(temp_dicom_dir_path)
            if os.path.isdir(os.path.join(temp_dicom_dir_path, subdir))]

        # Use the first (or only) subdirectory inside 'temp_dicom_dir', or None if there are none
        first_dicom_subdirectory = dicom_subdirectories[0] if dicom_subdirectories else None

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

            if os.path.exists(largest_slice):
                st.title("Largest Slice")
                st.image(largest_slice, use_column_width=True)

            # Display the video generated by the inference script
            if os.path.exists(video_path):
                st.title("Video")
                st.video(video_path, format="video/mp4")

            # Display the image generated by the inference script
            if os.path.exists(image_path):
                st.title("Diameter Graph")
                st.image(image_path, use_column_width=True)
        else:
            st.warning("No subdirectories or folders found inside 'temp_dicom_dir'.")
    else:
        st.warning("No subdirectories or folders found inside the 'outputs' directory.")